[custom op] Update kernel registry binding (#7141)

* [custom op] Update kernel registry binding

Split kernel registry creation from backend context creation

Signed-off-by: Vladimir Plazun <v.plazun@samsung.com>
* format fix

* fix unused parameter compile error

diff --git a/runtimes/neurun/backend/acl_cl/Backend.h b/runtimes/neurun/backend/acl_cl/Backend.h
index 68789e2..19536d3 100644 (file)
--- a/runtimes/neurun/backend/acl_cl/Backend.h
+++ b/runtimes/neurun/backend/acl_cl/Backend.h
@@ -42,14 +42,15 @@ public:
 
   std::shared_ptr<IConfig> config() const override { return _config; }
 
-  std::unique_ptr<BackendContext> newContext(const model::Operands &operands) const override
+  std::unique_ptr<BackendContext>
+  newContext(const model::Operands &operands,
+             const std::shared_ptr<custom::KernelRegistry> &) const override
   {
     auto tensor_builder = std::make_shared<TensorBuilder>(createMemoryManager());
-    auto kernel_registry = std::make_shared<custom::KernelRegistry>();
     return std::unique_ptr<BackendContext>{new BackendContext{
         this, tensor_builder, std::make_shared<ConstantInitializer>(operands, tensor_builder),
         std::make_shared<KernelGenerator>(operands, tensor_builder),
-        std::make_shared<ShapeFixer>(operands, tensor_builder), kernel_registry}};
+        std::make_shared<ShapeFixer>(operands, tensor_builder)}};
   }
 
 private:

diff --git a/runtimes/neurun/backend/acl_neon/Backend.h b/runtimes/neurun/backend/acl_neon/Backend.h
index a455db7..6cac3f3 100644 (file)
--- a/runtimes/neurun/backend/acl_neon/Backend.h
+++ b/runtimes/neurun/backend/acl_neon/Backend.h
@@ -42,14 +42,15 @@ public:
 
   std::shared_ptr<IConfig> config() const override { return _config; }
 
-  std::unique_ptr<BackendContext> newContext(const model::Operands &operands) const override
+  std::unique_ptr<BackendContext>
+  newContext(const model::Operands &operands,
+             const std::shared_ptr<custom::KernelRegistry> &) const override
   {
     auto tensor_builder = std::make_shared<TensorBuilder>(createMemoryManager());
-    auto kernel_registry = std::make_shared<custom::KernelRegistry>();
     return std::unique_ptr<BackendContext>{new BackendContext{
         this, tensor_builder, std::make_shared<ConstantInitializer>(operands, tensor_builder),
         std::make_shared<KernelGenerator>(operands, tensor_builder),
-        std::make_shared<ShapeFixer>(operands, tensor_builder), kernel_registry}};
+        std::make_shared<ShapeFixer>(operands, tensor_builder)}};
   }
 
 private:

diff --git a/runtimes/neurun/backend/cpu/Backend.h b/runtimes/neurun/backend/cpu/Backend.h
index bb02f0c..94624ee 100644 (file)
--- a/runtimes/neurun/backend/cpu/Backend.h
+++ b/runtimes/neurun/backend/cpu/Backend.h
@@ -41,14 +41,15 @@ public:
 
   std::shared_ptr<IConfig> config() const override { return _config; }
 
-  std::unique_ptr<BackendContext> newContext(const model::Operands &operands) const override
+  std::unique_ptr<BackendContext>
+  newContext(const model::Operands &operands,
+             const std::shared_ptr<custom::KernelRegistry> &registry) const override
   {
     auto tensor_builder = std::make_shared<TensorBuilder>();
-    auto kernel_registry = std::make_shared<custom::KernelRegistry>();
     return std::unique_ptr<BackendContext>{new BackendContext{
         this, tensor_builder, std::make_shared<ConstantInitializer>(operands, tensor_builder),
-        std::make_shared<KernelGenerator>(operands, tensor_builder, kernel_registry),
-        std::make_shared<ShapeFixer>(operands, tensor_builder), kernel_registry}};
+        std::make_shared<KernelGenerator>(operands, tensor_builder, registry),
+        std::make_shared<ShapeFixer>(operands, tensor_builder)}};
   }
 
 private:

