From e6ccc6168e4c80ffb0538c3e655005b47a9d8436 Mon Sep 17 00:00:00 2001 From: Eli Amesefe Date: Wed, 22 Jul 2026 21:02:53 -0700 Subject: [PATCH] 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