Introduce reduce mean operation for pure acl runtime (#2256)

* Rebase on master instead of merging

This commit does rebase on master instead of merging.

Signed-off-by: SungJin Choi <lotieye.choi@samsung.com>
* Fix and modify a typo

This commit fixes a typo
- axises => axis

Signed-off-by: SungJin Choi <lotieye.choi@samsung.com>

diff --git a/runtimes/pure_arm_compute/src/compilation.cc b/runtimes/pure_arm_compute/src/compilation.cc
index 37ee38d..a3c9354 100644 (file)
--- a/runtimes/pure_arm_compute/src/compilation.cc
+++ b/runtimes/pure_arm_compute/src/compilation.cc
@@ -22,6 +22,7 @@
 #include <arm_compute/runtime/CL/functions/CLCast.h>
 #include <arm_compute/runtime/CL/functions/CLDepthwiseConvolutionLayer.h>
 #include <arm_compute/runtime/CL/functions/CLDequantizationLayer.h>
+#include <arm_compute/runtime/CL/functions/CLReductionMean.h>
 
 #include "internal/arm_compute/Cast.h"
 #include "internal/arm_compute/kernel/View.h"
@@ -368,6 +369,7 @@ public:
   void visit(const ::internal::tflite::op::ReLU6::Node &node) override;
   void visit(const ::internal::tflite::op::Tanh::Node &node) override;
   void visit(const ::internal::tflite::op::Logistic::Node &node) override;
+  void visit(const ::internal::tflite::op::Mean::Node &node) override;
 
 private:
   const ::internal::tflite::operand::Set &_ctx;
@@ -2612,6 +2614,93 @@ void Planner::visit(const ::internal::tflite::op::Logistic::Node &node)
   _builder.addStage(stage);
 }
 
