Merge branch 'main' into dis

Author: xadupre

onnxscript/optimizer/_constant_folding.py

@@ -278,9 +278,18 @@ def _get_numpy_value(
         if size_limit is not None and const_value.size > size_limit:
             return None
         try:
-            # Reinterpret the array with `.view()` because some implementations of
-            # ir.TensorProtocol (e.g. PyTorch<=2.7) do not use ml_dtypes for bfloat16 etc.
-            array = const_value.numpy().view(const_value.dtype.numpy())
+            # Turn the constant value into a numpy array representation with the
+            # specifics of this conversion handled by the tensor type
+            array = const_value.numpy()
+            # Can/should not reinterpret strings via .view, resulting in
+            #   "TypeError: Cannot change data-type for array of references."
+            # There is also no reason to reinterpret strings, this is only
+            # relevant for some arithmetic types
+            if const_value.dtype != ir.DataType.STRING:
+                # Reinterpret the array with `.view()` because some
+                # implementations  of ir.TensorProtocol (e.g. PyTorch<=2.7) do
+                # not use ml_dtypes for bfloat16 etc.
+                array = array.view(const_value.dtype.numpy())
         except FileNotFoundError:
             # External data is not available.
             logger.warning(
@@ -344,6 +353,33 @@ def _get_int_attribute(node: ir.Node, name: str, default: int | None = None) ->
     return default
 
 
+@register("Add")
+def add(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
+    """Propagate symbolic dim values."""
+
+    def get_dim_value(input_index):
+        input = _get_input(node, input_index)
+        if input is None:
+            return None
+        shape_value: ir.Shape | None = state.get_shape_value(input)
+        if shape_value is None or len(shape_value) != 1:
+            return None
+        dim: int | ir.SymbolicDim = shape_value[0]
+        return dim if isinstance(dim, int) else dim.value
+
+    dim0 = get_dim_value(0)
+    dim1 = get_dim_value(1)
+    if dim0 is None or dim1 is None:
+        return None
+    if isinstance(dim0, int) and isinstance(dim1, int):
+        result_dim_value: int | ir.SymbolicDim = dim0 + dim1
+    else:
+        result_dim_value = ir.SymbolicDim(f"{dim0}+{dim1}")
+    output = _get_output(node, 0)
+    if output is not None:
+        state.set_sym_value(output, ir.Shape([result_dim_value]))
+
+
 @register("Abs")
 def abs(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
     """Replace an Abs node by Identity when applicable.
@@ -392,9 +428,26 @@ def gather(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
     return None
 
 
+def _propagate_shape_value(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
+    """Propagates symbolic shape value of input 0 to output 0.
+
+    Applies to ops like Reshape/Squeeze/Unsqueeze where the shape of the tensor may change
+    but the values in the tensor remain the same.
+    """
+    input = _get_input(node, 0)
+    input_shape_value = state.get_shape_value(input)
+    output = _get_output(node, 0)
+    if output is not None and input_shape_value is not None:
+        state.set_sym_value(output, input_shape_value)
+    return None
+
+
 @register("Reshape")
 def reshape(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
-    """Replace a Reshape node by Identity when applicable."""
+    """Replace a Reshape node by Identity when applicable.
+
+    Also propagate symbolic shape values.
+    """
     input = _get_input(node, 0)
     shape = _get_input(node, 1)
     if input is None or shape is None:
@@ -404,12 +457,18 @@ def reshape(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
     shape_value = state.get_shape_value(shape)
 
     if shape_value is None or input_shape is None:
-        return None
+        return _propagate_shape_value(node, op, state)
 
     # No need to check for special values like -1, 0, etc. here
     if _same_shape(input_shape, shape_value):
         return op.Identity(input)
-    return None
+    return _propagate_shape_value(node, op, state)
+
+
+@register("Squeeze")
+def squeeze(node: ir.Node, op, state: OptimizerState) -> ReturnValue:
+    """Propagate symbolic shape values."""
+    return _propagate_shape_value(node, op, state)
 
 
 @register("Cast")

onnxscript/rewriter/onnx_fusions/_rotary_embedding.py

@@ -30,7 +30,7 @@ def _rotate_half_pattern(op, x, start1, end1, start2, end2):
 
 class RotaryEmbedding23Fusion(pattern.RewriteRuleClassBase):
     def __init__(self):
-        super().__init__(name="RotaryEmbedding23", as_function=True)
+        super().__init__(name="RotaryEmbedding23")
 
     def pattern(self, op, x, cos, sin, start1, end1, start2, end2):
         return x * cos + _rotate_half_pattern(op, x, start1, end1, start2, end2) * sin