From e6688936a2d77f2c07d30528fd20acfaa240e7b2 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 14 Aug 2026 19:14:08 -0700 Subject: [PATCH] Cortex-M: lower subtraction, and stop miscompiling add with alpha quantized_add carries a multiplier per operand, so subtraction is the same kernel with the second one negated. quantize_multiplier_aot always lands in [2^30, 2^31), so the negation cannot saturate; the first requantize stage rounds asymmetrically under negation, but the second absorbs it and the int8 output is identical across 400 scale and zero-point configurations. This exposed an existing bug. alpha scales the second operand, quantized_add has nowhere to put it, and nothing checked, so torch.add(x, y, alpha=2) computed x + y and returned wrong values silently. The quantizer now declines those nodes and they stay fp32. It has to be the quantizer: by the time the lowering runs FoldAndAnnotateQParamsPass has removed the dq nodes, so declining there would leave an fp32 add over raw int8. The lowering raises instead. Subtraction also makes an existing overflow reachable, where the kernel's int32 left shift wraps and saturates to the opposite rail. No multiplier split avoids it -- shrinking max_scale_2x buys a shift and costs the same headroom -- so the lowering refuses the configuration, which needs operands that nearly cancel. Two gaps left open: ActivationFusionPass does not fold relu into sub, and rsub with alpha is still wrong because Arm's pass drops alpha before annotation. Authored with Claude Code. --- .../cortex_m/passes/aten_to_cortex_m_pass.py | 41 +++++ .../cortex_m/quantizer/pattern_checkers.py | 5 +- .../cortex_m/quantizer/quantizer_support.py | 2 + backends/cortex_m/test/build_test_runner.sh | 1 + backends/cortex_m/test/ops/test_add.py | 37 +++- backends/cortex_m/test/ops/test_sub.py | 168 ++++++++++++++++++ 6 files changed, 245 insertions(+), 9 deletions(-) create mode 100644 backends/cortex_m/test/ops/test_sub.py diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py index 08484ce3b38..1215dd62eb1 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -901,6 +901,7 @@ def _get_dequantize_per_tensor_replacement( @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.add.Tensor) +@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.sub.Tensor) def _get_add_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: @@ -908,6 +909,17 @@ def _get_add_replacement( if not _has_qparams(node): return None + # CortexMAddMulCheck declines alpha, so reaching here means some other + # quantizer annotated the node -- by which point FoldAndAnnotateQParamsPass + # has removed the dq pair and returning None would leave an fp32 add over + # raw int8. + if node.kwargs.get("alpha", 1) != 1: + raise RuntimeError( + f"{node.target} carries alpha; quantized_add cannot express it." + ) + + is_sub = node.target is exir_ops.edge.aten.sub.Tensor + scale1 = node.meta["input_qparams"][0].scale zero_point1 = node.meta["input_qparams"][0].zp scale2 = node.meta["input_qparams"][1].scale @@ -918,10 +930,39 @@ def _get_add_replacement( max_scale_2x = 2 * max(scale1, scale2) input1_mult, input1_shift = quantize_multiplier_aot(scale1 / max_scale_2x) input2_mult, input2_shift = quantize_multiplier_aot(scale2 / max_scale_2x) + if is_sub: + # quantized_add carries a multiplier per operand, so subtraction is the + # same kernel with the second one negated. + input2_mult = -input2_mult output_mult, output_shift = quantize_multiplier_aot( max_scale_2x / (output_scale * (1 << SHIFT_INT8)) ) + # A positive output shift takes arm_nn_requantize's left-shift branch, which + # evaluates val * (1 << shift) in int32, and past a point the summed operands + # wrap there and saturate to the opposite rail. Re-splitting the scaling does + # not help -- the product below is invariant under the choice of + # max_scale_2x -- so refuse instead. + if output_shift > 0: + qparams1 = node.meta["input_qparams"][0] + qparams2 = node.meta["input_qparams"][1] + span1 = (qparams1.qmin - zero_point1, qparams1.qmax - zero_point1) + span2 = (qparams2.qmin - zero_point2, qparams2.qmax - zero_point2) + sign = -1 if is_sub else 1 + worst_case_sum = ( + max(abs(d1 * scale1 + sign * d2 * scale2) for d1 in span1 for d2 in span2) + * (1 << SHIFT_INT8) + / max_scale_2x + ) + if worst_case_sum * (2**output_shift) >= 2**31: + raise RuntimeError( + f"{node.target}: an output scale of {output_scale} against " + f"operand scales of {scale1} and {scale2} needs a " + "requantization the int32 kernel cannot hold. The operands " + "nearly cancel over the calibration set; calibrate on data " + "where they do not." + ) + activation_min = node.meta["output_qparams"][0].qmin activation_max = node.meta["output_qparams"][0].qmax diff --git a/backends/cortex_m/quantizer/pattern_checkers.py b/backends/cortex_m/quantizer/pattern_checkers.py index cc89715b537..039556f6fd7 100644 --- a/backends/cortex_m/quantizer/pattern_checkers.py +++ b/backends/cortex_m/quantizer/pattern_checkers.py @@ -29,9 +29,12 @@ class CortexMAddMulCheck(PatternCheck): @classmethod def check_pattern(cls, pattern): """ - Checks that the pattern does not perform unsupported broadcasting. + Checks that the pattern does not perform unsupported broadcasting, and + that add/sub carry no alpha, which quantized_add has nowhere to put. """ for node in pattern: + if node.kwargs.get("alpha", 1) != 1: + return False if len(node.all_input_nodes) == 2: t1 = get_first_fake_tensor(node.all_input_nodes[0]) t2 = get_first_fake_tensor(node.all_input_nodes[1]) diff --git a/backends/cortex_m/quantizer/quantizer_support.py b/backends/cortex_m/quantizer/quantizer_support.py index aaaf6414d06..6d949f09e87 100644 --- a/backends/cortex_m/quantizer/quantizer_support.py +++ b/backends/cortex_m/quantizer/quantizer_support.py @@ -26,6 +26,8 @@ (torch.ops.aten.add.Tensor, torch.ops.aten.clamp.default): CortexMAddMulCheck, (torch.ops.aten.add.Tensor, torch.ops.aten.clamp_.default): CortexMAddMulCheck, (torch.ops.aten.add_.Tensor,): CortexMAddMulCheck, + (torch.ops.aten.sub.Tensor,): CortexMAddMulCheck, + (torch.ops.aten.sub_.Tensor,): CortexMAddMulCheck, (torch.ops.aten.mul.Tensor,): CortexMAddMulCheck, (torch.ops.aten.mul_.Tensor,): CortexMAddMulCheck, (torch.ops.aten.hardswish.default,): CortexMAddMulCheck, # lowers to mul diff --git a/backends/cortex_m/test/build_test_runner.sh b/backends/cortex_m/test/build_test_runner.sh index 4d8502ec59e..a38c6d53256 100755 --- a/backends/cortex_m/test/build_test_runner.sh +++ b/backends/cortex_m/test/build_test_runner.sh @@ -42,6 +42,7 @@ join_by_comma() { ops_list=( aten::add.out + aten::sub.out aten::clamp.out aten::mul.out aten::convolution.out diff --git a/backends/cortex_m/test/ops/test_add.py b/backends/cortex_m/test/ops/test_add.py index 8e64ab4f132..a57389afba8 100644 --- a/backends/cortex_m/test/ops/test_add.py +++ b/backends/cortex_m/test/ops/test_add.py @@ -59,20 +59,36 @@ class CortexMTensorAdd(Model): } -class CortexMAlphaAdd(ModelAlpha): +class CortexMIntAlphaAdd(ModelAlpha): + """An integer alpha is the case that was silently miscompiled. + + ModelAlpha(0.5) below is caught anyway, by "alpha argument of type float + cannot be safely cast", so it never exercised the silent path. + + The boundary quant/dequant pairs are pinned so that a quantizer which + stopped annotating anything at all would not read as a successful decline. + """ + ops_before_transforms = { "executorch_exir_dialects_edge__ops_aten_add_Tensor": 1, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 3, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 3, } - ops_after_transforms = { - "executorch_exir_dialects_edge__ops_cortex_m_quantized_add_default": 1, - "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 2, - "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_quantized_add_default": 0, + "executorch_exir_dialects_edge__ops_aten_add_Tensor": 1, + "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 3, + "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 3, } +class CortexMAlphaAdd(ModelAlpha): + """A float alpha is declined by the same guard, and stays in fp32.""" + + ops_before_transforms = CortexMIntAlphaAdd.ops_before_transforms + ops_after_transforms = CortexMIntAlphaAdd.ops_after_transforms + + class CortexMAddReLU(torch.nn.Module): ops_before_transforms = { "executorch_exir_dialects_edge__ops_aten_add_Tensor": 1, @@ -186,6 +202,13 @@ def forward(self, x, y): ramp_tensor(-2, 2, (1, 8, 1, 1)), ), ), + "alpha_int": McuTestCase( + CortexMIntAlphaAdd(2), + ( + ramp_tensor(-10, 10, (4, 5)), + ramp_tensor(-20, 20, (4, 5)), + ), + ), "alpha": McuTestCase( CortexMAlphaAdd(0.5), ( @@ -224,9 +247,7 @@ def forward(self, x, y): } -xfails_implementation: dict[str, xfail_type] = { - "alpha": "Expecting kwargs for aten op IR to be empty - alpha arg not supported.", -} +xfails_implementation: dict[str, xfail_type] = {} xfails_dialect: dict[str, xfail_type] = xfails_implementation | { # Cortex-M quantizer will not quantize additions that require broadcasting # leading to the add op not being replaced by a cortex-m specific implementation diff --git a/backends/cortex_m/test/ops/test_sub.py b/backends/cortex_m/test/ops/test_sub.py new file mode 100644 index 00000000000..959d9486b9b --- /dev/null +++ b/backends/cortex_m/test/ops/test_sub.py @@ -0,0 +1,168 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + + +import pytest +import torch +from executorch.backends.arm.test.common import parametrize, xfail_type +from executorch.backends.cortex_m.test.tester import ( + CortexMTester, + McuTestCase, + ramp_tensor, +) + + +# Subtraction reuses quantized_add with the second operand's multiplier +# negated, so every lowered case here expects a quantized_add. Equal operands +# cannot see a dropped or misplaced negation, so the pairs are unlike +# everywhere except collapsed_output_scale, which is there for the overflow +# guard. +class CortexMTensorSub(torch.nn.Module): + ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_sub_Tensor": 1, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 3, + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 3, + } + ops_after_transforms = { + "executorch_exir_dialects_edge__ops_cortex_m_quantized_add_default": 1, + "executorch_exir_dialects_edge__ops_aten_sub_Tensor": 0, + } + + def forward(self, x, y): + return x - y + + +class CortexMInplaceSub(CortexMTensorSub): + def forward(self, x, y): + return x.sub_(y) + + +class CortexMAlphaSub(torch.nn.Module): + """alpha has nowhere to go in quantized_add, so this must stay in fp32. + + An integer alpha is the case that matters: a float one is rejected earlier + by the dtype cast, so it would fail even without the quantizer declining. + + The boundary quant/dequant pairs are pinned so that a quantizer which + stopped annotating anything at all would not read as a successful decline. + """ + + ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_sub_Tensor": 1, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 3, + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 3, + } + ops_after_transforms = { + "executorch_exir_dialects_edge__ops_cortex_m_quantized_add_default": 0, + "executorch_exir_dialects_edge__ops_aten_sub_Tensor": 1, + "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 3, + "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 3, + } + + def forward(self, x, y): + return torch.sub(x, y, alpha=2) + + +test_cases = { + "tensor": McuTestCase( + model=CortexMTensorSub(), + example_inputs=( + ramp_tensor(-10, 10, (4, 5)), + ramp_tensor(-2, 8, (4, 5)), + ), + ), + # Rank 4 in the layout the convolutions produce, which is where a rewrite + # that only handles the small ranks would go unnoticed. + "rank_4_channels_last": McuTestCase( + model=CortexMTensorSub(), + example_inputs=( + ramp_tensor(-5, 5, (2, 8, 4, 4)).to(memory_format=torch.channels_last), + ramp_tensor(-1, 9, (2, 8, 4, 4)).to(memory_format=torch.channels_last), + ), + ), + "rank_5": McuTestCase( + model=CortexMTensorSub(), + example_inputs=( + ramp_tensor(-5, 5, (2, 2, 2, 2, 2)), + ramp_tensor(-3, 1, (2, 2, 2, 2, 2)), + ), + ), + # One operand an order of magnitude wider than the other, so the two + # multipliers differ and swapping them changes the result. + "mismatched_scales": McuTestCase( + model=CortexMTensorSub(), + example_inputs=( + ramp_tensor(-100, 100, (32,)), + ramp_tensor(-3, 7, (32,)), + ), + ), + # Cancelling operands: the output scale collapses to the observer's + # eps floor and the kernel's requantization runs one shift below the point + # at which its int32 left shift would wrap. Sits just inside what the + # overflow guard permits. + "collapsed_output_scale": McuTestCase( + model=CortexMTensorSub(), + example_inputs=( + ramp_tensor(-100, 100, (32,)), + ramp_tensor(-100, 100, (32,)), + ), + ), + "inplace": McuTestCase( + model=CortexMInplaceSub(), + example_inputs=( + ramp_tensor(-10, 10, (4, 5)), + ramp_tensor(-2, 8, (4, 5)), + ), + ), + "alpha_int": McuTestCase( + model=CortexMAlphaSub(), + example_inputs=( + ramp_tensor(-10, 10, (4, 5)), + ramp_tensor(-2, 8, (4, 5)), + ), + ), +} + +xfails: dict[str, xfail_type] = {} + + +def test_cancelling_operands_are_refused(cortex_m_target): + """Operands that cancel leave an output scale the kernel cannot reach. + + `collapsed_output_scale` above is the same shape one shift lower, where the + kernel still holds; widening the operands is what pushes it over. + """ + tester = CortexMTester( + CortexMTensorSub(), + (ramp_tensor(-150, 150, (32,)), ramp_tensor(-150, 150, (32,))), + target_config=cortex_m_target, + ) + tester.quantize().export().to_edge() + with pytest.raises(Exception) as raised: + tester.run_passes() + # The pass manager wraps whatever a pass raises, so check the cause. + cause = raised.value + while cause.__cause__ is not None: + cause = cause.__cause__ + assert "the int32 kernel cannot hold" in str(cause), cause + + +@parametrize("test_case", test_cases, xfails=xfails) +def test_dialect_sub(test_case, cortex_m_target): + tester = CortexMTester( + test_case.model, test_case.example_inputs, target_config=cortex_m_target + ) + tester.test_dialect( + test_case.model.ops_before_transforms, test_case.model.ops_after_transforms + ) + + +@parametrize("test_case", test_cases, xfails=xfails) +def test_implementation_sub(test_case, cortex_m_target): + tester = CortexMTester( + test_case.model, test_case.example_inputs, target_config=cortex_m_target + ) + tester.test_implementation()