From 60ae655de0d6c9b820dcd6f944576335faa31837 Mon Sep 17 00:00:00 2001 From: amacharla15 Date: Mon, 14 Sep 2026 01:55:24 -0700 Subject: [PATCH 1/5] Cortex-M: fix quantized lowering for decomposed attention --- .../cortex_m/passes/cortex_m_pass_manager.py | 9 +- .../passes/cortex_m_quantized_pass.py | 79 ++++++++++++ .../passes/materialize_mul_broadcast_pass.py | 74 ++++++++++++ .../passes/replace_quant_nodes_pass.py | 6 +- .../passes/replace_scalar_with_tensor_pass.py | 17 +++ .../cortex_m/test/models/test_transformer.py | 113 ++++++++++++++++++ 6 files changed, 295 insertions(+), 3 deletions(-) create mode 100644 backends/cortex_m/passes/cortex_m_quantized_pass.py create mode 100644 backends/cortex_m/passes/materialize_mul_broadcast_pass.py create mode 100644 backends/cortex_m/passes/replace_scalar_with_tensor_pass.py create mode 100644 backends/cortex_m/test/models/test_transformer.py diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index a15a7262660..e5c9ed691fb 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -8,6 +8,7 @@ from typing import Any, Optional, Type from executorch.backends.arm._passes import ( + ComputeConstantOpsAOTPass, DeduplicateGetAttrPass, FoldAndAnnotateQParamsPass, ScalarsToAttributePass, @@ -41,9 +42,11 @@ CortexMReplaceOpsWithChannelsLastVariants, ValidateCortexMExplicitLayoutPass, ) +from .materialize_mul_broadcast_pass import MaterializeMulBroadcastPass from .matmul_to_bmm_pass import MatmulToBmmPass from .quantized_clamp_activation_pass import QuantizedClampActivationPass from .replace_quant_nodes_pass import ReplaceQuantNodesPass +from .replace_scalar_with_tensor_pass import CortexMReplaceScalarWithTensorArgPass PassClass = Type[ExportPass] @@ -53,7 +56,10 @@ class CortexMPassManager(PassManager): # Run before folding so qparams attach to max_pool2d values, not tuple + getitem. RemoveGetItemPass, FoldAndAnnotateQParamsPass, - ReplaceScalarWithTensorArgPass, + # Materialize constant ops after qparams are attached so quantized + # scalar constants use the Cortex-M input dtype instead of remaining float. + ComputeConstantOpsAOTPass, + CortexMReplaceScalarWithTensorArgPass, ReplaceQuantNodesPass, ActivationFusionPass, QuantizedClampActivationPass, @@ -84,6 +90,7 @@ class CortexMPassManager(PassManager): pass_list_transform_for_annotation: list[PassClass] = [ ScalarsToAttributePass, ReplaceScalarWithTensorArgPass, + MaterializeMulBroadcastPass, ClampHardswishPass, DecomposeMeanPass, MatmulToBmmPass, diff --git a/backends/cortex_m/passes/cortex_m_quantized_pass.py b/backends/cortex_m/passes/cortex_m_quantized_pass.py new file mode 100644 index 00000000000..082ee44ec82 --- /dev/null +++ b/backends/cortex_m/passes/cortex_m_quantized_pass.py @@ -0,0 +1,79 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# 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 Any + +import torch + +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue +from torch.utils import _pytree as pytree + + +class CortexMQuantizedPass(ExportPass): + """Base pass for retracing folded Cortex-M quantized graphs. + + Cortex-M's quantized linear kernel includes requantization and therefore + produces the dtype described by its output qparams, normally int8. + + Once Q/DQ nodes have been folded, retracing ``aten.linear`` through the + regular PyTorch fake kernel would incorrectly recreate an int32 output. + Preserve the Cortex-M output dtype instead. + """ + + @staticmethod + def _is_quantized_meta(meta: NodeMetadata) -> bool: + input_qparams = meta.data.get("input_qparams", {}) + output_qparams = meta.data.get("output_qparams", {}) + return bool(input_qparams) and bool(output_qparams) + + def _call_quantized_op_without_fake_kernel( + self, + op, + args: tuple[ProxyValue, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + old_val = meta.data["val"] + output_qparams = meta.data.get("output_qparams", {}) + dtype = ( + next(iter(output_qparams.values())).dtype + if output_qparams + else old_val.dtype + ) + + result_data = torch.empty_like(old_val, dtype=dtype) + + args_proxy, kwargs_proxy = pytree.tree_map_only( + ProxyValue, + lambda value: value.proxy, + (args, kwargs), + ) + + result_proxy = self.tracer.create_proxy( + "call_function", + op, + args_proxy, + kwargs_proxy, + ) + result_proxy.node.meta.update(meta.data) + self.tracer.set_metadata(result_proxy.node, result_data) + + return ProxyValue(result_data, result_proxy) + + def call_operator(self, op, args, kwargs, meta): + if ( + op == exir_ops.edge.aten.linear.default + and isinstance(meta, NodeMetadata) + and self._is_quantized_meta(meta) + ): + return self._call_quantized_op_without_fake_kernel( + op, + args, + kwargs, + meta, + ) + + return super().call_operator(op, args, kwargs, meta) diff --git a/backends/cortex_m/passes/materialize_mul_broadcast_pass.py b/backends/cortex_m/passes/materialize_mul_broadcast_pass.py new file mode 100644 index 00000000000..0fecf0c28f9 --- /dev/null +++ b/backends/cortex_m/passes/materialize_mul_broadcast_pass.py @@ -0,0 +1,74 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor +from executorch.backends.cortex_m.passes.passes_utils import ( + is_channel_broadcast, + is_channels_last, +) +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import GraphModule, Node + + +class MaterializeMulBroadcastPass(ExportPass): + """Materialize MUL broadcasts unsupported by the Cortex-M kernel. + + ``cortex_m.quantized_mul`` supports equal-sized tensors and the backend's + channels-last channel-broadcast form. Other valid PyTorch broadcasting + forms are expanded before quantization so that the resulting MUL can be + annotated and lowered to the Cortex-M kernel. + """ + + def call(self, graph_module: GraphModule) -> PassResult: + modified = False + + for node in list(graph_module.graph.nodes): + if node.op != "call_function" or node.target != torch.ops.aten.mul.Tensor: + continue + + lhs, rhs = node.args[:2] + + if not isinstance(lhs, Node) or not isinstance(rhs, Node): + continue + + lhs_tensor = get_first_fake_tensor(lhs) + rhs_tensor = get_first_fake_tensor(rhs) + output_tensor = get_first_fake_tensor(node) + + if lhs_tensor.shape == rhs_tensor.shape: + continue + + if is_channel_broadcast(lhs_tensor, rhs_tensor) and is_channels_last( + lhs_tensor + ): + continue + + output_shape = list(output_tensor.shape) + new_args = list(node.args) + + for index, arg in enumerate((lhs, rhs)): + tensor = get_first_fake_tensor(arg) + + if tensor.shape == output_tensor.shape: + continue + + with graph_module.graph.inserting_before(node): + expanded = graph_module.graph.call_function( + torch.ops.aten.expand_copy.default, + args=(arg, output_shape), + ) + + new_args[index] = expanded + modified = True + + node.args = tuple(new_args) + + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + + return PassResult(graph_module, modified) diff --git a/backends/cortex_m/passes/replace_quant_nodes_pass.py b/backends/cortex_m/passes/replace_quant_nodes_pass.py index e8ba16d11c0..e08d2bcf998 100644 --- a/backends/cortex_m/passes/replace_quant_nodes_pass.py +++ b/backends/cortex_m/passes/replace_quant_nodes_pass.py @@ -11,10 +11,12 @@ import torch from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue +from executorch.exir.pass_base import NodeMetadata, ProxyValue +from .cortex_m_quantized_pass import CortexMQuantizedPass -class ReplaceQuantNodesPass(ExportPass): + +class ReplaceQuantNodesPass(CortexMQuantizedPass): """ Replace quantize and dequantize nodes with the corresponding cortex_m.quantize_per_tensor and cortex_m.dequantize_per_tensor nodes. diff --git a/backends/cortex_m/passes/replace_scalar_with_tensor_pass.py b/backends/cortex_m/passes/replace_scalar_with_tensor_pass.py new file mode 100644 index 00000000000..5c954044fdf --- /dev/null +++ b/backends/cortex_m/passes/replace_scalar_with_tensor_pass.py @@ -0,0 +1,17 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.backends.transforms.replace_scalar_with_tensor import ( + ReplaceScalarWithTensorArgPass, +) + +from .cortex_m_quantized_pass import CortexMQuantizedPass + + +class CortexMReplaceScalarWithTensorArgPass( + ReplaceScalarWithTensorArgPass, + CortexMQuantizedPass, +): + """Scalar-to-tensor replacement with Cortex-M quantized retracing.""" diff --git a/backends/cortex_m/test/models/test_transformer.py b/backends/cortex_m/test/models/test_transformer.py new file mode 100644 index 00000000000..ae02ff69a4b --- /dev/null +++ b/backends/cortex_m/test/models/test_transformer.py @@ -0,0 +1,113 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from executorch.backends.arm._passes import ( + ConstantFoldingPass, + DecomposeLayerNormPass, + DecomposeSDPAWithRegularSoftmaxPass, +) +from executorch.backends.cortex_m.test.tester import CortexMQuantize, CortexMTester +from torch.export import export +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +class TransformerModel(torch.nn.Module): + """Small Transformer used to exercise decomposed attention lowering.""" + + def __init__(self): + super().__init__() + self.transformer = torch.nn.Transformer( + d_model=32, + nhead=2, + num_encoder_layers=1, + num_decoder_layers=1, + dim_feedforward=64, + dropout=0.0, + batch_first=False, + ) + + def forward(self, src, tgt): + return self.transformer(src, tgt) + + +class TransformerDecompositionQuantize(CortexMQuantize): + """Apply the decomposition workflow before Cortex-M PT2E quantization.""" + + def run(self, artifact, inputs): + assert inputs is not None + + captured_graph = export( + artifact, + inputs, + strict=True, + ).module() + + captured_graph = ConstantFoldingPass().call(captured_graph).graph_module + captured_graph = ( + DecomposeSDPAWithRegularSoftmaxPass().call(captured_graph).graph_module + ) + captured_graph = DecomposeLayerNormPass().call(captured_graph).graph_module + + prepared = prepare_pt2e( + captured_graph, + self.quantizer, + ) + + with torch.no_grad(): + prepared(*inputs) + + self.converted_graph = convert_pt2e( + prepared, + fold_quantize=self.fold_quantize, + ) + + +def test_transformer_decomposed_attention_lowering(cortex_m_target): + """Decomposed attention ops should lower to Cortex-M quantized kernels.""" + + torch.manual_seed(0) + + model = TransformerModel().eval() + inputs = ( + torch.randn(8, 1, 32), + torch.randn(8, 1, 32), + ) + + tester = CortexMTester( + model, + inputs, + target_config=cortex_m_target, + ) + + tester.quantize(TransformerDecompositionQuantize()) + tester.export() + tester.to_edge() + tester.run_passes() + + # Decomposition details may change with PyTorch, so require Cortex-M + # lowering without depending on exact generated-op counts. + tester.check( + [ + "executorch_exir_dialects_edge__ops_cortex_m_quantized_batch_matmul_default", + "executorch_exir_dialects_edge__ops_cortex_m_softmax_default", + "executorch_exir_dialects_edge__ops_cortex_m_quantized_mul_default", + ] + ) + + tester.check_not( + [ + "executorch_exir_dialects_edge__ops_aten_bmm_default", + "executorch_exir_dialects_edge__ops_aten_mul_Tensor", + "executorch_exir_dialects_edge__ops_aten__softmax_default", + "executorch_exir_dialects_edge__ops_aten__safe_softmax_default", + ] + ) + + tester.run_method_and_compare_outputs( + inputs=inputs, + qtol=2, + ) From e1171a60a28abd29c1962259adb4647548a8f64d Mon Sep 17 00:00:00 2001 From: amacharla15 Date: Mon, 14 Sep 2026 06:19:11 -0700 Subject: [PATCH 2/5] Cortex-M: narrow decomposed attention lowering fixes --- .../cortex_m/passes/cortex_m_pass_manager.py | 9 ++-- .../passes/materialize_mul_broadcast_pass.py | 32 ++++++----- ...aterialize_quantized_mul_constants_pass.py | 53 +++++++++++++++++++ .../cortex_m/test/models/test_transformer.py | 1 - 4 files changed, 78 insertions(+), 17 deletions(-) create mode 100644 backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index e5c9ed691fb..bf3652956c1 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -8,7 +8,6 @@ from typing import Any, Optional, Type from executorch.backends.arm._passes import ( - ComputeConstantOpsAOTPass, DeduplicateGetAttrPass, FoldAndAnnotateQParamsPass, ScalarsToAttributePass, @@ -43,6 +42,9 @@ ValidateCortexMExplicitLayoutPass, ) from .materialize_mul_broadcast_pass import MaterializeMulBroadcastPass +from .materialize_quantized_mul_constants_pass import ( + MaterializeQuantizedMulConstantsPass, +) from .matmul_to_bmm_pass import MatmulToBmmPass from .quantized_clamp_activation_pass import QuantizedClampActivationPass from .replace_quant_nodes_pass import ReplaceQuantNodesPass @@ -56,9 +58,8 @@ class CortexMPassManager(PassManager): # Run before folding so qparams attach to max_pool2d values, not tuple + getitem. RemoveGetItemPass, FoldAndAnnotateQParamsPass, - # Materialize constant ops after qparams are attached so quantized - # scalar constants use the Cortex-M input dtype instead of remaining float. - ComputeConstantOpsAOTPass, + # Materialize only quantized MUL constants after qparams are attached. + MaterializeQuantizedMulConstantsPass, CortexMReplaceScalarWithTensorArgPass, ReplaceQuantNodesPass, ActivationFusionPass, diff --git a/backends/cortex_m/passes/materialize_mul_broadcast_pass.py b/backends/cortex_m/passes/materialize_mul_broadcast_pass.py index 0fecf0c28f9..f817e0eeb0e 100644 --- a/backends/cortex_m/passes/materialize_mul_broadcast_pass.py +++ b/backends/cortex_m/passes/materialize_mul_broadcast_pass.py @@ -14,14 +14,21 @@ class MaterializeMulBroadcastPass(ExportPass): - """Materialize MUL broadcasts unsupported by the Cortex-M kernel. + """Materialize scalar MUL broadcasts unsupported by the Cortex-M kernel. - ``cortex_m.quantized_mul`` supports equal-sized tensors and the backend's - channels-last channel-broadcast form. Other valid PyTorch broadcasting - forms are expanded before quantization so that the resulting MUL can be - annotated and lowered to the Cortex-M kernel. + ``ReplaceScalarWithTensorArgPass`` represents scalar multiplication as a + rank-1 ``aten.full`` tensor. When that broadcast is unsupported by the + Cortex-M quantized MUL kernel, materialize the scalar constant directly at + the MUL output shape. Leave general tensor broadcasting unchanged so those + graphs keep their existing fallback behavior. """ + @staticmethod + def _is_scalar_full(node: Node) -> bool: + if node.op != "call_function" or node.target != torch.ops.aten.full.default: + return False + return get_first_fake_tensor(node).numel() == 1 + def call(self, graph_module: GraphModule) -> PassResult: modified = False @@ -30,7 +37,6 @@ def call(self, graph_module: GraphModule) -> PassResult: continue lhs, rhs = node.args[:2] - if not isinstance(lhs, Node) or not isinstance(rhs, Node): continue @@ -51,17 +57,19 @@ def call(self, graph_module: GraphModule) -> PassResult: for index, arg in enumerate((lhs, rhs)): tensor = get_first_fake_tensor(arg) - - if tensor.shape == output_tensor.shape: + if tensor.shape == output_tensor.shape or not self._is_scalar_full(arg): continue + full_args = list(arg.args) + full_args[0] = output_shape with graph_module.graph.inserting_before(node): - expanded = graph_module.graph.call_function( - torch.ops.aten.expand_copy.default, - args=(arg, output_shape), + materialized = graph_module.graph.call_function( + torch.ops.aten.full.default, + args=tuple(full_args), + kwargs=dict(arg.kwargs), ) - new_args[index] = expanded + new_args[index] = materialized modified = True node.args = tuple(new_args) diff --git a/backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py b/backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py new file mode 100644 index 00000000000..2e17c44bc24 --- /dev/null +++ b/backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py @@ -0,0 +1,53 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.backends.arm._passes import ComputeConstantOpsAOTPass +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import Node + + +class MaterializeQuantizedMulConstantsPass(ComputeConstantOpsAOTPass): + """Precompute only quantized scalar constants that feed Cortex-M MUL. + + The generic constant-materialization pass also folds unrelated operators + such as ``arange``. Cortex-M only needs the ``aten.full`` constants created + for scalar MULs here, after qparams have been attached, so keep the scope + limited to those nodes. + """ + + targeted_ops = [exir_ops.edge.aten.full.default] + + _passthrough_ops = { + exir_ops.edge.aten.clone.default, + exir_ops.edge.dim_order_ops._clone_dim_order.default, + } + + @classmethod + def _feeds_quantized_mul(cls, node: Node) -> bool: + worklist = list(node.users) + visited: set[Node] = set() + + while worklist: + user = worklist.pop() + if user in visited: + continue + visited.add(user) + + if user.target == exir_ops.edge.aten.mul.Tensor: + return bool(user.meta.get("input_qparams")) and bool( + user.meta.get("output_qparams") + ) + + if user.target in cls._passthrough_ops: + worklist.extend(user.users) + + return False + + def compute_node_aot(self, node: Node) -> bool: + if not node.meta.get("output_qparams"): + return False + if not self._feeds_quantized_mul(node): + return False + return super().compute_node_aot(node) diff --git a/backends/cortex_m/test/models/test_transformer.py b/backends/cortex_m/test/models/test_transformer.py index ae02ff69a4b..f5dceb26768 100644 --- a/backends/cortex_m/test/models/test_transformer.py +++ b/backends/cortex_m/test/models/test_transformer.py @@ -101,7 +101,6 @@ def test_transformer_decomposed_attention_lowering(cortex_m_target): tester.check_not( [ "executorch_exir_dialects_edge__ops_aten_bmm_default", - "executorch_exir_dialects_edge__ops_aten_mul_Tensor", "executorch_exir_dialects_edge__ops_aten__softmax_default", "executorch_exir_dialects_edge__ops_aten__safe_softmax_default", ] From 3bc33a411d371faaa8bedc66163d2d733fc1b1c2 Mon Sep 17 00:00:00 2001 From: amacharla15 Date: Mon, 14 Sep 2026 08:32:23 -0700 Subject: [PATCH 3/5] Cortex-M: address decomposed attention review feedback --- ...aterialize_quantized_mul_constants_pass.py | 25 ++++---- .../cortex_m/test/models/test_transformer.py | 50 ++++++++++++++++ ...est_materialize_quantized_mul_constants.py | 59 +++++++++++++++++++ 3 files changed, 124 insertions(+), 10 deletions(-) create mode 100644 backends/cortex_m/test/test_materialize_quantized_mul_constants.py diff --git a/backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py b/backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py index 2e17c44bc24..1ad5c7ce2ec 100644 --- a/backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py +++ b/backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py @@ -25,14 +25,18 @@ class MaterializeQuantizedMulConstantsPass(ComputeConstantOpsAOTPass): } @classmethod - def _feeds_quantized_mul(cls, node: Node) -> bool: - worklist = list(node.users) + def _feeds_single_quantized_mul(cls, node: Node) -> bool: + current = node visited: set[Node] = set() - while worklist: - user = worklist.pop() + while True: + users = list(current.users) + if len(users) != 1: + return False + + user = users[0] if user in visited: - continue + return False visited.add(user) if user.target == exir_ops.edge.aten.mul.Tensor: @@ -40,14 +44,15 @@ def _feeds_quantized_mul(cls, node: Node) -> bool: user.meta.get("output_qparams") ) - if user.target in cls._passthrough_ops: - worklist.extend(user.users) + if user.target not in cls._passthrough_ops: + return False - return False + current = user def compute_node_aot(self, node: Node) -> bool: - if not node.meta.get("output_qparams"): + output_qparams = node.meta.get("output_qparams", {}) + if set(output_qparams) != {0}: return False - if not self._feeds_quantized_mul(node): + if not self._feeds_single_quantized_mul(node): return False return super().compute_node_aot(node) diff --git a/backends/cortex_m/test/models/test_transformer.py b/backends/cortex_m/test/models/test_transformer.py index f5dceb26768..cb013068936 100644 --- a/backends/cortex_m/test/models/test_transformer.py +++ b/backends/cortex_m/test/models/test_transformer.py @@ -15,6 +15,13 @@ from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e +class AttentionModel(torch.nn.Module): + """Minimal SDPA model used to isolate attention-scale MUL lowering.""" + + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention(q, k, v) + + class TransformerModel(torch.nn.Module): """Small Transformer used to exercise decomposed attention lowering.""" @@ -66,6 +73,49 @@ def run(self, artifact, inputs): ) +def test_decomposed_attention_scale_mul_lowering(cortex_m_target): + """The two SDPA scaling MULs should lower to Cortex-M quantized MUL.""" + + torch.manual_seed(0) + + model = AttentionModel().eval() + inputs = ( + torch.randn(1, 2, 4, 16), + torch.randn(1, 2, 4, 16), + torch.randn(1, 2, 4, 16), + ) + + tester = CortexMTester( + model, + inputs, + target_config=cortex_m_target, + ) + + tester.quantize(TransformerDecompositionQuantize()) + tester.export() + tester.to_edge() + tester.run_passes() + + tester.check_count( + { + "executorch_exir_dialects_edge__ops_cortex_m_quantized_mul_default": 2, + "executorch_exir_dialects_edge__ops_cortex_m_quantized_batch_matmul_default": 2, + "executorch_exir_dialects_edge__ops_cortex_m_softmax_default": 1, + } + ) + + tester.check_not( + [ + "executorch_exir_dialects_edge__ops_aten_mul_Tensor", + ] + ) + + tester.run_method_and_compare_outputs( + inputs=inputs, + qtol=2, + ) + + def test_transformer_decomposed_attention_lowering(cortex_m_target): """Decomposed attention ops should lower to Cortex-M quantized kernels.""" diff --git a/backends/cortex_m/test/test_materialize_quantized_mul_constants.py b/backends/cortex_m/test/test_materialize_quantized_mul_constants.py new file mode 100644 index 00000000000..25dbf3fe0fe --- /dev/null +++ b/backends/cortex_m/test/test_materialize_quantized_mul_constants.py @@ -0,0 +1,59 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.backends.cortex_m.passes.materialize_quantized_mul_constants_pass import ( + MaterializeQuantizedMulConstantsPass, +) +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import Graph + + +def _mark_quantized(node): + node.meta["input_qparams"] = {0: object()} + node.meta["output_qparams"] = {0: object()} + + +def test_accepts_single_quantized_mul_path(): + graph = Graph() + + full = graph.call_function( + exir_ops.edge.aten.full.default, + args=((1,), 0.5), + ) + clone = graph.call_function( + exir_ops.edge.aten.clone.default, + args=(full,), + ) + other = graph.placeholder("other") + mul = graph.call_function( + exir_ops.edge.aten.mul.Tensor, + args=(other, clone), + ) + _mark_quantized(mul) + + assert MaterializeQuantizedMulConstantsPass._feeds_single_quantized_mul(full) + + +def test_rejects_shared_constant_with_unrelated_consumer(): + graph = Graph() + + full = graph.call_function( + exir_ops.edge.aten.full.default, + args=((1,), 0.5), + ) + other = graph.placeholder("other") + + mul = graph.call_function( + exir_ops.edge.aten.mul.Tensor, + args=(other, full), + ) + _mark_quantized(mul) + + graph.call_function( + exir_ops.edge.aten.add.Tensor, + args=(other, full), + ) + + assert not MaterializeQuantizedMulConstantsPass._feeds_single_quantized_mul(full) From 59d3be34f76d0dedd956cc56beb21900ba55de11 Mon Sep 17 00:00:00 2001 From: amacharla15 Date: Mon, 14 Sep 2026 09:06:23 -0700 Subject: [PATCH 4/5] Cortex-M: add new passes to Buck targets --- backends/cortex_m/passes/BUCK | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index 4d758703cb4..38f6ec2727c 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -10,6 +10,18 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") +fbcode_target(_kind = runtime.python_library, + name="cortex_m_quantized_pass", + srcs=[ + "cortex_m_quantized_pass.py", + ], + deps=[ + "//caffe2:torch", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], +) + fbcode_target(_kind = runtime.python_library, name="replace_quant_nodes_pass", srcs=[ @@ -21,6 +33,7 @@ fbcode_target(_kind = runtime.python_library, "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", "//executorch/backends/cortex_m/ops:ops", + "//executorch/backends/cortex_m/passes:cortex_m_quantized_pass", "//executorch/backends/cortex_m/passes:passes_utils", ], ) @@ -37,7 +50,10 @@ fbcode_target(_kind = runtime.python_library, "decompose_mean_pass.py", "explicit_layout_pass.py", "matmul_to_bmm_pass.py", + "materialize_mul_broadcast_pass.py", + "materialize_quantized_mul_constants_pass.py", "quantized_clamp_activation_pass.py", + "replace_scalar_with_tensor_pass.py", "scratch_buffer_sizes.py", ], deps=[ @@ -46,6 +62,7 @@ fbcode_target(_kind = runtime.python_library, "//executorch/backends/cortex_m:cmsis_nn", "//executorch/backends/cortex_m:target_config", "//executorch/backends/cortex_m/ops:ops", + "//executorch/backends/cortex_m/passes:cortex_m_quantized_pass", "//executorch/backends/cortex_m/passes:passes_utils", "//executorch/backends/cortex_m/passes:replace_quant_nodes_pass", "//executorch/backends/transforms:aten_to_dialect_pass", From c09864c17849acaa0b6e4fbbd832ee3ca8424b7a Mon Sep 17 00:00:00 2001 From: amacharla15 Date: Mon, 14 Sep 2026 09:37:41 -0700 Subject: [PATCH 5/5] Cortex-M: preserve quantized BMM dtype during retracing --- .../passes/cortex_m_quantized_pass.py | 6 +- .../cortex_m/test/models/test_transformer.py | 58 +++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/backends/cortex_m/passes/cortex_m_quantized_pass.py b/backends/cortex_m/passes/cortex_m_quantized_pass.py index 082ee44ec82..6b9ff98041c 100644 --- a/backends/cortex_m/passes/cortex_m_quantized_pass.py +++ b/backends/cortex_m/passes/cortex_m_quantized_pass.py @@ -65,7 +65,11 @@ def _call_quantized_op_without_fake_kernel( def call_operator(self, op, args, kwargs, meta): if ( - op == exir_ops.edge.aten.linear.default + op + in { + exir_ops.edge.aten.linear.default, + exir_ops.edge.aten.bmm.default, + } and isinstance(meta, NodeMetadata) and self._is_quantized_meta(meta) ): diff --git a/backends/cortex_m/test/models/test_transformer.py b/backends/cortex_m/test/models/test_transformer.py index cb013068936..5612fc7e01e 100644 --- a/backends/cortex_m/test/models/test_transformer.py +++ b/backends/cortex_m/test/models/test_transformer.py @@ -9,8 +9,16 @@ ConstantFoldingPass, DecomposeLayerNormPass, DecomposeSDPAWithRegularSoftmaxPass, + FoldAndAnnotateQParamsPass, +) +from executorch.backends.cortex_m.passes.replace_scalar_with_tensor_pass import ( + CortexMReplaceScalarWithTensorArgPass, ) from executorch.backends.cortex_m.test.tester import CortexMQuantize, CortexMTester +from executorch.backends.test.harness.stages import StageType +from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.program._program import _transform from torch.export import export from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e @@ -73,6 +81,56 @@ def run(self, artifact, inputs): ) +def test_quantized_bmm_retrace_preserves_output_dtype(cortex_m_target): + """Retracing folded quantized BMM must preserve its requantized dtype.""" + + torch.manual_seed(0) + + model = AttentionModel().eval() + inputs = ( + torch.randn(1, 2, 4, 16), + torch.randn(1, 2, 4, 16), + torch.randn(1, 2, 4, 16), + ) + + tester = CortexMTester( + model, + inputs, + target_config=cortex_m_target, + ) + + tester.quantize(TransformerDecompositionQuantize()) + tester.export() + tester.to_edge() + + exported_program = tester.get_artifact(StageType.TO_EDGE).exported_program() + + exported_program = _transform( + exported_program, + RemoveGetItemPass(), + ) + exported_program = _transform( + exported_program, + FoldAndAnnotateQParamsPass(exported_program), + ) + exported_program = _transform( + exported_program, + CortexMReplaceScalarWithTensorArgPass(), + ) + + bmm_nodes = [ + node + for node in exported_program.graph_module.graph.nodes + if node.target == exir_ops.edge.aten.bmm.default + ] + + assert len(bmm_nodes) == 2 + + for node in bmm_nodes: + output_qparams = node.meta["output_qparams"] + assert node.meta["val"].dtype == output_qparams[0].dtype + + def test_decomposed_attention_scale_mul_lowering(cortex_m_target): """The two SDPA scaling MULs should lower to Cortex-M quantized MUL."""