Arm backend: Sort passes in transform_for_annotation_pipeline (#15790)

The passes listed in ArmPassManager.transform_for_annotation_pipeline
can feel a bit arbitrary because there is no clearly intended structure
or pattern being applied there. Restructure the list into clearly
labelled blocks to make the code easier to read and maintain.

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/arm_pass_manager.py

@@ -291,9 +291,14 @@ def transform_to_backend_pipeline(
             )
 
     def transform_for_annotation_pipeline(self, graph_module: GraphModule):
+        # Preprocessing passes
+
         self.add_pass(
             RemoveGraphAssertsPass()
         )  # ConvertInt64ConstOpsToInt32Pass requires this pass to remove the assertation in Graph
+
+        # Transformation passes (pre scalar -> tensor)
+
         self.add_pass(ConvertInt64ConstOpsToInt32Pass())
         self.add_pass(ConvertInt64OutputOpsToInt32Pass())
         self.add_pass(InsertInt32CastsAfterInt64PlaceholdersPass())
@@ -304,12 +309,18 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
         self.add_pass(CastBoolToInt8Pass())
         self.add_pass(DecomposeSignPass())
         self.add_pass(DecomposeAddmmPass())
-        self.add_pass(ReplaceScalarWithTensorByProfilePass())
         self.add_pass(DecomposeRemainderPass())
         self.add_pass(DecomposeFloorDividePass())
         self.add_pass(DecomposeDivTensorModePass())
-        self.add_pass(DecomposeAddSubAlphaPass())
+
+        # Scalars -> tensors
+
+        self.add_pass(ReplaceScalarWithTensorByProfilePass())
         self.add_pass(ScalarsToAttributePass())
+
+        # Transformation passes (post scalar removal)
+
+        self.add_pass(DecomposeAddSubAlphaPass())
         self.add_pass(DecomposeGroupNormPass())
         self.add_pass(DecomposeLayerNormPass())
         self.add_pass(DecomposeVarPass())
@@ -323,16 +334,16 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
         self.add_pass(DecomposeSqrtPass())
         self.add_pass(DecomposeSiluPass())
         self.add_pass(DecomposeAvgPool2d())
-
         if self.tosa_spec.is_U55_subset:
             # Numerically stable softmax uses amax which is not supported on Ethos-U55
             self.add_pass(DecomposeSoftmaxUnstablePass())
         else:
             self.add_pass(DecomposeSoftmaxPass())
-
         self.add_pass(ConvertMinMaxPass())
-        self.add_pass(ReplaceInfValues())
 
+        # Postprocessing passes
+
+        self.add_pass(ReplaceInfValues())
         if not self.tosa_spec.is_U55_subset:
             # Uses where which is not supported on Ethos-U55
             self.add_pass(DecomposeMaskedFill())

backends/arm/_passes/decompose_remainder_pass.py

@@ -3,7 +3,7 @@
 # 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 typing import Dict, Set, Type
 
 import torch
 from executorch.backends.arm._passes import ArmPass
@@ -17,46 +17,50 @@
 
 Op = OpOverload | EdgeOpOverload
 
-
-def _get_remainder_decomposition_ops(op: Op) -> tuple[Op, Op, Op]:
-    """
-    Returns the (div_mode_op, mul_op, sub_op) needed to lower the provided
-    remainder operator. The concrete ops depend on whether the remainder op is
-    the aten or edge variant.
-    """
-    if op == exir_ops.edge.aten.remainder.Tensor:
-        return (
-            exir_ops.edge.aten.div.Tensor_mode,
-            exir_ops.edge.aten.mul.Tensor,
-            exir_ops.edge.aten.sub.Tensor,
-        )
-    if op == torch.ops.aten.remainder.Tensor:
-        return (
-            torch.ops.aten.div.Tensor_mode,
-            torch.ops.aten.mul.Tensor,
-            torch.ops.aten.sub.Tensor,
-        )
-    raise RuntimeError(f"Can't get remainder decomposition ops for op {op}")
+_decomposition_ops: Dict[Op, tuple[Op, Op, Op]] = {
+    exir_ops.edge.aten.remainder.Scalar: (
+        exir_ops.edge.aten.div.Tensor_mode,
+        exir_ops.edge.aten.mul.Scalar,
+        exir_ops.edge.aten.sub.Tensor,
+    ),
+    torch.ops.aten.remainder.Tensor: (
+        torch.ops.aten.div.Tensor_mode,
+        torch.ops.aten.mul.Tensor,
+        torch.ops.aten.sub.Tensor,
+    ),
+    torch.ops.aten.remainder.Scalar: (
+        torch.ops.aten.div.Tensor_mode,
+        torch.ops.aten.mul.Scalar,
+        torch.ops.aten.sub.Tensor,
+    ),
+    exir_ops.edge.aten.remainder.Tensor: (
+        exir_ops.edge.aten.div.Tensor_mode,
+        exir_ops.edge.aten.mul.Tensor,
+        exir_ops.edge.aten.sub.Tensor,
+    ),
+}
 
 
 class DecomposeRemainderPass(ArmPass):
     """
     Decompose the remainder operation into primitive arithmetic:
         remainder(x, y) -> x - floor_div(x, y) * y
-    where floor_div(x, y) == div(x, y, rounding_mode=\"floor\").
+    where floor_div(x, y) == div(x, y, rounding_mode="floor").
     """
 
     _passes_required_after: Set[Type[ExportPass]] = {DecomposeDivTensorModePass}
 
     def call_operator(self, op, args, kwargs, meta, updated=False):
         supported_ops = (
+            exir_ops.edge.aten.remainder.Scalar,
             exir_ops.edge.aten.remainder.Tensor,
+            torch.ops.aten.remainder.Scalar,
             torch.ops.aten.remainder.Tensor,
         )
         if op not in supported_ops:
             return super().call_operator(op, args, kwargs, meta, updated)
 
-        div_op, mul_op, sub_op = _get_remainder_decomposition_ops(op)
+        div_op, mul_op, sub_op = _decomposition_ops[op]
         x, y = args[0], args[1]
 
         floor_div = super().call_operator(