[GPU] Add new operation GatherElements to IE clDNN plugin

inference-engine/src/cldnn_engine/cldnn_primitives_list.hpp

@@ -204,6 +204,7 @@ REGISTER_FACTORY(v5, Loop);
 // ------------------------------ Supported v6 ops ------------------------------ //
 REGISTER_FACTORY(v6, CTCGreedyDecoderSeqLen);
 REGISTER_FACTORY(v6, MVN);
+REGISTER_FACTORY(v6, GatherElements);
 
 // ------------------------------ Supported v7 ops ------------------------------ //
 REGISTER_FACTORY(v7, Gather);

inference-engine/src/cldnn_engine/ops/gather_elements.cpp

@@ -0,0 +1,66 @@
+// Copyright (C) 2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "cldnn_program.h"
+#include "cldnn_common_utils.h"
+
+#include "ngraph/op/gather_elements.hpp"
+#include "ngraph/op/constant.hpp"
+
+#include "cldnn/primitives/gather_elements.hpp"
+
+namespace CLDNNPlugin {
+
+static cldnn::gather_elements::gather_elements_axis GetGatherElementsAxis(int axis, unsigned rank) {
+    if (axis < 0)
+        axis += rank;
+    if (axis < 0 || axis >= rank)
+        IE_THROW() << "GatherElements axis is not correspond to number of dimensions";
+
+    // Difference in dimension ordering between IE and clDNN,
+    // reverse spatial dimensions after batch and feature.
+    unsigned cldnn_axis = axis;
+    if (axis >= 2) {
+        auto spatial_axis = axis - 2;
+        // Default and minimum number of dimensions is 4
+        auto spatial_size = std::max(rank, 4u) - 2;
+        cldnn_axis = spatial_size - spatial_axis - 1 + 2;
+    }
+
+    switch (cldnn_axis) {
+        case 0: return cldnn::gather_elements::gather_elements_axis::along_b;
+        case 1: return cldnn::gather_elements::gather_elements_axis::along_f;
+        case 2: return cldnn::gather_elements::gather_elements_axis::along_x;
+        case 3: return cldnn::gather_elements::gather_elements_axis::along_y;
+        case 4: return cldnn::gather_elements::gather_elements_axis::along_z;
+        case 5: return cldnn::gather_elements::gather_elements_axis::along_w;
+        default: IE_THROW() << "Unsupported GatherElements axis: " << axis;
+    }
+    return cldnn::gather_elements::gather_elements_axis::along_f;  // shouldn't get here
+}
+
+void CreateGatherElementsOp(Program& p, const std::shared_ptr<ngraph::op::v6::GatherElements>& op) {
+    p.ValidateInputs(op, {2});
+    auto inputPrimitives = p.GetInputPrimitiveIDs(op);
+    std::string layerName = layer_type_name_ID(op);
+
+    size_t rank = op->get_input_shape(0).size();
+    int32_t axis = static_cast<int32_t>(op->get_axis());
+
+    auto outLayout = DefaultFormatForDims(op->get_output_shape(0).size());
+
+    auto primitive = cldnn::gather_elements(layerName,
+                                            inputPrimitives[0],
+                                            inputPrimitives[1],
+                                            outLayout,
+                                            CldnnTensorFromIEDims(op->get_output_shape(0)),
+                                            GetGatherElementsAxis(axis, rank));
+
+    p.AddPrimitive(primitive);
+    p.AddPrimitiveToProfiler(op);
+}
+
+REGISTER_FACTORY_IMPL(v6, GatherElements);
+
+}  // namespace CLDNNPlugin

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

@@ -4,7 +4,8 @@
 
 #include <vector>
 
-#include "shared_test_classes/single_layer/gather_elements.hpp"
+#include "single_layer_tests/gather_elements.hpp"
+#include "common_test_utils/test_constants.hpp"
 
 using namespace LayerTestsDefinitions;
 
@@ -73,4 +74,4 @@ INSTANTIATE_TEST_SUITE_P(smoke_set5, GatherElementsLayerTest,
                             ::testing::ValuesIn(iPrecisions),
                             ::testing::Values(CommonTestUtils::DEVICE_CPU)),
                         GatherElementsLayerTest::getTestCaseName);
-}  // namespace
+}  // namespace

inference-engine/tests/functional/plugin/gpu/shared_tests_instances/single_layer_tests/gather_elements.cpp

