Add reproduction test case for incorrect slice rewrite and add potential fix (#2478)

This adds a reproduction of the rule introduced in
f42c2bb leading to an incorrect rewrite
of the graph.

The original rule does not consider the step parameter, which can
influence the result of a `Slice` to be the identity even when input and
output shape are equivalent.

The potential fix seems to be to not apply the rule on `step != 1`,
therefore the second commit adds this to the original rule
implementation.

---------

Co-authored-by: Justin Chu <[email protected]>

Author: sebimarkgraf

onnxscript/rewriter/collapse_slices.py

@@ -5,6 +5,7 @@
 import logging
 
 from onnxscript import ir
+from onnxscript.rewriter._ir_utils import is_singleton_value
 from onnxscript.rewriter._rewrite_rule import RewriteRule, RewriteRuleSet
 
 logger = logging.getLogger(__name__)
@@ -76,10 +77,14 @@ def _potential_redundant_slice(op, data, starts, ends, axes, steps):
     return op.Slice(data, starts, ends, axes, steps, _outputs=["slice_output"])
 
 
-def _same_shape(op, data: ir.Value, slice_output: ir.Value, **_):
+def _same_shape(op, data: ir.Value, slice_output: ir.Value, steps: ir.Value, **_):
     """Check if the shape of the slice output is the same as the data."""
     if data.shape is None or slice_output.shape is None:
         return False
+
+    if not is_singleton_value(steps, 1):
+        return False
+
     return data.shape == slice_output.shape
 
 

onnxscript/rewriter/collapse_slices_test.py

@@ -100,3 +100,22 @@ def test_slice_equal_dynamic_shape(self):
         model = ir.serde.deserialize_model(model_proto)
         count = collapse_slices.rules.apply_to_model(model)
         self.assertEqual(count, 1)
+
+    def test_slice_equal_dynamic_shape_but_step_reverse(self):
+        model_proto = onnx.parser.parse_model(
+            f"""
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[L, M, N] data) => (float[L, M, N] output)
+            {{
+                starts = Constant<value: tensor = int64[1] {{0}}>()
+                ends = Constant<value: tensor = int64[1] {{{9}}}>()
+                axes = Constant<value: tensor = int64[1] {{0}}>()
+                steps = Constant<value: tensor = int64[1] {{-1}}>()
+                output = Slice (data, starts, ends, axes, steps)
+            }}
+        """
+        )
+        model = ir.serde.deserialize_model(model_proto)
+        count = collapse_slices.rules.apply_to_model(model)
+        # Should not change the output shape if we did not use the default step of 1
+        self.assertEqual(count, 0)