Update parse slice (#847)

* init reverseOp branch: ref op + ref test. WIP

* first passing basic test

* cleanup

* additional axis implementation

* additional test

* ref op implementation vec to int for axis

* ref op test change for axis

* initial gpu files and test

* updates to implementation and test

* fixed some issues

* clang format

* cleanup

* formatting

* removing comments

* changes to parse_slice.cpp debug copy

* cleanup + additional axis  for reverse instruction

* formatting

* remove local size, back to default

* update tests: replace with std functions

* multiple axis for reverse op

* fix a build error

* clang format

* changes to parse_slice.cpp debug copy

* cleanup + additional axis  for reverse instruction

* formatting

* axes update to parse slice

* typo

* more tests

* fix a bug for the reverse device function

* clang format

* fix a bug

* clang format

* ref test updates, multiaxis

* formatting

* formatting, cleanup bool op

* casting for tidy warning

* tidy fix

* remove bool, add steps, check only negative axis

* clang-format

* step op for parse slice

* cleanup & format

* missing axis for logsoftmax_nonstd_input_test

* updated onnx file for logsoftmax_nonstd_input_test

* updates to parse slice. tests for slice+reverse, slice+step+reverse

* removing tests for slice+step+reverse as step requires normalization, will move it to other branch. removed related lines and tests

* duplicate test removal

* some refinement of the code

* clang format

* undefined behavior fix

* undef behavior v2

* formatting

* formatting & updates

* change to parse slice

* update to parse_slice for undef/asan + test update

* formatting

* remove header, no if

* assertions + change the loop from axis to steps for logsoftmax test segfault
Co-authored-by: Shucai Xiao <Shucai.Xiao@amd.com>
Co-authored-by: mvermeulen <5479696+mvermeulen@users.noreply.github.com>

src/onnx/parse_slice.cpp

@@ -20,17 +20,18 @@ struct parse_slice : op_parser<parse_slice>
     {
         op::slice op;
 
+        std::vector<int64_t> steps;
+
         // slice can have up to 5 inputs, we first check the 5th one
         // to decide whether MIGRAPHX can handle this slice
         if(args.size() == 5)
         {
             migraphx::argument step_arg = args.back()->eval();
             check_arg_empty(step_arg, "PARSE_SLICE: cannot handle variable steps for slice");
-            std::vector<int> steps;
             step_arg.visit([&](auto s) { steps.assign(s.begin(), s.end()); });
-            if(!std::all_of(steps.begin(), steps.end(), [](auto s) { return s == 1; }))
+            if(!std::all_of(steps.begin(), steps.end(), [](auto s) { return abs(s) == 1; }))
             {
-                MIGRAPHX_THROW("PARSE_SLICE: cannot handle step other than 1");
+                MIGRAPHX_THROW("PARSE_SLICE: cannot handle step other than 1 or -1");
             }
         }
 
@@ -77,7 +78,29 @@ struct parse_slice : op_parser<parse_slice>
             op.axes = axes;
         }
 
-        return info.add_instruction(op, args[0]);
+        std::vector<int64_t> raxes;
+
+        assert(steps.empty() or steps.size() == op.axes.size());
+        assert(op.axes.size() == op.starts.size());
+        assert(op.axes.size() == op.ends.size());
+
+        for(auto i : range(steps.size()))
+        {
+            if(steps[i] >= 0)
+                continue;
+            op.starts[i] += 1;
+            if(op.starts[i] == 0)
+                op.starts[i] = INT_MAX;
+            op.ends[i] += 1;
+            raxes.push_back(op.axes[i]);
+            std::swap(op.starts[i], op.ends[i]);
+        }
+
+        auto ins = info.add_instruction(op, args[0]);
+        if(not raxes.empty())
+            return info.add_instruction(make_op("reverse", {{"axes", raxes}}), ins);
+        else
+            return ins;
     }
 };
 

test/onnx/gen_onnx.py

