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