[neurun] Enable Equal, SquaredDifference, TopKV2 ops (#4738)

* [neurun] Enable Equal, SquaredDifference, TopKV2 ops

This commit enables `Equal`, `SquaredDifference`, `TopKV2` ops for `acl_cl`.

Signed-off-by: sjsujinkim <sjsujin.kim@samsung.com>

* Fix build break

11 files changed, 502 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
diff --git a/runtimes/neurun/backend/acl_cl/StageGenerator.h b/runtimes/neurun/backend/acl_cl/StageGenerator.h
index 3a05104c9..5dad5c707 100644
--- a/runtimes/neurun/backend/acl_cl/StageGenerator.h
+++ b/runtimes/neurun/backend/acl_cl/StageGenerator.h
@@ -76,6 +76,9 @@ public:
   virtual void visit(const model::operation::SQRTNode &) override;
   virtual void visit(const model::operation::LogicalOrNode &) override;
   virtual void visit(const model::operation::LogicalNotNode &) override;
+  virtual void visit(const model::operation::EqualNode &) override;
+  virtual void visit(const model::operation::SquaredDifferenceNode &) override;
+  virtual void visit(const model::operation::TopKV2Node &) override;
 
 private:
   const neurun::model::operand::Set &_ctx;
diff --git a/runtimes/neurun/core/include/model/operation/EqualNode.h b/runtimes/neurun/core/include/model/operation/EqualNode.h
new file mode 100644
index 000000000..a1792f1a5
--- /dev/null
+++ b/runtimes/neurun/core/include/model/operation/EqualNode.h
@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef __NEURUN_MODEL_OPERATION_EQUAL_NODE_H__
+#define __NEURUN_MODEL_OPERATION_EQUAL_NODE_H__
+
+#include "model/operation/Node.h"
+
+namespace neurun
+{
+namespace model
+{
+namespace operation
+{
+
+class EqualNode : public model::operation::Node
+{
+public:
+  enum Input
+  {
+    INPUT0 = 0,
+    INPUT1 = 1
+  };
+
+public:
+  EqualNode(const operand::IndexSet &inputs, const operand::IndexSet &outputs);
+
+public:
+  virtual void accept(NodeVisitor &&) const override;
+  virtual std::string getName() const override { return "Equal"; }
+};
+
+} // namespace operation
+} // namespace model
+} // namespace neurun
+
+#endif // __NEURUN_MODEL_OPERATION_EQUAL_NODE_H__
diff --git a/runtimes/neurun/core/include/model/operation/Node.Include.h b/runtimes/neurun/core/include/model/operation/Node.Include.h
index 2cefc1a79..3b2350743 100644
--- a/runtimes/neurun/core/include/model/operation/Node.Include.h
+++ b/runtimes/neurun/core/include/model/operation/Node.Include.h
@@ -39,6 +39,7 @@
 #include "ExpNode.h"
 #include "ReduceMaxNode.h"
 #include "NotEqualNode.h"
+#include "EqualNode.h"
 #include "LogicalAndNode.h"
 #include "LogicalOrNode.h"
 #include "LogicalNotNode.h"
@@ -56,3 +57,5 @@
 #include "PReLUNode.h"
 #include "TransposeConvNode.h"
 #include "SQRTNode.h"
+#include "SquaredDifferenceNode.h"
+#include "TopKV2Node.h"
diff --git a/runtimes/neurun/core/include/model/operation/Op.lst b/runtimes/neurun/core/include/model/operation/Op.lst
index d26074ab9..dea9cb5ef 100644
--- a/runtimes/neurun/core/include/model/operation/Op.lst
+++ b/runtimes/neurun/core/include/model/operation/Op.lst
@@ -60,5 +60,8 @@ OP(HashtableLookupNode     , true)
 OP(PReLUNode               , true)
 OP(TransposeConvNode       , true)
 OP(SQRTNode                , true)
+OP(EqualNode               , true)
+OP(SquaredDifferenceNode   , true)
+OP(TopKV2Node              , true)
 OP(PermuteNode             , false)
 
diff --git a/runtimes/neurun/core/include/model/operation/SquaredDifferenceNode.h b/runtimes/neurun/core/include/model/operation/SquaredDifferenceNode.h
new file mode 100644
index 000000000..4f7fb6a25
--- /dev/null
+++ b/runtimes/neurun/core/include/model/operation/SquaredDifferenceNode.h
@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef __NEURUN_MODEL_OPERATION_SQUARED_DIFFERENCE_NODE_H__
+#define __NEURUN_MODEL_OPERATION_SQUARED_DIFFERENCE_NODE_H__
+
+#include "model/operation/Node.h"
+
+namespace neurun
+{
+namespace model
+{
+namespace operation
+{
+
+class SquaredDifferenceNode : public model::operation::Node
+{
+public:
+  enum Input
+  {
+    LHS = 0,
+    RHS
+  };
+
+public:
+  SquaredDifferenceNode(const operand::IndexSet &inputs, const operand::IndexSet &outputs);
+
+public:
+  virtual void accept(NodeVisitor &&) const override;
+  virtual std::string getName() const override { return "SquaredDifference"; }
+};
+
+} // namespace operation
+} // namespace model
+} // namespace neurun
+
+#endif // __NEURUN_MODEL_OPERATION_SQUARED_DIFFERENCE_NODE_H__
diff --git a/runtimes/neurun/core/include/model/operation/TopKV2Node.h b/runtimes/neurun/core/include/model/operation/TopKV2Node.h
new file mode 100644
index 000000000..eb99a51fc
--- /dev/null
+++ b/runtimes/neurun/core/include/model/operation/TopKV2Node.h
@@ -0,0 +1,68 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef __NEURUN_MODEL_OPERATION_TOPK_V2_H__
+#define __NEURUN_MODEL_OPERATION_TOPK_V2_H__
+
+#include <memory>
+
+#include "model/operation/Node.h"
+
+namespace neurun
+{
+namespace model
+{
+namespace operation
+{
+
+class TopKV2Node : public model::operation::Node
+{
+public:
+  enum Input
+  {
+    INPUT
+  };
+
+  enum Output
+  {
+    OUTPUT_VALUES = 0,
+    OUTPUT_INDICES,
+  };
+
+  struct Param
+  {
+    operand::Index k_index;
+  };
+
+public:
+  TopKV2Node(const operand::IndexSet &inputs, const operand::IndexSet &outputs, const Param &param);
+
+public:
+  virtual void accept(NodeVisitor &&) const override;
+  virtual std::string getName() const override { return "TopKV2"; }
+
+public:
+  const Param &param() const { return _param; }
+
+private:
+  Param _param;
+};
+
+} // namespace operation
+} // namespace model
+} // namespace neurun
+
+#endif // __NEURUN_MODEL_OPERATION_TOPK_V2_H__
diff --git a/runtimes/neurun/core/src/model/operation/EqualNode.cc b/runtimes/neurun/core/src/model/operation/EqualNode.cc
new file mode 100644
index 000000000..9004b03a8
--- /dev/null
+++ b/runtimes/neurun/core/src/model/operation/EqualNode.cc
@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "model/operation/EqualNode.h"
+
+#include <cassert>
+
+#include "model/operation/NodeVisitor.h"
+
+namespace neurun
+{
+namespace model
+{
+namespace operation
+{
+
+void EqualNode::accept(NodeVisitor &&v) const { v.visit(*this); }
+
+EqualNode::EqualNode(const operand::IndexSet &inputs, const operand::IndexSet &outputs)
+    : model::operation::Node{OperandConstraint::createExact(2u), inputs, outputs}
+{
+}
+
+} // namespace operation
+} // namespace model
+} // namespace neurun
diff --git a/runtimes/neurun/core/src/model/operation/SquaredDifferenceNode.cc b/runtimes/neurun/core/src/model/operation/SquaredDifferenceNode.cc
new file mode 100644
index 000000000..a039fdf0b
--- /dev/null
+++ b/runtimes/neurun/core/src/model/operation/SquaredDifferenceNode.cc
@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "model/operation/SquaredDifferenceNode.h"
+
+#include <cassert>
+
+#include "model/operation/NodeVisitor.h"
+
+namespace neurun
+{
+namespace model
+{
+namespace operation
+{
+
+void SquaredDifferenceNode::accept(NodeVisitor &&v) const { v.visit(*this); }
+
+SquaredDifferenceNode::SquaredDifferenceNode(const operand::IndexSet &inputs,
+                                             const operand::IndexSet &outputs)
+    : model::operation::Node{OperandConstraint::createExact(2u), inputs, outputs}
+{
+}
+
+} // namespace operation
+} // namespace model
+} // namespace neurun
diff --git a/runtimes/neurun/core/src/model/operation/TopKV2Node.cc b/runtimes/neurun/core/src/model/operation/TopKV2Node.cc
new file mode 100644
index 000000000..c57fe9bce
--- /dev/null
+++ b/runtimes/neurun/core/src/model/operation/TopKV2Node.cc
@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "model/operation/TopKV2Node.h"
+
+#include <cassert>
+
+#include "model/operation/NodeVisitor.h"
+
+namespace neurun
+{
+namespace model
+{
+namespace operation
+{
+
+void TopKV2Node::accept(NodeVisitor &&v) const { v.visit(*this); }
+
+TopKV2Node::TopKV2Node(const operand::IndexSet &inputs, const operand::IndexSet &outputs,
+                       const Param &param)
+    : model::operation::Node{OperandConstraint::createExact(1u), inputs, outputs}, _param{param}
+{
+}
+
+} // namespace operation
+} // namespace model
+} // namespace neurun
diff --git a/runtimes/neurun/frontend/nnapi/wrapper/OperationFactory.cc b/runtimes/neurun/frontend/nnapi/wrapper/OperationFactory.cc
index 3256f446e..7846fa89c 100644
--- a/runtimes/neurun/frontend/nnapi/wrapper/OperationFactory.cc
+++ b/runtimes/neurun/frontend/nnapi/wrapper/OperationFactory.cc
@@ -909,6 +909,55 @@ OperationFactory::OperationFactory()
 
     return new operation::LogicalNotNode{inputs, outputs};
   };
