[GPU] Einsum with repeated labels and ellipsis support (#11615)

* Einsum test helper

* Einsum single layer tests

* Add Einsum decomposition with repeated labels and ellipsis support
to GPU transformations pipeline

Co-authored-by: Oleksii Khovan <okhovan@lohika.com>

src/tests/functional/plugin/gpu/shared_tests_instances/single_layer_tests/einsum.cpp

@@ -0,0 +1,42 @@
+// Copyright (C) 2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <vector>
+
+#include "single_layer_tests/einsum.hpp"
+
+using namespace ngraph::helpers;
+using namespace LayerTestsDefinitions;
+
+namespace {
+const std::vector<InferenceEngine::Precision> precisions = {
+        InferenceEngine::Precision::FP32,
+        InferenceEngine::Precision::FP16
+};
+
+const std::vector<EinsumEquationWithInput> equationsWithInput = {
+    { "ij->ji", {{{1, 2}}} }, // transpose 2d
+    { "ijk->kij", { {1, 2, 3} } }, // transpose 3d
+    { "ij->i", { {2, 3} } }, // reduce
+    { "ab,cd->abcd", { { 1, 2}, {3, 4} } }, // no reduction
+    { "ab,bc->ac", { {2, 3}, {3, 2} } }, // matrix multiplication
+    { "ab,bcd,bc->ca", { {2, 4}, {4, 3, 1}, {4, 3} } },  // multiple multiplications
+    { "kii->ki", { {1, 3, 3} } }, // diagonal
+    { "abbac,bad->ad", { {2, 3, 3, 2, 4}, {3, 2, 1} } }, // diagonal and multiplication with repeated labels
+    { "a...->...a", { {2, 2, 3} } }, // transpose with ellipsis
+    { "a...->...", { {2, 2, 3} } }, // reduce with ellipsis
+    { "ab...,...->ab...", { {2, 2, 3}, {1} } }, // multiply by scalar
+    { "a...j,j...->a...", { {1, 1, 4, 3}, {3, 4, 2, 1} } } // complex multiplication
+};
+
+const auto params = ::testing::Combine(
+        ::testing::ValuesIn(precisions),
+        ::testing::ValuesIn(equationsWithInput),
+        ::testing::Values(CommonTestUtils::DEVICE_GPU));
+
+INSTANTIATE_TEST_SUITE_P(smoke_Einsum, EinsumLayerTest,
+        params,
+        EinsumLayerTest::getTestCaseName);
+
+}  // namespace

src/tests/functional/plugin/shared/include/single_layer_tests/einsum.hpp

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

src/tests/functional/shared_test_classes/include/shared_test_classes/single_layer/einsum.hpp

@@ -0,0 +1,35 @@
+// Copyright (C) 2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <string>
+#include <tuple>
+#include <vector>
+
+#include "shared_test_classes/base/layer_test_utils.hpp"
+
+namespace LayerTestsDefinitions {
+
+typedef std::tuple<
+        std::string,                       // Equation
+        std::vector<std::vector<size_t>>   // Input shapes
+> EinsumEquationWithInput;
+
+typedef std::tuple<
+        InferenceEngine::Precision,         // Input precision
+        EinsumEquationWithInput,            // Equation with corresponding input shapes
+        std::string                         // Device name
+> EinsumLayerTestParamsSet;
+
+class EinsumLayerTest : public testing::WithParamInterface<EinsumLayerTestParamsSet>,
+                              virtual public LayerTestsUtils::LayerTestsCommon {
+public:
+    static std::string getTestCaseName(const testing::TestParamInfo<EinsumLayerTestParamsSet>& obj);
+
+protected:
+    void SetUp() override;
+};
+
+}  // namespace LayerTestsDefinitions

src/tests/functional/shared_test_classes/src/single_layer/einsum.cpp

@@ -0,0 +1,44 @@
+// Copyright (C) 2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "shared_test_classes/single_layer/einsum.hpp"
+#include "ngraph_functions/builders.hpp"
+
+namespace LayerTestsDefinitions {
+
+std::string EinsumLayerTest::getTestCaseName(const testing::TestParamInfo<EinsumLayerTestParamsSet>& obj) {
+    InferenceEngine::Precision precision;
+    EinsumEquationWithInput equationWithInput;
+    std::string targetDevice;
+    std::tie(precision, equationWithInput, targetDevice) = obj.param;
+    std::string equation;
+    std::vector<InferenceEngine::SizeVector> inputShapes;
+    std::tie(equation, inputShapes) = equationWithInput;
+
+    std::ostringstream result;
+    result << "PRC=" << precision.name() << "_";
+    result << "IS=" << CommonTestUtils::vec2str(inputShapes) << "_";
+    result << "Eq=" << equation << "_";
+    result << "trgDev=" << targetDevice;
+    return result.str();
+}
+
+void EinsumLayerTest::SetUp() {
+    InferenceEngine::Precision precision;
+    EinsumEquationWithInput equationWithInput;
+    std::tie(precision, equationWithInput, targetDevice) = this->GetParam();
+    std::string equation;
+    std::vector<InferenceEngine::SizeVector> inputShapes;
+    std::tie(equation, inputShapes) = equationWithInput;
+
+    const auto ngPrc = FuncTestUtils::PrecisionUtils::convertIE2nGraphPrc(precision);
+    const auto params = ngraph::builder::makeParams(ngPrc, inputShapes);
+    const auto paramOuts = ngraph::helpers::convert2OutputVector(ngraph::helpers::castOps2Nodes<ngraph::op::Parameter>(params));
+
+    const std::shared_ptr<ngraph::Node> einsum = ngraph::builder::makeEinsum(paramOuts, equation);
+    const ngraph::ResultVector results{std::make_shared<ngraph::opset3::Result>(einsum)};
+    function = std::make_shared<ngraph::Function>(results, params, "einsum");
+}
+
+}  // namespace LayerTestsDefinitions

src/tests/ngraph_helpers/ngraph_functions/include/ngraph_functions/builders.hpp

@@ -587,5 +587,8 @@ std::shared_ptr<ngraph::Node> makeDFT(const ngraph::Output<Node> &dataNode,
                                       const std::vector<int64_t> &axes,
                                       const std::vector<int64_t> &signalSize,
                                       const ngraph::helpers::DFTOpType opType);
+
+std::shared_ptr<ngraph::Node> makeEinsum(const OutputVector& inputs,
+                                         const std::string& equation);
 }  // namespace builder
 }  // namespace ngraph

src/tests/ngraph_helpers/ngraph_functions/src/einsum.cpp

@@ -0,0 +1,20 @@
+// Copyright (C) 2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <memory>
+#include <string>
+
+#include "ngraph_functions/builders.hpp"
+
+namespace ngraph {
+namespace builder {
+
+    std::shared_ptr<ngraph::Node> makeEinsum(const OutputVector& inputs,
+                                             const std::string& equation) {
+        std::shared_ptr<ngraph::Node> einsum = std::make_shared<ngraph::opset7::Einsum>(inputs, equation);
+        return einsum;
+    }
+
+}  // namespace builder
+}  // namespace ngraph