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()