From 3d14f28b3911a9a236d100d171cf208ffd9b1239 Mon Sep 17 00:00:00 2001 From: Prashant Rawat Date: Fri, 23 Oct 2026 12:35:34 -0700 Subject: [PATCH 1/2] Validate layout copies on qparams, not fake-tensor dtype (#24071) Summary: `ValidateCortexMExplicitLayoutPass` requires every NHWC layout copy to move an int8 tensor, and reads `node.meta["val"].dtype` to decide. In `explicit_layout_pass_list` that validator runs before `AtenToCortexMPass`, so at validation time only the activations that `ActivationFusionPass` has already folded into the convolution clamp -- relu, clamp, hardtanh -- carry an int8 fake tensor. The LUT activations (tanh, sigmoid, which become `cortex_m::quantized_activation`) and softmax (which becomes `cortex_m::softmax`) are still `aten.*` with an fp32 fake tensor, even though `FoldAndAnnotateQParamsPass` has already attached their int8 qparams. That fp32 propagates into any convolution fed by one, and the layout copy on its output trips the assertion. The explicit-layout path validates rather than falling back, so this is a hard abort, not a coverage loss -- and explicit layout is the only mode that supports Conv1d. A two-node graph reproduces it: `Conv1d(relu(x))` lowers, `Conv1d(tanh(x))` raises RuntimeError: Cortex-M layout copy channels_last_permute_copy_default must move an int8 tensor. out of `_CortexMLoweringPass`. `LutActivationThenConv1d` in `test_explicit_layout_pipeline.py` is that graph. Judge the condition the kernels will actually see. `produces_int8` accepts an int8 fake tensor as before, and otherwise reads the folded `output_qparams[0]`, which is in place well before this point. Graphs that were rejected lower unchanged: on a 31-convolution ECAPA-TDNN speaker model, 0 of 31 convolutions reached a CMSIS-NN kernel before and 31 of 31 do after, with 0.000e+00 max absolute difference against the quantized eager module. Folded qparams on their own are the wrong bar, though, and this is what the first version of the diff got wrong. They say what dtype a node was annotated for, not that anything acts on the annotation, and a node nothing rewrites keeps an fp32 kernel in the middle of an int8 chain. So `produces_int8` also asks whether a rewrite is coming: either `AtenToCortexMPass` has a substitution registered for the target, or `QuantizedClampActivationPass` has already put the node on int8 terms. It lives in `aten_to_cortex_m_pass.py`, which owns the substitution registry and is therefore the module that can answer the question. `QuantizedClampActivationPass` is the pass that makes the second arm necessary, and it had the same premature read. It quantizes the bounds of a clamp that will see int8, and decided that from the producer's fake tensor, so a clamp behind a LUT activation was skipped -- harmless while the validator rejected those graphs outright, silently wrong once it stopped. `sigmoid -> relu -> Conv1d` left `aten.relu` clamping at 0 instead of the zero point, erasing everything below it, where the relu should have been a no-op; `tanh -> clamp(-0.5, 0.5) -> Conv2d` kept fp32 limits and an fp32 value on the tensor `cortex_m::quantized_conv2d_nhwc` reads as int8. Both fail an output comparison against the quantized eager module, so this was wrong numerics, not a wrong annotation. It now asks `produces_int8`, seven passes ahead of the validator, and records the clamps it quantized. That record is a node predicate in `cortex_m/utils.py`, alongside `skips_quantized_max_pool2d`, which uses the same `meta["custom"]` convention. A clamp never becomes a `cortex_m::` kernel -- it stays on the portable one and reads int8 only because its bounds were rewritten into codes -- so no other pass can re-derive the decision, and a clamp the pass declined (bounds that are not compile-time scalars, output qparams that are not per-tensor) is indistinguishable from one it never saw. With the decision recorded, a declined clamp is rejected by the validator instead of being lowered. `has_qparams` moves to `cortex_m/utils.py` as part of this. It was private to `aten_to_cortex_m_pass`; utils is already where this backend keeps its node predicates. Its ten existing call sites change in name only. The error message now also names the producing node. The layout copy is synthesized by this pass and appears in no graph the caller holds, so its own name gives nothing to go on. Differential Revision: D126708891 --- .../cortex_m/passes/explicit_layout_pass.py | 14 +- .../passes/quantized_clamp_activation_pass.py | 13 +- .../test/test_explicit_layout_pipeline.py | 173 ++++++++++++++++++ backends/cortex_m/utils.py | 20 ++ 4 files changed, 208 insertions(+), 12 deletions(-) diff --git a/backends/cortex_m/passes/explicit_layout_pass.py b/backends/cortex_m/passes/explicit_layout_pass.py index 659b67d729e..f1063bd8eb6 100644 --- a/backends/cortex_m/passes/explicit_layout_pass.py +++ b/backends/cortex_m/passes/explicit_layout_pass.py @@ -10,6 +10,7 @@ import torch from executorch.backends.cortex_m.utils import ( coerce_int_pair, + lowered_output_dtype, skips_quantized_max_pool2d, ) from executorch.backends.transforms.canonicalize_view_copy_permute_pass import ( @@ -134,12 +135,15 @@ def call(self, graph_module: GraphModule) -> PassResult: input_value = ( input_node.meta.get("val") if isinstance(input_node, Node) else None ) - if ( - not isinstance(input_value, torch.Tensor) - or input_value.dtype != torch.int8 - ): + input_dtype = ( + lowered_output_dtype(input_node) + if isinstance(input_node, Node) + else None + ) + if not isinstance(input_value, torch.Tensor) or input_dtype != torch.int8: raise RuntimeError( - f"Cortex-M layout copy {node.name} must move an int8 tensor." + f"Cortex-M layout copy {node.name} must move an int8 tensor, " + f"but its input {input_node} is {input_dtype}." ) rank = input_value.dim() if ( diff --git a/backends/cortex_m/passes/quantized_clamp_activation_pass.py b/backends/cortex_m/passes/quantized_clamp_activation_pass.py index 429d1ebb445..59f8f1cadb1 100644 --- a/backends/cortex_m/passes/quantized_clamp_activation_pass.py +++ b/backends/cortex_m/passes/quantized_clamp_activation_pass.py @@ -7,11 +7,14 @@ from typing import Any import torch -from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( get_output_qparams, ) -from executorch.backends.cortex_m.utils import get_activation_bounds, quantize_val +from executorch.backends.cortex_m.utils import ( + get_activation_bounds, + lowered_output_dtype, + quantize_val, +) from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass from torch.fx import GraphModule, Node @@ -71,11 +74,7 @@ def _is_quantized_int8_activation(self, node: Node) -> bool: input_node = node.args[0] if len(node.args) > 0 else None if not isinstance(input_node, Node): return False - try: - tensor = get_first_fake_tensor(input_node) - except Exception: - return False - if tensor is None or tensor.dtype != torch.int8: + if lowered_output_dtype(input_node) != torch.int8: return False try: diff --git a/backends/cortex_m/test/test_explicit_layout_pipeline.py b/backends/cortex_m/test/test_explicit_layout_pipeline.py index 30528fff9ed..c956f6945e9 100644 --- a/backends/cortex_m/test/test_explicit_layout_pipeline.py +++ b/backends/cortex_m/test/test_explicit_layout_pipeline.py @@ -20,6 +20,7 @@ CortexMRunPasses, CortexMTester, ) +from executorch.backends.cortex_m.utils import lowered_output_dtype from executorch.backends.test.harness.stages import Quantize, StageType from executorch.backends.transforms.remove_unused_constants_pass import ( RemoveUnusedConstantsPass, @@ -67,6 +68,69 @@ def forward(self, x): ) +class LutActivationThenConv1d(torch.nn.Module): + """A LUT activation feeding a convolution. + + tanh lowers to ``cortex_m::quantized_activation``, but only in + ``AtenToCortexMPass``, which runs after the layout validator. Until then + it keeps an fp32 fake tensor and passes that dtype on to the convolution, + so validating on ``val`` rejects this graph even though it lowers. + An activation that ``ActivationFusionPass`` folds into the convolution + clamp, such as relu, never exposes the gap. + """ + + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv1d(8, 8, 1) + + def forward(self, x): + return self.conv(torch.tanh(x)) + + +class LutActivationViewThenConv1d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv1d(8, 8, 1) + + def forward(self, x): + return self.conv(torch.tanh(x).view(1, 8, 16)) + + +class LutActivationReluThenConv1d(torch.nn.Module): + """A relu between a LUT activation and a convolution. + + Nothing fuses this relu. ``ActivationFusionPass`` only folds an + activation that follows a convolution, so it reaches + ``QuantizedClampActivationPass``, which has to quantize its bound + against the sigmoid's qparams. Left alone on int8 codes, ``aten.relu`` + clamps at 0 rather than at the zero point and erases everything below + it -- and it is meant to be a no-op here, sigmoid being positive. + """ + + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv1d(8, 8, 1) + + def forward(self, x): + return self.conv(torch.relu(torch.sigmoid(x))) + + +class LutActivationClampThenConv2d(torch.nn.Module): + """A two-sided clamp between a LUT activation and a convolution. + + The same gap as ``LutActivationReluThenConv1d``, with bounds of its + own: unquantized, ``aten.clamp`` holds fp32 limits against int8 codes + and keeps an fp32 value on a tensor the convolution reads as int8. + """ + + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 1) + + def forward(self, x): + return self.conv(torch.clamp(torch.tanh(x), -0.5, 0.5)) + + class NHWCPaddedConv(torch.nn.Module): def __init__(self, channels, shared_pad): super().__init__() @@ -84,6 +148,12 @@ def _count(exported_program, target) -> int: return sum(node.target == target for node in exported_program.graph.nodes) +def _only(exported_program, target) -> Node: + nodes = [node for node in exported_program.graph.nodes if node.target == target] + assert len(nodes) == 1, f"expected one {target}, found {len(nodes)}" + return nodes[0] + + def _run_explicit_layout_pass_manager(tester: CortexMTester) -> CortexMTester: target_config = CortexMTargetConfig(cpu=CortexM.M55) tester.run_passes( @@ -372,6 +442,109 @@ def test_implementation_transpose_conv2d_strided_pointwise(hardtanh): tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) +def test_explicit_layout_accepts_conv_fed_by_lut_activation(): + """A convolution behind a LUT activation lowers and keeps its numerics.""" + torch.manual_seed(0) + inputs = (torch.randn(1, 8, 16),) + tester = _run_explicit_layout_passes( + CortexMTester(LutActivationThenConv1d().eval(), inputs) + ) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + assert _count(program, exir_ops.edge.aten.convolution.default) == 0 + assert _count(program, exir_ops.edge.cortex_m.quantized_activation.default) == 1 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize( + "model,inputs,bounds", + [ + (LutActivationReluThenConv1d, (torch.randn(1, 8, 16),), (int, type(None))), + (LutActivationClampThenConv2d, (torch.randn(1, 4, 8, 8),), (int, int)), + ], +) +def test_clamp_between_lut_activation_and_conv_is_quantized(model, inputs, bounds): + """A clamp-like activation inside an int8 chain runs on int8 terms. + + The convolution reaches a CMSIS-NN kernel either way. What the fake + tensor costs is the clamp between it and the LUT activation: left in + float, its bounds mean nothing against int8 codes and the tensor the + convolution reads is still declared fp32. + """ + torch.manual_seed(0) + tester = _run_explicit_layout_passes(CortexMTester(model().eval(), inputs)) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + assert _count(program, exir_ops.edge.cortex_m.quantized_activation.default) == 1 + clamp = _only(program, exir_ops.edge.aten.clamp.default) + assert tuple(type(bound) for bound in clamp.args[1:]) == bounds + assert clamp.meta["val"].dtype == torch.int8 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize( + "model,inputs,bounds", + [ + (LutActivationReluThenConv1d, (torch.randn(1, 8, 16),), (int, type(None))), + (LutActivationClampThenConv2d, (torch.randn(1, 4, 8, 8),), (int, int)), + ], +) +def test_clamp_between_lut_activation_and_conv_is_quantized_on_legacy_passes( + model, inputs, bounds +): + """The same clamp on the default pass list, which has no validator. + + Here the clamp is the only thing standing between the LUT activation + and the rest of the graph: nothing rejects an unquantized one, so it + reaches the program. `aten.relu` left alone cuts at code 0 while the + zero point is -128, and `aten.clamp` left on float bounds promotes + the tensor back to fp32 in front of a dequantize that reads int8. + """ + torch.manual_seed(0) + tester = CortexMTester(model().eval(), inputs) + tester.quantize().export().to_edge().run_passes() + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_activation.default) == 1 + assert _count(program, exir_ops.edge.aten.relu.default) == 0 + clamp = _only(program, exir_ops.edge.aten.clamp.default) + assert tuple(type(bound) for bound in clamp.args[1:]) == bounds + assert clamp.meta["val"].dtype == torch.int8 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +def test_explicit_layout_accepts_conv_fed_by_view_of_lut_activation(): + """A dtype-agnostic op between the LUT activation and the conv stays int8.""" + torch.manual_seed(0) + inputs = (torch.randn(1, 128),) + tester = _run_explicit_layout_passes( + CortexMTester(LutActivationViewThenConv1d().eval(), inputs) + ) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + assert _count(program, exir_ops.edge.cortex_m.quantized_activation.default) == 1 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +def test_lowered_output_dtype_reads_first_output_of_multi_output_node(): + node = torch.fx.Graph().call_function( + exir_ops.edge.aten.max_pool2d_with_indices.default + ) + node.meta["val"] = ( + torch.empty(1, dtype=torch.int8), + torch.empty(1, dtype=torch.int64), + ) + assert lowered_output_dtype(node) == torch.int8 + + +def test_lowered_output_dtype_declines_a_non_node_argument(): + """Substitutions read it straight off ``node.args``, which holds literals.""" + assert lowered_output_dtype(5.0) is None + + def test_explicit_layout_rejects_unsupported_spatial_operator(): tester = CortexMTester(UnsupportedAvgPool(), (torch.randn(1, 3, 8, 8),)) diff --git a/backends/cortex_m/utils.py b/backends/cortex_m/utils.py index b8d9134c0a4..3dbb5c4dca2 100644 --- a/backends/cortex_m/utils.py +++ b/backends/cortex_m/utils.py @@ -14,6 +14,7 @@ from executorch.exir.dialects.edge._ops import EdgeOpOverload from torch.fx import Node +from torch.fx.node import Argument # L-shift value used in CMSIS-NN for int8 operations SHIFT_INT8 = 20 @@ -198,6 +199,25 @@ def skips_quantized_max_pool2d(node: Node) -> bool: ) +def lowered_output_dtype(node: Argument) -> torch.dtype | None: + """The dtype of this node's first output on device. + + Folded output qparams win over the fake tensor, which stays fp32 on + ops such as tanh and softmax until AtenToCortexMPass rewrites them. + An argument that is not a node, or carries neither, reports None, so + a caller gating a substitution on int8 declines it rather than raising. + """ + if not isinstance(node, Node): + return None + qparams = node.meta.get("output_qparams", {}).get(0) + if qparams is not None: + return qparams.dtype + value = node.meta.get("val") + if isinstance(value, (tuple, list)): + value = value[0] if value else None + return value.dtype if isinstance(value, torch.Tensor) else None + + def is_qualified_int8_node(args) -> bool: try: if len(args) < 6: From eb39e44b4aae410fa7c328399a7d0176fe7794a6 Mon Sep 17 00:00:00 2001 From: Prashant Rawat Date: Fri, 23 Oct 2026 12:35:34 -0700 Subject: [PATCH 2/2] Substitute on qparams, not the input fake tensor (#24121) Summary: `_get_permute_replacement`, `_get_minimum_replacement` and `_get_maximum_replacement` choose their kernel from the input fake tensor. That tensor is stale inside `AtenToCortexMPass`: `AtenToDialectPass` gives each new `cortex_m` node the replaced node's meta dict, so an input already rewritten from `aten.tanh` still reads fp32 until the retrace at the end of the pass. Read `lowered_output_dtype` instead, as the layout validator and clamp pass now do. Rewritten nodes keep their folded `output_qparams`, so it reports int8 for them; a `cortex_m` dequantize has none and still reports float. The cost was coverage, not numerics. `Conv1d(tanh(x))` lowered its input-side layout move to a portable `aten::permute_copy` and its output-side move to `cortex_m::transpose`, and `minimum(tanh(x), y)` / `maximum(tanh(x), y)` stayed on portable kernels. That breaks the explicit-layout guarantee that every layout move is a CMSIS transpose, and a selective build that registers only the expected kernels fails to load. Differential Revision: D127046193 --- .../cortex_m/passes/aten_to_cortex_m_pass.py | 13 ++-- .../test/test_explicit_layout_pipeline.py | 65 +++++++++++++++++++ 2 files changed, 73 insertions(+), 5 deletions(-) diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py index b9339b838b8..7ae1ecb174a 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -24,6 +24,7 @@ from executorch.backends.cortex_m.utils import ( build_activation_lut, is_foldable_alpha, + lowered_output_dtype, quantize_multiplier_aot, quantize_val, SHIFT_INT8, @@ -1209,8 +1210,7 @@ def _get_minimum_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: del dialect_pass - input_tensor = _get_input_tensor_data(node) - if input_tensor.dtype not in (torch.int8, torch.int32): + if lowered_output_dtype(node.args[0]) not in (torch.int8, torch.int32): return None return DialectNodeSpec(exir_ops.edge.cortex_m.minimum.default, node.args) @@ -1220,8 +1220,7 @@ def _get_maximum_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: del dialect_pass - input_tensor = _get_input_tensor_data(node) - if input_tensor.dtype != torch.int8: + if lowered_output_dtype(node.args[0]) != torch.int8: return None return DialectNodeSpec(exir_ops.edge.cortex_m.maximum.default, node.args) @@ -1241,8 +1240,12 @@ def _get_permute_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: del dialect_pass + if lowered_output_dtype(node.args[0]) != torch.int8: + return None + # Shape, unlike dtype, survives AtenToDialectPass: a substitution keeps the + # replaced node's meta and computes the same function over it. input_tensor = _get_input_tensor_data(node) - if input_tensor.dtype != torch.int8 or not 1 <= input_tensor.dim() <= 4: + if not 1 <= input_tensor.dim() <= 4: return None return _transpose_spec(node, input_tensor) diff --git a/backends/cortex_m/test/test_explicit_layout_pipeline.py b/backends/cortex_m/test/test_explicit_layout_pipeline.py index c956f6945e9..dbdef44c31c 100644 --- a/backends/cortex_m/test/test_explicit_layout_pipeline.py +++ b/backends/cortex_m/test/test_explicit_layout_pipeline.py @@ -131,6 +131,24 @@ def forward(self, x): return self.conv(torch.clamp(torch.tanh(x), -0.5, 0.5)) +class LutActivationThenElementwise(torch.nn.Module): + """An elementwise binary op fed by a LUT activation. + + ``minimum`` and ``maximum`` pick their kernel from the same fake + tensor the input-side layout permute does, and there is no layout + copy here to involve the validator -- the convolution is what drags + one in. So this is the gap on its own: a graph that is int8 start to + finish and still leaves one op on the portable kernel. + """ + + def __init__(self, fn) -> None: + super().__init__() + self.fn = fn + + def forward(self, x, y): + return self.fn(torch.tanh(x), y) + + class NHWCPaddedConv(torch.nn.Module): def __init__(self, channels, shared_pad): super().__init__() @@ -457,6 +475,53 @@ def test_explicit_layout_accepts_conv_fed_by_lut_activation(): tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) +def test_layout_permute_behind_lut_activation_lowers_to_transpose(): + """Every layout move in an int8 chain reaches a CMSIS-NN transpose. + + The convolution has a layout move on each side. The one on its output + is synthesized as ``channels_last.permute_copy``, whose substitution + reads no dtype; the one on its input is an ``aten.permute_copy``, and + its substitution read the same fake tensor the layout validator used + to. Behind a LUT activation that tensor is fp32, so the input-side + move was left on the portable kernel while the output-side move + converted. Numerics are unaffected -- portable ``permute_copy`` + computes the right thing on int8 -- so what this costs is the + invariant: a selective build registering only the kernels the graph + is supposed to need fails to load. + """ + torch.manual_seed(0) + inputs = (torch.randn(1, 8, 16),) + tester = _run_explicit_layout_passes( + CortexMTester(LutActivationThenConv1d().eval(), inputs) + ) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + assert _count(program, exir_ops.edge.aten.permute_copy.default) == 0 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize( + "fn,cortex_m_op,aten_op", + [ + (torch.minimum, exir_ops.edge.cortex_m.minimum, exir_ops.edge.aten.minimum), + (torch.maximum, exir_ops.edge.cortex_m.maximum, exir_ops.edge.aten.maximum), + ], +) +def test_elementwise_behind_lut_activation_lowers_to_cortex_m(fn, cortex_m_op, aten_op): + """An elementwise binary op behind a LUT activation reaches CMSIS-NN.""" + torch.manual_seed(0) + inputs = (torch.randn(1, 4, 8), torch.randn(1, 4, 8)) + tester = _run_explicit_layout_passes( + CortexMTester(LutActivationThenElementwise(fn).eval(), inputs) + ) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, cortex_m_op.default) == 1 + assert _count(program, aten_op.default) == 0 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + @pytest.mark.parametrize( "model,inputs,bounds", [