Add support for uneven Split operations (#2254)

Since opset version 18, the Split operator allows splitting into unevenly sized outputs when the split input is not present, and a num_outputs attribute has been introduced.

src/onnx/parse_split.cpp

@@ -68,13 +68,34 @@ struct parse_split : op_parser<parse_split>
         // no split attribute, input is equally divided
         else
         {
-            if((lens[tuned_axis] % info.num_outputs) != 0)
+            std::size_t num_outputs = info.num_outputs;
+            // the num_outputs attribute seems to be redundant since we already have
+            // node_info::num_outputs, but we can still perform an error check
+            if(contains(info.attributes, "num_outputs"))
             {
-                MIGRAPHX_THROW("PARSE_SPLIT: input cannot be equally divided into " +
-                               std::to_string(info.num_outputs) + " splits!");
+                num_outputs =
+                    parser.parse_value(info.attributes.at("num_outputs")).at<std::size_t>();
+                if(num_outputs != info.num_outputs)
+                {
+                    MIGRAPHX_THROW("PARSE_SPLIT: num_outputs attribute " +
+                                   std::to_string(num_outputs) +
+                                   " doesn't match actual number of outputs " +
+                                   std::to_string(info.num_outputs) + "!");
+                }
+            }
+
+            if(lens[tuned_axis] % num_outputs == 0)
+            {
+                std::size_t chunk_size = lens[tuned_axis] / num_outputs;
+                vec_splits.resize(num_outputs, chunk_size);
+            }
+            else
+            {
+                std::size_t chunk_size      = lens[tuned_axis] / num_outputs + 1;
+                std::size_t last_chunk_size = lens[tuned_axis] - chunk_size * (num_outputs - 1);
+                vec_splits.resize(num_outputs - 1, chunk_size);
+                vec_splits.push_back(last_chunk_size);
             }
-            auto dl = lens[tuned_axis] / info.num_outputs;
-            vec_splits.resize(info.num_outputs, dl);
         }
 
         if(std::accumulate(vec_splits.begin(), vec_splits.end(), int64_t(0)) !=

test/onnx/gen_onnx.py

@@ -8080,6 +8080,42 @@ def split_test_no_attribute():
     return ([const_node, node], [x], [y1, y2, y3, y4])
 
 
+@onnx_test()
+def split_test_uneven():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [12, 15])
+    y1 = helper.make_tensor_value_info('y1', TensorProto.FLOAT, [3, 15])
+    y2 = helper.make_tensor_value_info('y2', TensorProto.FLOAT, [3, 15])
+    y3 = helper.make_tensor_value_info('y3', TensorProto.FLOAT, [3, 15])
+    y4 = helper.make_tensor_value_info('y4', TensorProto.FLOAT, [3, 15])
+    y5 = helper.make_tensor_value_info('y5', TensorProto.FLOAT, [0, 15])
+
+    node = onnx.helper.make_node(
+        'Split',
+        inputs=['x'],
+        outputs=['y1', 'y2', 'y3', 'y4', 'y5'],
+    )
+
+    return ([node], [x], [y1, y2, y3, y4, y5])
+
+
+@onnx_test()
+def split_test_uneven_num_outputs():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [11, 15])
+    y1 = helper.make_tensor_value_info('y1', TensorProto.FLOAT, [3, 15])
+    y2 = helper.make_tensor_value_info('y2', TensorProto.FLOAT, [3, 15])
+    y3 = helper.make_tensor_value_info('y3', TensorProto.FLOAT, [3, 15])
+    y4 = helper.make_tensor_value_info('y4', TensorProto.FLOAT, [2, 15])
+
+    node = onnx.helper.make_node(
+        'Split',
+        inputs=['x'],
+        outputs=['y1', 'y2', 'y3', 'y4'],
+        num_outputs=4,
+    )
+
+    return ([node], [x], [y1, y2, y3, y4])
+
+
 @onnx_test()
 def split_test_no_attribute_invalid_split():
     x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [300, 15])
@@ -8139,6 +8175,24 @@ def split_test_no_attribute_invalid_input_split():
     return ([node], [x], [y1, y2, y3])
 
 
+@onnx_test()
+def split_test_invalid_num_outputs():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [11, 15])
+    y1 = helper.make_tensor_value_info('y1', TensorProto.FLOAT, [3, 15])
+    y2 = helper.make_tensor_value_info('y2', TensorProto.FLOAT, [3, 15])
+    y3 = helper.make_tensor_value_info('y3', TensorProto.FLOAT, [3, 15])
+    y4 = helper.make_tensor_value_info('y4', TensorProto.FLOAT, [2, 15])
+
+    node = onnx.helper.make_node(
+        'Split',
+        inputs=['x'],
+        outputs=['y1', 'y2', 'y3', 'y4'],
+        num_outputs=5,
+    )
+
+    return ([node], [x], [y1, y2, y3, y4])
+
+
 @onnx_test()
 def sqrt_test():
     x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10, 15])

