Step op for parse_slice.cpp (#858)

* slice+step+reverse: parsing + onnxfile + test

* cleanup

* updates for step operator, abs axis

* test updates

* updates to tests

* slice+reverse verify test

* verify test for slice+step+reverse

* clang format

* reverse with lens

* simplify normalize_compute_shape

* step op: normalization for axes

* clang format

* change of order: fixing wrong functionality in some cases

* test for slice,reverse,step

* clang format
Co-authored-by: Paul Fultz II <pfultz2@yahoo.com>

src/include/migraphx/op/reverse.hpp

@@ -39,9 +39,7 @@ struct reverse
 
     shape normalize_compute_shape(std::vector<shape> inputs) const
     {
-        auto lens = inputs[0].lens();
-        auto type = inputs[0].type();
-        return shape{type, lens};
+        return inputs[0].with_lens(inputs[0].lens());
     }
 
     argument compute(const shape& s, std::vector<argument> args) const

src/include/migraphx/op/step.hpp

@@ -7,6 +7,7 @@
 #include <migraphx/argument.hpp>
 #include <migraphx/functional.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/op/normalize_attribute.hpp>
 #include <cmath>
 #include <utility>
 
@@ -25,8 +26,15 @@ struct step
         return pack(f(self.axes, "axes"), f(self.steps, "steps"));
     }
 
+    value attributes() const
+    {
+        value normalize;
+        normalize["axes"] = value::array{normalize_attribute::include_min};
+        return {{"normalize_axes", normalize}};
+    }
+
     std::string name() const { return "step"; }
-    shape compute_shape(std::vector<shape> inputs) const
+    shape normalize_compute_shape(std::vector<shape> inputs) const
     {
         check_shapes{inputs, *this}.has(1);
         auto input   = inputs.at(0);

src/onnx/parse_slice.cpp

@@ -29,10 +29,6 @@ struct parse_slice : op_parser<parse_slice>
             migraphx::argument step_arg = args.back()->eval();
             check_arg_empty(step_arg, "PARSE_SLICE: cannot handle variable steps for slice");
             step_arg.visit([&](auto s) { steps.assign(s.begin(), s.end()); });
-            if(!std::all_of(steps.begin(), steps.end(), [](auto s) { return abs(s) == 1; }))
-            {
-                MIGRAPHX_THROW("PARSE_SLICE: cannot handle step other than 1 or -1");
-            }
         }
 
         if(args.size() >= 4)
@@ -98,7 +94,16 @@ struct parse_slice : op_parser<parse_slice>
 
         auto ins = info.add_instruction(op, args[0]);
         if(not raxes.empty())
-            return info.add_instruction(make_op("reverse", {{"axes", raxes}}), ins);
+            ins = info.add_instruction(make_op("reverse", {{"axes", raxes}}), ins);
+        if(std::any_of(steps.begin(), steps.end(), [](auto s) { return std::abs(s) != 1; }))
+        {
+            std::vector<int64_t> nsteps;
+            std::transform(steps.begin(), steps.end(), std::back_inserter(nsteps), [](auto s) {
+                return std::abs(s);
+            });
+            return ins = info.add_instruction(
+                       make_op("step", {{"axes", op.axes}, {"steps", nsteps}}), ins);
+        }
         else
             return ins;
     }

test/onnx/gen_onnx.py

@@ -3584,7 +3584,7 @@ def slice_5arg_reverse_test():
                                 outputs=['arg_axis'],
                                 value=axis_tensor)
 
-    end = np.array([-1, -1])
+    end = np.array([-5, -1])
     end_tensor = helper.make_tensor(name="end",
                                     data_type=TensorProto.INT32,
                                     dims=end.shape,
@@ -3594,7 +3594,60 @@ def slice_5arg_reverse_test():
                                outputs=['arg_end'],
                                value=end_tensor)
 
