Arm: Enable rsub lowering (#8525)

Enable rsub lowering

Update ScalarsToAttrbibutePass to replace rsub with sub.

Signed-off-by: Oscar Andersson <[email protected]>

Author: oscarandersson8218

backends/arm/_passes/scalars_to_attribute_pass.py

@@ -76,5 +76,17 @@ def call(self, graph_module: GraphModule) -> PassResult:
                     new_args.append(get_attr_node)
             n.args = tuple(new_args)
 
+            # Replace rsub.Scalar with sub.Tensor as retracing will fail otherwise
+            if n.target == torch.ops.aten.rsub.Scalar:
+                with graph_module.graph.inserting_after(n):
+                    reversed_args = (n.args[1], n.args[0])
+                    sub = graph_module.graph.create_node(
+                        "call_function", torch.ops.aten.sub.Tensor, reversed_args, {}
+                    )
+                    n.replace_all_uses_with(sub)
+                    sub.meta["val"] = n.meta["val"]
+                graph_module.graph.erase_node(n)
+
         graph_module.recompile()
+        graph_module = super().call(graph_module).graph_module
         return PassResult(graph_module, True)

backends/arm/test/ops/test_scalars.py

@@ -14,11 +14,6 @@
 """
 Summary of non-working cases.
 MI:
-    Any case with int scalar: A to_copy is inserted to cast the value which we don't partition.
-        This makes the constant end up outside our partition and the input to the delegate becomes
-        a to_copy placeholder. In ArmTester, the placeholder is then interpreted as an input.
-        Potential fix: partition int -> float to_copy-ops in ArmBackend.
-        # MLETORCH-407
     Op(scalar, tensor):
         One issue is that lift_constant_tensor_pass looks for a fake_tensor in the meta of the first
         node which does not work the first node is a scalar.
@@ -27,17 +22,12 @@
         somewhere in _transform in the to_edge step. This makes ArmPartitioner miss tagging the
         data in tag_constant_data.
         # MLETORCH-408
-
-BI:
-    sub(Scalar, Tensor) becomes rsub, which either fails since the scalar does not become an attribute
-        in scalars_to_attribute_pass, or, if added to targeted_ops in that pass, fails since rsub expects a
-        Scalar.
-        Potential fix: Create pass to convert rsub.Scalar to sub.Tensor
+    Sub or inplace-sub with an integer input.
 """
 
 
 class TestScalars(unittest.TestCase):
-    """Tests various scalar cases for for"""
+    """Tests various scalar cases"""
 
     class Add(torch.nn.Module):
         def forward(self, x, y):
@@ -133,13 +123,10 @@ def forward(self, x):
                 scalar = dtype[1]
                 tensor_scalar_tests.append((test_name + "_ts", op[1], tensor, scalar))
 
-                # Don't add (scalar, tensor) test case for inplace and .Scalar ops.
-                if op[0][-1] == "_" or op[0][-6:] == "Scalar":
+                # Don't add (scalar, tensor) test case for .Scalar ops.
+                if op[0][-6:] == "Scalar":
                     continue
 
-                # sub(scalar, tensor) does not work in any case.
-                if op[0][0:3] == "Sub":
-                    continue
                 tensor_scalar_tests.append((test_name + "_st", op[1], scalar, tensor))
 
     tensor_const_tests = []
@@ -182,8 +169,8 @@ def _test_add_tosa_BI_pipeline(self, module: torch.nn.Module, test_data: tuple):
     def test_MI(self, test_name: str, op: torch.nn.Module, x, y):
         expected_exception = None
         if any(token in test_name for token in ("Sub_int", "Sub__int")):
-            expected_exception = (AssertionError, ValueError)
-        elif test_name.endswith("_st"):
+            expected_exception = AssertionError
+        if test_name.endswith("_st"):
             expected_exception = AttributeError
 
         if expected_exception:
@@ -204,5 +191,13 @@ def test_MI_const(self, test_name: str, op: torch.nn.Module, x):
     def test_BI(self, test_name: str, op: torch.nn.Module, x, y):
         self._test_add_tosa_BI_pipeline(op, (x, y))
 
+    # op(Scalar float, tensor) works if the scalar is constant.
+    @parameterized.expand(tensor_const_tests)
+    def test_BI_const(self, test_name: str, op: torch.nn.Module, x):
+        self._test_add_tosa_BI_pipeline(op, (x,))
+
     def test_shift_sub_inplace_tosa_MI(self):
         self._test_add_tosa_MI_pipeline(self.ShiftInplaceSub(), (torch.IntTensor(5),))
+
+    def test_shift_sub_inplace_tosa_BI(self):
+        self._test_add_tosa_BI_pipeline(self.ShiftInplaceSub(), (torch.IntTensor(5),))