[nnc] Fix reduceOp reduction dimensions handling (#2494)

- fix shape inference
- fix interpreter index computation

Signed-off-by: Efimov Alexander <a.efimov@samsung.com>

diff --git a/contrib/nnc/include/core/modelIR/ir_dot_node_info.h b/contrib/nnc/include/core/modelIR/ir_dot_node_info.h
index 620aba3..c6e97f7 100644 (file)
--- a/contrib/nnc/include/core/modelIR/ir_dot_node_info.h
+++ b/contrib/nnc/include/core/modelIR/ir_dot_node_info.h
@@ -18,7 +18,7 @@
 #define NNCC_IR_NODE_DOT_BUILDER_H
 
 #include "core/modelIR/Shape.h"
-#include "core/modelIR/operations/common.h"
+#include "core/modelIR/operations/CommonProps.h"
 #include "core/modelIR/operations/PoolOp.h"
 
 namespace nnc

diff --git a/contrib/nnc/include/core/modelIR/operations/common.h b/contrib/nnc/include/core/modelIR/operations/CommonProps.h
similarity index 86%
rename from contrib/nnc/include/core/modelIR/operations/common.h
rename to contrib/nnc/include/core/modelIR/operations/CommonProps.h
index 42a9992..a4ce61b 100644 (file)
--- a/contrib/nnc/include/core/modelIR/operations/common.h
+++ b/contrib/nnc/include/core/modelIR/operations/CommonProps.h
@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 
-#ifndef _NNC_CORE_IR_MODEL_COMMON_H_
-#define _NNC_CORE_IR_MODEL_COMMON_H_
+#ifndef _NNC_CORE_IR_MODEL_COMMON_PROPS_H_
+#define _NNC_CORE_IR_MODEL_COMMON_PROPS_H_
 
 namespace nnc {
 namespace mir {
@@ -31,4 +31,4 @@ enum class PaddingType {
 } // namespace mir
 } // namespace nnc
 
-#endif //_NNC_CORE_IR_MODEL_COMMOND_H_
+#endif //_NNC_CORE_IR_MODEL_COMMON_PROPS_H_

diff --git a/contrib/nnc/include/core/modelIR/operations/Conv2DOp.h b/contrib/nnc/include/core/modelIR/operations/Conv2DOp.h
index 6fa821a..f17e0de 100644 (file)
--- a/contrib/nnc/include/core/modelIR/operations/Conv2DOp.h
+++ b/contrib/nnc/include/core/modelIR/operations/Conv2DOp.h
@@ -20,7 +20,7 @@
 #include <vector>
 
 #include "core/modelIR/Operation.h"
-#include "core/modelIR/operations/common.h"
+#include "core/modelIR/operations/CommonProps.h"
 #include "core/modelIR/TensorVariant.h"
 
 namespace nnc {

diff --git a/contrib/nnc/include/core/modelIR/operations/Deconv2DOp.h b/contrib/nnc/include/core/modelIR/operations/Deconv2DOp.h
index 0be89fd..a91aec0 100644 (file)
--- a/contrib/nnc/include/core/modelIR/operations/Deconv2DOp.h
+++ b/contrib/nnc/include/core/modelIR/operations/Deconv2DOp.h
@@ -18,7 +18,7 @@
 #define _NNC_CORE_IR_MODEL_DECONV_2D_H_
 
 #include "core/modelIR/Operation.h"
-#include "core/modelIR/operations/common.h"
+#include "core/modelIR/operations/CommonProps.h"
 #include "core/modelIR/TensorVariant.h"
 
 namespace nnc {

diff --git a/contrib/nnc/include/core/modelIR/operations/DepthwiseConv2DOp.h b/contrib/nnc/include/core/modelIR/operations/DepthwiseConv2DOp.h
index eb0f078..dd99400 100644 (file)
--- a/contrib/nnc/include/core/modelIR/operations/DepthwiseConv2DOp.h
+++ b/contrib/nnc/include/core/modelIR/operations/DepthwiseConv2DOp.h
@@ -21,7 +21,7 @@
 
 #include "core/modelIR/Operation.h"
 #include "core/modelIR/TensorVariant.h"
-#include "core/modelIR/operations/common.h"
+#include "core/modelIR/operations/CommonProps.h"
 
 namespace nnc {
 namespace mir {

diff --git a/contrib/nnc/include/core/modelIR/operations/PoolOp.h b/contrib/nnc/include/core/modelIR/operations/PoolOp.h
index b8c5c7f..19e9f32 100644 (file)
--- a/contrib/nnc/include/core/modelIR/operations/PoolOp.h
+++ b/contrib/nnc/include/core/modelIR/operations/PoolOp.h
@@ -18,7 +18,7 @@
 #define _NNC_CORE_IR_MODEL_POOL_H_
 
 #include "core/modelIR/Operation.h"
-#include "core/modelIR/operations/common.h"
+#include "core/modelIR/operations/CommonProps.h"
 #include <vector>
 #include <cmath>
 

diff --git a/contrib/nnc/include/core/modelIR/operations/ReduceFOp.h b/contrib/nnc/include/core/modelIR/operations/ReduceFOp.h
index e247434..1bdce27 100644 (file)
--- a/contrib/nnc/include/core/modelIR/operations/ReduceFOp.h
+++ b/contrib/nnc/include/core/modelIR/operations/ReduceFOp.h
@@ -42,25 +42,31 @@ public:
             FuncType func_type)
       : Operation(Type::reduceF, {arg}), _reduceDims(reduce_dims), _keepDims(keep_dims),
         _funcType(func_type) {
+
     // Infer output shapes.
     const auto& input_shape = getInputShape(0);
     const auto& red_dims = getReductionDims();
     Shape output_shape;
+
     if (getKeepDims()) {
       output_shape = input_shape;
       for (auto red_axis: red_dims) {
         output_shape.dim(red_axis) = 1;
       }
     } else {
+      // This mask contains true for axis indexes that should be reduced
+      // for example, if we want to reduce 1 and 3 axes, with total number of dims 4
+      // mask will contain: [false, true, false, true]
+      std::vector<bool> reduce_axis_mask(input_shape.rank(), false);
+      for (auto axis: red_dims)
+        reduce_axis_mask[axis] = true;
+
+      // Actual shape inference
       std::vector<int32_t> out_dims;
       out_dims.reserve(input_shape.rank() - red_dims.size());
-      auto red_axis = red_dims.begin();
       for (int32_t axis_id = 0; axis_id < input_shape.rank(); axis_id++) {
-        if (axis_id == (*red_axis)) {
-          red_axis++;
-        } else {
+        if (!reduce_axis_mask[axis_id])
           out_dims.emplace_back(input_shape.dim(axis_id));
-        }
       }
       output_shape = Shape(out_dims);
     }

diff --git a/contrib/nnc/passes/caffe_frontend/caffe_op_creator.h b/contrib/nnc/passes/caffe_frontend/caffe_op_creator.h
index 55fac40..63da772 100644 (file)
--- a/contrib/nnc/passes/caffe_frontend/caffe_op_creator.h
+++ b/contrib/nnc/passes/caffe_frontend/caffe_op_creator.h
@@ -24,7 +24,7 @@
 
 #include "core/modelIR/Graph.h"
 #include "core/modelIR/TensorVariant.h"
-#include "core/modelIR/operations/common.h"
+#include "core/modelIR/operations/CommonProps.h"
 #include "core/modelIR/Shape.h"
 
 #include "caffe/proto/caffe.pb.h"

diff --git a/contrib/nnc/passes/interpreter/ops/Depthwise_conv_2D.h b/contrib/nnc/passes/interpreter/ops/Depthwise_conv_2D.h
index e4e7707..c0f0a17 100644 (file)
--- a/contrib/nnc/passes/interpreter/ops/Depthwise_conv_2D.h
+++ b/contrib/nnc/passes/interpreter/ops/Depthwise_conv_2D.h
@@ -19,7 +19,7 @@
 
 #include "OperationImpl.h"
 
-#include "core/modelIR/operations/common.h"
+#include "core/modelIR/operations/CommonProps.h"
 #include "core/modelIR/operations/DepthwiseConv2DOp.h"
 
 namespace nnc

diff --git a/contrib/nnc/passes/interpreter/ops/Pool.h b/contrib/nnc/passes/interpreter/ops/Pool.h
index 8b63533..d87d827 100644 (file)
--- a/contrib/nnc/passes/interpreter/ops/Pool.h
+++ b/contrib/nnc/passes/interpreter/ops/Pool.h
@@ -19,7 +19,7 @@
 
 #include "OperationImpl.h"
 #include "core/modelIR/operations/PoolOp.h"
-#include "core/modelIR/operations/common.h"
+#include "core/modelIR/operations/CommonProps.h"
 
 namespace nnc
 {

diff --git a/contrib/nnc/passes/interpreter/ops/Reduce.h b/contrib/nnc/passes/interpreter/ops/Reduce.h
index f4d2d96..f12d403 100644 (file)
--- a/contrib/nnc/passes/interpreter/ops/Reduce.h
+++ b/contrib/nnc/passes/interpreter/ops/Reduce.h
@@ -90,16 +90,22 @@ public:
     out_id.resize(_outputShape.rank());
     for (const mir::Index& input_id : mir::ShapeRange(_inShape)) {
       int32_t out_idx_id = 0;
-      int32_t red_dim = 0;
+
+      // This mask contains true for axis indexes that should be reduced
+      // for example, if we want to reduce 1 and 3 axes, with total number of dims 4
+      // mask will contain: [false, true, false, true]
+      std::vector<bool> reduce_axis_mask(_inShape.rank(), false);
+      for (auto axis: _reductionDims)
+        reduce_axis_mask[axis] = true;
+
       // change out id to point to the correct cell
-      for (int d = 0; d != _inShape.rank(); ++d) {
-        if (d == _reductionDims[red_dim]) {
-          red_dim++;
-          if (_keepDims)
-            out_id.at(out_idx_id++) = 0;
-          else
+      if (_keepDims) {
+        for (int32_t d = 0; d < _inShape.rank(); ++d)
+          out_id.at(out_idx_id++) = reduce_axis_mask[d] ? 0 : input_id.at(d);
+      } else {
+        for (int32_t d = 0; d < _inShape.rank(); ++d) {
+          if (reduce_axis_mask[d])
             continue;
-        } else {
           out_id.at(out_idx_id++) = input_id.at(d);
         }
       }

diff --git a/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.h b/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.h
index cf4268b..c0eee9e 100644 (file)
--- a/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.h
+++ b/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.h
@@ -23,7 +23,7 @@
 #include <memory>
 #include "core/modelIR/Graph.h"
 #include "core/modelIR/TensorVariant.h"
-#include "core/modelIR/operations/common.h"
+#include "core/modelIR/operations/CommonProps.h"
 #include "core/modelIR/operations/ElementwiseOp.h"
 #include "core/modelIR/Shape.h"
 #include "onnx/onnx.pb.h"

diff --git a/contrib/nnc/passes/tflite_frontend/tflite_op_creator.h b/contrib/nnc/passes/tflite_frontend/tflite_op_creator.h
index 5ada473..a8fdf30 100644 (file)
--- a/contrib/nnc/passes/tflite_frontend/tflite_op_creator.h
+++ b/contrib/nnc/passes/tflite_frontend/tflite_op_creator.h
@@ -28,7 +28,7 @@
 #include "core/modelIR/Scalar.h"
 #include "core/modelIR/Shape.h"
 
-#include "core/modelIR/operations/common.h"
+#include "core/modelIR/operations/CommonProps.h"
 #include "core/modelIR/operations/ReduceFOp.h"
 
 #include "schema_generated.h"

diff --git a/contrib/nnc/tests/interpreter/op_info_util.h b/contrib/nnc/tests/interpreter/op_info_util.h
index c7e1ec3..5531616 100644 (file)
--- a/contrib/nnc/tests/interpreter/op_info_util.h
+++ b/contrib/nnc/tests/interpreter/op_info_util.h
@@ -21,7 +21,7 @@
 #include <cassert>
 
 #include "core/modelIR/TensorVariant.h"
-#include "core/modelIR/operations/common.h"
+#include "core/modelIR/operations/CommonProps.h"
 #include "core/modelIR/operations/PoolOp.h"
 
 #include "op_info_generated.h"