@@ -2218,7 +2218,10 @@ def logsoftmax_nonstd_input_test():
                                   ends=[4, 4],
                                   outputs=['1'])
 
-    node1 = onnx.helper.make_node('LogSoftmax', inputs=['1'], outputs=['2'])
+    node1 = onnx.helper.make_node('LogSoftmax',
+                                  inputs=['1'],
+                                  outputs=['2'],
+                                  axis=-1)
 
     return ([node0, node1], [x], [z])
 
@@ -3429,6 +3432,59 @@ def slice_5arg_test():
     return ([arg_step, arg_axis, arg_end, arg_start, node], [x], [y])
 
 
+@onnx_test
+def slice_5arg_reverse_test():
+    step = np.array([-1, 1])
+    step_tensor = helper.make_tensor(name="step",
+                                     data_type=TensorProto.INT32,
+                                     dims=step.shape,
+                                     vals=step.astype(int))
+    arg_step = helper.make_node("Constant",
+                                inputs=[],
+                                outputs=['arg_step'],
+                                value=step_tensor)
+
+    axis = np.array([-1, -2])
+    axis_tensor = helper.make_tensor(name="axis",
+                                     data_type=TensorProto.INT32,
+                                     dims=axis.shape,
+                                     vals=axis.astype(int))
+    arg_axis = helper.make_node("Constant",
+                                inputs=[],
+                                outputs=['arg_axis'],
+                                value=axis_tensor)
+
+    end = np.array([-1, -1])
+    end_tensor = helper.make_tensor(name="end",
+                                    data_type=TensorProto.INT32,
+                                    dims=end.shape,
+                                    vals=end.astype(int))
+    arg_end = helper.make_node("Constant",
+                               inputs=[],
+                               outputs=['arg_end'],
+                               value=end_tensor)
+
+    start = np.array([-5, -3])
+    start_tensor = helper.make_tensor(name="start",
+                                      data_type=TensorProto.INT32,
+                                      dims=start.shape,
+                                      vals=start.astype(int))
+    arg_start = helper.make_node("Constant",
+                                 inputs=[],
+                                 outputs=['arg_start'],
+                                 value=start_tensor)
+
+    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [5, 5])
+    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [4, 2])
+
+    node = onnx.helper.make_node(
+        'Slice',
+        inputs=['0', 'arg_start', 'arg_end', 'arg_axis', 'arg_step'],
+        outputs=['1'])
+
+    return ([arg_step, arg_axis, arg_end, arg_start, node], [x], [y])
+
+
 @onnx_test
 def slice_max_end_test():
     x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [10, 20])

test/onnx/onnx_test.cpp

@@ -2989,6 +2989,26 @@ TEST_CASE(slice_5arg_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(slice_5arg_reverse_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    auto l0  = mm->add_parameter("0", migraphx::shape{migraphx::shape::float_type, {5, 5}});
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {-1, 1}});
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {-1, -2}});
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {-1, -1}});
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {-5, -3}});
+    auto slice_out = mm->add_instruction(
+        migraphx::make_op("slice", {{"axes", {-1, -2}}, {"starts", {0, -3}}, {"ends", {-4, -1}}}),
+        l0);
+    auto ret = mm->add_instruction(migraphx::make_op("reverse", {{"axes", {-1}}}), slice_out);
+    mm->add_return({ret});
+
+    auto prog = migraphx::parse_onnx("slice_5arg_reverse_test.onnx");
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(slice_max_end_test)
 {
     migraphx::program p;

test/onnx/slice_5arg_reverse_test.onnx

+slice_5arg_reverse_test:
+9arg_step"Constant*#
+value**

Bstep

+Barg_axis"Constant*,
+value* *

Baxis

+@arg_end"Constant*+
+value**

Bend

+D	arg_start"Constant*-
+value*!*

Bstart

+5
+
0
+	arg_start
+arg_end
+arg_axis
+arg_step
1"Sliceslice_5arg_reverse_testZ
+
0
+

+
+b
+
1
+

+
+B
\ No newline at end of file