Cortex-M: lower subtraction, and stop miscompiling add with alpha

backends/cortex_m/passes/aten_to_cortex_m_pass.py

@@ -901,13 +901,25 @@ 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:
     del dialect_pass
     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
 

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

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

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

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

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