[IE CLDNN] Implement EmbeddingBag operations

inference-engine/src/cldnn_engine/cldnn_program.cpp

@@ -63,6 +63,7 @@
 #include <api/grn.hpp>
 #include <api/ctc_greedy_decoder.hpp>
 #include <api/cum_sum.hpp>
+#include <api/embedding_bag.hpp>
 
 #include <chrono>
 #include <cmath>
@@ -598,6 +599,9 @@ Program::LayerType Program::LayerTypeFromStr(const std::string &str) {
         { "CTCGreedyDecoder", CTCGreedyDecoder },
         { "PriorBoxClustered", PriorBoxClustered },
         { "CumSum", CumSum },
+        { "EmbeddingBagPackedSum", EmbeddingBagPackedSum },
+        { "EmbeddingBagOffsetsSum", EmbeddingBagOffsetsSum },
+        { "EmbeddingSegmentsSum", EmbeddingSegmentsSum },
     };
     auto it = LayerNameToType.find(str);
     if (it != LayerNameToType.end())
@@ -1280,6 +1284,12 @@ void Program::CreateSingleLayerPrimitive(cldnn::topology& topology, InferenceEng
             break;
         case CumSum: CreateCumSumPrimitive(topology, layer);
             break;
+        case EmbeddingBagPackedSum: CreateEmbeddingBagPackedSumPrimitive(topology, layer);
+            break;
+        case EmbeddingBagOffsetsSum: CreateEmbeddingBagOffsetsSumPrimitive(topology, layer);
+            break;
+        case EmbeddingSegmentsSum: CreateEmbeddingSegmentsSumPrimitive(topology, layer);
+            break;
         default: THROW_CLDNN_EXCEPTION("Unknown Layer Type: " << layer->type);
     }
 }
@@ -4487,6 +4497,95 @@ void Program::CreatePriorBoxClusteredPrimitive(cldnn::topology& topology, Infere
     AddPrimitiveToProfiler(priorBoxLayerName, layer);
 }
 
