[custom op] Bind custom kernel registry to each backend context at creation time...

Make backend context hold information about registered custom kernels
This will be later used by `CustomOperationManager`

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

diff --git a/runtimes/neurun/backend/acl_cl/Backend.h b/runtimes/neurun/backend/acl_cl/Backend.h
index 2637d54..68789e2 100644 (file)
--- a/runtimes/neurun/backend/acl_cl/Backend.h
+++ b/runtimes/neurun/backend/acl_cl/Backend.h
@@ -26,6 +26,7 @@
 #include "KernelGenerator.h"
 #include "ShapeFixer.h"
 #include "MemoryManager.h"
+#include "backend/CustomKernelRegistry.h"
 
 namespace neurun
 {
@@ -44,10 +45,11 @@ public:
   std::unique_ptr<BackendContext> newContext(const model::Operands &operands) 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)}};
+        std::make_shared<ShapeFixer>(operands, tensor_builder), kernel_registry}};
   }
 
 private:

diff --git a/runtimes/neurun/backend/acl_neon/Backend.h b/runtimes/neurun/backend/acl_neon/Backend.h
index 3a2bc90..a455db7 100644 (file)
--- a/runtimes/neurun/backend/acl_neon/Backend.h
+++ b/runtimes/neurun/backend/acl_neon/Backend.h
@@ -26,6 +26,7 @@
 #include "KernelGenerator.h"
 #include "ShapeFixer.h"
 #include "MemoryManager.h"
+#include "backend/CustomKernelRegistry.h"
 
 namespace neurun
 {
@@ -44,10 +45,11 @@ public:
   std::unique_ptr<BackendContext> newContext(const model::Operands &operands) 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)}};
+        std::make_shared<ShapeFixer>(operands, tensor_builder), kernel_registry}};
   }
 
 private:

diff --git a/runtimes/neurun/backend/cpu/Backend.h b/runtimes/neurun/backend/cpu/Backend.h
index 33970e3..bb02f0c 100644 (file)
--- a/runtimes/neurun/backend/cpu/Backend.h
+++ b/runtimes/neurun/backend/cpu/Backend.h
@@ -25,6 +25,7 @@
 #include "ConstantInitializer.h"
 #include "KernelGenerator.h"
 #include "ShapeFixer.h"
+#include "backend/CustomKernelRegistry.h"
 
 namespace neurun
 {
@@ -43,10 +44,11 @@ public:
   std::unique_ptr<BackendContext> newContext(const model::Operands &operands) 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),
-        std::make_shared<ShapeFixer>(operands, tensor_builder)}};
+        std::make_shared<KernelGenerator>(operands, tensor_builder, kernel_registry),
+        std::make_shared<ShapeFixer>(operands, tensor_builder), kernel_registry}};
   }
 
 private:

diff --git a/runtimes/neurun/backend/cpu/KernelGenerator.cc b/runtimes/neurun/backend/cpu/KernelGenerator.cc
index 3b6a25b..f809daf 100644 (file)
--- a/runtimes/neurun/backend/cpu/KernelGenerator.cc
+++ b/runtimes/neurun/backend/cpu/KernelGenerator.cc
@@ -47,8 +47,9 @@ namespace cpu
 {
 
 KernelGenerator::KernelGenerator(const neurun::model::Operands &operand_ctx,
-                                 const std::shared_ptr<TensorBuilder> &tensor_builder)
-    : _ctx(operand_ctx), _tensor_builder(tensor_builder)
+                                 const std::shared_ptr<TensorBuilder> &tensor_builder,
+                                 const std::shared_ptr<custom::KernelRegistry> &kernel_registry)
+    : _ctx(operand_ctx), _tensor_builder(tensor_builder), _kernel_registry(kernel_registry)
 {
   // DO NOTHING
 }

diff --git a/runtimes/neurun/backend/cpu/KernelGenerator.h b/runtimes/neurun/backend/cpu/KernelGenerator.h
index 4c76661..d932658 100644 (file)
--- a/runtimes/neurun/backend/cpu/KernelGenerator.h
+++ b/runtimes/neurun/backend/cpu/KernelGenerator.h
@@ -20,6 +20,7 @@
 #include "backend/IKernelGenerator.h"
 #include "model/Operands.h"
 #include "operand/Tensor.h"
+#include "backend/CustomKernelRegistry.h"
 #include "TensorBuilder.h"
 
 namespace neurun
@@ -33,7 +34,8 @@ class KernelGenerator : public IKernelGenerator
 {
 public:
   KernelGenerator(const neurun::model::Operands &ctx,
-                  const std::shared_ptr<TensorBuilder> &tensor_builder);
+                  const std::shared_ptr<TensorBuilder> &tensor_builder,
+                  const std::shared_ptr<custom::KernelRegistry> &kernel_registry);
 
   using IKernelGenerator::visit;
 
@@ -53,6 +55,7 @@ public:
 private:
   const neurun::model::Operands &_ctx;
   std::shared_ptr<TensorBuilder> _tensor_builder;
+  std::shared_ptr<custom::KernelRegistry> _kernel_registry;
   model::Layout _current_subg_layout;
 };
 

diff --git a/runtimes/neurun/core/include/backend/Backend.h b/runtimes/neurun/core/include/backend/Backend.h
index 4b0ee08..382d655 100644 (file)
--- a/runtimes/neurun/core/include/backend/Backend.h
+++ b/runtimes/neurun/core/include/backend/Backend.h
@@ -33,6 +33,11 @@ class IKernelGenerator;
 class IShapeFixer;
 struct ITensorBuilder;
 
+namespace custom
+{
+class KernelRegistry;
+}
+
 class BackendContext
 {
 public:
@@ -41,6 +46,7 @@ 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