+
+  _map[ANEURALNETWORKS_EQUAL_EX] = [](const OperationFactory::Param &init_param) {
+    assert(init_param.input_count == 2 && init_param.output_count == 1);
+
+    operand::IndexSet outputs{init_param.outputs[0]};
+
+    // Each input should be interpreted as follows:
+    //
+    //  0 -> input0 Tensor Index
+    //  1 -> input1 Tensor Index
+    operand::IndexSet inputs{init_param.inputs[0], init_param.inputs[1]};
+
+    return new operation::EqualNode{inputs, outputs};
+  };
+
+  _map[ANEURALNETWORKS_SQUARED_DIFFERENCE_EX] = [](const OperationFactory::Param &init_param) {
+    assert(init_param.input_count == 2 && init_param.output_count == 1);
+
+    operand::IndexSet outputs{init_param.outputs[0]};
+
+    // Each input should be interpreted as follows:
+    //
+    //  0 -> LHS Tensor Index
+    //  1 -> RHS Tensor Index
+    operand::IndexSet inputs{init_param.inputs[0], init_param.inputs[1]};
+
+    return new operation::SquaredDifferenceNode{inputs, outputs};
+  };
+
+  _map[ANEURALNETWORKS_TOPK_V2_EX] = [](const OperationFactory::Param &init_param) {
+    assert(init_param.input_count == 2 && init_param.output_count == 2);
+
+    // Each output should be interpreted as follows:
+    //
+    //  0 -> Index for Output Values
+    //  1 -> Index for Output Indices
+    operand::IndexSet outputs{init_param.outputs[0], init_param.outputs[1]};
+
+    // Each input should be interpreted as follows:
+    //
+    //  0 -> Index for Input Data
+    //  1 -> Index for K
+    operand::IndexSet inputs{init_param.inputs[0]};
+
+    operation::TopKV2Node::Param param;
+    param.k_index = operand::Index{init_param.inputs[1]};
+
+    return new operation::TopKV2Node{inputs, outputs, param};
+  };
 }
 
 neurun::model::operation::Node *OperationFactory::create(ANeuralNetworksOperationType type,