+void Program::CreateEmbeddingBagPackedSumPrimitive(cldnn::topology& topology, InferenceEngine::CNNLayerPtr& layer) {
+    ValidateLayer(layer, {2, 3});
+
+    auto inputPrimitives = GetPrevLayersPrimitives(layer);
+    auto embeddingBag = as<InferenceEngine::GenericLayer*>(layer);
+
+    auto layerName = layer_type_name_ID(layer);
+    auto embeddingBagPrim = cldnn::embedding_bag(
+            layerName,
+            inputPrimitives,
+            cldnn::embedding_bag::packed_sum,
+            CldnnTensorFromIEDims(embeddingBag->outData[0]->getTensorDesc().getDims()));
+
+    topology.add(embeddingBagPrim);
+    AddPrimitiveToProfiler(layerName, layer);
+}
+
+void Program::CreateEmbeddingBagOffsetsSumPrimitive(cldnn::topology& topology, InferenceEngine::CNNLayerPtr& layer) {
+    ValidateLayer(layer, {3, 4, 5});
+
+    auto inputPrimitives = GetPrevLayersPrimitives(layer);
+    auto embeddingBag = as<InferenceEngine::GenericLayer*>(layer);
+
+    int32_t defaultIndex = -1;
+    if (inputPrimitives.size() > 3) {
+        auto defaultIndexInput = layer->insData[3].lock();
+        auto defaultIndexInputCreator = defaultIndexInput->getCreatorLayer().lock();
+        if (defaultIndexInputCreator->blobs.size() == 1) {
+            auto constantBlob = defaultIndexInputCreator->blobs.begin()->second;
+            auto defaultIndexPrecision = constantBlob->getTensorDesc().getPrecision();
+            if (defaultIndexPrecision == InferenceEngine::Precision::I32) {
+                auto data = constantBlob->buffer().as<int32_t*>();
+                defaultIndex = data[0];
+            } else {
+                THROW_IE_EXCEPTION << layer->name << "Incorrect EmbeddingBagOfsetsSum default_index precision";
+            }
+        }
+        inputPrimitives.erase(inputPrimitives.begin() + 3); // Remove "default_index"
+    }
+
+    auto layerName = layer_type_name_ID(layer);
+    auto embeddingBagPrim = cldnn::embedding_bag(
+            layerName,
+            inputPrimitives,
+            cldnn::embedding_bag::offsets_sum,
+            CldnnTensorFromIEDims(embeddingBag->outData[0]->getTensorDesc().getDims()),
+            defaultIndex);
+
+    topology.add(embeddingBagPrim);
+    AddPrimitiveToProfiler(layerName, layer);
+}
+
+void Program::CreateEmbeddingSegmentsSumPrimitive(cldnn::topology& topology, InferenceEngine::CNNLayerPtr& layer) {
+    ValidateLayer(layer, {4, 5, 6});
+
+    auto inputPrimitives = GetPrevLayersPrimitives(layer);
+    auto embeddingBag = as<InferenceEngine::GenericLayer*>(layer);
+
+    inputPrimitives.erase(inputPrimitives.begin() + 3); // Remove "num_segments"
+
+    int32_t defaultIndex = -1;
+    if (inputPrimitives.size() > 3) {
+        auto defaultIndexInput = layer->insData[4].lock();
+        auto defaultIndexInputCreator = defaultIndexInput->getCreatorLayer().lock();
+        if (defaultIndexInputCreator->blobs.size() == 1) {
+            auto constantBlob = defaultIndexInputCreator->blobs.begin()->second;
+            auto defaultIndexPrecision = constantBlob->getTensorDesc().getPrecision();
+            if (defaultIndexPrecision == InferenceEngine::Precision::I32) {
+                auto data = constantBlob->buffer().as<int32_t*>();
+                defaultIndex = data[0];
+            } else {
+                THROW_IE_EXCEPTION << layer->name << "Incorrect EmbeddingBagOfsetsSum default_index precision";
+            }
+        }
+        inputPrimitives.erase(inputPrimitives.begin() + 3); // Remove "default_index"
+    }
+
+    auto layerName = layer_type_name_ID(layer);
+    auto embeddingBagPrim = cldnn::embedding_bag(
+            layerName,
+            inputPrimitives,
+            cldnn::embedding_bag::segments_sum,
+            CldnnTensorFromIEDims(embeddingBag->outData[0]->getTensorDesc().getDims()),
+            defaultIndex);
+
+    topology.add(embeddingBagPrim);
+    AddPrimitiveToProfiler(layerName, layer);
+}
+
 bool Program::IsValidSplitConvMerge(const InferenceEngine::SplitLayer *splitLayer) const {
     if (splitLayer->outData.size() != 2) return false;  // split into 2
 

inference-engine/src/cldnn_engine/cldnn_program.h

@@ -217,6 +217,9 @@ class Program {
         CTCGreedyDecoder,
         PriorBoxClustered,
         CumSum,
+        EmbeddingBagPackedSum,
+        EmbeddingBagOffsetsSum,
+        EmbeddingSegmentsSum,
         NO_TYPE
     };
     using GenericBlobMap = std::map<cldnn::primitive_id, cldnn::primitive_id>;
@@ -373,6 +376,9 @@ class Program {
     void CreateCTCGreedyDecoderPrimitive(cldnn::topology& topology, InferenceEngine::CNNLayerPtr& layer);
     void CreatePriorBoxClusteredPrimitive(cldnn::topology& topology, InferenceEngine::CNNLayerPtr& layer);
     void CreateCumSumPrimitive(cldnn::topology& topology, InferenceEngine::CNNLayerPtr& layer);
+    void CreateEmbeddingBagPackedSumPrimitive(cldnn::topology& topology, InferenceEngine::CNNLayerPtr& layer);
+    void CreateEmbeddingBagOffsetsSumPrimitive(cldnn::topology& topology, InferenceEngine::CNNLayerPtr& layer);
+    void CreateEmbeddingSegmentsSumPrimitive(cldnn::topology& topology, InferenceEngine::CNNLayerPtr& layer);
 };
 
 }  // namespace CLDNNPlugin

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