-    start = np.array([-5, -3])
+    start = np.array([-1, -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_5arg_step_test():
+    step = np.array([-2, 2])
+    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([-5, -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([-1, -3])
     start_tensor = helper.make_tensor(name="start",
                                       data_type=TensorProto.INT32,
                                       dims=start.shape,

test/onnx/onnx_test.cpp

@@ -3314,10 +3314,11 @@ TEST_CASE(slice_5arg_reverse_test)
     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}});
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {-5, -1}});
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {-1, -3}});
     auto slice_out = mm->add_instruction(
-        migraphx::make_op("slice", {{"axes", {-1, -2}}, {"starts", {0, -3}}, {"ends", {-4, -1}}}),
+        migraphx::make_op("slice",
+                          {{"axes", {-1, -2}}, {"starts", {-4, -3}}, {"ends", {2147483647, -1}}}),
         l0);
     auto ret = mm->add_instruction(migraphx::make_op("reverse", {{"axes", {-1}}}), slice_out);
     mm->add_return({ret});
@@ -3327,6 +3328,30 @@ TEST_CASE(slice_5arg_reverse_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(slice_5arg_step_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}}, {-2, 2}});
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {-1, -2}});
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {-5, -1}});
+    mm->add_literal({{migraphx::shape::int32_type, {2}}, {-1, -3}});
+    auto slice_out = mm->add_instruction(
+        migraphx::make_op("slice",
+                          {{"axes", {-1, -2}}, {"starts", {-4, -3}}, {"ends", {2147483647, -1}}}),
+        l0);
+    auto reverse_out =
+        mm->add_instruction(migraphx::make_op("reverse", {{"axes", {-1}}}), slice_out);
+    auto step_out = mm->add_instruction(
+        migraphx::make_op("step", {{"axes", {-1, -2}}, {"steps", {2, 2}}}), reverse_out);
+    mm->add_return({step_out});
+
+    auto prog = migraphx::parse_onnx("slice_5arg_step_test.onnx");
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(slice_max_end_test)
 {
     migraphx::program p;

test/onnx/slice_5arg_reverse_test.onnx

@@ -4,9 +4,9 @@
 Barg_axis"Constant*,
 value* *

Baxis

 @arg_end"Constant*+
-value**

Bend

+value**

Bend

 D	arg_start"Constant*-
-value*!*

Bstart

+value*!*

Bstart

 5
 
0
 	arg_start

test/onnx/slice_5arg_step_test.onnx

+slice_5arg_step_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_step_testZ
+
0
+

+
+b
+
1
+

+
+B
\ No newline at end of file

test/onnx/slice_5arg_test.onnx

-slice_5arg_test:
+slice_5arg_test:
 0arg_step"Constant*
 value**

Bstep

 Barg_axis"Constant*,
@@ -20,4 +20,4 @@ D	arg_start"Constant*-
 
1
 

 
-B
\ No newline at end of file
+B
\ No newline at end of file

test/onnx/verify_onnx.cpp

@@ -408,6 +408,85 @@ TEST_CASE(selu_test)
     EXPECT(migraphx::verify_range(result_vector, gold));
 }
 
