Variadic Split

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

@@ -11,6 +11,22 @@
 #include "plaidml/op/op.h"
 
 using namespace plaidml;          // NOLINT[build/namespaces]
+using namespace InferenceEngine;  // NOLINT[build/namespaces]
+
+namespace {
+
+template <typename T>
+std::vector<T> cast_constant_operand(size_t operand_idx, ngraph::Node* layer) {
+  auto* ngraph_const = ngraph::as_type<ngraph::op::Constant>(layer->get_input_node_ptr(operand_idx));
+  if (ngraph_const) {
+    return ngraph_const->cast_vector<T>();
+  } else {
+    THROW_IE_EXCEPTION << "Dynamic padding not currently supported by PlaidML plugin; all of pads_begin, pads_end, "
+                          "and pads_value must be Constants";
+  }
+}
+
+}  // namespace
 
 namespace PlaidMLPlugin {
 
@@ -20,7 +36,7 @@ static OpRegistration reg("variadicsplit", [](const Context& ctx) {
   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 split_lengths = cast_constant_operand<int32_t>(2, ctx.layer);
 
   auto ndims = I.rank();
   std::vector<edsl::TensorDim> I_dims(ndims);
@@ -45,10 +61,10 @@ static OpRegistration reg("variadicsplit", [](const Context& ctx) {
       O_dims[axis] = placeholder;
       offset = offset + placeholder;
     } else {
-      O_dims[axis] = static_cast<edsl::TensorDim>(split);
+      O_dims[axis] = edsl::TensorDim(split);
       offset = offset + split;
     }
-    Os.push_back(plaidml::edsl::Contraction().outShape(O_dims).outAccess(O_idxs).assign(I(I_idxs)));
+    Os.push_back(edsl::Contraction().outShape(O_dims).outAccess(O_idxs).assign(I(I_idxs)));
   }
   return edsl::make_tuple(Os);
 });

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

@@ -17,12 +17,12 @@ namespace {
     };
 
     // Sum of elements numSplits = inputShapes[Axis]
-    const std::vector<std::vector<size_t>> numSplits = {
+    const std::vector<std::vector<int32_t>> numSplits = {
             {1, 16, 5, 8},
             {2, 19, 5, 4},
             {7, 13, 2, 8},
             {5, 8, 12, 5},
-            {4, 11, 6, 9}
+            {4, 11, -1, 9}
     };
 
     INSTANTIATE_TEST_CASE_P(NumSplitsCheck, VariadicSplitLayerTest,

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

@@ -17,12 +17,12 @@ namespace {
     };
 
     // Sum of elements numSplits = inputShapes[Axis]
-    const std::vector<std::vector<size_t>> numSplits = {
+    const std::vector<std::vector<int32_t>> numSplits = {
             {1, 16, 5, 8},
             {2, 19, 5, 4},
             {7, 13, 2, 8},
             {5, 8, 12, 5},
-            {4, 11, 6, 9}
+            {4, 11, -1, 9}
     };
 
     INSTANTIATE_TEST_CASE_P(NumSplitsCheck, VariadicSplitLayerTest,

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

@@ -15,7 +15,7 @@
 namespace LayerTestsDefinitions {
 
     typedef std::tuple<
-            std::vector<size_t>,            // Num splits
+            std::vector<int32_t>,           // Num splits
             size_t,                         // Axis
             InferenceEngine::Precision,     // Net precision
             std::vector<size_t>,            // Input shapes

inference-engine/tests/functional/plugin/shared/src/single_layer_tests/variadic_split.cpp

@@ -23,7 +23,7 @@ namespace LayerTestsDefinitions {
 
     std::string VariadicSplitLayerTest::getTestCaseName(testing::TestParamInfo<VariadicSplitParams> obj) {
         size_t axis;
-        std::vector<size_t> numSplits;
+        std::vector<int32_t> numSplits;
         InferenceEngine::Precision netPrecision;
         InferenceEngine::SizeVector inputShapes;
         std::string targetDevice;
@@ -41,7 +41,8 @@ namespace LayerTestsDefinitions {
     void VariadicSplitLayerTest::SetUp() {
         SetRefMode(LayerTestsUtils::RefMode::CONSTANT_FOLDING);
         size_t axis;
-        std::vector<size_t> inputShape, numSplits;
+        std::vector<size_t> inputShape;
+        std::vector<int32_t> numSplits;
         InferenceEngine::Precision netPrecision;
         std::tie(numSplits, axis, netPrecision, inputShape, targetDevice) = this->GetParam();
         auto ngPrc = FuncTestUtils::PrecisionUtils::convertIE2nGraphPrc(netPrecision);

inference-engine/tests/ngraph_functions/include/ngraph_functions/builders.hpp

@@ -185,7 +185,7 @@ std::shared_ptr<ngraph::Node> makeSplit(const ngraph::Output<Node> &in,
                                         size_t axis);
 
 std::shared_ptr<ngraph::Node> makeVariadicSplit(const ngraph::Output<Node> &in,
-                                                const std::vector<size_t> numSplits,
+                                                const std::vector<int32_t> numSplits,
                                                 size_t axis);
 
 std::shared_ptr<ngraph::Node> makeActivation(const ngraph::Output<Node> &in,

inference-engine/tests/ngraph_functions/src/variadic_split.cpp

@@ -10,11 +10,11 @@
 namespace ngraph {
 namespace builder {
         std::shared_ptr<ngraph::Node> makeVariadicSplit(const ngraph::Output<Node> &in,
-                                                        const std::vector<size_t> numSplits,
+                                                        const std::vector<int32_t> numSplits,
                                                         size_t axis) {
             auto splitAxisOp = std::make_shared<ngraph::opset3::Constant>(element::u64, ngraph::Shape{},
                                                                           std::vector<size_t>{axis});
-            auto numSplit = std::make_shared<ngraph::opset3::Constant>(element::u64, ngraph::Shape{numSplits.size()},
+            auto numSplit = std::make_shared<ngraph::opset3::Constant>(element::i64, ngraph::Shape{numSplits.size()},
                                                                        numSplits);
             auto VariadicSplitNode = std::make_shared<ngraph::opset3::VariadicSplit>(in, splitAxisOp, numSplit);
             return VariadicSplitNode;