@@ -0,0 +1,227 @@
+// Copyright (C) 2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+// //
+
+#include <vector>
+#include <ngraph/opsets/opset6.hpp>
+
+#include "single_layer_tests/gather_elements.hpp"
+#include "common_test_utils/test_constants.hpp"
+
+using namespace LayerTestsDefinitions;
+using namespace ngraph::opset6;
+
+namespace {
+
+const std::vector<InferenceEngine::Precision> inputPrecisions = {
+        InferenceEngine::Precision::FP32,
+        InferenceEngine::Precision::FP16,
+        InferenceEngine::Precision::I32,
+};
+
+const std::vector<InferenceEngine::Precision> idxPrecisions = {
+        InferenceEngine::Precision::I32,
+        InferenceEngine::Precision::I64,
+};
+
+INSTANTIATE_TEST_CASE_P(smoke_set1, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>({2, 2})),
+        ::testing::Values(std::vector<size_t>({2, 2})),
+        ::testing::ValuesIn(std::vector<int>({-1, 0, 1})),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_set2, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>({2, 2, 1})),
+        ::testing::Values(std::vector<size_t>({4, 2, 1})),
+        ::testing::ValuesIn(std::vector<int>({0, -3})),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_set3, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>({2, 2, 3, 5})),
+        ::testing::Values(std::vector<size_t>({2, 2, 3, 7})),
+        ::testing::Values(3, -1),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_set4, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>({3, 2, 3, 8})),
+        ::testing::Values(std::vector<size_t>({2, 2, 3, 8})),
+        ::testing::Values(0, -4),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_set5, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>({3, 2, 3, 4, 8})),
+        ::testing::Values(std::vector<size_t>({3, 2, 3, 5, 8})),
+        ::testing::Values(3, -2),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank4axis0, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{7, 7, 8, 4}),
+        ::testing::Values(std::vector<size_t>{2, 7, 8, 4}),
+        ::testing::Values(0),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank4axis1, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{6, 1, 8, 4}),
+        ::testing::Values(std::vector<size_t>{6, 8, 8, 4}),
+        ::testing::Values(1, -3),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank4axis2, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{6, 7, 4, 4}),
+        ::testing::Values(std::vector<size_t>{6, 7, 2, 4}),
+        ::testing::Values(2, -2),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank4axis3, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{6, 5, 8, 7}),
+        ::testing::Values(std::vector<size_t>{6, 5, 8, 7}),
+        ::testing::Values(3, -1),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank5axis0, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{2, 3, 9, 4, 9}),
+        ::testing::Values(std::vector<size_t>{1, 3, 9, 4, 9}),
+        ::testing::Values(0),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank5axis1, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{2, 3, 5, 4, 7}),
+        ::testing::Values(std::vector<size_t>{2, 9, 5, 4, 7}),
+        ::testing::Values(1, -4),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank5axis2, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{1, 2, 6, 8, 9}),
+        ::testing::Values(std::vector<size_t>{1, 2, 6, 8, 9}),
+        ::testing::Values(2, -3),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank5axis3, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{2, 2, 4, 7, 7}),
+        ::testing::Values(std::vector<size_t>{2, 2, 4, 3, 7}),
+        ::testing::Values(3, -2),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank5axis4, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{1, 3, 9, 3, 2}),
+        ::testing::Values(std::vector<size_t>{1, 3, 9, 3, 9}),
+        ::testing::Values(4, -1),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank6axis0, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{3, 3, 2, 4, 4, 3}),
+        ::testing::Values(std::vector<size_t>{7, 3, 2, 4, 4, 3}),
+        ::testing::Values(0),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank6axis1, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{1, 6, 2, 3, 5, 9}),
+        ::testing::Values(std::vector<size_t>{1, 6, 2, 3, 5, 9}),
+        ::testing::Values(1, -5),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank6axis2, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{2, 3, 9, 7, 2, 1}),
+        ::testing::Values(std::vector<size_t>{2, 3, 5, 7, 2, 1}),
+        ::testing::Values(2, -4),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank6axis3, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{1, 3, 4, 5, 1, 3}),
+        ::testing::Values(std::vector<size_t>{1, 3, 4, 4, 1, 3}),
+        ::testing::Values(3, -3),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank6axis4, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{1, 3, 2, 4, 3, 3}),
+        ::testing::Values(std::vector<size_t>{1, 3, 2, 4, 6, 3}),
+        ::testing::Values(4, -2),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+INSTANTIATE_TEST_CASE_P(smoke_GatherElements_rank6axis5, GatherElementsLayerTest,
+    ::testing::Combine(
+        ::testing::Values(std::vector<size_t>{2, 1, 7, 8, 1, 6}),
+        ::testing::Values(std::vector<size_t>{2, 1, 7, 8, 1, 5}),
+        ::testing::Values(5, -1),
+        ::testing::ValuesIn(inputPrecisions),
+        ::testing::ValuesIn(idxPrecisions),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    GatherElementsLayerTest::getTestCaseName);
+
+}  // namespace

inference-engine/tests/functional/plugin/shared/include/single_layer_tests/gather_elements.hpp

@@ -0,0 +1,15 @@
+// Copyright (C) 2018-2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "shared_test_classes/single_layer/gather_elements.hpp"
+
+namespace LayerTestsDefinitions {
+
+TEST_P(GatherElementsLayerTest, CompareWithRefs) {
+    Run();
+}
+
+}  // namespace LayerTestsDefinitions

inference-engine/tests/functional/shared_test_classes/src/single_layer/gather_elements.cpp

@@ -48,7 +48,4 @@ void GatherElementsLayerTest::SetUp() {
     function = std::make_shared<ngraph::Function>(results, params, "gatherEl");
 }
 
-TEST_P(GatherElementsLayerTest, CompareWithRefs) {
-    Run();
-}
 }  // namespace LayerTestsDefinitions