[nnkit] Introduce nni toolchain (#366)

This commit introduces 'nni' toolchain which loads backend/action
modules, and executes them.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>

diff --git a/contrib/nnkit/tools/nni/CMakeLists.txt b/contrib/nnkit/tools/nni/CMakeLists.txt
index 800a422..73edd02 100644 (file)
--- a/contrib/nnkit/tools/nni/CMakeLists.txt
+++ b/contrib/nnkit/tools/nni/CMakeLists.txt
@@ -1,3 +1,6 @@
 file(GLOB_RECURSE SOURCES "*.cpp")
 
 add_executable(nni ${SOURCES})
+target_link_libraries(nni nnkit_intf_action)
+target_link_libraries(nni nnkit_intf_backend)
+target_link_libraries(nni dl)

diff --git a/contrib/nnkit/tools/nni/nni.cpp b/contrib/nnkit/tools/nni/nni.cpp
index aee9e2e..5bd2a73 100644 (file)
--- a/contrib/nnkit/tools/nni/nni.cpp
+++ b/contrib/nnkit/tools/nni/nni.cpp
@@ -1,4 +1,340 @@
+#include <nnkit/CmdlineArguments.h>
+
+#include <string>
+#include <vector>
+
+// TODO Extract this helper class
+class VectorArguments : public nnkit::CmdlineArguments
+{
+public:
+  uint32_t size(void) const { return _args.size(); }
+
+public:
+  const char *at(uint32_t nth) const { return _args.at(nth).c_str(); }
+
+public:
+  VectorArguments &append(const std::string &arg)
+  {
+    _args.emplace_back(arg);
+    return (*this);
+  }
+
+private:
+  std::vector<std::string> _args;
+};
+
+namespace
+{
+
+class Section
+{
+public:
+  Section() = default;
+
+public:
+  const nnkit::CmdlineArguments &args(void) const { return _args; }
+
+public:
+  void append(const std::string &arg) { _args.append(arg); }
+
+private:
+  VectorArguments _args;
+};
+
+}
+
+// TODO Extract Backend-related helpers
+#include <nnkit/Backend.h>
+
+#include <memory>
+
+#include <dlfcn.h>
+#include <assert.h>
+
+namespace
+{
+
+class BackendBinder
+{
+private:
+  typedef std::unique_ptr<nnkit::Backend> (*Entry)(const nnkit::CmdlineArguments &);
+
+public:
+  BackendBinder(const std::string &path)
+  {
+    // NOTE Some backend (such as tflite) needs multithreading support (std::thread).
+    //
+    //      std::thread in libstdc++.so includes weak symbols for pthread_XXX functions,
+    //      and these weak symbols should be overridden by strong symbols in libpthread.so.
+    //      If not, std::thread will not work correctly.
+    //
+    //      RTLD_GLOBAL flag is necessary to allow weak symbols to be overridden.
+    _handle = dlopen(path.c_str(), RTLD_LAZY | RTLD_GLOBAL);
+    assert(_handle != nullptr);
+
+    _entry = reinterpret_cast<Entry>(dlsym(_handle, "make_backend"));
+    assert(_entry != nullptr);
+  }
+
+public:
+  // Copy is not allowed to avoid double close
+  BackendBinder(const BackendBinder &) = delete;
+  BackendBinder(BackendBinder &&binder)
+  {
+    // Handle is transferd from 'binder' instance into this instance.
+    _handle = binder._handle;
+    _entry = binder._entry;
+
+    binder._handle = nullptr;
+    binder._entry = nullptr;
+  }
+
+public:
+  ~BackendBinder()
+  {
+    if (_handle != nullptr)
+    {
+      dlclose(_handle);
+    }
+  }
+
+public:
+  std::unique_ptr<nnkit::Backend> make(const nnkit::CmdlineArguments &args) const
+  {
+    return _entry(args);
+  }
+
+private:
+  void *_handle;
+  Entry _entry;
+};
+
+}
+
+namespace
+{
+
+class BackendSection : public Section
+{
+public:
+  BackendSection(const std::string &path) : _binder{path}
+  {
+    // DO NOTHING
+  }
+
+public:
+  std::unique_ptr<nnkit::Backend> initialize(void) const
+  {
+    return _binder.make(args());
+  }
+
+private:
+  BackendBinder _binder;
+};
+
+}
+
+// TODO Extract Action-related helpers
+#include <nnkit/Action.h>
+
+#include <memory>
+
+#include <dlfcn.h>
+#include <assert.h>
+
+namespace
+{
+
+class ActionBinder
+{
+private:
+  typedef std::unique_ptr<nnkit::Action> (*Entry)(const nnkit::CmdlineArguments &);
+
+public:
+  ActionBinder(const std::string &path)
+  {
+    // Q: Do we need RTLD_GLOBAL here?
+    _handle = dlopen(path.c_str(), RTLD_LAZY);
+    assert(_handle != nullptr);
+
+    _entry = reinterpret_cast<Entry>(dlsym(_handle, "make_action"));
+    assert(_entry != nullptr);
+  }
+
+public:
+  // Copy is not allowed to avoid double close
+  ActionBinder(const ActionBinder &) = delete;
+  ActionBinder(ActionBinder &&binder)
+  {
+    // Handle is transferd from 'binder' instance into this instance.
+    _handle = binder._handle;
+    _entry = binder._entry;
+
+    binder._handle = nullptr;
+    binder._entry = nullptr;
+  }
+
+public:
+  ~ActionBinder()
+  {
+    if (_handle)
+    {
+      dlclose(_handle);
+    }
+  }
+
+public:
+  std::unique_ptr<nnkit::Action> make(const nnkit::CmdlineArguments &args) const
+  {
+    return _entry(args);
+  }
+
+private:
+  void *_handle;
+  Entry _entry;
+};
+
+}
+
+namespace
+{
+
+class ActionSection : public Section
+{
+public:
+  ActionSection(const std::string &path) : _binder{path}
+  {
+    // DO NOTHING
+  }
+
+public:
+  std::unique_ptr<nnkit::Action> initialize(void) const
+  {
+    return _binder.make(args());
+  }
+
+private:
+  ActionBinder _binder;
+};
+
+}
+
+#include <nncc/foundation/Memory.h>
+
+#include <map>
+#include <iostream>
+
 int main(int argc, char **argv)
 {
+  // Usage:
+  //  [Command] --backend [Backend module path] --backend-arg ... --backend-arg ...
+  //            --pre [Action module path] --action-arg ... --action-arg ...
+  //            --post [Action module path] --action-arg ... --action-arg ...
+  // TODO Show usage message on invalid arguments
+
+  // 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 = nncc::foundation::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);
+  };
+
+  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);
+  }
+  
+  // 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;
 }