Create slice and reshape improvements

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: justinchuby

onnxscript/optimizer/_constant_folding.py

@@ -196,6 +196,10 @@ def set_sym_value(self, value: ir.Value, sym_value: SymbolicValue) -> None:
 
     def get_shape_value(self, value: ir.Value | None) -> ir.Shape | None:
         const_value = _get_numpy_value(value, ir.DataType.INT64, size_limit=10)
+        if const_value is None:
+            # Reshape accepts shape input of INT32 type as well, so we also check for INT32 here
+            # This is common for tflite models
+            const_value = _get_numpy_value(value, ir.DataType.INT32, size_limit=10)
         if const_value is not None:
             if const_value.ndim == 1:
                 return ir.Shape(const_value.tolist())

onnxscript/rewriter/__init__.py

@@ -42,6 +42,7 @@
     _fuse_batchnorm,
     _fuse_pad_into_conv,
     _fuse_relus_clips,
+    _materialize_reshape_shape,
     _min_max_to_clip,
     _no_op,
     _redundant_scatter_nd,
@@ -54,6 +55,7 @@
     *_broadcast_to_matmul.rules,
     *_cast_constant_of_shape.rules,
     *_collapse_slices.rules,
+    *_materialize_reshape_shape.rules,
     *_min_max_to_clip.rules,
     *_fuse_relus_clips.rules,
     *_basic_rules.basic_optimization_rules(),

onnxscript/rewriter/rules/common/__init__.py

@@ -25,6 +25,7 @@
     "max_min_rule",
     "gemm_to_matmul_add_rule",
     "matmul_add_to_gemm_rule",
+    "materialize_reshape_shape_rule",
     "mul_by_1_rule",
     "no_op_cast_rule",
     "no_op_dynamic_scatter_nd_rule",
@@ -107,6 +108,9 @@
     transpose_ab_matmul_add_to_gemm_rule,
     transpose_b_matmul_add_to_gemm_rule,
 )
+from onnxscript.rewriter.rules.common._materialize_reshape_shape import (
+    materialize_reshape_shape_rule,
+)
 from onnxscript.rewriter.rules.common._min_max_to_clip import (
     max_max_rule,
     max_min_rule,

onnxscript/rewriter/rules/common/_collapse_slices.py

@@ -82,7 +82,12 @@ def _same_shape(op, data: ir.Value, slice_output: ir.Value, steps: ir.Value, **_
     if data.shape is None or slice_output.shape is None:
         return False
 
-    if not _ir_utils.is_singleton_value(steps, 1):
+    # All steps must be 1
+    steps_np = _ir_utils.get_numpy_value(steps)
+    if steps_np is not None:
+        if not all(s == 1 for s in steps_np.flatten()):
+            return False
+    elif not _ir_utils.is_singleton_value(steps, 1):
         return False
 
     return _ir_utils.same_shape(data.shape, slice_output.shape)

onnxscript/rewriter/rules/common/_materialize_reshape_shape.py

@@ -0,0 +1,63 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Materialize Reshape shape input from known output shape.
+
+When symbolic shape inference has been run, a Reshape node may have a known
+output shape even though its shape input is computed dynamically (e.g., via a
+Shape → Cast → Split → Concat chain).  This rule replaces the shape input
+with a concrete constant, allowing the dynamic chain to become dead code and
+be removed by unused-node elimination.
+
+- Fully static output shape → constant with exact dims.
+- Exactly one symbolic dim → replace it with ``-1`` (Reshape infers it).
+"""
+
+from __future__ import annotations
+
+from onnxscript import ir
+from onnxscript.rewriter import _ir_utils as ir_utils
+from onnxscript.rewriter._basics import MatchResult
+from onnxscript.rewriter._rewrite_rule import RewriteRuleClassBase, RewriteRuleSet
+
+
+class MaterializeReshapeShape(RewriteRuleClassBase):
+    """Replace a dynamic Reshape shape input with a constant when output shape is known."""
+
+    def pattern(self, op, data, shape):
+        return op.Reshape(data, shape)
+
+    def check(self, context, data: ir.Value, shape: ir.Value) -> MatchResult:
+        check_result = MatchResult()
+
+        # Shape input must not already be a constant
+        if ir_utils.get_numpy_value(shape) is not None:
+            return check_result.fail("Shape input is already a constant.")
+
+        output = context.output_values[0]
+        if output.shape is None:
+            return check_result.fail("Output shape is not known.")
+
+        dims = list(output.shape)
+        sym_count = sum(1 for d in dims if not isinstance(d, int))
+
+        if sym_count == 0:
+            self._new_dims = [int(d) for d in dims]
+        elif sym_count == 1:
+            self._new_dims = [-1 if not isinstance(d, int) else int(d) for d in dims]
+        else:
+            return check_result.fail(
+                f"Output shape has {sym_count} symbolic dims, cannot materialize."
+            )
+        return check_result
+
+    def rewrite(self, op, data: ir.Value, shape: ir.Value):
+        new_shape = op.Constant(
+            value=ir.tensor(self._new_dims, dtype=ir.DataType.INT64),
+        )
+        return op.Reshape(data, new_shape)
+
+
+materialize_reshape_shape_rule = MaterializeReshapeShape.rule()
+
+rules = RewriteRuleSet([materialize_reshape_shape_rule])
+