Variadic Split

inference-engine/src/plaidml_plugin/ops/variadic_split.cpp

@@ -0,0 +1,56 @@
+// Copyright (C) 2020 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "plaidml_ops.hpp"
+#include "plaidml_util.hpp"
+
+#include "ngraph/opsets/opset.hpp"
+#include "ngraph/opsets/opset1.hpp"
+
+#include "plaidml/op/op.h"
+
+using namespace plaidml;          // NOLINT[build/namespaces]
+
+namespace PlaidMLPlugin {
+
+static OpRegistration reg("variadicsplit", [](const Context& ctx) {
+  IE_ASSERT(ctx.operands.size() == 3);
+  auto I = ctx.operands.at(0);
+  auto axes = get_axis_vector_from_constant_operand(1, ctx.layer);
+  IE_ASSERT(axes.size() == 1);
+  auto axis = axes[0];
+  auto split_lengths = get_shape_from_constant_operand(2, ctx.layer);
+
+  auto ndims = I.rank();
+  std::vector<edsl::TensorDim> I_dims(ndims);
+  std::vector<edsl::TensorIndex> I_idxs(ndims);
+  std::vector<edsl::Tensor> Os;
+  I.bind_dims(I_dims);
+  auto O_dims = I_dims;
+
+  size_t split_size = 0;
+  for (auto split : split_lengths) {
+      if (split != -1) {
+        split_size += split;
+      }
+  }
+  auto placeholder = I_dims[axis] - split_size;
+
+  edsl::TensorDim offset(0);
+  for (auto split : split_lengths) {
+    auto O_idxs = I_idxs;
+    O_idxs[axis] = I_idxs[axis] - offset;
+    if (split == -1) {
+      O_dims[axis] = placeholder;
+      offset = offset + placeholder;
+    } else {
+      O_dims[axis] = static_cast<edsl::TensorDim>(split);
+      offset = offset + split;
+    }
+    Os.push_back(plaidml::edsl::Contraction().outShape(O_dims).outAccess(O_idxs).assign(I(I_idxs)));
+  }
+  return edsl::make_tuple(Os);
+});
+
+}  // namespace PlaidMLPlugin

inference-engine/tests/functional/plugin/plaidml/shared_tests_instances/single_layer_tests/variadic_split.cpp

@@ -0,0 +1,36 @@
+// Copyright (C) 2019 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <vector>
+
+#include "single_layer_tests/variadic_split.hpp"
+#include "common_test_utils/test_constants.hpp"
+
+using namespace LayerTestsDefinitions;
+
+namespace {
+
+    const std::vector<InferenceEngine::Precision> netPrecisions = {
+            InferenceEngine::Precision::FP32,
+            // InferenceEngine::Precision::FP16
+    };
+
+    // Sum of elements numSplits = inputShapes[Axis]
+    const std::vector<std::vector<size_t>> numSplits = {
+            {1, 16, 5, 8},
+            {2, 19, 5, 4},
+            {7, 13, 2, 8},
+            {5, 8, 12, 5},
+            {4, 11, 6, 9}
+    };
+
+    INSTANTIATE_TEST_CASE_P(NumSplitsCheck, VariadicSplitLayerTest,
+            ::testing::Combine(
+            ::testing::ValuesIn(numSplits),
+            ::testing::Values(0, 1, 2, 3),
+            ::testing::ValuesIn(netPrecisions),
+            ::testing::Values(std::vector<size_t>({30, 30, 30, 30})),
+            ::testing::Values(CommonTestUtils::DEVICE_PLAIDML)),
+            VariadicSplitLayerTest::getTestCaseName);
+}  // namespace