From fa7dd6c8eada5393b34b9d32797d9ae481551ef5 Mon Sep 17 00:00:00 2001
From: Eli Amesefe <eliamesefe@meta.com>
Date: Wed, 22 Jul 2026 19:24:46 -0700
Subject: [PATCH] Arm backend: regression test for equal non-simple IFM scale
 Add on Ethos-U55 (#21151)

Summary:

Add EqualScaleAdd (x + x) tests on Ethos-U55 and Ethos-U85. When both Add/Sub
operands share one equal, non-simple (non-power-of-2) scale, the Ethos-U55
equal-scale code path writes both scales into 16-bit OPA/OPB registers and the
wide value can overflow to zero, silently dropping an operand (x + x returns x,
i.e. the output is halved; softmax's Sub(x, max) loses an operand). Correct
lowering keeps the operand as an explicit rescale so the FVP result matches the
quantized reference.

The check is Ethos-U55-codegen-specific (it reproduces on the FVP, not on the
TOSA reference path). Non-power-of-2 input magnitudes force the non-simple equal
scale that triggers the overflow.

Reviewed By: rascani

Differential Revision: D113307975
---
 backends/arm/test/ops/test_add.py | 60 +++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/backends/arm/test/ops/test_add.py b/backends/arm/test/ops/test_add.py
index 1a66b818280..464cc5e12fe 100644
--- a/backends/arm/test/ops/test_add.py
+++ b/backends/arm/test/ops/test_add.py
@@ -239,6 +239,66 @@ def test_add_tensor_tosa_INT_2(test_data: input_t2):
     pipeline.run()
 
 
+# --- Equal non-simple IFM scale Add on Ethos-U55 -------------------------------
+# When both operands of an Add/Sub share one *equal* non-simple (non-power-of-2)
+# scale, the Ethos-U55 equal-scale code path writes both scales into 16-bit OPA/OPB
+# registers; the wide (pre-narrowing) value can exceed 16 bits and overflow to
+# zero, silently dropping an operand -- ``x + x`` then returns ``x`` (output
+# halved), and softmax's ``Sub(x, max)`` loses an operand. Correct lowering keeps
+# the operand as an explicit rescale so the on-device (FVP) result matches the
+# quantized reference.
+#
+# This is Ethos-U55-codegen-specific: it only reproduces on the FVP (the TOSA
+# reference path does not exhibit the register overflow).
+class EqualScaleAdd(torch.nn.Module):
+    """``x + x`` -> both Add operands share one (equal) IFM scale."""
+
+    def forward(self, x: torch.Tensor):
+        return x + x
+
+    # Non-power-of-2 magnitudes force a non-simple activation scale (max/127 with a
+    # non-power-of-2 max), i.e. exactly the equal-scale case that overflows the
+    # Ethos-U55 OPA/OPB registers.
+    test_data: dict[str, input_t1] = {
+        "equal_nonsimple_scale": lambda: (
+            torch.tensor([0.3, 1.7, 2.9, 5.1, 6.3, 6.9]),
+        ),
+        "equal_nonsimple_scale_4d": lambda: (0.7 * torch.randn(1, 3, 8, 8) + 3.1,),
+    }
+
+
+@common.parametrize("test_data", EqualScaleAdd.test_data)
+@common.XfailIfNoCorstone300
+def test_add_equal_scale_u55_INT(test_data: input_t1):
+    """Equal non-simple IFM scales must not drop an Add operand on Ethos-U55.
+
+    If the equal-scale operand is dropped, the FVP output is halved and the
+    comparison against the quantized reference fails.
+
+    """
+    pipeline = EthosU55PipelineINT[input_t1](
+        EqualScaleAdd(),
+        test_data(),
+        aten_op,
+        exir_op,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", EqualScaleAdd.test_data)
+@common.XfailIfNoCorstone320
+def test_add_equal_scale_u85_INT(test_data: input_t1):
+    """U85 companion: U85 has an advanced-mode dual-IFM rescale so it is unaffected;
+    this guards against a future regression of the shared Add/Sub scaling path."""
+    pipeline = EthosU85PipelineINT[input_t1](
+        EqualScaleAdd(),
+        test_data(),
+        aten_op,
+        exir_op,
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", Add2.test_data)
 @common.XfailIfNoCorstone300
 def test_add_tensor_u55_INT_2(test_data: input_t2):