Scatter update ref impl revision (#7476)

* add template tests for scatter_update

* rm CT from ref impl

* fix namings

* minor code refactor

* fix style

* change ref impl

* add lambda func

* new tests

* remove CT from reference implementation, add template tests

* fix cpu IE

* refactor template tests

* split template tests

inference-engine/tests/functional/plugin/cpu/shared_tests_instances/single_layer_tests/scatter_update.cpp

@@ -25,6 +25,9 @@ const std::vector<InferenceEngine::Precision> idxPrecisions = {
 
 // map<inputShape, map<indicesShape, axis>>
 std::map<std::vector<size_t>, std::map<std::vector<size_t>, std::vector<int>>> axesShapeInShape {
+    {{10, 16, 12, 15}, {{{2, 2, 2}, {0, 1, 2, 3}}, {{2, 4}, {0, 1, 2, 3}}, {{8}, {0, 1, 2, 3}}}},
+    {{10, 9, 10, 9, 10}, {{{8}, {0, 1, 2, 3, 4}}, {{4, 2}, {0, 1, 2, 3, 4}}}},
+    {{10, 9, 10, 9, 10, 12}, {{{8}, {0, 1, 2, 3, 4, 5}}}},
     {{10, 16, 12, 15}, {{{2, 4}, {0, 1, 2, 3}}, {{8}, {-1, -2, -3, -4}}}},
     {{10, 9, 10, 9, 10}, {{{8}, {-3, -1, 0, 2, 4}}, {{4, 2}, {-2, 2}}}},
 };
@@ -43,4 +46,4 @@ const auto ScatterUpdateCase = ::testing::Combine(
 
 INSTANTIATE_TEST_SUITE_P(smoke_ScatterUpdate, ScatterUpdateLayerTest, ScatterUpdateCase, ScatterUpdateLayerTest::getTestCaseName);
 
-}  // namespace
+}  // namespace

ngraph/core/reference/include/ngraph/runtime/reference/scatter_update.hpp

@@ -4,13 +4,32 @@
 
 #pragma once
 
+#include <numeric>
+
 #include "ngraph/check.hpp"
 #include "ngraph/coordinate_transform.hpp"
 #include "ngraph/shape.hpp"
+#include "ngraph/util.hpp"
 
 namespace ngraph {
 namespace runtime {
 namespace reference {
+static const CoordinateTransformBasic get_target_shape(const Shape& data_shape,
+                                                       const Coordinate& start_corner,
+                                                       const Coordinate& end_corner) {
+    const auto m_n_axes = data_shape.size();
+    Shape target_shape;
+    target_shape.reserve(m_n_axes);
+    AxisVector axis_order(m_n_axes);
+    std::iota(axis_order.begin(), axis_order.end(), 0);
+    const Strides strides(m_n_axes, 1);
+    for (size_t axis = 0; axis < m_n_axes; axis++) {
+        target_shape.push_back(
+            ceil_div(end_corner[axis_order[axis]] - start_corner[axis_order[axis]], strides[axis_order[axis]]));
+    }
+    return target_shape;
+}
+
 void scatter_update(const char* input_data,
                     const int64_t* indices,
                     const char* updates,
@@ -36,43 +55,57 @@ void scatter_update(const char* input_data,
     //     for d_coord in slice data[..., i_idx, ...],
     //         u_coord in slice updates[..., i_coord, ...]
     //          data[index(d_coord)] = updates[index(u_coord)]
-
-    NGRAPH_SUPPRESS_DEPRECATED_START
-    CoordinateTransform indices_transform{indices_shape};
-    CoordinateTransform data_transform{data_shape};
+    CoordinateTransformBasic indices_transform{indices_shape};
+    const auto indices_in_strides = row_major_strides(indices_shape);
 
     size_t indices_ndim = indices_shape.size();
     size_t updates_ndim = updates_shape.size();
+    size_t data_ndim = data_shape.size();
+
+    const auto size_after_axis = shape_size(Shape(data_shape.begin() + axis + 1, data_shape.end()));
+    int num_axis_jumps{0};
+    int num_unary_moves{0};
+    for (size_t i = axis + 1; i < updates_ndim; ++i) {
+        const auto updates_size_after_axis = shape_size(Shape(updates_shape.begin() + i, updates_shape.end()));
+        if (updates_size_after_axis > size_after_axis)
+            ++num_axis_jumps;
+        if (updates_shape[i] == 1)
+            ++num_unary_moves;
+    }
+
+    if (!num_axis_jumps)
+        num_axis_jumps = updates_ndim - data_ndim;
+
+    auto updates_axis_dim = axis + num_axis_jumps + num_unary_moves;
+
+    if (updates_axis_dim >= updates_ndim)
+        updates_axis_dim = updates_ndim - 1;
 
-    // Create an outer CoordinateTransform for "update", which would allow to
-    // iterate only over "indices" dimensions:
-    // set to "1" all non-indices dimensions
-    // updates[1, ..., 1, m, n, ..., p, 1, 1,..., 1]
     Coordinate updates_indices_start_corner(updates_ndim, 0);
     Coordinate updates_indices_end_corner(updates_ndim, 1);
+
     for (size_t i = 0; i < indices_ndim; ++i) {
         updates_indices_end_corner[axis + i] = updates_shape[axis + i];
     }
-    CoordinateTransform updates_indices_transform(updates_shape,
-                                                  updates_indices_start_corner,
-                                                  updates_indices_end_corner);
-    // Is needed to simultaneously iterate over updates coordinates while
-    // iterating over indices.
+
+    const auto updates_indices_transform =
+        get_target_shape(updates_shape, updates_indices_start_corner, updates_indices_end_corner);
     auto updates_indices_coord_iter = updates_indices_transform.begin();
 
+    int iteration{0};
     for (const Coordinate& indices_cord : indices_transform) {
-        const size_t indices_idx = indices_transform.index(indices_cord);
+        const size_t indices_idx =
+            std::inner_product(indices_cord.begin(), indices_cord.end(), indices_in_strides.begin(), 0);
         int64_t slice_index = indices[indices_idx];
 
-        // Define the extent of coordinates which will be updated.
         Coordinate out_start_corner(data_shape.size(), 0);
         Coordinate out_end_corner(data_shape);
         out_start_corner[axis] = static_cast<size_t>(slice_index);
         out_end_corner[axis] = out_start_corner[axis] + 1;
-        CoordinateTransform out_transform(data_shape, out_start_corner, out_end_corner);
 
-        // Define the CoordinateTransform for updates coordinates.
-        // All except indices-dimensions.
+        const auto out_transform = get_target_shape(data_shape, out_start_corner, out_end_corner);
+        const auto out_transform_in_strides = row_major_strides(data_shape);
+
         if (updates_indices_coord_iter == updates_indices_transform.end())
             break;
         Coordinate updates_update_start_corner = *updates_indices_coord_iter;
@@ -80,27 +113,32 @@ void scatter_update(const char* input_data,
         for (size_t i = 0; i < indices_ndim; ++i) {
             updates_update_end_corner[axis + i] = updates_update_start_corner[axis + i] + 1;
         }
-        // The m, n, .., p symbols stand for values at those axes.
-        // The m+1 means value at axis m plus 1.
-        // udpates_shape (start): [ 0, ...,  0, m  , n  , ... p  ,  0, ...,  0]
-        // updates_shape (end):   [-1, ..., -1, m+1, n+1, ... p+1, -1, ..., -1]
-        CoordinateTransform updates_update_transform(updates_shape,
-                                                     updates_update_start_corner,
-                                                     updates_update_end_corner);
+
+        const auto updates_update_transform =
+            get_target_shape(updates_shape, updates_update_start_corner, updates_update_end_corner);
+        const auto updates_update_in_strides = row_major_strides(updates_shape);
         auto updates_update_coord_iter = updates_update_transform.begin();
+
         for (const Coordinate& out_cord : out_transform) {
             if (updates_update_coord_iter == updates_update_transform.end())
                 break;
-            const auto src_idx = updates_update_transform.index(*updates_update_coord_iter) * elem_size;
-            std::copy(updates + src_idx,
-                      updates + (src_idx + elem_size),
-                      out_buf + out_transform.index(out_cord) * elem_size);
+            Coordinate update_cord = *updates_update_coord_iter;
+            Coordinate out_coord = out_cord;
+            out_coord.at(axis) = slice_index;
+            update_cord.at(updates_axis_dim) += iteration;
+            const auto data_idx =
+                std::inner_product(out_coord.begin(), out_coord.end(), out_transform_in_strides.begin(), 0);
+            const auto updates_idx =
+                std::inner_product(update_cord.begin(), update_cord.end(), updates_update_in_strides.begin(), 0) *
+                elem_size;
+
+            std::copy(updates + updates_idx, updates + (updates_idx + elem_size), out_buf + data_idx * elem_size);
             updates_update_coord_iter++;
         }
         updates_indices_coord_iter++;
+        iteration++;
     }
-    NGRAPH_SUPPRESS_DEPRECATED_END
 }
 }  // namespace reference
 }  // namespace runtime
-}  // namespace ngraph
+}  // namespace ngraph

ngraph/core/src/op/scatter_update.cpp

@@ -109,4 +109,4 @@ bool op::v3::ScatterUpdate::has_evaluate() const {
         break;
     }
     return false;
-}
+}