nnkit-run.cpp File Reference

#include <nnkit/CmdlineArguments.h>
#include <nnkit/VectorArguments.h>
#include <nnkit/BackendPlugin.h>
#include <nnkit/Action.h>
#include <memory>
#include <dlfcn.h>
#include <assert.h>
#include <map>
#include <iostream>

Go to the source code of this file.

Functions
int	main (int argc, char **argv)

Function Documentation

main()

int main	(	int	argc,
		char **	argv
	)

Definition at line 148 of file nnkit-run.cpp.

{
  // Usage:
  //  [Command] --backend [Backend module path] --backend-arg ... --backend-arg ...
  //            --pre [Action module path] --pre-arg ... --pre-arg ...
  //            --post [Action module path] --post-arg ... --post-arg ...
 
  // Argument sections
  //
  // NOTE Command-line arguments should include one backend section, and may include multiple
  //      pre/post action sections.
  struct Sections
  {
    std::unique_ptr<BackendSection> backend;
    std::vector<ActionSection> pre;
    std::vector<ActionSection> post;
  };
 
  Sections sections;
 
  // Simple argument parser (based on map)
  std::map<std::string, std::function<void(const std::string &arg)>> argparse;
 
  argparse["--backend"] = [&sections](const std::string &tag) {
    sections.backend = std::make_unique<BackendSection>(tag);
  };
 
  argparse["--backend-arg"] = [&sections](const std::string &arg) {
    sections.backend->append(arg);
  };
 
  argparse["--pre"] = [&sections](const std::string &tag) { sections.pre.emplace_back(tag); };
 
  argparse["--pre-arg"] = [&sections](const std::string &arg) { sections.pre.back().append(arg); };
 
  argparse["--post"] = [&sections](const std::string &tag) { sections.post.emplace_back(tag); };
 
  argparse["--post-arg"] = [&sections](const std::string &arg) {
    sections.post.back().append(arg);
  };
 
  if (argc < 2)
  {
    std::cerr << "Usage:" << std::endl
              << "[Command] --backend [Backend module path] "
              << "--backend-arg [Backend argument] ..." << std::endl
              << "          --pre [Pre-Action module path] "
              << "--pre-arg [Pre-Action argument] ..." << std::endl
              << "          --post [Post-Action module path] "
              << "--post-arg [Post-Action argument] ..." << std::endl;
    return 255;
  }
 
  for (int n = 1; n < argc; n += 2)
  {
    const std::string tag{argv[n]};
    const std::string arg{argv[n + 1]};
 
    auto it = argparse.find(tag);
 
    if (it == argparse.end())
    {
      std::cerr << "Option '" << tag << "' is not supported" << std::endl;
      return 255;
    }
 
    it->second(arg);
  }
 
  // we need a backend
  if (sections.backend == nullptr)
  {
    std::cerr << "Error: Backend is required. Provide with [--backend]" << std::endl;
    return 255;
  }
 
  // Initialize a backend
  auto backend = sections.backend->initialize();
 
  // Initialize pre actions
  std::vector<std::unique_ptr<nnkit::Action>> pre_actions;
 
  for (const auto &section : sections.pre)
  {
    pre_actions.emplace_back(section.initialize());
  }
 
  // Initialize post actions
  std::vector<std::unique_ptr<nnkit::Action>> post_actions;
 
  for (const auto &section : sections.post)
  {
    post_actions.emplace_back(section.initialize());
  }
 
  //
  // Run inference
  //
  backend->prepare([&pre_actions](nnkit::TensorContext &ctx) {
    // Run pre-actions on prepared tensor context
    for (auto &action : pre_actions)
    {
      action->run(ctx);
    }
  });
 
  backend->run();
 
  backend->teardown([&post_actions](nnkit::TensorContext &ctx) {
    // Run post-actions before teardown
    for (auto &action : post_actions)
    {
      action->run(ctx);
    }
  });
 
  return 0;
}