[NGraph] Add scatterNDUpdate and scatterUpdate reference implementations

openvinotoolkit · Jul 30, 2020 · cd4f520 · cd4f520
1 parent 77345a7
commit cd4f520
Show file tree

Hide file tree

Showing 10 changed files with 236 additions and 199 deletions.
diff --git a/inference-engine/src/mkldnn_plugin/nodes/list_tbl.hpp b/inference-engine/src/mkldnn_plugin/nodes/list_tbl.hpp
@@ -51,7 +51,6 @@ MKLDNN_EXTENSION_NODE(FillImpl, Fill);
 MKLDNN_EXTENSION_NODE(UniqueImpl, Unique);
 MKLDNN_EXTENSION_NODE(PSROIPoolingImpl, PSROIPooling);
 MKLDNN_EXTENSION_NODE(DepthToSpaceImpl, DepthToSpace);
-MKLDNN_EXTENSION_NODE(ScatterImpl, ScatterUpdate);
 MKLDNN_EXTENSION_NODE(OneHotImpl, OneHot);
 MKLDNN_EXTENSION_NODE(BroadcastImpl, Broadcast);
 MKLDNN_EXTENSION_NODE(ExperimentalSparseWeightedReduceImpl, ExperimentalSparseWeightedSum);

diff --git a/inference-engine/src/mkldnn_plugin/nodes/scatter.cpp b/inference-engine/src/mkldnn_plugin/nodes/scatter.cpp
diff --git a/...sts/functional/plugin/cpu/shared_tests_instances/single_layer_tests/scatter_ND_update.cpp b/...sts/functional/plugin/cpu/shared_tests_instances/single_layer_tests/scatter_ND_update.cpp
@@ -36,7 +36,6 @@ const auto ScatterNDUpdateCases = ::testing::Combine(
         ::testing::Values(CommonTestUtils::DEVICE_CPU)
 );
 
-// open after ops support in ngraph merged
-// INSTANTIATE_TEST_CASE_P(ScatterNDUpdate, ScatterNDUpdateLayerTest, ScatterNDUpdateCases, ScatterNDUpdateLayerTest::getTestCaseName);
+INSTANTIATE_TEST_CASE_P(ScatterNDUpdate, ScatterNDUpdateLayerTest, ScatterNDUpdateCases, ScatterNDUpdateLayerTest::getTestCaseName);
 
-}  // namespace
+}  // namespace
diff --git a/.../tests/functional/plugin/cpu/shared_tests_instances/single_layer_tests/scatter_update.cpp b/.../tests/functional/plugin/cpu/shared_tests_instances/single_layer_tests/scatter_update.cpp
@@ -41,7 +41,6 @@ const auto ScatterUpdateCase = ::testing::Combine(
         ::testing::Values(CommonTestUtils::DEVICE_CPU)
 );
 
-// open after ngraph reference implementation merged
-// INSTANTIATE_TEST_CASE_P(ScatterUpdate, ScatterUpdateLayerTest, ScatterUpdateCase, ScatterUpdateLayerTest::getTestCaseName);
+INSTANTIATE_TEST_CASE_P(ScatterUpdate, ScatterUpdateLayerTest, ScatterUpdateCase, ScatterUpdateLayerTest::getTestCaseName);
 
-}  // namespace
+}  // namespace
diff --git a/inference-engine/tests/ngraph_functions/include/ngraph_functions/builders.hpp b/inference-engine/tests/ngraph_functions/include/ngraph_functions/builders.hpp
@@ -10,6 +10,7 @@
 #include <ngraph/opsets/opset1.hpp>
 #include <ngraph/opsets/opset2.hpp>
 #include <ngraph/opsets/opset3.hpp>
+#include <ngraph/opsets/opset4.hpp>
 
 #include "ngraph_functions/utils/data_utils.hpp"
 

diff --git a/inference-engine/tests/ngraph_functions/src/scatter_ND_update.cpp b/inference-engine/tests/ngraph_functions/src/scatter_ND_update.cpp
@@ -13,10 +13,8 @@ std::shared_ptr<ngraph::Node> makeScatterNDUpdate(const ngraph::Output<Node> &in
                                                   const std::vector<size_t>& indices,
                                                   const ngraph::Output<Node> &update) {
     auto indicesNode = std::make_shared<ngraph::opset1::Constant>(indicesType, indicesShape, indices);
-    // blocked by ngraph merge
-    // auto dtsNode = std::make_shared<ngraph::opset3::ScatterNDUpdate>(in, indicesNode, update);
-    // return dtsNode;
-    return nullptr;
+    auto dtsNode = std::make_shared<ngraph::opset4::ScatterNDUpdate>(in, indicesNode, update);
+    return dtsNode;
 }
 
 }  // namespace builder

