Repository: microsoft/onnxscript

Issue:
Issue in gemm_to_matmul_add
Error:

```
  File "/home/xadupre/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "<eval_with_key>.4", line 5, in forward
    fused_0 = self.fused_0(primals_3, primals_1, primals_2);  primals_3 = primals_1 = primals_2 = None
  File "/home/xadupre/.local/lib/python3.10/site-packages/torch/fx/graph_module.py", line 735, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
  File "/home/xadupre/.local/lib/python3.10/site-packages/torch/onnx/_internal/onnxruntime.py", line 931, in _ort_acclerated_call
    onnx_model = optimizer.optimize(onnx_model)
  File "/home/xadupre/github/onnxscript/onnxscript/optimizer/__init__.py", line 80, in optimize
    model = rewriter.rewrite(
  File "/home/xadupre/github/onnxscript/onnxscript/rewriter/__init__.py", line 38, in rewrite
    count = pattern_rewrite_rules.apply_to_model(model_ir)
  File "/home/xadupre/github/onnxscript/onnxscript/rewriter/pattern.py", line 881, in apply_to_model
    count = self._apply_to_graph_or_function(model, model.graph)
  File "/home/xadupre/github/onnxscript/onnxscript/rewriter/pattern.py", line 871, in _apply_to_graph_or_function
    delta = rule.try_rewrite(model, graph_or_function, node)
  File "/home/xadupre/github/onnxscript/onnxscript/rewriter/pattern.py", line 743, in try_rewrite
    match = self.matches(node, model)
  File "/home/xadupre/github/onnxscript/onnxscript/rewriter/pattern.py", line 734, in matches
    and not self._condition_function(**match.bindings)
  File "/home/xadupre/github/onnxscript/onnxscript/rewriter/broadcast_to_matmul.py", line 114, in check_if_need_reshape
    broadcast_matmul_output_shape.pop(-2)
IndexError: pop index out of range
```

To reproduce:

```python
        import torch

        def _make_aot_ort(dynamic: bool = False) -> tuple:
            import onnxruntime
            from torch.onnx import (
                _OrtBackend as OrtBackend,
                _OrtBackendOptions as OrtBackendOptions,
                ExportOptions,
            )
            export_options = ExportOptions(dynamic_shapes=dynamic)
            options = OrtBackendOptions(export_options=export_options)
            ort_backend = OrtBackend(options=options)
            return ort_backend

        class Linear(torch.nn.Module):
            def __init__(self):
                super(Linear, self).__init__()
                self.linear = torch.nn.Linear(128, 10)
                self.activation = torch.nn.ReLU()

            def forward(self, *inputs):
                input = self.linear(inputs[0])
                input = self.activation(input)
                return input

        model = Linear()
        model.train()
        loss_fn = torch.nn.MSELoss()

        input = torch.randn((64, 128), requires_grad=True)
        labels = torch.randn((64, 10), requires_grad=True)

        compiled_model = torch.compile(model, backend=_make_aot_ort())
        output = compiled_model(*input)
        loss = loss_fn(output, labels)
        loss.backward()
```


Discussion:

Resolving patch:
diff --git a/docs/tutorial/rewriter/examples/broadcast_matmul.py b/docs/tutorial/rewriter/examples/broadcast_matmul.py
--- a/docs/tutorial/rewriter/examples/broadcast_matmul.py
+++ b/docs/tutorial/rewriter/examples/broadcast_matmul.py
@@ -64,7 +64,7 @@ def matmul_pattern(op, input_a: ir.Value, input_b: ir.Value, **_):
 # =====================
 
 
