Cortex-M: quantize aten.matmul by rewriting it to bmm before annotation

backends/cortex_m/passes/cortex_m_pass_manager.py

@@ -27,6 +27,7 @@
 from .clamp_hardswish_pass import ClampHardswishPass
 from .decompose_hardswish_pass import DecomposeHardswishPass
 from .decompose_mean_pass import DecomposeMeanPass
+from .matmul_to_bmm_pass import MatmulToBmmPass
 from .quantized_clamp_activation_pass import QuantizedClampActivationPass
 from .replace_quant_nodes_pass import ReplaceQuantNodesPass
 
@@ -51,6 +52,7 @@ class CortexMPassManager(PassManager):
         ReplaceScalarWithTensorArgPass,
         ClampHardswishPass,
         DecomposeMeanPass,
+        MatmulToBmmPass,
         DeduplicateGetAttrPass,
     ]
 

backends/cortex_m/passes/matmul_to_bmm_pass.py

@@ -0,0 +1,81 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import cast, Dict
+
+import torch
+from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue
+
+from torch._ops import OpOverload
+from torch.fx.node import Argument
+
+
+class MatmulToBmmPass(ExportPass):
+    """
+    Rewrites ``aten.matmul.default`` into ``aten.bmm.default`` so the cortex_m
+    annotator and ``quantized_batch_matmul`` lowering pick it up.
+
+    ``a @ b`` / ``torch.matmul`` is captured as ``aten.matmul.default`` at
+    annotation time and only decomposes to ``bmm`` at ``to_edge`` -- after
+    quantization -- so it would otherwise never receive qparams. This pass runs
+    before annotation and normalizes matmul to bmm:
+
+    - rank-3 @ rank-3 with matching batch dims: replaced directly with
+      ``aten.bmm.default``.
+    - rank>3 (e.g. attention [B,H,N,d]@[B,H,d,N]) with matching leading batch
+      dims: the leading batch dims are folded into a single batch dim (reshape to
+      3D), ``aten.bmm.default`` is applied, and the result is reshaped back to the
+      original leading dims. This is required because the cortex_m bmm checker
+      rejects non-rank-3.
+    - everything else is left unchanged, including rank-2 (mm / linear territory)
+      and broadcasting matmuls whose batch dims differ (``aten.bmm`` requires
+      equal batch dims, so rewriting those would crash).
+    """
+
+    def call_operator(
+        self,
+        op: OpOverload,
+        args: tuple[Argument, ...],
+        kwargs: Dict[str, Argument],
+        meta: NodeMetadata,
+    ) -> ProxyValue:
+        if op != torch.ops.aten.matmul.default:
+            return super().call_operator(op, args, kwargs, meta)
+
+        lhs = cast(ProxyValue, args[0])
+        rhs = cast(ProxyValue, args[1])
+        lhs_shape = lhs.to_tensor().shape
+        rhs_shape = rhs.to_tensor().shape
+        lhs_rank = len(lhs_shape)
+        rhs_rank = len(rhs_shape)
+
+        if lhs_rank == 3 and rhs_rank == 3 and lhs_shape[:-2] == rhs_shape[:-2]:
+            return super().call_operator(
+                torch.ops.aten.bmm.default, (lhs, rhs), {}, meta
+            )
+
+        if lhs_rank == rhs_rank and lhs_rank > 3 and lhs_shape[:-2] == rhs_shape[:-2]:
+            batch_dims = list(lhs_shape[:-2])
+            m, k = lhs_shape[-2], lhs_shape[-1]
+            n = rhs_shape[-1]
+
+            lhs_3d = super().call_operator(
+                torch.ops.aten.reshape.default, (lhs, [-1, m, k]), {}, meta
+            )
+            rhs_3d = super().call_operator(
+                torch.ops.aten.reshape.default, (rhs, [-1, k, n]), {}, meta
+            )
+            bmm_out = super().call_operator(
+                torch.ops.aten.bmm.default, (lhs_3d, rhs_3d), {}, meta
+            )
+            return super().call_operator(
+                torch.ops.aten.reshape.default,
+                (bmm_out, batch_dims + [m, n]),
+                {},
+                meta,
+            )
+
+        return super().call_operator(op, args, kwargs, meta)

