Arm backend: Make remaining passes inherit from ArmPass (#15764)

cc @freddan80 @per @zingo @oscarandersson8218 @digantdesai

Signed-off-by: Martin Lindström <Martin.Lindstroem@arm.com>
Co-authored-by: Martin Lindström <Martin.Lindstroem@arm.com>

Author: martinlsm

backends/arm/_passes/__init__.py

@@ -64,6 +64,7 @@
 from .decompose_ne_pass import DecomposeNotEqualPass  # noqa
 from .decompose_remainder_pass import DecomposeRemainderPass  # noqa
 from .decompose_round_pass import DecomposeRoundPass  # noqa
+from .decompose_sdpa_pass import DecomposeScaledDotProductAttention  # noqa
 from .decompose_select import DecomposeSelectPass  # noqa
 from .decompose_sign_pass import DecomposeSignPass  # noqa
 from .decompose_silu_pass import DecomposeSiluPass  # noqa
@@ -83,6 +84,7 @@
 from .fuse_duplicate_users_pass import FuseDuplicateUsersPass  # noqa
 from .fuse_equal_placeholders_pass import FuseEqualPlaceholdersPass  # noqa
 from .fuse_quantized_activation_pass import FuseQuantizedActivationPass  # noqa
+from .fuse_view_copy_transform_pass import FuseViewCopyTransformPass  # noqa
 from .insert_int32_casts_after_int64_placeholders import (  # noqa
     InsertInt32CastsAfterInt64PlaceholdersPass,
 )
@@ -91,6 +93,8 @@
 from .match_arg_dtype_pass import MatchArgDtypePass  # noqa
 from .match_arg_ranks_pass import MatchArgRanksPass  # noqa
 from .mm_to_bmm_pass import ConvertMmToBmmPass  # noqa
+from .remove_getitem_pass import RemoveGetItemPass  # noqa
+from .remove_graph_asserts_pass import RemoveGraphAssertsPass  # noqa
 from .remove_noop_pass import RemoveNoopPass  # noqa
 from .replace_scalar_with_tensor_pass import (  # noqa
     ReplaceScalarWithTensorByProfilePass,

backends/arm/_passes/arm_pass_manager.py

@@ -66,6 +66,7 @@
     DecomposeNotEqualPass,
     DecomposeRemainderPass,
     DecomposeRoundPass,
+    DecomposeScaledDotProductAttention,
     DecomposeSelectPass,
     DecomposeSignPass,
     DecomposeSiluPass,
@@ -82,13 +83,16 @@
     FuseDuplicateUsersPass,
     FuseEqualPlaceholdersPass,
     FuseQuantizedActivationPass,
+    FuseViewCopyTransformPass,
     InsertInt32CastsAfterInt64PlaceholdersPass,
     InsertRescaleInt32Pass,
     InsertRescalePass,
     InsertTableOpsPass,
     MatchArgDtypePass,
     MatchArgRanksPass,
     QuantizeOperatorArguments,
+    RemoveGetItemPass,
+    RemoveGraphAssertsPass,
     RemoveNoopPass,
     ReplaceInfValues,
     ReplaceScalarWithTensorByProfilePass,
@@ -107,14 +111,8 @@
     TosaLoweringContext,
     TosaSpecification,
 )
-from executorch.backends.transforms.decompose_sdpa import (
-    DecomposeScaledDotProductAttention,
-)
-from executorch.backends.transforms.fuse_view_copy import FuseViewCopyTransform
-from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass
 from executorch.exir import ExportedProgram
 from executorch.exir.pass_manager import PassManager
-from executorch.exir.passes.remove_graph_asserts_pass import RemoveGraphAssertsPass
 from torch.fx import GraphModule
 from torch.fx.passes.infra.pass_base import PassResult
 from torch.nn.modules import Module
@@ -258,7 +256,7 @@ def _tosa_pipeline(
         self.add_pass(CastToInt32Pass())
         self.add_pass(BroadcastArgsPass())
         self.add_pass(ConvertPermuteSingletonToViewPass())
-        self.add_pass(FuseViewCopyTransform())
+        self.add_pass(FuseViewCopyTransformPass())
         self.add_pass(DecomposeConv2dWithInt16ActivationPass())
         self.add_pass(DecomposeSumPass())
         self.add_pass(InsertTableOpsPass(exported_program))

backends/arm/_passes/convert_permute_singleton_to_view_pass.py

@@ -6,6 +6,8 @@
 
 from typing import Sequence, Set, Tuple, Type
 
+from executorch.backends.arm._passes.arm_pass import ArmPass
+
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
@@ -18,7 +20,7 @@
 )
 
 
-class ConvertPermuteSingletonToViewPass(ExportPass):
+class ConvertPermuteSingletonToViewPass(ArmPass):
     """Replace permutations that only move singleton axes with a reshape.
 
     Examples:

backends/arm/_passes/convert_squeezes_to_view.py

@@ -1,5 +1,4 @@
 # Copyright 2025 Arm Limited and/or its 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.
@@ -8,9 +7,9 @@
 from typing import Set, Type
 
 from executorch.backends.arm._passes import ArmPass
-
-from executorch.backends.transforms.fuse_view_copy import FuseViewCopyTransform
-
+from executorch.backends.arm._passes.fuse_view_copy_transform_pass import (
+    FuseViewCopyTransformPass,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
@@ -20,7 +19,7 @@ class ConvertSqueezesToViewPass(ArmPass):
     Replaces squeeze/unsqueeze operators with view. These are simply special cases of the view op, so removing them gives us less cases to handle in the node visitiors.
     """
 
-    _passes_required_after: Set[Type[ExportPass]] = {FuseViewCopyTransform}
+    _passes_required_after: Set[Type[ExportPass]] = {FuseViewCopyTransformPass}
 
     def call_operator(self, op, args, kwargs, meta):
         if op not in [

backends/arm/_passes/decompose_embedding_pass.py

@@ -10,7 +10,9 @@
 
 import torch
 from executorch.backends.arm._passes import ArmPass
-from executorch.backends.transforms.fuse_view_copy import FuseViewCopyTransform
+from executorch.backends.arm._passes.fuse_view_copy_transform_pass import (
+    FuseViewCopyTransformPass,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 
@@ -33,7 +35,7 @@ class DecomposeEmbeddingPass(ArmPass):
          i = indices is expected to be int32 before this pass
     """
 
-    _passes_required_after: Set[Type[ExportPass]] = {FuseViewCopyTransform}
+    _passes_required_after: Set[Type[ExportPass]] = {FuseViewCopyTransformPass}
 
     aten_ops = (torch.ops.aten.embedding.default,)
     edge_ops = (exir_ops.edge.aten.embedding.default,)

backends/arm/_passes/decompose_int16_activation_conv2d_pass.py

@@ -4,24 +4,27 @@
 # LICENSE file in the root directory of this source tree.
 
 
-from typing import cast
+from typing import cast, Set, Type
 
 import torch
+from executorch.backends.arm._passes.arm_pass import ArmPass
 from executorch.backends.arm._passes.quant_args import QuantArgs
 
 from executorch.backends.arm.tosa.specification import get_context_spec, Tosa_1_00
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
 
-class DecomposeConv2dWithInt16ActivationPass(ExportPass):
+class DecomposeConv2dWithInt16ActivationPass(ArmPass):
     """
     This pass decomposes a convolution with input dtype int16 and bias
     into a convolution without bias followed by an addition of the bias
     since the TOSA op requires the bias to be int48 which is hard to represent
     in torch. Instead rescale the int48 output to int16 and add the bias in int16.
     """
 
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
     def call_operator(self, op, args, kwargs, meta):
         if op != exir_ops.edge.aten.convolution.default:
             return super().call_operator(op, args, kwargs, meta)

backends/arm/_passes/decompose_sdpa_pass.py

@@ -0,0 +1,16 @@
+# 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.
+
+from typing import Set, Type
+
+from executorch.backends.arm._passes.arm_pass import ArmPass
+from executorch.backends.transforms import decompose_sdpa
+from executorch.exir.pass_base import ExportPass
+
+
+class DecomposeScaledDotProductAttention(
+    ArmPass, decompose_sdpa.DecomposeScaledDotProductAttention
+):
+    _passes_required_after: Set[Type[ExportPass]] = set()

backends/arm/_passes/fuse_view_copy_transform_pass.py

@@ -0,0 +1,14 @@
+# 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.
+
+from typing import Set, Type
+
+from executorch.backends.arm._passes.arm_pass import ArmPass
+from executorch.backends.transforms.fuse_view_copy import FuseViewCopyTransform
+from executorch.exir.pass_base import ExportPass
+
+
+class FuseViewCopyTransformPass(ArmPass, FuseViewCopyTransform):
+    _passes_required_after: Set[Type[ExportPass]] = set()

backends/arm/_passes/remove_getitem_pass.py

@@ -0,0 +1,14 @@
+# 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.
+
+from typing import Set, Type
+
+from executorch.backends.arm._passes.arm_pass import ArmPass
+from executorch.backends.transforms import remove_getitem_op
+from executorch.exir.pass_base import ExportPass
+
+
+class RemoveGetItemPass(ArmPass, remove_getitem_op.RemoveGetItemPass):
+    _passes_required_after: Set[Type[ExportPass]] = set()

backends/arm/_passes/remove_graph_asserts_pass.py

@@ -0,0 +1,14 @@
+# 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.
+
+from typing import Set, Type
+
+from executorch.backends.arm._passes.arm_pass import ArmPass
+from executorch.exir.pass_base import ExportPass
+from executorch.exir.passes import remove_graph_asserts_pass
+
+
+class RemoveGraphAssertsPass(remove_graph_asserts_pass.RemoveGraphAssertsPass, ArmPass):
+    _passes_required_after: Set[Type[ExportPass]] = set()