-def check_if_need_reshape(
+def check_if_not_need_reshape(
     input_a: ir.Value, input_b: ir.Value, shape_c: ir.Value, **_
 ) -> bool:
     """If matmul broadcasting is enough, then we don't need the reshapes.
@@ -184,7 +184,7 @@ def apply_rewrite(model):
     two_reshapes_matmul_reshape_rule = pattern.RewriteRule(
         two_reshapes_matmul_reshape_pattern,  # target pattern
         matmul_pattern,  # replacement pattern
-        check_if_need_reshape,  # match_condition function
+        check_if_not_need_reshape,  # match_condition function
     )
     # Create a Rewrite Rule Set
     rewrite_rule_set = pattern.RewriteRuleSet([two_reshapes_matmul_reshape_rule])
diff --git a/onnxscript/rewriter/broadcast_to_matmul.py b/onnxscript/rewriter/broadcast_to_matmul.py
--- a/onnxscript/rewriter/broadcast_to_matmul.py
+++ b/onnxscript/rewriter/broadcast_to_matmul.py
@@ -11,7 +11,7 @@
 
 
 # condition to check if we need to replace the pattern
-def check_if_need_reshape(input_a, input_b, shape_c, **_) -> bool:
+def check_if_not_need_reshape(input_a, input_b, shape_c, **_) -> bool:
     """If matmul broadcasting is enough, then we don't need the reshapes.
 
     To validate this, we need to check the following:
@@ -108,9 +108,12 @@ def check_if_need_reshape(input_a, input_b, shape_c, **_) -> bool:
     broadcast_matmul_output_shape = (
         longer_shape[: -len(shorter_shape)] + broadcast_matmul_output_shape
     )
-    if mimic_matmul_broadcast_behavior and dim_b == 2:
+    if mimic_matmul_broadcast_behavior and dim_b == 2 and input_b_shape[-1] == 1:
+        # If input_b is expanded to 2-D, then we need to remove the last dimension
         broadcast_matmul_output_shape = broadcast_matmul_output_shape[:-1]
-    if mimic_matmul_broadcast_behavior and dim_a == 2:
+    if mimic_matmul_broadcast_behavior and dim_a == 2 and input_a_shape[0] == 1:
+        # If input_a is expanded to 2-D, then we need to remove the first dimension
+        # of input_a, which would be the -2nd dimension of the output shape.
         broadcast_matmul_output_shape.pop(-2)
     if shape_c != broadcast_matmul_output_shape:
         logger.info(
@@ -149,14 +152,14 @@ def one_reshape_matmul_reshape_pattern(input_a, input_b, shape_a, shape_c):
 two_reshapes_matmul_reshape_rule = pattern.RewriteRule(
     two_reshapes_matmul_reshape_pattern,
     matmul,
-    check_if_need_reshape,
+    check_if_not_need_reshape,
 )
 one_reshape_matmul_reshape_rule = pattern.RewriteRule(
     one_reshape_matmul_reshape_pattern,
     matmul,
-    # We can use the same check_if_need_reshape function for both the rules,
+    # We can use the same check_if_not_need_reshape function for both the rules,
     # as one_reshape_matmul_reshape_pattern is a subset of two_reshapes_matmul_reshape_pattern.
-    check_if_need_reshape,
+    check_if_not_need_reshape,
 )
 
 # NOTE: The order of the rules is important. Larger pattern should be checked first.
diff --git a/onnxscript/rewriter/gemm_to_matmul_add.py b/onnxscript/rewriter/gemm_to_matmul_add.py
--- a/onnxscript/rewriter/gemm_to_matmul_add.py
+++ b/onnxscript/rewriter/gemm_to_matmul_add.py
@@ -1,5 +1,5 @@
 from onnxscript.rewriter import pattern
-from onnxscript.rewriter.broadcast_to_matmul import check_if_need_reshape
+from onnxscript.rewriter.broadcast_to_matmul import check_if_not_need_reshape
 
 op = pattern.onnxop
 
@@ -18,4 +18,4 @@ def matmul_add(op, input_a, input_b, input_c, **_):
     return op.Add(matmul, input_c)
 
 
-rule = pattern.RewriteRule(reshape_gemm_reshape_pattern, matmul_add, check_if_need_reshape)
+rule = pattern.RewriteRule(reshape_gemm_reshape_pattern, matmul_add, check_if_not_need_reshape)


Test patch:
diff --git a/onnxscript/rewriter/gemm_to_matmul_add_test.py b/onnxscript/rewriter/gemm_to_matmul_add_test.py
--- a/onnxscript/rewriter/gemm_to_matmul_add_test.py
+++ b/onnxscript/rewriter/gemm_to_matmul_add_test.py
@@ -163,7 +163,7 @@ def test_reshape_gemm_reshape_replace_when_first_input_is_one_dimension_and_broa
         self.assertEqual(model.graph[2].op_type, "MatMul")
         self.assertEqual(model.graph[3].op_type, "Add")
 
-    def test_reshape_gemm_reshape_replace_when_first_input_is_one_dimension_and_not_broadcastable(
+    def test_reshape_gemm_reshape_remain_when_first_input_is_one_dimension_and_not_broadcastable(
         self,
     ):
         model_proto = onnx.parser.parse_model(
@@ -207,7 +207,7 @@ def test_reshape_gemm_reshape_replace_when_second_input_is_one_dimension_and_bro
         self.assertEqual(model.graph[2].op_type, "MatMul")
         self.assertEqual(model.graph[3].op_type, "Add")
 
-    def test_reshape_gemm_reshape_replace_when_second_input_is_one_dimension_and_not_broadcastable(
+    def test_reshape_gemm_reshape_remain_when_second_input_is_one_dimension_and_not_broadcastable(
         self,
     ):
         model_proto = onnx.parser.parse_model(
@@ -228,6 +228,48 @@ def test_reshape_gemm_reshape_replace_when_second_input_is_one_dimension_and_not
         self.assertEqual(count, 0)
         self.assertEqual(len(model.graph), 5)
 
+    def test_reshape_gemm_reshape_replaces_when_inputs_are_two_dimensional_and_broadcastable(
+        self,
+    ):
+        model_proto = onnx.parser.parse_model(
+            """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[3, 5] input_x, float[5, 10] input_y, float[3, 10] input_z) => (float[3, 10] output)
+            {
+                shape_a = Constant<value: tensor = int64[2] {3, 5}>()
+                reshape_x = Reshape (input_x, shape_a)
+                gemm = Gemm<alpha=1.0, beta=1.0> (reshape_x, input_y, input_z)
+                shape_d = Constant<value: tensor = int64[2] {3, 10}>()
+                output = Reshape (gemm, shape_d)
+            }
+        """
+        )
+        model = ir.serde.deserialize_model(model_proto)
+        replacement_count = gemm_to_matmul_add.rule.apply_to_model(model)
+        self.assertEqual(replacement_count, 1)
+        self.assertEqual(len(model.graph), 4)
+
+    def test_reshape_gemm_reshape_remain_when_inputs_are_two_dimension_and_not_broadcastable(
+        self,
+    ):
+        model_proto = onnx.parser.parse_model(
+            """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[5, 3] input_x, float[5, 10] input_y, float[3, 10] input_z) => (float[3, 10] output)
+            {
+                shape_a = Constant<value: tensor = int64[2] {3, 5}>()
+                reshape_x = Reshape (input_x, shape_a)
+                gemm = Gemm<alpha=1.0, beta=1.0> (reshape_x, input_y, input_z)
+                shape_d = Constant<value: tensor = int64[2] {3, 10}>()
+                output = Reshape (gemm, shape_d)
+            }
+        """
+        )
+        model = ir.serde.deserialize_model(model_proto)
+        count = gemm_to_matmul_add.rule.apply_to_model(model)
+        self.assertEqual(count, 0)
+        self.assertEqual(len(model.graph), 5)
+
     def test_reshape_gemm_reshape_remain_when_output_is_not_matmul_broadcasted(
         self,
     ):