diff --git a/ngraph/test/runtime/interpreter/int_executable.hpp b/ngraph/test/runtime/interpreter/int_executable.hpp
@@ -95,6 +95,9 @@
 #include "ngraph/state/uniform_rng_state.hpp"
 #include "op/avg_pool.hpp"
 
+#include "reference/scatter_nd_update.hpp"
+#include "reference/scatter_update.hpp"
+
 namespace ngraph
 {
     namespace runtime
@@ -1160,6 +1163,81 @@ class INTERPRETER_BACKEND_API ngraph::runtime::interpreter::INTExecutable : publ
             }
             break;
         }
+        case OP_TYPEID::ScatterNDUpdate_v3:
+        {
+            const op::ScatterNDUpdate* scatterNDUpd =
+                static_cast<const op::v3::ScatterNDUpdate*>(&node);
+            auto idxType = scatterNDUpd->get_input_element_type(1);
+            if (idxType == element::i32)
+            {
+                reference::scatterNdUpdate<T, int32_t>(args[0]->get_data_ptr<const T>(),
+                                                       args[1]->get_data_ptr<const int32_t>(),
+                                                       args[2]->get_data_ptr<const T>(),
+                                                       out[0]->get_data_ptr<T>(),
+                                                       node.get_input_shape(0),
+                                                       node.get_input_shape(1),
+                                                       node.get_input_shape(2));
+            }
+            else if (idxType == element::i64)
+            {
+                reference::scatterNdUpdate<T, int64_t>(args[0]->get_data_ptr<const T>(),
+                                                       args[1]->get_data_ptr<const int64_t>(),
+                                                       args[2]->get_data_ptr<const T>(),
+                                                       out[0]->get_data_ptr<T>(),
+                                                       node.get_input_shape(0),
+                                                       node.get_input_shape(1),
+                                                       node.get_input_shape(2));
+            }
+            else
+            {
+                throw ngraph_error(
+                    "ScatterNDUpdate layer support only i32 and i64 'indices' input precision!");
+            }
+
+            break;
+        }
+        case OP_TYPEID::ScatterUpdate_v3:
+        {
+            const op::v3::ScatterUpdate* scatterUpd =
+                static_cast<const op::v3::ScatterUpdate*>(&node);
+
+            if (scatterUpd->get_input_element_type(3) != element::i64)
+                throw ngraph_error(
+                    "ScatterNDUpdate layer support only i64 'axis' input precision!");
+
+            auto idxType = scatterUpd->get_input_element_type(1);
+            if (idxType == element::i32)
+            {
+                reference::scatterUpdate<T, int32_t, int64_t>(
+                    args[0]->get_data_ptr<const T>(),
+                    args[1]->get_data_ptr<const int32_t>(),
+                    args[2]->get_data_ptr<const T>(),
+                    args[3]->get_data_ptr<const int64_t>(),
+                    out[0]->get_data_ptr<T>(),
+                    node.get_input_shape(0),
+                    node.get_input_shape(1),
+                    node.get_input_shape(2));
+            }
+            else if (idxType == element::i64)
+            {
+                reference::scatterUpdate<T, int64_t, int64_t>(
+                    args[0]->get_data_ptr<const T>(),
+                    args[1]->get_data_ptr<const int64_t>(),
+                    args[2]->get_data_ptr<const T>(),
+                    args[3]->get_data_ptr<const int64_t>(),
+                    out[0]->get_data_ptr<T>(),
+                    node.get_input_shape(0),
+                    node.get_input_shape(1),
+                    node.get_input_shape(2));
+            }
+            else
+            {
+                throw ngraph_error(
+                    "ScatterUpdate layer support only i32 and i64 'indices' input precision!");
+            }
+
+            break;
+        }
 
         // Fused Ops are not supported in interpreter. They need to be decomposed before execution
         case OP_TYPEID::DepthToSpace:

diff --git a/ngraph/test/runtime/interpreter/opset_int_tbl.hpp b/ngraph/test/runtime/interpreter/opset_int_tbl.hpp
@@ -33,4 +33,6 @@ NGRAPH_OP(EmbeddingSegmentsSum, op::v3)
 NGRAPH_OP(ExtractImagePatches, op::v3)
 NGRAPH_OP(ShapeOf, op::v3)
 NGRAPH_OP(NonZero, op::v3)
+NGRAPH_OP(ScatterNDUpdate, op::v3)
+NGRAPH_OP(ScatterUpdate, op::v3)
 #undef ID_SUFFIX