Arm backend: normalize inplace arithmetic ops. (#18231)

Just turn them into normal out-of-place ops.
This solves some inplace scalar xfails.

Author: Erik-Lundell

backends/arm/_passes/__init__.py

@@ -146,6 +146,7 @@
     RewriteHighRankSingletonPermutePass,
 )
 from .rewrite_index_put_pass import RewriteIndexPutPass  # noqa
+from .rewrite_inplace_arithmetic_pass import RewriteInplaceArithmeticPass  # noqa
 from .rewrite_le_lt_to_ge_gt_pass import RewriteLeLtToGeGtPass  # noqa
 from .rewrite_matmul import RewriteMatmulPass  # noqa
 from .rewrite_pad import RewritePadPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -125,6 +125,7 @@
     RewriteConvPass,
     RewriteHighRankSingletonPermutePass,
     RewriteIndexPutPass,
+    RewriteInplaceArithmeticPass,
     RewriteLeLtToGeGtPass,
     RewriteMatmulPass,
     RewritePadPass,
@@ -553,6 +554,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                 DecomposeFloorDividePass(tfa_pass=True),
                 DecomposeDivTensorModePass(tfa_pass=True),
                 DecomposeWhereScalarOtherPass(tfa_pass=True),
+                RewriteInplaceArithmeticPass(tfa_pass=True),
             ]
         )
 

backends/arm/_passes/rewrite_inplace_arithmetic_pass.py

@@ -0,0 +1,39 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Set, Type
+
+import torch
+
+from executorch.backends.arm._passes import ArmPass
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+OP_MAP = {
+    torch.ops.aten.add_.Tensor: torch.ops.aten.add.Tensor,
+    torch.ops.aten.sub_.Tensor: torch.ops.aten.sub.Tensor,
+    torch.ops.aten.mul_.Tensor: torch.ops.aten.mul.Tensor,
+    torch.ops.aten.div_.Tensor: torch.ops.aten.div.Tensor,
+    exir_ops.edge.aten.add_.Tensor: exir_ops.edge.aten.add.Tensor,
+    exir_ops.edge.aten.sub_.Tensor: exir_ops.edge.aten.sub.Tensor,
+    exir_ops.edge.aten.mul_.Tensor: exir_ops.edge.aten.mul.Tensor,
+    exir_ops.edge.aten.div_.Tensor: exir_ops.edge.aten.div.Tensor,
+}
+
+
+class RewriteInplaceArithmeticPass(ArmPass):
+    """Rewrite inplace arithmetic ops into functional equivalents."""
+
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
+    def call_operator(self, op, args, kwargs, meta):
+        if not self.allowed_to_transform(meta):
+            return super().call_operator(op, args, kwargs, meta)
+
+        target_op = OP_MAP.get(op)
+        if target_op is None:
+            return super().call_operator(op, args, kwargs, meta)
+
+        return super().call_operator(target_op, args, kwargs, meta, updated=True)

backends/arm/_passes/scalars_to_attribute_pass.py

@@ -27,14 +27,10 @@ class ScalarsToAttributePass(ArmPass):
 
     targeted_ops = [
         torch.ops.aten.add.Tensor,
-        torch.ops.aten.add_.Tensor,
         torch.ops.aten.sub.Tensor,
-        torch.ops.aten.sub_.Tensor,
         torch.ops.aten.rsub.Scalar,
         torch.ops.aten.mul.Tensor,
-        torch.ops.aten.mul_.Tensor,
         torch.ops.aten.div.Tensor,
-        torch.ops.aten.div_.Tensor,
     ]
 
     def _convert_scalar_args(

backends/arm/quantizer/quantization_annotator.py

@@ -703,13 +703,10 @@ def any_or_hardtanh_min_zero(n: Node):
         quant_properties.quant_output = _QuantProperty(0, output_act_qspec)
     elif node.target in (
         torch.ops.aten.add.Tensor,
-        torch.ops.aten.add_.Tensor,
         torch.ops.aten.sub.Tensor,
-        torch.ops.aten.sub_.Tensor,
         torch.ops.aten.mm.default,
         torch.ops.aten.bmm.default,
         torch.ops.aten.mul.Tensor,
-        torch.ops.aten.mul_.Tensor,
     ):
         quant_properties.quant_inputs = [
             _QuantProperty(0, input_act_qspec),

backends/arm/test/ops/test_scalars.py

@@ -168,13 +168,6 @@ def forward(self, x):
     "float_r4_st": "MLETORCH-408: Arithmetic ops can't handle scalars first",
 }
 
-int_inplace_xfails = {
-    "int_r1_ts": "MLETORCH-1708: Numerical error in TFA/quantization",
-    "int_r4_ts": "MLETORCH-1708: Numerical error in TFA/quantization",
-    "float_r1_ts": "MLETORCH-1708: Numerical error in TFA/quantization",
-    "float_r4_ts": "MLETORCH-1708: Numerical error in TFA/quantization",
-}
-
 
 # ADD FP ------------------------------------------------------
 @common.parametrize("test_data", tensor_scalar_tests, xfails=xfails)
@@ -215,9 +208,7 @@ def test_add_tensor_tosa_INT_scalar(test_data):
     pipeline.run()
 
 
-@common.parametrize(
-    "test_data", tensor_scalar_tests, xfails=int_inplace_xfails, strict=False
-)
+@common.parametrize("test_data", tensor_scalar_tests)
 def test_add_tensor_tosa_INT_inplace(test_data):
     """Tests inplace add with one scalar input."""
     pipeline = TosaPipelineINT[input_t1](AddInplace(), test_data, aten_op=[])
@@ -285,9 +276,7 @@ def test_sub_tensor_tosa_INT_scalar(test_data):
     pipeline.run()
 
 
-@common.parametrize(
-    "test_data", tensor_scalar_tests, xfails=int_inplace_xfails, strict=False
-)
+@common.parametrize("test_data", tensor_scalar_tests)
 def test_sub_tensor_tosa_INT_inplace(test_data):
     """Tests inplace sub with one scalar input."""
     pipeline = TosaPipelineINT[input_t1](SubInplace(), test_data, aten_op=[])
@@ -344,9 +333,7 @@ def test_mul_tensor_tosa_INT_scalar(test_data):
     pipeline.run()
 
 
-@common.parametrize(
-    "test_data", tensor_scalar_tests, xfails=int_inplace_xfails, strict=False
-)
+@common.parametrize("test_data", tensor_scalar_tests)
 def test_mul_tensor_tosa_INT_inplace(test_data):
     """Tests inplace mul with one scalar input."""
     pipeline = TosaPipelineINT[input_t1](MulInplace(), test_data, aten_op=[])