1 files changed, 157 insertions, 0 deletions
diff --git a/runtimes/neurun/backend/acl_cl/StageGenerator.cc b/runtimes/neurun/backend/acl_cl/StageGenerator.cc
index cc09d564f..73530aa51 100644
--- a/runtimes/neurun/backend/acl_cl/StageGenerator.cc
+++ b/runtimes/neurun/backend/acl_cl/StageGenerator.cc
@@ -2458,6 +2458,163 @@ void StageGenerator::visit(const model::operation::LogicalNotNode &node)
   });
 }
 
+void StageGenerator::visit(const model::operation::EqualNode &node)
+{
+  const auto output_index{node.getOutputs().at(0)};
+  const auto input0_index{node.getInputs().at(model::operation::EqualNode::Input::INPUT0)};
+  const auto input1_index{node.getInputs().at(model::operation::EqualNode::Input::INPUT1)};
+
+  if (!(_ctx.at(input0_index).shape() == _ctx.at(input1_index).shape()))
+  {
+    const auto broadcast_rank =
+        std::max(_ctx.at(input0_index).shape().rank(), _ctx.at(input1_index).shape().rank());
+    const_cast<neurun::model::operand::Shape &>(_ctx.at(input0_index).shape())
+        .extendRank(broadcast_rank);
+    const_cast<neurun::model::operand::Shape &>(_ctx.at(input1_index).shape())
+        .extendRank(broadcast_rank);
+  }
+
+  // Construct operation parameters
+  struct Param
+  {
+    model::operand::Index output_index;
+    model::operand::Index input0_index;
+    model::operand::Index input1_index;
+  };
+
+  Param param;
+
+  param.output_index = output_index;
+  param.input0_index = input0_index;
+  param.input1_index = input1_index;
+
+  auto tensors = _tensor_builder;
+
+  returnStage([tensors, param](IExecutionBuilder &builder) {
+    auto output_alloc = tensors->at(param.output_index).get();
+    auto input0_alloc = tensors->at(param.input0_index).get();
+    auto input1_alloc = tensors->at(param.input1_index).get();
+
+    std::unique_ptr<::arm_compute::IFunction> fn;
+
+    auto l = make_layer<::arm_compute::CLComparison>();
+
+    l->configure(input0_alloc->handle(), input1_alloc->handle(), output_alloc->handle(),
+                 ::arm_compute::ComparisonOperation::Equal);
+
+    fn = std::move(l);
+
+    auto acl_fn = make_cl_function(std::move(fn));
+
+    builder.append(std::move(acl_fn));
+  });
+}
+
+void StageGenerator::visit(const model::operation::SquaredDifferenceNode &node)
+{
+  const auto ofm_index{node.getOutputs().at(0)};
+  const auto lhs_index{node.getInputs().at(model::operation::SquaredDifferenceNode::Input::LHS)};
+  const auto rhs_index{node.getInputs().at(model::operation::SquaredDifferenceNode::Input::RHS)};
+
+  if (!(_ctx.at(lhs_index).shape() == _ctx.at(rhs_index).shape()))
+  {
+    const auto broadcast_rank =
+        std::max(_ctx.at(lhs_index).shape().rank(), _ctx.at(rhs_index).shape().rank());
+    const_cast<neurun::model::operand::Shape &>(_ctx.at(lhs_index).shape())
+        .extendRank(broadcast_rank);
+    const_cast<neurun::model::operand::Shape &>(_ctx.at(rhs_index).shape())
+        .extendRank(broadcast_rank);
+  }
+
+  // Construct operation parameters
+  struct Param
+  {
+    model::operand::Index ofm_index;
+    model::operand::Index lhs_index;
+    model::operand::Index rhs_index;
+  };
+
+  Param param;
+
+  param.ofm_index = ofm_index;
+  param.lhs_index = lhs_index;
+  param.rhs_index = rhs_index;
+
+  auto tensors = _tensor_builder;
+
+  returnStage([tensors, param](IExecutionBuilder &builder) {
+    auto ofm_alloc = tensors->at(param.ofm_index).get();
+    auto lhs_alloc = tensors->at(param.lhs_index).get();
+    auto rhs_alloc = tensors->at(param.rhs_index).get();
+
+    std::unique_ptr<::arm_compute::IFunction> fn;
+
+    auto l = make_layer<::arm_compute::CLSquaredDifference>();
+
+    l->configure(lhs_alloc->handle(), rhs_alloc->handle(), ofm_alloc->handle());
+
+    fn = std::move(l);
+
+    auto acl_fn = make_cl_function(std::move(fn));
+
+    builder.append(std::move(acl_fn));
+  });
+}
+
+void StageGenerator::visit(const model::operation::TopKV2Node &node)
+{
+  const auto outputValues_index{
+      node.getOutputs().at(model::operation::TopKV2Node::Output::OUTPUT_VALUES)};
+  const auto outputIndices_index{
+      node.getOutputs().at(model::operation::TopKV2Node::Output::OUTPUT_INDICES)};
+
+  const auto inputData_index{node.getInputs().at(model::operation::TopKV2Node::Input::INPUT)};
+  const auto k_index{node.param().k_index};
+
+  // Currently, we only support the vector input.
+  assert(_ctx.at(inputData_index).shape().rank() == 1 ||
+         _ctx.at(inputData_index).shape().rank() == 2);
+
+  const int32_t k = _ctx.at(k_index).asScalar<int32_t>();
+
+  // Construct operation parameters
+  struct Param
+  {
+    model::operand::Index outputValues_index;
+    model::operand::Index outputIndices_index;
+
+    model::operand::Index inputData_index;
+    int32_t k;
+  };
+
+  Param param;
+
+  param.outputValues_index = outputValues_index;
+  param.outputIndices_index = outputIndices_index;
+  param.inputData_index = inputData_index;
+  param.k = k;
+
+  auto tensors = _tensor_builder;
+
+  returnStage([tensors, param](IExecutionBuilder &builder) {
+    auto values_alloc = tensors->at(param.outputValues_index).get();
+    auto indices_alloc = tensors->at(param.outputIndices_index).get();
+    auto input_alloc = tensors->at(param.inputData_index).get();
+
+    std::unique_ptr<::arm_compute::IFunction> fn;
+
+    auto l = make_layer<::arm_compute::CLTopKV2>();
+
+    l->configure(input_alloc->handle(), param.k, values_alloc->handle(), indices_alloc->handle());
+
+    fn = std::move(l);
+
+    auto acl_fn = make_cl_function(std::move(fn));
+
+    builder.append(std::move(acl_fn));
+  });
+}
+
 } // namespace acl_cl
 } // namespace backend
 } // namespace neurun