Arm backend: Clean up shift support (#9573)

- Handle lshift.Tensor and rshift.Tensor
- Convert *.Scalar to *.Tensor
- Test Scalar and Tensor cases with multiple dtypes
- Move cast logic from node visitor to pass

Signed-off-by: Erik Lundell <[email protected]>

Author: Erik-Lundell

backends/arm/_passes/__init__.py

@@ -7,7 +7,8 @@
 from . import arm_pass_utils  # noqa
 from .annotate_channels_last_dim_order_pass import AnnotateChannelsLastDimOrder  # noqa
 from .annotate_decomposed_matmul import AnnotateDecomposedMatmulPass  # noqa
-from .cast_int64_pass import CastInt64ToInt32Pass  # noqa
+from .cast_int64_pass import CastInt64BuffersToInt32Pass  # noqa
+from .cast_to_int32_pass import CastToInt32Pass  # noqa
 from .conv1d_unsqueeze_pass import Conv1dUnsqueezePass  # noqa
 from .convert_any_default_dim_dims_pass import ConvertAnyDefaultDimDimsPass  # noqa
 from .convert_expand_copy_to_repeat import ConvertExpandCopyToRepeatPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -10,7 +10,8 @@
 from executorch.backends.arm._passes import (
     AnnotateChannelsLastDimOrder,
     AnnotateDecomposedMatmulPass,
-    CastInt64ToInt32Pass,
+    CastInt64BuffersToInt32Pass,
+    CastToInt32Pass,
     ComputeConstantOpsAOT,
     Conv1dUnsqueezePass,
     ConvertAnyDefaultDimDimsPass,
@@ -80,6 +81,8 @@ def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(ConvertToClampPass())
         self.add_pass(ConvertMinMaxPass())
         self.add_pass(ConvertAnyDefaultDimDimsPass())
+        if isinstance(self.tosa_spec, Tosa_0_80) and self.tosa_spec.is_U55_subset:
+            self.add_pass(CastToInt32Pass())
 
         self.add_pass(ReplaceScalarWithTensorArgPass())
         self.add_pass(AnnotateDecomposedMatmulPass())
@@ -94,7 +97,7 @@ def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(SizeAdjustConv2DPass())
         self.add_pass(ConvertExpandCopyToRepeatPass())
         self.add_pass(UnsqueezeBeforeRepeatPass())
-        self.add_pass(CastInt64ToInt32Pass(exported_program))
+        self.add_pass(CastInt64BuffersToInt32Pass(exported_program))
         self.add_pass(KeepDimsFalseToSqueezePass())
         self.add_pass(Conv1dUnsqueezePass(exported_program))
         self.add_pass(DecomposeSelectPass())
@@ -141,7 +144,7 @@ def _tosa_080_MI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(SizeAdjustConv2DPass())
         self.add_pass(ConvertExpandCopyToRepeatPass())
         self.add_pass(UnsqueezeBeforeRepeatPass())
-        self.add_pass(CastInt64ToInt32Pass(exported_program))
+        self.add_pass(CastInt64BuffersToInt32Pass(exported_program))
         self.add_pass(KeepDimsFalseToSqueezePass())
         self.add_pass(Conv1dUnsqueezePass(exported_program))
         self.add_pass(DecomposeSelectPass())

backends/arm/_passes/cast_int64_pass.py

@@ -15,13 +15,13 @@
 logger.setLevel(logging.WARNING)
 
 
-class CastInt64ToInt32Pass(ExportPass):
+class CastInt64BuffersToInt32Pass(ExportPass):
     """
     Cast int64 buffers to int32 if the int64 data is in int32 range.
     """
 
     def __init__(self, exported_program: torch.export.ExportedProgram):
-        super(CastInt64ToInt32Pass, self).__init__()
+        super(CastInt64BuffersToInt32Pass, self).__init__()
         self.exported_program = exported_program
 
     def _assert_within_int32(self, tensor: torch.Tensor, node: torch.fx.Node):

backends/arm/_passes/cast_to_int32_pass.py

@@ -0,0 +1,54 @@
+# Copyright 2025 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.
+
+import torch
+
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+
+class CastToInt32Pass(ExportPass):
+    """Casts the input to int32 if it is not already and casts back the output to the original input dtype."""
+
+    targeted_ops = {
+        exir_ops.edge.aten.bitwise_left_shift.Tensor,
+        exir_ops.edge.aten.bitwise_right_shift.Tensor,
+    }
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in self.targeted_ops:
+            return super().call_operator(op, args, kwargs, meta)
+
+        new_args: list = []
+        did_cast = False
+        for arg in args:
+            if arg.data.dtype != torch.int32:
+                new_args.append(
+                    super().call_operator(
+                        exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+                        (arg,),
+                        {"dtype": torch.int32},
+                        meta,
+                    )
+                )
+                did_cast = True
+            else:
+                new_args.append(arg)
+
+        output = super().call_operator(
+            op,
+            tuple(new_args),
+            {},
+            meta,
+        )
+
+        if did_cast:
+            output = super().call_operator(
+                exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+                (output,),
+                {"dtype": args[0].data.dtype},
+                meta,
+            )
+        return output

backends/arm/_passes/match_arg_ranks_pass.py

@@ -45,6 +45,8 @@ def __init__(self, exported_program):
         exir_ops.edge.aten.sub.Tensor,
         exir_ops.edge.aten.mul.Tensor,
         exir_ops.edge.aten.div.Tensor,
+        exir_ops.edge.aten.bitwise_right_shift.Tensor,
+        exir_ops.edge.aten.bitwise_left_shift.Tensor,
     ]
 
     def _match_op_rank(self, graph_module, node, arg, max_rank):

backends/arm/operator_support/right_shift_support.py

@@ -22,7 +22,10 @@
 
 @register_tosa_support_check
 class RightShiftSupported(SupportedTOSAOperatorCheck):
-    targets = [exir_ops.edge.aten.__rshift__.Scalar]
+    targets = [
+        exir_ops.edge.aten.bitwise_right_shift.Tensor,
+        exir_ops.edge.aten.__rshift__.Scalar,
+    ]
 
     tosa_specs = [
         TosaSpecification.create_from_string("TOSA-0.80+BI"),

backends/arm/operator_support/tosa_supported_operators.py

@@ -205,6 +205,8 @@ def is_node_supported(
             exir_ops.edge.aten.amin.default,
             exir_ops.edge.aten.eye.default,
             exir_ops.edge.aten.linspace.default,
+            exir_ops.edge.aten.bitwise_left_shift.Tensor,
+            exir_ops.edge.aten.__lshift__.Scalar,
             torch.ops.aten.scalar_tensor.default,
         ]
 

backends/arm/operators/__init__.py

@@ -35,7 +35,7 @@
     op_reciprocal,
     op_repeat,
     op_rescale,
-    op_rshift,
+    op_rshift_tensor,
     op_rsqrt,
     op_sigmoid,
     op_slice,

backends/arm/operators/op_rshift.py

@@ -0,0 +1,46 @@
+# Copyright 2024-2025 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.
+
+# pyre-unsafe
+
+from typing import List
+
+import serializer.tosa_serializer as ts  # type: ignore
+import torch
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from executorch.backends.arm.tosa_specification import Tosa_0_80
+from serializer.tosa_serializer import TosaOp
+
+
+@register_node_visitor
+class RshiftVisitor(NodeVisitor):
+    target = "aten.bitwise_right_shift.Tensor"
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+
+        attr = ts.TosaSerializerAttribute()
+        round = False
+        if isinstance(self.tosa_spec, Tosa_0_80) and self.tosa_spec.is_U55_subset:
+            # U55 only supports INT32 and round == True
+            # TODO MLETORCH-525 Emulate round == False with different decomposition
+            round = True
+        attr.ArithmeticRightShiftAttribute(round=round)
+
+        tosa_graph.addOperator(
+            TosaOp.Op().ARITHMETIC_RIGHT_SHIFT,
+            [inputs[0].name, inputs[1].name],
+            [output.name],
+            attr,
+        )