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 2a30e38b520..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 @@ -10,6 +10,7 @@ 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, ) @@ -194,3 +195,89 @@ 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 + + +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. + + 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