+TEST_CASE(slice_test)
+{
+    migraphx::program p = migraphx::parse_onnx("slice_test.onnx");
+    p.compile(migraphx::ref::target{});
+
+    migraphx::shape sh_data{migraphx::shape::float_type, {3, 2}};
+    std::vector<float> data = {0, 1, 2, 3, 4, 5};
+
+    migraphx::parameter_map pp;
+    pp["0"] = migraphx::argument(sh_data, data.data());
+
+    auto result = p.eval(pp).back();
+    std::vector<float> result_vector;
+    result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold = {2, 3};
+
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
+TEST_CASE(slice_5arg_test)
+{
+    migraphx::program p = migraphx::parse_onnx("slice_5arg_test.onnx");
+    p.compile(migraphx::ref::target{});
+
+    migraphx::shape sh_data{migraphx::shape::float_type, {5, 5}}; // start
+    std::vector<float> data = {0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12,
+                               13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24};
+
+    migraphx::parameter_map pp;
+    pp["0"] = migraphx::argument(sh_data, data.data());
+
+    auto result = p.eval(pp).back();
+    std::vector<float> result_vector;
+    result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });
+
+    std::vector<float> gold = {10, 11, 12, 13, 15, 16, 17, 18};
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
+TEST_CASE(slice_reverse_test)
+{
+    migraphx::program p = migraphx::parse_onnx("slice_5arg_reverse_test.onnx");
+    p.compile(migraphx::ref::target{});
+
+    migraphx::shape sh_data{migraphx::shape::float_type, {5, 5}}; // start
+    std::vector<float> data = {0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12,
+                               13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24};
+
+    migraphx::parameter_map pp;
+    pp["0"] = migraphx::argument(sh_data, data.data());
+
+    auto result = p.eval(pp).back();
+    std::vector<float> result_vector;
+    result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });
+
+    std::vector<float> gold = {14, 13, 12, 11, 19, 18, 17, 16};
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
+TEST_CASE(slice_step_test)
+{
+    migraphx::program p = migraphx::parse_onnx("slice_5arg_step_test.onnx");
+    p.compile(migraphx::ref::target{});
+
+    migraphx::shape sh_data{migraphx::shape::float_type, {5, 5}}; // start
+    std::vector<float> data = {0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12,
+                               13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24};
+
+    migraphx::parameter_map pp;
+    pp["0"] = migraphx::argument(sh_data, data.data());
+
+    auto result = p.eval(pp).back();
+    std::vector<float> result_vector;
+    result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });
+
+    std::vector<float> gold = {14, 12};
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
 TEST_CASE(upsample_test)
 {
     migraphx::program p = migraphx::parse_onnx("upsample_test.onnx");

test/verify/test_slice_reverse.cpp

+
+#include "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_slice_reverse : verify_program<test_slice_reverse>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        migraphx::shape s{migraphx::shape::int32_type, {3, 5}};
+        auto x         = mm->add_parameter("x", s);
+        auto slice_out = mm->add_instruction(
+            migraphx::make_op("slice", {{"axes", {0, 1}}, {"starts", {0, 2}}, {"ends", {2, -1}}}),
+            x);
+        mm->add_instruction(migraphx::make_op("reverse", {{"axes", {0}}}), slice_out);
+
+        return p;
+    }
+};

test/verify/test_slice_reverse_step.cpp

+
+#include "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_slice_reverse_step : verify_program<test_slice_reverse_step>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        migraphx::shape s{migraphx::shape::int32_type, {7, 5}};
+        auto x         = mm->add_parameter("x", s);
+        auto slice_out = mm->add_instruction(
+            migraphx::make_op("slice", {{"axes", {0, 1}}, {"starts", {0, 2}}, {"ends", {2, -1}}}),
+            x);
+        auto step_out =
+            mm->add_instruction(migraphx::make_op("reverse", {{"axes", {0, 1}}}), slice_out);
+        mm->add_instruction(migraphx::make_op("step", {{"axes", {0, 1}}, {"steps", {2, 2}}}),
+                            step_out);
+        return p;
+    }
+};

test/verify/test_slice_step_reverse.cpp

+
+#include "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_slice_step_reverse : verify_program<test_slice_step_reverse>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        migraphx::shape s{migraphx::shape::int32_type, {7, 5}};
+        auto x         = mm->add_parameter("x", s);
+        auto slice_out = mm->add_instruction(
+            migraphx::make_op("slice", {{"axes", {0, 1}}, {"starts", {0, 2}}, {"ends", {2, -1}}}),
+            x);
+        auto step_out = mm->add_instruction(
+            migraphx::make_op("step", {{"axes", {0, 1}}, {"steps", {2, 2}}}), slice_out);
+        mm->add_instruction(migraphx::make_op("reverse", {{"axes", {0}}}), step_out);
+
+        return p;
+    }
+};