[mir] Add ReduceMean operation (#6994)

This is a future replacement of `Reduce` operation.

Signed-off-by: Sergei Barannikov <s.barannikov@samsung.com>

diff --git a/compiler/mir/CMakeLists.txt b/compiler/mir/CMakeLists.txt
index 2711d4a..bc265e5 100644 (file)
--- a/compiler/mir/CMakeLists.txt
+++ b/compiler/mir/CMakeLists.txt
@@ -8,6 +8,7 @@ set(MIR_SOURCES
     src/ops/GatherOp.cpp
     src/ops/PadOp.cpp
     src/ops/PoolOp.cpp
+    src/ops/ReduceMeanOp.cpp
     src/ops/SqueezeOp.cpp
     src/ops/SliceOp.cpp
     src/ops/TransposeOp.cpp

diff --git a/compiler/mir/include/mir/IrDotDumper.h b/compiler/mir/include/mir/IrDotDumper.h
index 3c44277..653f453 100644 (file)
--- a/compiler/mir/include/mir/IrDotDumper.h
+++ b/compiler/mir/include/mir/IrDotDumper.h
@@ -45,6 +45,7 @@ public:
   void visit(ops::PadOp &op) override;
   void visit(ops::PoolOp &op) override;
   void visit(ops::ReduceOp &op) override;
+  void visit(ops::ReduceMeanOp &op) override;
   void visit(ops::ReluOp &op) override;
   void visit(ops::ReshapeOp &op) override;
   void visit(ops::ResizeOp &op) override;

diff --git a/compiler/mir/include/mir/OpDefs.h b/compiler/mir/include/mir/OpDefs.h
index 7d73d50..5da3050 100644 (file)
--- a/compiler/mir/include/mir/OpDefs.h
+++ b/compiler/mir/include/mir/OpDefs.h
@@ -37,6 +37,7 @@
 #include "mir/ops/PadOp.h"
 #include "mir/ops/PoolOp.h"
 #include "mir/ops/ReduceOp.h"
+#include "mir/ops/ReduceMeanOp.h"
 #include "mir/ops/ReluOp.h"
 #include "mir/ops/ReshapeOp.h"
 #include "mir/ops/ResizeOp.h"

diff --git a/compiler/mir/include/mir/Operations.inc b/compiler/mir/include/mir/Operations.inc
index e0acb21..9867c6b 100644 (file)
--- a/compiler/mir/include/mir/Operations.inc
+++ b/compiler/mir/include/mir/Operations.inc
@@ -37,6 +37,7 @@ HANDLE_OP(output, OutputOp)
 HANDLE_OP(pad, PadOp)
 HANDLE_OP(pool, PoolOp)
 HANDLE_OP(reduce, ReduceOp)
+HANDLE_OP(reduceMean, ReduceMeanOp)
 HANDLE_OP(ReLU, ReluOp)
 HANDLE_OP(reshape, ReshapeOp)
 HANDLE_OP(resizeIm, ResizeOp)

diff --git a/compiler/mir/include/mir/ops/ReduceMeanOp.h b/compiler/mir/include/mir/ops/ReduceMeanOp.h
new file mode 100644 (file)
index 0000000..fc7ebc1
--- /dev/null
+++ b/compiler/mir/include/mir/ops/ReduceMeanOp.h
@@ -0,0 +1,57 @@
+/*
+ * 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 _MIR_OPS_REDUCE_MEAN_OP_H_
+#define _MIR_OPS_REDUCE_MEAN_OP_H_
+
+#include "mir/Operation.h"
+#include <vector>
+
+namespace mir
+{
+namespace ops
+{
+
+class ReduceMeanOp : public Operation
+{
+public:
+  ReduceMeanOp(Output *arg, std::vector<int> reduction_dims, bool keep_dims)
+      : Operation(Type::reduceMean, {arg}), _reduction_dims(std::move(reduction_dims)),
+        _keep_dims(keep_dims)
+  {
+    inferOutputShapes();
+  }
+
+  Operation *copyWithInputs(const std::vector<Output *> &inputs) override
+  {
+    return new ReduceMeanOp(inputs[0], _reduction_dims, _keep_dims);
+  }
+
+  const std::vector<int> &getReductionDims() const { return _reduction_dims; };
+
+  bool getKeepDims() const { return _keep_dims; };
+
+private:
+  void inferOutputShapes();
+
+  std::vector<int> _reduction_dims;
+  bool _keep_dims;
+};
+
+} // namespace ops
+} // namespace mir
+
+#endif //_MIR_OPS_REDUCE_MEAN_OP_H_

diff --git a/compiler/mir/src/IrDotDumper.cpp b/compiler/mir/src/IrDotDumper.cpp
index a683107..7cf234b 100644 (file)
--- a/compiler/mir/src/IrDotDumper.cpp
+++ b/compiler/mir/src/IrDotDumper.cpp
@@ -15,7 +15,6 @@
  */
 
 #include "mir/IrDotDumper.h"
-
 #include "mir/OpDefs.h"
 
 #include <map>
@@ -262,6 +261,19 @@ void IrDotDumper::visit(ops::ReduceOp &op)
   _dot_builder.updateWithOp(&op, node_info);
 }
 
+void IrDotDumper::visit(ops::ReduceMeanOp &op)
+{
+  auto node_info =
+      DotIrNodeInfo()
+          .withType("ReduceMeanOp", op.getName())
+          .withInShapes(getInputShapes(op))
+          .withOutShapes(getOutputShapes(op))
+          .withShape("Reduction dims", Shape(op.getReductionDims())) // appropriated shape to dims
+          .withMisc("Keep dims", op.getKeepDims());
+
+  _dot_builder.updateWithOp(&op, node_info);
+}
+
 void IrDotDumper::visit(ops::ResizeOp &op)
 {
   static const std::map<ops::ResizeOp::ResizeMethod, const char *> modes{

diff --git a/compiler/mir/src/ops/ReduceMeanOp.cpp b/compiler/mir/src/ops/ReduceMeanOp.cpp
new file mode 100644 (file)
index 0000000..8fe23b3
--- /dev/null
+++ b/compiler/mir/src/ops/ReduceMeanOp.cpp
@@ -0,0 +1,61 @@
+/*
+ * 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 "mir/ops/ReduceMeanOp.h"
+
+namespace mir
+{
+namespace ops
+{
+
+void ReduceMeanOp::inferOutputShapes()
+{
+  const auto &input_shape = getInputShape(0);
+  const auto &reduction_dims = getReductionDims();
+  Shape output_shape;
+
+  if (getKeepDims())
+  {
+    output_shape = input_shape;
+    for (const int dim : reduction_dims)
+    {
+      output_shape.dim(dim) = 1;
+    }
+  }
+  else
+  {
+    // This mask contains 'true' for dimension indices that should be reduced.
+    // for example, if we want to reduce 1 and 3 dimensions with total number of dimensions 4,
+    // the mask will contain: [false, true, false, true].
+    std::vector<bool> reduction_dims_mask(input_shape.rank(), false);
+    for (auto axis : reduction_dims)
+      reduction_dims_mask[axis] = true;
+
+    std::vector<std::int32_t> out_dims;
+    out_dims.reserve(input_shape.rank() - reduction_dims.size());
+    for (int axis_id = 0; axis_id < input_shape.rank(); axis_id++)
+    {
+      if (!reduction_dims_mask[axis_id])
+        out_dims.emplace_back(input_shape.dim(axis_id));
+    }
+    output_shape = Shape(out_dims);
+  }
+
+  setOutputShape(0, output_shape);
+}
+
+} // namespace ops
+} // namespace mir

diff --git a/compiler/mir/unittests/ShapeInference.cpp b/compiler/mir/unittests/ShapeInference.cpp
index 3d47dbd..3a5459d 100644 (file)
--- a/compiler/mir/unittests/ShapeInference.cpp
+++ b/compiler/mir/unittests/ShapeInference.cpp
@@ -19,7 +19,7 @@
 #include "mir/ops/ReshapeOp.h"
 #include "mir/ops/ResizeOp.h"
 #include "mir/ops/SqueezeOp.h"
-#include "mir/ops/ReduceOp.h"
+#include "mir/ops/ReduceMeanOp.h"
 #include "mir/Shape.h"
 
 #include <vector>
@@ -78,8 +78,8 @@ TEST(ShapeInferenceTest, ReduceChangeRank)
 
   auto input = g.create<ops::InputOp>("input", Shape{10, 2, 10, 9});
 
-  auto n = g.create<ops::ReduceOp>("reduce", input->getOutput(0), std::vector<int32_t>{1, 3}, false,
-                                   ops::ReduceOp::FuncType::mean);
+  auto n =
+      g.create<ops::ReduceMeanOp>("reduce", input->getOutput(0), std::vector<int32_t>{1, 3}, false);
 
   ASSERT_EQ(resultShape, n->getOutputShape(0));
 }