test/onnx/onnx_test.cpp

@@ -7780,6 +7780,46 @@ TEST_CASE(split_test_default)
     EXPECT(p == prog);
 }
 
+TEST_CASE(split_test_uneven)
+{
+    migraphx::program p;
+    auto* mm   = p.get_main_module();
+    auto input = mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, {12, 15}});
+    auto r1    = mm->add_instruction(
+        migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {3}}}), input);
+    auto r2 = mm->add_instruction(
+        migraphx::make_op("slice", {{"axes", {0}}, {"starts", {3}}, {"ends", {6}}}), input);
+    auto r3 = mm->add_instruction(
+        migraphx::make_op("slice", {{"axes", {0}}, {"starts", {6}}, {"ends", {9}}}), input);
+    auto r4 = mm->add_instruction(
+        migraphx::make_op("slice", {{"axes", {0}}, {"starts", {9}}, {"ends", {12}}}), input);
+    auto r5 = mm->add_instruction(
+        migraphx::make_op("slice", {{"axes", {0}}, {"starts", {12}}, {"ends", {12}}}), input);
+    mm->add_return({r1, r2, r3, r4, r5});
+
+    auto prog = migraphx::parse_onnx("split_test_uneven.onnx");
+    EXPECT(p == prog);
+}
+
+TEST_CASE(split_test_uneven_num_outputs)
+{
+    migraphx::program p;
+    auto* mm   = p.get_main_module();
+    auto input = mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, {11, 15}});
+    auto r1    = mm->add_instruction(
+        migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {3}}}), input);
+    auto r2 = mm->add_instruction(
+        migraphx::make_op("slice", {{"axes", {0}}, {"starts", {3}}, {"ends", {6}}}), input);
+    auto r3 = mm->add_instruction(
+        migraphx::make_op("slice", {{"axes", {0}}, {"starts", {6}}, {"ends", {9}}}), input);
+    auto r4 = mm->add_instruction(
+        migraphx::make_op("slice", {{"axes", {0}}, {"starts", {9}}, {"ends", {11}}}), input);
+    mm->add_return({r1, r2, r3, r4});
+
+    auto prog = migraphx::parse_onnx("split_test_uneven_num_outputs.onnx");
+    EXPECT(p == prog);
+}
+
 TEST_CASE(split_test_no_attribute_invalid_split)
 {
     EXPECT(
@@ -7797,6 +7837,11 @@ TEST_CASE(split_test_no_attribute_invalid_input_split)
         [&] { migraphx::parse_onnx("split_test_no_attribute_invalid_input_split.onnx"); }));
 }
 
+TEST_CASE(split_test_invalid_num_outputs)
+{
+    EXPECT(test::throws([&] { migraphx::parse_onnx("split_test_invalid_num_outputs.onnx"); }));
+}
+
 TEST_CASE(sqrt_test)
 {
     migraphx::program p;

test/onnx/split_test_invalid_num_outputs.onnx

+	split_test_invalid_num_outputs:

+.
+
xy1y2y3y4"Split*
+num_outputs
split_test_invalid_num_outputsZ
+
x
+

+
+b
+y1
+

+
+b
+y2
+

+
+b
+y3
+

+
+b
+y4
+

+
+B
\ No newline at end of file

test/onnx/split_test_uneven_num_outputs.onnx

+	split_test_uneven_num_outputs:

+.
+
xy1y2y3y4"Split*
+num_outputs
split_test_uneven_num_outputsZ
+
x
+

+
+b
+y1
+

+
+b
+y2
+

+
+b
+y3
+

+
+b
+y4
+

+
+B
\ No newline at end of file

test/py/onnx_backend_test.py

@@ -835,10 +835,6 @@ def disabled_tests_onnx_1_13_0(backend_test):
     backend_test.exclude(r'test_scatter_elements_with_reduction_max_cpu')
     backend_test.exclude(r'test_scatter_elements_with_reduction_min_cpu')
 
-    # The following tests fail due to the CastLike operator being unsupported
-    backend_test.exclude(r'test_split_1d_uneven_split_opset18_cpu')
-    backend_test.exclude(r'test_split_2d_uneven_split_opset18_cpu')
-
 
 def disabled_tests_onnx_1_14_0(backend_test):
     # fails