Arm backend: skip FoldScalarMulIntoConv when conv output feeds carried state (#21234)

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

backends/arm/_passes/fold_scalar_mul_into_conv_pass.py

@@ -10,14 +10,15 @@
 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,
     delete_constant_placeholder,
 )
 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

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