Improve redundant slice removal (#2441)

Improve the optimization for removal of redundant slices. (It doesn't
currently handle dynamic shapes.) The optimization is fairly simple, and
eliminates the slice when the input and output shapes are same.

---------

Signed-off-by: Ganesan Ramalingam <grama@microsoft.com>

Author: gramalingam

onnxscript/rewriter/collapse_slices.py

@@ -73,7 +73,14 @@ def _identity_to_itself(op, data, **_):
 
 def _potential_redundant_slice(op, data, starts, ends, axes, steps):
     """To identify a slice op"""
-    return op.Slice(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, **_):
+    """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
+    return data.shape == slice_output.shape
 
 
 # Register the rewrite rules
@@ -83,5 +90,13 @@ def _potential_redundant_slice(op, data, starts, ends, axes, steps):
     _check_if_redundant_slice,
 )
 
-# NOTE: The order of the rules is important. Larger pattern should be checked first.
-rules = RewriteRuleSet([remove_redundant_slice])
+remove_redundant_slice2 = RewriteRule(
+    _potential_redundant_slice,
+    _identity_to_itself,
+    _same_shape,
+)
+
+# NOTE: The second rule subsumes the first one. So, we may be able to remove the first one,
+# provided shape-inference is run before the rewriter and computes the shape of the slice output.
+
+rules = RewriteRuleSet([remove_redundant_slice, remove_redundant_slice2])

onnxscript/rewriter/collapse_slices_test.py

@@ -65,11 +65,11 @@ def test_slice_is_redundant_when_ends_reaches_int64_max(self):
             (np.random.rand(512, 16, 112).astype(np.float32),),
         )
 
-    def test_slice_pattern_is_not_matched_when_input_is_dynamic(self):
+    def test_slice_unequal_dynamic_shape(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)
+            agraph (float[L, M, N] data) => (float[P, M, N] output)
             {{
                 starts = Constant<value: tensor = int64[1] {{0}}>()
                 ends = Constant<value: tensor = int64[1] {{{9}}}>()
@@ -82,3 +82,21 @@ def test_slice_pattern_is_not_matched_when_input_is_dynamic(self):
         model = ir.serde.deserialize_model(model_proto)
         count = collapse_slices.rules.apply_to_model(model)
         self.assertEqual(count, 0)
+
+    def test_slice_equal_dynamic_shape(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)
+        self.assertEqual(count, 1)