From af5fba787f0f4dc812575b41f1652debda251fbb Mon Sep 17 00:00:00 2001
From: Daria Mityagina <daria.mityagina@intel.com>
Date: Wed, 8 Dec 2021 13:44:06 +0300
Subject: [PATCH] [VPU] Add new operation "GatherND-8" (#8745)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

In the MyriadX plugin, GatherND-5 will be called instead of GatherND-8, since the non-default batch_dims parameter values ​​are not plugin use cases.
---
 .../transformations/convert_gatherND8.hpp     | 17 ++++++++
 .../transformations/convert_gatherND8.cpp     | 41 +++++++++++++++++++
 .../src/frontend/frontend.cpp                 |  2 +
 .../single_layer_tests/gather_nd.cpp          | 23 +++++++++++
 4 files changed, 83 insertions(+)
 create mode 100644 inference-engine/src/vpu/common/include/vpu/ngraph/transformations/convert_gatherND8.hpp
 create mode 100644 inference-engine/src/vpu/common/src/ngraph/transformations/convert_gatherND8.cpp

diff --git a/inference-engine/src/vpu/common/include/vpu/ngraph/transformations/convert_gatherND8.hpp b/inference-engine/src/vpu/common/include/vpu/ngraph/transformations/convert_gatherND8.hpp
new file mode 100644
index 00000000000..a0581648f2d
--- /dev/null
+++ b/inference-engine/src/vpu/common/include/vpu/ngraph/transformations/convert_gatherND8.hpp
@@ -0,0 +1,17 @@
+// Copyright (C) 2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <transformations_visibility.hpp>
+#include <ngraph/pass/graph_rewrite.hpp>
+
+namespace vpu {
+
+class ConvertGatherND8ToGatherND5 : public ngraph::pass::MatcherPass {
+public:
+    NGRAPH_RTTI_DECLARATION;
+    ConvertGatherND8ToGatherND5();
+};
+}
diff --git a/inference-engine/src/vpu/common/src/ngraph/transformations/convert_gatherND8.cpp b/inference-engine/src/vpu/common/src/ngraph/transformations/convert_gatherND8.cpp
new file mode 100644
index 00000000000..1351e72ea20
--- /dev/null
+++ b/inference-engine/src/vpu/common/src/ngraph/transformations/convert_gatherND8.cpp
@@ -0,0 +1,41 @@
+// Copyright (C) 2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "vpu/ngraph/transformations/convert_gatherND8.hpp"
+#include <ngraph/opsets/opset5.hpp>
+#include <ngraph/opsets/opset8.hpp>
+#include <ngraph/rt_info.hpp>
+#include <ngraph/pattern/op/wrap_type.hpp>
+
+NGRAPH_RTTI_DEFINITION(vpu::ConvertGatherND8ToGatherND5, "ConvertGatherND8ToGatherND5", 0);
+
+namespace vpu {
+
+ConvertGatherND8ToGatherND5::ConvertGatherND8ToGatherND5() {
+    auto gather_nd_v8_pattern = ngraph::pattern::wrap_type<ngraph::opset8::GatherND>();
+
+    ngraph::matcher_pass_callback callback = [=](ngraph::pattern::Matcher& m) {
+        auto gather_nd_v8_node = std::dynamic_pointer_cast<ngraph::opset8::GatherND>(m.get_match_root());
+        if (!gather_nd_v8_node) {
+            return false;
+        }
+        if (gather_nd_v8_node->get_batch_dims() == 0) {
+            auto gather_nd_v5_node = std::make_shared<ngraph::opset5::GatherND>(gather_nd_v8_node->input_value(0),
+                                                                gather_nd_v8_node->input_value(1),
+                                                                gather_nd_v8_node->get_batch_dims());
+
+            gather_nd_v5_node->set_friendly_name(gather_nd_v8_node->get_friendly_name());
+            ngraph::copy_runtime_info(gather_nd_v8_node, gather_nd_v5_node);
+            ngraph::replace_node(gather_nd_v8_node, gather_nd_v5_node);
+            return true;
+        } else {
+            return false;
+        }
+    };
+
+    auto m = std::make_shared<ngraph::pattern::Matcher>(gather_nd_v8_pattern, "ConvertGatherND8ToGatherND5");
+    register_matcher(m, callback);
+}
+
+}  // namespace vpu
diff --git a/inference-engine/src/vpu/graph_transformer/src/frontend/frontend.cpp b/inference-engine/src/vpu/graph_transformer/src/frontend/frontend.cpp
index cd68c0b4bdb..447ba461280 100644
--- a/inference-engine/src/vpu/graph_transformer/src/frontend/frontend.cpp
+++ b/inference-engine/src/vpu/graph_transformer/src/frontend/frontend.cpp
@@ -34,6 +34,7 @@
 #include <transformations/op_conversions/hswish_decomposition.hpp>
 #include <transformations/op_conversions/simplify_ctc_greedy_decoder_seq_len.hpp>
 #include <transformations/op_conversions/convert_gather_downgrade.hpp>
+#include <vpu/ngraph/transformations/convert_gatherND8.hpp>
 #include <transformations/convert_precision.hpp>
 #include <legacy/transformations/convert_opset1_to_legacy/convert_opset1_to_legacy.hpp>
 #include <transformations/common_optimizations/common_optimizations.hpp>
@@ -181,6 +182,7 @@ ie::CNNNetwork FrontEnd::convertNetwork(ie::CNNNetwork& network) {
     manager.register_pass<ngraph::pass::ConvertNMS3ToNMS5>();
     manager.register_pass<ngraph::pass::ConvertNMS4ToNMS5>();
     manager.register_pass<ngraph::pass::ConvertGather7ToGather1>();
+    manager.register_pass<vpu::ConvertGatherND8ToGatherND5>();
     manager.register_pass<vpu::MergeGatherGatherElements>();
     manager.register_pass<ngraph::pass::CommonOptimizations>();
 
diff --git a/inference-engine/tests/functional/plugin/myriad/shared_tests_instances/single_layer_tests/gather_nd.cpp b/inference-engine/tests/functional/plugin/myriad/shared_tests_instances/single_layer_tests/gather_nd.cpp
index dc01dde5d87..8e7ebbc396f 100644
--- a/inference-engine/tests/functional/plugin/myriad/shared_tests_instances/single_layer_tests/gather_nd.cpp
+++ b/inference-engine/tests/functional/plugin/myriad/shared_tests_instances/single_layer_tests/gather_nd.cpp
@@ -39,6 +39,18 @@ const std::vector<GatherNDParamsSubset> layerParams = {
     GatherNDParamsSubset{{2, 2, 2, 2}, {2, 2, 1}, 2},
 };
 
+const std::vector<GatherNDParamsSubset> layerParamsND8 = {
+    GatherNDParamsSubset{{500, 256, 10, 15}, {25, 125, 3}, 0},
+    GatherNDParamsSubset{{3, 3}, {2, 2}, 0},
+    GatherNDParamsSubset{{5, 3}, {2, 1}, 0},
+    GatherNDParamsSubset{{5, 3, 4}, {2, 2}, 0},
+    GatherNDParamsSubset{{6, 3, 4}, {2, 1, 2}, 0},
+    GatherNDParamsSubset{{5, 2, 6, 8}, {1}, 0},
+    GatherNDParamsSubset{{6, 6, 9, 7}, {2}, 0},
+    GatherNDParamsSubset{{2, 4, 9, 4}, {3}, 0},
+    GatherNDParamsSubset{{5, 2, 3, 7}, {4}, 0},
+};
+
 INSTANTIATE_TEST_SUITE_P(
     smoke_GatherND,
     GatherNDLayerTest,
@@ -50,4 +62,15 @@ INSTANTIATE_TEST_SUITE_P(
         testing::Values<Config>({{InferenceEngine::MYRIAD_DETECT_NETWORK_BATCH, CONFIG_VALUE(NO)}})),
     GatherNDLayerTest::getTestCaseName);
 
+INSTANTIATE_TEST_SUITE_P(
+    smoke_GatherND,
+    GatherND8LayerTest,
+    testing::Combine(
+        testing::ValuesIn(layerParamsND8),
+        testing::ValuesIn(netPrecisions),
+        testing::ValuesIn(indicesPrecisions),
+        testing::Values(CommonTestUtils::DEVICE_MYRIAD),
+        testing::Values<Config>({{InferenceEngine::MYRIAD_DETECT_NETWORK_BATCH, CONFIG_VALUE(NO)}})),
+    GatherND8LayerTest::getTestCaseName);
+
 }  // namespace