diff --git a/runtimes/neurun/core/include/backend/Backend.h b/runtimes/neurun/core/include/backend/Backend.h
index 382d655..e8bfac2 100644 (file)
--- a/runtimes/neurun/core/include/backend/Backend.h
+++ b/runtimes/neurun/core/include/backend/Backend.h
@@ -46,7 +46,6 @@ public:
   std::shared_ptr<IConstantInitializer> constant_initializer;
   std::shared_ptr<IKernelGenerator> kernel_gen;
   std::shared_ptr<IShapeFixer> shape_fixer;
-  std::shared_ptr<custom::KernelRegistry> _custom_kernel_registry;
 };
 
 class Backend
@@ -54,7 +53,10 @@ class Backend
 public:
   virtual ~Backend() = default;
   virtual std::shared_ptr<neurun::backend::IConfig> config() const = 0;
-  virtual std::unique_ptr<BackendContext> newContext(const model::Operands &operands) const = 0;
+
+  virtual std::unique_ptr<BackendContext>
+  newContext(const model::Operands &operands,
+             const std::shared_ptr<custom::KernelRegistry> &registry) const = 0;
 };
 
 } // namespace backend

diff --git a/runtimes/neurun/core/include/graph/Graph.h b/runtimes/neurun/core/include/graph/Graph.h
index 13b29b9..b3e6d54 100644 (file)
--- a/runtimes/neurun/core/include/graph/Graph.h
+++ b/runtimes/neurun/core/include/graph/Graph.h
@@ -54,6 +54,17 @@ class BackendResolver;
 
 namespace neurun
 {
+namespace backend
+{
+namespace custom
+{
+class KernelRegistry;
+} // namespace neurun
+} // namespace backend
+} // namespace neurun
+
+namespace neurun
+{
 namespace graph
 {
 
@@ -131,6 +142,21 @@ public:
 private:
   void initializeUseDef();
 
+  // Custom operations support
+public:
+  void bindKernelRegistry(const std::shared_ptr<backend::custom::KernelRegistry> &registry)
+  {
+    _kernel_registry = registry;
+  }
+
+  const std::shared_ptr<backend::custom::KernelRegistry> &getKernelRegistry() const
+  {
+    return _kernel_registry;
+  }
+
+private:
+  std::shared_ptr<backend::custom::KernelRegistry> _kernel_registry;
+
   // Accessors
 public:
   const model::OperandIndexSequence &getInputs() const { return _model->inputs; }

diff --git a/runtimes/neurun/core/src/compiler/BackendResolver.h b/runtimes/neurun/core/src/compiler/BackendResolver.h
index c152d78..248ef2f 100644 (file)
--- a/runtimes/neurun/core/src/compiler/BackendResolver.h
+++ b/runtimes/neurun/core/src/compiler/BackendResolver.h
@@ -35,11 +35,12 @@ class BackendResolver
 {
 public:
   BackendResolver(const model::Operands &operands,
-                  const std::vector<const backend::Backend *> &backends)
+                  const std::vector<const backend::Backend *> &backends,
+                  const std::shared_ptr<backend::custom::KernelRegistry> &registry)
   {
     for (const auto backend : backends)
     {
-      _context_manager.emplace(backend, backend->newContext(operands));
+      _context_manager.emplace(backend, backend->newContext(operands, registry));
     }
   }
 

diff --git a/runtimes/neurun/core/src/compiler/Compiler.cc b/runtimes/neurun/core/src/compiler/Compiler.cc
index 467ae81..8b07dd1 100644 (file)
--- a/runtimes/neurun/core/src/compiler/Compiler.cc
+++ b/runtimes/neurun/core/src/compiler/Compiler.cc
@@ -56,7 +56,8 @@ void Compiler::compile(void)
   if (util::getConfigBool(util::config::USE_SCHEDULER))
   {
     auto scheduler =
-        compiler::Scheduler(_graph->operands(), backend::BackendManager::instance().getAll());
+        compiler::Scheduler(_graph->operands(), backend::BackendManager::instance().getAll(),
+                            _graph->getKernelRegistry());
     br = scheduler.schedule(*_graph);
     indexed_ranks = scheduler.getIndexedRanks();
   }

diff --git a/runtimes/neurun/core/src/compiler/ManualScheduler.cc b/runtimes/neurun/core/src/compiler/ManualScheduler.cc
index 7dd1491..c7de7b2 100644 (file)
--- a/runtimes/neurun/core/src/compiler/ManualScheduler.cc
+++ b/runtimes/neurun/core/src/compiler/ManualScheduler.cc
@@ -30,7 +30,7 @@ namespace compiler
 std::unique_ptr<BackendResolver> ManualScheduler::schedule(const graph::Graph &graph)
 {
   auto backend_resolver = nnfw::cpp14::make_unique<compiler::BackendResolver>(
-      graph.operands(), backend::BackendManager::instance().getAll());
+      graph.operands(), backend::BackendManager::instance().getAll(), graph.getKernelRegistry());
 
   // 1. Backend for All operations
   const auto backend_all_str = util::getConfigString(util::config::OP_BACKEND_ALLOPS);

diff --git a/runtimes/neurun/core/src/compiler/Scheduler.h b/runtimes/neurun/core/src/compiler/Scheduler.h
index 80d6eaf..130c81d 100644 (file)
--- a/runtimes/neurun/core/src/compiler/Scheduler.h
+++ b/runtimes/neurun/core/src/compiler/Scheduler.h
@@ -47,13 +47,14 @@ public:
    * @param[in] model Graph model
    * @param[in] backend_resolver backend resolver
    */
-  Scheduler(const neurun::model::Operands &operands, std::vector<const backend::Backend *> backends)
+  Scheduler(const neurun::model::Operands &operands, std::vector<const backend::Backend *> backends,
+            const std::shared_ptr<backend::custom::KernelRegistry> &registry)
       : _is_supported{}, _backends_avail_time{}, _ops_eft{},
         _op_to_rank{std::make_shared<model::OperationIndexMap<int64_t>>()},
         _all_backends(std::move(backends))
   {
     _backend_resolver =
-        nnfw::cpp14::make_unique<compiler::BackendResolver>(operands, _all_backends);
+        nnfw::cpp14::make_unique<compiler::BackendResolver>(operands, _all_backends, registry);
     _exec_time = nnfw::cpp14::make_unique<backend::ExecTime>(_all_backends);
 
     // Find cpu backend

diff --git a/runtimes/neurun/test/core/backend/ExecTime.test.cc b/runtimes/neurun/test/core/backend/ExecTime.test.cc
index ff67f65..0409c07 100644 (file)
--- a/runtimes/neurun/test/core/backend/ExecTime.test.cc
+++ b/runtimes/neurun/test/core/backend/ExecTime.test.cc
@@ -38,7 +38,9 @@ struct MockBackend : public ::neurun::backend::Backend
   {
     return std::make_shared<MockConfig>();
   }
-  std::unique_ptr<BackendContext> newContext(const model::Operands &) const override
+  std::unique_ptr<BackendContext>
+  newContext(const model::Operands &,
+             const std::shared_ptr<backend::custom::KernelRegistry> &) const override
   {
     return nullptr;
   }

diff --git a/runtimes/neurun/test/core/compiler/Scheduler.cc b/runtimes/neurun/test/core/compiler/Scheduler.cc
index 9fab9d7..44431b7 100644 (file)
--- a/runtimes/neurun/test/core/compiler/Scheduler.cc
+++ b/runtimes/neurun/test/core/compiler/Scheduler.cc
@@ -51,7 +51,9 @@ struct MockConfigCPU : public IConfig
 struct MockBackendCPU : public Backend
 {
   std::shared_ptr<IConfig> config() const override { return std::make_shared<MockConfigCPU>(); }
-  std::unique_ptr<BackendContext> newContext(const Operands &) const override
+  std::unique_ptr<BackendContext>
+  newContext(const Operands &,
+             const std::shared_ptr<backend::custom::KernelRegistry> &) const override
   {
     return std::unique_ptr<BackendContext>(
         new BackendContext{this, nullptr, nullptr, nullptr, nullptr});
@@ -68,7 +70,9 @@ struct MockConfigGPU : public IConfig
 struct MockBackendGPU : public Backend
 {
   std::shared_ptr<IConfig> config() const override { return std::make_shared<MockConfigGPU>(); }
-  std::unique_ptr<BackendContext> newContext(const Operands &) const override
+  std::unique_ptr<BackendContext>
+  newContext(const Operands &,
+             const std::shared_ptr<backend::custom::KernelRegistry> &) const override
   {
     return std::unique_ptr<BackendContext>(
         new BackendContext{this, nullptr, nullptr, nullptr, nullptr});
@@ -85,7 +89,9 @@ struct MockConfigNPU : public IConfig
 struct MockBackendNPU : public Backend
 {
   std::shared_ptr<IConfig> config() const override { return std::make_shared<MockConfigNPU>(); }
-  std::unique_ptr<BackendContext> newContext(const Operands &) const override
+  std::unique_ptr<BackendContext>
+  newContext(const Operands &,
+             const std::shared_ptr<backend::custom::KernelRegistry> &) const override
   {
     return std::unique_ptr<BackendContext>(
         new BackendContext{this, nullptr, nullptr, nullptr, nullptr});
@@ -349,7 +355,7 @@ TEST_P(SchedulerTestWithExecutorParam, straight_graph_known_exec_time)
     et.uploadOperationsExecTime();
 
     // Test scheduler
-    auto scheduler = compiler::Scheduler(graph->operands(), _mock_backends);
+    auto scheduler = compiler::Scheduler(graph->operands(), _mock_backends, nullptr);
     const auto br = scheduler.schedule(*graph);
     ASSERT_EQ(br->getBackend(add_op_idx)->config()->id(), "cpu");
     ASSERT_EQ(br->getBackend(sub_op_idx)->config()->id(), "gpu");
@@ -363,7 +369,7 @@ TEST_P(SchedulerTestWithExecutorParam, straight_graph_known_exec_time)
     setPermutationsExecutionTime(_mock_backends, OPERAND_SIZE, 1e5);
 
     // Test scheduler
-    auto scheduler = compiler::Scheduler(graph->operands(), _mock_backends);
+    auto scheduler = compiler::Scheduler(graph->operands(), _mock_backends, nullptr);
     const auto br = scheduler.schedule(*graph);
     ASSERT_EQ(br->getBackend(add_op_idx)->config()->id(), "cpu");
     ASSERT_EQ(br->getBackend(sub_op_idx)->config()->id(), "cpu");
@@ -402,7 +408,7 @@ TEST_P(SchedulerTestWithExecutorParam, branched_graph_known_exec_time)
     et.uploadOperationsExecTime();
 
     // Test scheduler
-    auto scheduler = compiler::Scheduler(graph->operands(), _mock_backends);
+    auto scheduler = compiler::Scheduler(graph->operands(), _mock_backends, nullptr);
     const auto br = scheduler.schedule(*graph);
 
     std::string branch1_expected_backend("npu"), branch2_expected_backend("npu");
@@ -437,7 +443,7 @@ TEST_P(SchedulerTestWithExecutorParam, branched_graph_known_exec_time)
     et.uploadOperationsExecTime();
 
     // Test scheduler
-    auto scheduler = compiler::Scheduler(graph->operands(), _mock_backends);
+    auto scheduler = compiler::Scheduler(graph->operands(), _mock_backends, nullptr);
     const auto br = scheduler.schedule(*graph);
     ASSERT_EQ(br->getBackend(add_op_idx)->config()->id(), "npu");
     ASSERT_EQ(br->getBackend(mul1_op_idx)->config()->id(), "npu");