Arm backend: Fold scalar mul into convolution by oscarandersson8218 · Pull Request #20838 · pytorch/executorch

oscarandersson8218 · 2026-07-10T13:27:01Z

Fold constant scalar or channel-wise multiplication after convolution into the convolution weights and bias. This removes post-conv mul ops when the scale is a scalar or is broadcastable only across output channels.

Handle both pre-edge ATen conv2d graphs and edge aten.convolution graphs, and skip folding when the convolution output has multiple users or the scale does not represent channel scaling.

cc @digantdesai @freddan80 @per @zingo @mansnils @Sebastian-Larsson @robell @rascani

Fold constant scalar or channel-wise multiplication after convolution into the convolution weights and bias. This removes post-conv mul ops when the scale is a scalar or is broadcastable only across output channels. Handle both pre-edge ATen conv2d graphs and edge aten.convolution graphs, and skip folding when the convolution output has multiple users or the scale does not represent channel scaling. Signed-off-by: Oscar Andersson <oscar.andersson@arm.com> Change-Id: Ic630bbe69473821aeda66e9dc844d33669586a1d

pytorch-bot · 2026-07-10T13:27:06Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20838

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit b8519c3 with merge base 3d9936e ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

…hazard (#21152) Summary: FoldScalarMulIntoConv #20838 rewrites conv(x) * scale to 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: activation qparams are calibrated and frozen on the pre-fold conv output, which has a uniform range. Folding a per-channel scale into the weights moves the scaling inside the convolution, so the real conv output spans the scaled dynamic range while it is still requantized with the stale pre-fold scale. With a wide-dynamic-range scale this clips / loses precision and the quantized result diverges from the reference where the mul is kept as a separate, well-scaled op. The same requant shift degrades models that carry a quantized value across steps, where it compounds. Add ConvMulWideDynamicRangeScale and an xfail test on the bit-accurate TOSA INT reference documenting the hazard. It will start passing (flagging that the xfail should be removed) once the fold is made quantization-aware or is skipped for the affected models. Differential Revision: D113310822

…hazard (#21152) Summary: FoldScalarMulIntoConv #20838 rewrites conv(x) * scale to 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: activation qparams are calibrated and frozen on the pre-fold conv output, which has a uniform range. Folding a per-channel scale into the weights moves the scaling inside the convolution, so the real conv output spans the scaled dynamic range while it is still requantized with the stale pre-fold scale. With a wide-dynamic-range scale this clips / loses precision and the quantized result diverges from the reference where the mul is kept as a separate, well-scaled op. The same requant shift degrades models that carry a quantized value across steps, where it compounds. Add ConvMulWideDynamicRangeScale and an xfail test on the bit-accurate TOSA INT reference documenting the hazard. It will start passing (flagging that the xfail should be removed) once the fold is made quantization-aware or is skipped for the affected models. Reviewed By: rascani Differential Revision: D113310822

…d state (#21234) Summary: FoldScalarMulIntoConv #20838 rewrites conv(x) * scale to a convolution with scaled weights. This is exact in floating point but moves the scaling inside the convolution, changing the conv output's quantization (its activation range and requantization). That is harmless for a plain feed-forward output, but if 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) -- the shifted requantization drifts the persisted quantized state and the error compounds over the sequence. Make the fold state-aware: skip it when the conv output transitively reaches a BUFFER_MUTATION / USER_INPUT_MUTATION output in the exported program's graph signature. Plain feed-forward convolutions are unaffected and still fold. Add StatefulConvMul (conv output * scale written into a mutable buffer) and assert the scalar mul survives. Differential Revision: D113310883

…hazard (#21152) Summary: FoldScalarMulIntoConv #20838 rewrites conv(x) * scale to 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: activation qparams are calibrated and frozen on the pre-fold conv output, which has a uniform range. Folding a per-channel scale into the weights moves the scaling inside the convolution, so the real conv output spans the scaled dynamic range while it is still requantized with the stale pre-fold scale. With a wide-dynamic-range scale this clips / loses precision and the quantized result diverges from the reference where the mul is kept as a separate, well-scaled op. The same requant shift degrades models that carry a quantized value across steps, where it compounds. Add ConvMulWideDynamicRangeScale and an xfail test on the bit-accurate TOSA INT reference documenting the hazard. It will start passing (flagging that the xfail should be removed) once the fold is made quantization-aware or is skipped for the affected models. Reviewed By: rascani Differential Revision: D113310822

…d state (#21234) Summary: FoldScalarMulIntoConv #20838 rewrites conv(x) * scale to a convolution with scaled weights. This is exact in floating point but moves the scaling inside the convolution, changing the conv output's quantization (its activation range and requantization). That is harmless for a plain feed-forward output, but if 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) -- the shifted requantization drifts the persisted quantized state and the error compounds over the sequence. Make the fold state-aware: skip it when the conv output transitively reaches a BUFFER_MUTATION / USER_INPUT_MUTATION output in the exported program's graph signature. Plain feed-forward convolutions are unaffected and still fold. Add StatefulConvMul (conv output * scale written into a mutable buffer) and assert the scalar mul survives. Differential Revision: D113310883

…hazard (#21152) Summary: FoldScalarMulIntoConv #20838 rewrites conv(x) * scale to 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: activation qparams are calibrated and frozen on the pre-fold conv output, which has a uniform range. Folding a per-channel scale into the weights moves the scaling inside the convolution, so the real conv output spans the scaled dynamic range while it is still requantized with the stale pre-fold scale. With a wide-dynamic-range scale this clips / loses precision and the quantized result diverges from the reference where the mul is kept as a separate, well-scaled op. The same requant shift degrades models that carry a quantized value across steps, where it compounds. Add ConvMulWideDynamicRangeScale and an xfail test on the bit-accurate TOSA INT reference documenting the hazard. It will start passing (flagging that the xfail should be removed) once the fold is made quantization-aware or is skipped for the affected models. Reviewed By: rascani Differential Revision: D113310822

…d state (#21234) Summary: FoldScalarMulIntoConv #20838 rewrites conv(x) * scale to a convolution with scaled weights. This is exact in floating point but moves the scaling inside the convolution, changing the conv output's quantization (its activation range and requantization). That is harmless for a plain feed-forward output, but if 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) -- the shifted requantization drifts the persisted quantized state and the error compounds over the sequence. Make the fold state-aware: skip it when the conv output transitively reaches a BUFFER_MUTATION / USER_INPUT_MUTATION output in the exported program's graph signature. Plain feed-forward convolutions are unaffected and still fold. Add StatefulConvMul (conv output * scale written into a mutable buffer) and assert the scalar mul survives. Differential Revision: D113310883

…hazard (#21152) Summary: FoldScalarMulIntoConv (#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

…d state (#21234) Summary: Follow-up to D113310822, which documents (as a strict xfail) that FoldScalarMulIntoConv (#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