backends/cortex_m/test/ops/test_batch_matmul.py

@@ -48,6 +48,46 @@ def forward(self, lhs):
         return torch.bmm(lhs, self.rhs)
 
 
+class CortexMMatmul(torch.nn.Module):
+    """``torch.matmul`` is captured as ``aten.matmul.default`` at annotation time
+    and only decomposes to ``bmm`` at ``to_edge`` -- after quantization -- so it
+    would never receive qparams. The pre-annotation matmul->bmm rewrite makes it
+    lower to ``cortex_m.quantized_batch_matmul`` for both rank-3 @ rank-3 and
+    rank>3 (leading batch dims folded to 3D and reshaped back)."""
+
+    ops_before_transforms = {
+        "executorch_exir_dialects_edge__ops_aten_bmm_default": 1,
+    }
+
+    ops_after_transforms = {
+        "executorch_exir_dialects_edge__ops_cortex_m_quantized_batch_matmul_default": 1,
+        "executorch_exir_dialects_edge__ops_aten_bmm_default": 0,
+        "executorch_exir_dialects_edge__ops_aten_matmul_default": 0,
+    }
+
+    def forward(self, lhs, rhs):
+        return torch.matmul(lhs, rhs)
+
+
+class CortexMMatmulBroadcast(torch.nn.Module):
+    """Broadcasting rank-3 matmul whose batch dims differ ([1,4,8] @ [2,8,4] ->
+    [2,4,4]). ``aten.bmm`` requires equal batch dims, so this must NOT be
+    rewritten to bmm; it falls through, stays ``aten.matmul`` -> fp32, and is not
+    lowered (no cortex_m_quantized_batch_matmul)."""
+
+    ops_before_transforms = {
+        "executorch_exir_dialects_edge__ops_aten_bmm_default": 1,
+    }
+
+    ops_after_transforms = {
+        "executorch_exir_dialects_edge__ops_cortex_m_quantized_batch_matmul_default": 0,
+        "executorch_exir_dialects_edge__ops_aten_bmm_default": 1,
+    }
+
+    def forward(self, lhs, rhs):
+        return torch.matmul(lhs, rhs)
+
+
 test_cases = {
     "bmm_small": McuTestCase(
         CortexMBmm(),
@@ -84,6 +124,22 @@ def forward(self, lhs):
 }
 
 
+matmul_test_cases = {
+    "matmul_rank3": McuTestCase(
+        CortexMMatmul(),
+        (torch.randn(2, 4, 8), torch.randn(2, 8, 4)),
+    ),
+    "matmul_rank4": McuTestCase(
+        CortexMMatmul(),
+        (torch.randn(1, 4, 16, 8), torch.randn(1, 4, 8, 16)),
+    ),
+    "matmul_broadcast_rank3": McuTestCase(
+        CortexMMatmulBroadcast(),
+        (torch.randn(1, 4, 8), torch.randn(2, 8, 4)),
+    ),
+}
+
+
 @parametrize("test_case", test_cases)
 def test_dialect_batch_matmul(test_case, cortex_m_target):
     tester = CortexMTester(
@@ -108,6 +164,18 @@ def test_dialect_batch_matmul_const_rhs(test_case, cortex_m_target):
     )
 
 
+@parametrize("test_case", matmul_test_cases)
+def test_dialect_matmul(test_case, cortex_m_target):
+    tester = CortexMTester(
+        test_case.model, test_case.example_inputs, target_config=cortex_m_target
+    )
+    tester.test_dialect(
+        test_case.model.ops_before_transforms,
+        test_case.model.ops_after_transforms,
+        qtol=1,
+    )
+
+
 @parametrize("test_case", test_cases)
 def test_implementation_batch_matmul(test_case, cortex_m_target):
     tester = CortexMTester(