@@ -0,0 +1,59 @@
+// Copyright (C) 2020 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "single_layer_tests/embedding_bag_offsets_sum.hpp"
+
+#include <vector>
+
+#include "common_test_utils/test_constants.hpp"
+
+using namespace LayerTestsDefinitions;
+
+namespace {
+
+const std::vector<InferenceEngine::Precision> netPrecisions = {
+    InferenceEngine::Precision::FP32,
+    InferenceEngine::Precision::FP16
+};
+
+const std::vector<InferenceEngine::Precision> indPrecisions = {
+    InferenceEngine::Precision::I64,
+    InferenceEngine::Precision::I32
+};
+
+const std::vector<std::vector<size_t>> emb_table_shape = {
+    {5, 6},
+    {10, 35},
+    {5, 4, 16}
+};
+
+const std::vector<std::vector<size_t>> indices = {
+    {0, 1, 2, 2, 3},
+    {4, 4, 3, 1, 0},
+    {1, 2, 1, 2, 1, 2, 1, 2, 1, 2}
+};
+
+const std::vector<std::vector<size_t>> offsets = {
+    {0, 2},
+    {0, 0, 2, 2},
+    {2, 4}
+};
+
+const std::vector<size_t> default_index = {0, 4};
+const std::vector<bool> with_weights = {false, true};
+const std::vector<bool> with_default_index = {false, true};
+
+const auto embBagOffsetSumArgSet = ::testing::Combine(
+    ::testing::ValuesIn(emb_table_shape), ::testing::ValuesIn(indices),
+    ::testing::ValuesIn(offsets), ::testing::ValuesIn(default_index),
+    ::testing::ValuesIn(with_weights), ::testing::ValuesIn(with_default_index));
+
+INSTANTIATE_TEST_CASE_P(
+    smoke, EmbeddingBagOffsetsSumLayerTest,
+    ::testing::Combine(embBagOffsetSumArgSet,
+                       ::testing::ValuesIn(netPrecisions),
+                       ::testing::ValuesIn(indPrecisions),
+                       ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    EmbeddingBagOffsetsSumLayerTest::getTestCaseName);
+}  // namespace

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

@@ -0,0 +1,49 @@
+// Copyright (C) 2020 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "single_layer_tests/embedding_bag_packed_sum.hpp"
+
+#include <vector>
+
+#include "common_test_utils/test_constants.hpp"
+
+using namespace LayerTestsDefinitions;
+
+namespace {
+
+const std::vector<InferenceEngine::Precision> netPrecisions = {
+    InferenceEngine::Precision::FP32,
+    InferenceEngine::Precision::FP16,
+};
+
+const std::vector<InferenceEngine::Precision> indPrecisions = {
+    InferenceEngine::Precision::I64,
+    InferenceEngine::Precision::I32
+};
+
+const std::vector<std::vector<size_t>> emb_table_shape = {
+    {5, 6},
+    {10, 35},
+    {5, 4, 16}
+};
+
+const std::vector<std::vector<std::vector<size_t>>> indices = {
+    {{0, 1}, {2, 2}, {3, 4}},
+    {{4, 4, 3}, {1, 0, 2}},
+    {{1, 2, 1, 2}, {1, 2, 1, 2}}
+};
+const std::vector<bool> with_weights = {false, true};
+
+const auto embBagPackedSumArgSet = ::testing::Combine(
+    ::testing::ValuesIn(emb_table_shape), ::testing::ValuesIn(indices),
+    ::testing::ValuesIn(with_weights));
+
+INSTANTIATE_TEST_CASE_P(
+    smoke, EmbeddingBagPackedSumLayerTest,
+    ::testing::Combine(embBagPackedSumArgSet,
+                       ::testing::ValuesIn(netPrecisions),
+                       ::testing::ValuesIn(indPrecisions),
+                       ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    EmbeddingBagPackedSumLayerTest::getTestCaseName);
+}  // namespace

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

@@ -0,0 +1,54 @@
+// Copyright (C) 2020 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <vector>
+
+#include "single_layer_tests/embedding_segments_sum.hpp"
+#include "common_test_utils/test_constants.hpp"
+
+using namespace LayerTestsDefinitions;
+
+namespace {
+
+const std::vector<InferenceEngine::Precision> netPrecisions = {
+    InferenceEngine::Precision::FP32,
+    InferenceEngine::Precision::FP16
+};
+
+const std::vector<InferenceEngine::Precision> indPrecisions = {
+    InferenceEngine::Precision::I64,
+    InferenceEngine::Precision::I32
+};
+
+const std::vector<std::vector<size_t>> emb_table_shape = {
+    {5, 6},
+    {10, 35},
+    {5, 4, 16}
+};
+const std::vector<std::vector<size_t>> indices = {
+    {0, 1, 2, 2, 3},
+    {4, 4, 3, 1, 2}
+};
+const std::vector<std::vector<size_t>> segment_ids = {
+    {0, 1, 2, 3, 4},
+    {0, 0, 2, 2, 4}
+};
+const std::vector<size_t> num_segments = {5, 7};
+const std::vector<size_t> default_index = {0, 4};
+const std::vector<bool> with_weights = {false, true};
+const std::vector<bool> with_default_index = {false, true};
+
+const auto embSegmentsSumArgSet = ::testing::Combine(
+    ::testing::ValuesIn(emb_table_shape), ::testing::ValuesIn(indices),
+    ::testing::ValuesIn(segment_ids), ::testing::ValuesIn(num_segments),
+    ::testing::ValuesIn(default_index), ::testing::ValuesIn(with_weights),
+    ::testing::ValuesIn(with_default_index));
+
+INSTANTIATE_TEST_CASE_P(
+    smoke, EmbeddingSegmentsSumLayerTest,
+    ::testing::Combine(embSegmentsSumArgSet, ::testing::ValuesIn(netPrecisions),
+                       ::testing::ValuesIn(indPrecisions),
+                       ::testing::Values(CommonTestUtils::DEVICE_GPU)),
+    EmbeddingSegmentsSumLayerTest::getTestCaseName);
+}  // namespace

inference-engine/thirdparty/clDNN/api/embedding_bag.hpp

@@ -0,0 +1,64 @@
+/*
+// Copyright (c) 2020 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+*/
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+#pragma once
+#include "primitive.hpp"
+
+namespace cldnn {
+/// @addtogroup cpp_api C++ API
+/// @{
+/// @addtogroup cpp_topology Network Topology
+/// @{
+/// @addtogroup cpp_primitives Primitives
+/// @{
+
+/// @brief Computes sums of "bags" of embeddings, without instantiating the intermediate embeddings.
+/// @details For each index in `indices` this operator gets values from `data` embedding table and sums all values belonging to each bag.
+struct embedding_bag : public primitive_base<embedding_bag> {
+    CLDNN_DECLARE_PRIMITIVE(embedding_bag)
+
+    /// @brief Select type of embedding_bag operation
+    enum embedding_bag_type {
+        packed_sum,
+        offsets_sum,
+        segments_sum
+    };
+
+    /// @brief Constructs embedding_bag primitive.
+    /// @param id This primitive id.
+    /// @param inputs Vector with different inputs.
+    /// @param output_shape Tensor with shape of output layout
+    /// @param default_index default index in embedding table to fill empty "bags"
+    embedding_bag(const primitive_id& id,
+                  const std::vector<primitive_id>& inputs,
+                  const embedding_bag_type& type,
+                  const tensor& output_shape,
+                  const int32_t default_index = -1,
+                  const padding& output_padding = padding())
+        : primitive_base(id, inputs, output_padding), type(type), output_shape(output_shape), default_index(default_index) {}
+
+    /// @brief Type of EmbeddingBag operation
+    embedding_bag_type type;
+    /// @brief Shape of output layout
+    tensor output_shape;
+    /// @brief Default index
+    int32_t default_index;
+};
+/// @}
+/// @}
+/// @}
+}  // namespace cldnn