From 72d52ff25e7168cfbb2e3f74ea803e450df14a85 Mon Sep 17 00:00:00 2001 From: Eli Amesefe Date: Wed, 22 Jul 2026 21:04:36 -0700 Subject: [PATCH 1/2] Arm backend: xfail test for FoldScalarMulIntoConv quantized-accuracy hazard (#21152) Summary: FoldScalarMulIntoConv (https://github.com/pytorch/executorch/pull/20838) rewrites conv(x) * scale into a convolution with scaled weights. This is exact in floating point (already covered by the existing tests in this file), but it is not safe under quantization when the scaled conv output feeds a carried/stateful quantization boundary -- a value written back into a mutable buffer / input that persists across invocations (recurrent state). Activation qparams are calibrated and frozen on the pre-fold conv output; folding the scale into the weights moves the scaling inside the convolution, so the conv output is requantized with the stale pre-fold scale. That shifted requantization drifts the persisted quantized state and the error compounds over the sequence. Add StatefulConvMul (conv output * scale written into a mutable buffer) and an xfail unit test (test_does_not_fold_when_output_feeds_stateful_buffer) documenting the hazard: the fold currently fires and drops the mul across the buffer boundary. The test is marked xfail (strict); the state-aware guard in D113310883 makes it pass, and the xfail is removed there. Reviewed By: rascani Differential Revision: D113310822 --- .../test_fold_scalar_mul_into_conv_pass.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/backends/arm/test/passes/test_fold_scalar_mul_into_conv_pass.py b/backends/arm/test/passes/test_fold_scalar_mul_into_conv_pass.py index 2a30e38b520..511e1404782 100644 --- a/backends/arm/test/passes/test_fold_scalar_mul_into_conv_pass.py +++ b/backends/arm/test/passes/test_fold_scalar_mul_into_conv_pass.py @@ -6,6 +6,7 @@ from collections import Counter from typing import Tuple +import pytest import torch from executorch.backends.arm._passes.arm_pass_manager import ( _ExportedProgramGraphPassAdapter, @@ -194,3 +195,53 @@ def test_does_not_fold_when_conv_has_multiple_users() -> None: mul_count = mul_tensor_count + mul_scalar_count assert mul_count == 1 torch.testing.assert_close(module(*inputs), transformed.module()(*inputs)) + + +class StatefulConvMul(torch.nn.Module): + """Conv output * scale is written into a mutable buffer (carried state). + + ``state`` is read at the start and written at the end of ``forward``, so + the (scaled) conv output crosses a carried/stateful quantization boundary. + Folding the scale into the conv weights would shift that requantization, so + the fold must be skipped here. + + """ + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 3, 3, padding=1) + self.register_buffer("state", torch.zeros(1, 3, 8, 8)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = self.conv(x + self.state) * 0.25 # type: ignore[operator] + self.state.copy_(y) # type: ignore[operator] + return y + + +# FoldScalarMulIntoConv is not yet carried-state-aware at this point in the stack, +# so the test below FAILS here (the fold happens and removes the mul). It is marked +# xfail. +@pytest.mark.xfail( + reason="FoldScalarMulIntoConv carried-state guard is added in the following diff", + strict=True, +) +def test_does_not_fold_when_output_feeds_stateful_buffer() -> None: + """The scale must survive when the conv output crosses a carried-state + boundary. + + Folding would move the scale into the conv weights and shift the conv output + requantization, drifting the quantized value persisted in the mutable + buffer. + + """ + module = StatefulConvMul().eval() + inputs = (torch.randn(1, 3, 8, 8),) + + transformed = _run_pass(module, inputs) + counts = _op_counts(transformed) + + assert counts[exir_ops.edge.aten.convolution.default] == 1 + mul_count = ( + counts[exir_ops.edge.aten.mul.Tensor] + counts[exir_ops.edge.aten.mul.Scalar] + ) + assert mul_count == 1 From b4c979384051e4d4e10ecb175ace30cd6d2e52a3 Mon Sep 17 00:00:00 2001 From: Eli Amesefe Date: Wed, 22 Jul 2026 21:04:36 -0700 Subject: [PATCH 2/2] Arm backend: skip FoldScalarMulIntoConv when conv output feeds carried state (#21234) Summary: Follow-up to D113310822, which documents (as a strict xfail) that FoldScalarMulIntoConv (https://github.com/pytorch/executorch/pull/20838) is unsafe when the scaled conv output feeds a carried/stateful quantization boundary: the fold moves the scaling inside the convolution, shifting the conv output requant, which drifts a quantized value persisted across invocations (recurrent state) and compounds over the sequence. Make the fold state-aware and skip it in that case. Two carried-state signals are handled: - Structural: the conv output transitively reaches a BUFFER_MUTATION / USER_INPUT_MUTATION output in the exported program graph signature (mutable buffer / input mutation). - Explicit: a new mark_carried_quant_state / is_carried_quant_state marker (arm_pass_utils) a model can set on the (conv * scale) output. Functional recurrent state (state_in arg + state_out return, shared quantization) has no structural signal before quantization, so it needs an explicit marker. Plain feed-forward convolutions are unaffected and still fold. Un-xfails test_does_not_fold_when_output_feeds_stateful_buffer (the buffer-mutation case introduced in D113310822) and adds test_does_not_fold_when_output_marked_carried_state (the explicit-marker case). Differential Revision: D113310883 --- backends/arm/_passes/arm_pass_utils.py | 21 ++++++ .../_passes/fold_scalar_mul_into_conv_pass.py | 75 ++++++++++++++++++- .../test_fold_scalar_mul_into_conv_pass.py | 52 +++++++++++-- 3 files changed, 139 insertions(+), 9 deletions(-) diff --git a/backends/arm/_passes/arm_pass_utils.py b/backends/arm/_passes/arm_pass_utils.py index 85d4ed514ea..d8fac246ca6 100644 --- a/backends/arm/_passes/arm_pass_utils.py +++ b/backends/arm/_passes/arm_pass_utils.py @@ -40,6 +40,27 @@ _Dim = int | torch.SymInt +# Meta key marking a value as *carried quantized state*: a tensor whose +# quantization is shared across invocations (e.g. recurrent state fed back frame +# to frame). Such a value's quantization must stay stable, so requantization- +# perturbing rewrites (e.g. FoldScalarMulIntoConvPass) must not fold into a +# convolution whose output reaches it. The marker is model-supplied because +# functional carried state has no structural signal before quantization runs. +CARRIED_QUANT_STATE_META_KEY = "carried_quant_state" + + +def mark_carried_quant_state(node: torch.fx.Node) -> None: + """Mark ``node``'s output as carried quantized state (see the meta key + doc). + """ + node.meta[CARRIED_QUANT_STATE_META_KEY] = True + + +def is_carried_quant_state(node: torch.fx.Node) -> bool: + """True if ``node`` was marked as carried quantized state.""" + return bool(node.meta.get(CARRIED_QUANT_STATE_META_KEY, False)) + + def is_submodule_node(node: torch.fx.Node): if node.op not in ("get_attr", "placeholder"): return False diff --git a/backends/arm/_passes/fold_scalar_mul_into_conv_pass.py b/backends/arm/_passes/fold_scalar_mul_into_conv_pass.py index 56df57d3b7a..51435ac338b 100644 --- a/backends/arm/_passes/fold_scalar_mul_into_conv_pass.py +++ b/backends/arm/_passes/fold_scalar_mul_into_conv_pass.py @@ -10,6 +10,7 @@ from executorch.backends.arm._passes.arm_pass_utils import ( get_constant_placeholder_kind, get_param_tensor, + is_carried_quant_state, ) from executorch.backends.transforms.utils import ( create_constant_placeholder, @@ -17,7 +18,7 @@ ) from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult -from torch.export.graph_signature import InputKind +from torch.export.graph_signature import InputKind, OutputKind from torch.fx import GraphModule, Node from torch.fx.node import Argument @@ -33,6 +34,23 @@ class FoldScalarMulIntoConvPass(ArmPass): broadcastable over the convolution output with only the channel dimension non-unit. + The rewrite is exact in floating point, but it moves the scaling *inside* + the convolution and therefore changes the convolution output's quantization + (its activation range, and hence its requantization). That is harmless for a + plain feed-forward output, but if the (scaled) convolution output feeds a + *carried/stateful quantization boundary* -- a value whose quantization is + shared across invocations (recurrent state) -- the shifted requantization + drifts the carried state and the error compounds over the sequence. To stay + safe, the fold is skipped when the convolution output transitively reaches + such a boundary, identified as either: + + * a value written into a mutable buffer / input (``BUFFER_MUTATION`` / + ``USER_INPUT_MUTATION`` in the graph signature), which is structurally + visible; or + * a value the model marked with ``mark_carried_quant_state`` -- needed for + *functional* carried state (state passed in and returned as plain I/O), + which has no structural signal before quantization runs. + """ _passes_required_after: Set[Type[ExportPass]] = set() @@ -64,6 +82,14 @@ def call(self, graph_module: GraphModule) -> PassResult: graph = graph_module.graph modified = False constant_placeholders_to_delete: set[Node] = set() + # Carried/stateful quant boundaries: values written to a mutable buffer / + # input (structurally visible), plus any value the model explicitly marked + # as carried quantized state (functional recurrent state, which has no + # structural signal before quantization). + state_boundary_names = self._state_mutation_value_names() + state_boundary_names |= { + node.name for node in graph.nodes if is_carried_quant_state(node) + } for node in list(graph.nodes): if not self.allowed_to_transform(node.meta): @@ -72,6 +98,12 @@ def call(self, graph_module: GraphModule) -> PassResult: if match is None: continue + # Do not fold when the (scaled) conv output feeds a carried/stateful + # quant boundary: moving the scale into the weights shifts the conv + # output requantization, which drifts persisted quantized state. + if self._feeds_state_boundary(node, state_boundary_names): + continue + conv_node, scale, scale_node = match folded_placeholders = self._fold_scale_into_conv( graph, conv_node, node, scale @@ -95,6 +127,47 @@ def call(self, graph_module: GraphModule) -> PassResult: return PassResult(graph_module, modified) + def _state_mutation_value_names(self) -> Set[str]: + """Names of graph output values that write a carried/stateful boundary. + + These are the values bound to ``BUFFER_MUTATION`` / ``USER_INPUT_MUTATION`` + outputs in the exported program's graph signature -- i.e. mutable buffers + or inputs whose value persists across invocations (recurrent state). When + the graph signature is unavailable (``exported_program is None``) no state + boundary can be identified and the set is empty. + + """ + if self.exported_program is None: + return set() + names: Set[str] = set() + for spec in self.exported_program.graph_signature.output_specs: + if spec.kind in ( + OutputKind.BUFFER_MUTATION, + OutputKind.USER_INPUT_MUTATION, + ): + name = getattr(spec.arg, "name", None) + if name is not None: + names.add(name) + return names + + def _feeds_state_boundary(self, node: Node, state_names: Set[str]) -> bool: + """True if ``node``'s output transitively reaches a state-mutation + value. + """ + if not state_names: + return False + seen: Set[Node] = set() + stack = [node] + while stack: + current = stack.pop() + if current in seen: + continue + seen.add(current) + if current.name in state_names: + return True + stack.extend(current.users) + return False + def _match_mul_after_conv(self, node: Node) -> Optional[ConvScaleMatch]: if node.op != "call_function": return None diff --git a/backends/arm/test/passes/test_fold_scalar_mul_into_conv_pass.py b/backends/arm/test/passes/test_fold_scalar_mul_into_conv_pass.py index 511e1404782..fc93a0990bb 100644 --- a/backends/arm/test/passes/test_fold_scalar_mul_into_conv_pass.py +++ b/backends/arm/test/passes/test_fold_scalar_mul_into_conv_pass.py @@ -6,11 +6,11 @@ from collections import Counter from typing import Tuple -import pytest import torch from executorch.backends.arm._passes.arm_pass_manager import ( _ExportedProgramGraphPassAdapter, ) +from executorch.backends.arm._passes.arm_pass_utils import mark_carried_quant_state from executorch.backends.arm._passes.fold_scalar_mul_into_conv_pass import ( FoldScalarMulIntoConvPass, ) @@ -218,13 +218,49 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y -# FoldScalarMulIntoConv is not yet carried-state-aware at this point in the stack, -# so the test below FAILS here (the fold happens and removes the mul). It is marked -# xfail. -@pytest.mark.xfail( - reason="FoldScalarMulIntoConv carried-state guard is added in the following diff", - strict=True, -) +def test_does_not_fold_when_output_marked_carried_state() -> None: + """A model-supplied carried-state marker must prevent the fold. + + ConvMulScalar folds by default; marking the (conv * scalar) output as + carried quantized state -- the mechanism for functional recurrent state, + which has no structural signal before quantization -- must make the pass + skip it. + + """ + module = ConvMulScalar().eval() + inputs = (torch.randn(2, 3, 8, 8),) + edge_program = to_edge( + torch.export.export(module, inputs, strict=True), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + edge_exported_program = edge_program.exported_program() + + marked = False + for node in edge_exported_program.graph_module.graph.nodes: + if node.op == "call_function" and node.target in ( + exir_ops.edge.aten.mul.Tensor, + exir_ops.edge.aten.mul.Scalar, + ): + mark_carried_quant_state(node) + marked = True + assert marked, "test setup: expected a mul node to mark" + + transformed = edge_program.transform( + [ + _ExportedProgramGraphPassAdapter( + FoldScalarMulIntoConvPass(edge_exported_program) + ) + ] + ).exported_program() + counts = _op_counts(transformed) + + assert counts[exir_ops.edge.aten.convolution.default] == 1 + mul_count = ( + counts[exir_ops.edge.aten.mul.Tensor] + counts[exir_ops.edge.aten.mul.Scalar] + ) + assert mul_count == 1 # fold skipped because the output is marked carried state + + def test_does_not_fold_when_output_feeds_stateful_buffer() -> None: """The scale must survive when the conv output crosses a carried-state boundary.