+// Reduce Mean
+void Planner::visit(const ::internal::tflite::op::Mean::Node &node)
+{
+  VERBOSE(Mean) << "Configure Mean operation" << std::endl;
+
+  const ::internal::tflite::operand::Index ofm_index{node.param().ofm_index};
+  const ::internal::tflite::operand::Index ifm_index{node.param().ifm_index};
+  const ::internal::tflite::operand::Index axis_index{node.param().axis_index};
+  const ::internal::tflite::operand::Index keep_dims_index{node.param().keep_dims_index};
+  const int keep_dims = _ctx.at(keep_dims_index).asScalar<int>();
+
+  // Set shape constraints
+  _builder.addShapeConstr(
+      ofm_index, asTensorInfo(_ctx.at(ofm_index).shape().asTensor(), _ctx.at(ofm_index).type()));
+  _builder.addShapeConstr(
+      ifm_index, asTensorInfo(_ctx.at(ifm_index).shape().asTensor(), _ctx.at(ifm_index).type()));
+  _builder.addShapeConstr(
+      axis_index, asTensorInfo(_ctx.at(axis_index).shape().asVector(), _ctx.at(axis_index).type()));
+
+  // TODO keep_dims==0
+  assert(keep_dims != 0);
+
+  // Set axis
+  // TODO Other axis (Axis for width and height are currently supported.)
+  // TODO Other ranks (Rank 4 is currently supported.)
+  assert(_ctx.at(ifm_index).shape().rank() == 4);
+
+  std::vector<uint32_t> axis;
+  {
+    const auto axis_base = _ctx.at(axis_index).data().base();
+    const auto axis_type = _ctx.at(axis_index).type();
+    const auto axis_size = _ctx.at(axis_index).shape().asVector();
+
+    // NHWC type -> WHCN type
+    if (_ctx.at(ofm_index).shape().asTensor().rank() == 4)
+    {
+      for (uint32_t n = 0; n < axis_size; ++n)
+      {
+        const ::arm_compute::Coordinates coordinate{n};
+        const int32_t *from = reinterpret_cast<const int32_t *>(axis_base) + n;
+        if (*from == 1)
+        {
+          axis.push_back(1); // h
+        }
+        else if (*from == 2)
+        {
+          axis.push_back(0); // w
+        }
+        else if (*from < 0)
+        {
+          // Nothing to do
+        }
+        else
+        {
+          throw std::runtime_error{"Not supported axis"};
+        }
+      }
+    }
+  }
+
+  struct Param
+  {
+    int ofm_index;
+    int ifm_index;
+    std::vector<uint32_t> axis;
+  };
+
+  Param param;
+
+  param.ofm_index = ofm_index.asInt();
+  param.ifm_index = ifm_index.asInt();
+  param.axis = axis;
+
+  auto stage = [param](const IAllocationContext &ctx, IExecutionBuilder &builder) {
+    auto ofm_alloc = ctx.at(::internal::tflite::operand::Index{param.ofm_index});
+    auto ifm_alloc = ctx.at(::internal::tflite::operand::Index{param.ifm_index});
+
+    auto fn = nnfw::make_unique<::arm_compute::CLReductionMean>();
+
+    fn->configure(ifm_alloc, ofm_alloc, param.axis);
+
+    builder.append("Mean", std::move(fn));
+  };
+
+  _builder.addStage(stage);
+}
+
 class AllocationContext final : public IAllocationContext
 {
 public:

diff --git a/runtimes/pure_arm_compute/src/internal/op/Mean.cc b/runtimes/pure_arm_compute/src/internal/op/Mean.cc
new file mode 100644 (file)
index 0000000..76c5f8f
--- /dev/null
+++ b/runtimes/pure_arm_compute/src/internal/op/Mean.cc
@@ -0,0 +1,51 @@
+#include "internal/op/Mean.h"
+#include "internal/op/NodeVisitor.h"
+
+#include <cassert>
+
+namespace internal
+{
+namespace tflite
+{
+namespace op
+{
+namespace Mean
+{
+
+void Node::accept(NodeVisitor &&v) const { v.visit(*this); }
+
+} // namespace Mean
+} // namespace op
+} // namespace tflite
+} // namespace internal
+
+namespace internal
+{
+namespace tflite
+{
+namespace op
+{
+namespace Mean
+{
+
+Param::Param(uint32_t inputCount, const uint32_t *inputs, uint32_t outputCount,
+             const uint32_t *outputs)
+{
+  assert(inputCount == 3 && outputCount == 1);
+
+  ofm_index = outputs[0];
+
+  // Each input should be interpreted as follows:
+  //
+  //  0 -> ifm Tensor Index
+  //  1 -> axis Tensor Index
+  //  2 -> keep_dims Index
+  ifm_index = inputs[0];
+  axis_index = inputs[1];
+  keep_dims_index = inputs[2];
+}
+
+} // namespace Mean
+} // namespace op
+} // namespace tflite
+} // namespace internal

diff --git a/runtimes/pure_arm_compute/src/internal/op/Mean.h b/runtimes/pure_arm_compute/src/internal/op/Mean.h
new file mode 100644 (file)
index 0000000..af3f125
--- /dev/null
+++ b/runtimes/pure_arm_compute/src/internal/op/Mean.h
@@ -0,0 +1,55 @@
+#ifndef __INTERNAL_OP_MEAN_H__
+#define __INTERNAL_OP_MEAN_H__
+
+#include "internal/op/Node.h"
+
+#include <cstdint>
+
+namespace internal
+{
+namespace tflite
+{
+namespace op
+{
+namespace Mean
+{
+
+struct Param
+{
+  int32_t ofm_index; // output
+
+  int32_t ifm_index;       // input
+  int32_t axis_index;      // axis
+  int32_t keep_dims_index; // keep_dims
+
+  Param() = default;
+  Param(uint32_t inputCount, const uint32_t *inputs, uint32_t outputCount, const uint32_t *outputs);
+};
+
+class Node final : public op::Node
+{
+public:
+  Node(const Param &param) : _param(param)
+  {
+    // DO NOTHING
+  }
+
+public:
+  virtual ~Node() = default;
+
+public:
+  const Param &param(void) const { return _param; }
+
+public:
+  void accept(NodeVisitor &&) const override;
+
+private:
+  const Param _param;
+};
+
+} // namespace Mean
+} // namespace op
+} // namespace tflite
+} // namespace internal
+
+#endif // __INTERNAL_OP_MEAN_H__

diff --git a/runtimes/pure_arm_compute/src/internal/op/NodeVisitor.h b/runtimes/pure_arm_compute/src/internal/op/NodeVisitor.h
index f68fd19..cf24d98 100644 (file)
--- a/runtimes/pure_arm_compute/src/internal/op/NodeVisitor.h
+++ b/runtimes/pure_arm_compute/src/internal/op/NodeVisitor.h
@@ -26,6 +26,7 @@
 #include "internal/op/Tanh.h"
 #include "internal/op/Squeeze.h"
 #include "internal/op/Logistic.h"
+#include "internal/op/Mean.h"
 
 namespace internal
 {
@@ -67,6 +68,7 @@ struct NodeVisitor
   virtual void visit(const Tanh::Node &) = 0;
   virtual void visit(const Squeeze::Node &) = 0;
   virtual void visit(const Logistic::Node &) = 0;
+  virtual void visit(const Mean::Node &) = 0;
 };
 
 } // namespace op

diff --git a/runtimes/pure_arm_compute/src/model.cc b/runtimes/pure_arm_compute/src/model.cc
index 5313df5..a7c6e52 100644 (file)
--- a/runtimes/pure_arm_compute/src/model.cc
+++ b/runtimes/pure_arm_compute/src/model.cc
@@ -449,6 +449,16 @@ int ANeuralNetworksModel_addOperation(ANeuralNetworksModel *model,
 
       break;
     }
+    case ANEURALNETWORKS_MEAN:
+    {
+      using internal::tflite::op::Mean::Param;
+      using internal::tflite::op::Mean::Node;
+
+      auto &operations = model->deref().operations();
+      operations.emplace_back<Node>(Param{inputCount, inputs, outputCount, outputs});
+
+      break;
+    }
     default:
       throw std::runtime_error{"Not supported operation"};
   };