[custom op] Build custom kernels (#6964)

* [custom op] Build custom kernels

Update cpu backend's `KernelGenerator` to build custom operation kernels

* Fix allocations in custom node kernel generator

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

2 files changed, 42 insertions, 0 deletions

diff --git a/runtimes/neurun/backend/cpu/KernelGenerator.cc b/runtimes/neurun/backend/cpu/KernelGenerator.cc
index 225428a26..8e7c2c7d4 100644
--- a/runtimes/neurun/backend/cpu/KernelGenerator.cc
+++ b/runtimes/neurun/backend/cpu/KernelGenerator.cc
@@ -409,6 +409,47 @@ void KernelGenerator::visit(const model::operation::PermuteNode &node)
   _execution_builder->append(std::move(fn));
 }
 
+void KernelGenerator::visit(const model::operation::CustomNode &node)
+{
+  auto get_type_info = [this](const model::Operand &operand) -> custom::TypeInfo {
+    auto backendShape = ::neurun::backend::cpu::kernel::getShape(operand, _current_subg_layout);
+
+    custom::Shape shape(backendShape.dimensions.size());
+    for (size_t d = 0; d < backendShape.dimensions.size(); ++d)
+    {
+      shape.dim(d) = backendShape.dimensions[d];
+    }
+
+    return {shape, backendShape.type};
+  };
+
+  auto fill_op_info = [&](const model::OperandIndexSequence &opSeq,
+                          std::vector<custom::TypeInfo> &types, std::vector<void *> &allocs) {
+    for (auto &idx : opSeq)
+    {
+      const auto &operand = _ctx.at(idx);
+      // TODO make sure using `_current_subg_layout` is correct for custom operations
+      types.emplace_back(get_type_info(operand));
+      auto in_alloc = _tensor_builder->at(idx).get();
+      allocs.emplace_back(in_alloc);
+    }
+  };
+
+  custom::Kernel::CustomKernelConfigParams params{};
+
+  fill_op_info(node.getInputs(), params.input_types, params.input_allocations);
+  fill_op_info(node.getOutputs(), params.output_types, params.output_allocations);
+
+  params.userdata = node.userdata().data;
+  params.userdata_size = node.userdata().size;
+
+  auto fn = _kernel_registry->buildKernelForOp(node.id());
+
+  fn->configure(std::move(params));
+
+  _execution_builder->append(std::move(fn));
+}
+
 } // namespace cpu
 } // namespace backend
 } // namespace neurun
diff --git a/runtimes/neurun/backend/cpu/KernelGenerator.h b/runtimes/neurun/backend/cpu/KernelGenerator.h
index d93265821..178540e3e 100644
--- a/runtimes/neurun/backend/cpu/KernelGenerator.h
+++ b/runtimes/neurun/backend/cpu/KernelGenerator.h
@@ -51,6 +51,7 @@ public:
   void visit(const model::operation::SoftmaxNode &) override;
   void visit(const model::operation::AddNode &) override;
   void visit(const model::operation::PermuteNode &) override;
+  void visit(const model::operation::CustomNode &node) override;
 
 private:
   const neurun::model::Operands &_ctx;