Arm backend: Sort passes in _tosa_pipeline into blocks

The passes listed in `ArmPassManager._tosa_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.

Signed-off-by: Martin Lindström <[email protected]>
Change-Id: Iadf37cda2c7a88cad80bf363062d38d492206be7

Author: Martin Lindström

backends/arm/_passes/arm_pass_manager.py

@@ -159,7 +159,12 @@ def _transform(self, graph_module: GraphModule):
     def _tosa_pipeline(
         self, exported_program: ExportedProgram, graph_module: GraphModule
     ) -> GraphModule:
+        # Preprocessing passes
+
         self.add_pass(AnnotateOutputDimOrderPass())
+
+        # Node transformation passes (pre q/dq folding)
+
         self.add_pass(FuseQuantizedActivationPass())
         self.add_pass(RemoveGetItemPass())
         self.add_pass(ConvertToClampPass())
@@ -174,8 +179,19 @@ def _tosa_pipeline(
         self.add_pass(ConvertELUParamsPass())
         self.add_pass(ConvertSplitToSlicePass())
         self.add_pass(QuantizeOperatorArguments())
+
+        # Fold Q/DQ nodes, insert INT8/INT32 rescales.
+
         self.add_pass(FoldAndAnnotateQParamsPass(exported_program))  # type: ignore[call-arg]
         self.add_pass(FuseDuplicateUsersPass())
+        # TODO: DecomposeLinearPass should run after InsertRescaleInt32Pass or
+        # before FoldAndAnnotateQParamsPass but is unable to at the moment.
+        # Ticket: MLETORCH-1539
+        self.add_pass(DecomposeLinearPass())
+        self.add_pass(InsertRescaleInt32Pass())
+
+        # Node transformation passes (post q/dq folding)
+
         self.add_pass(DecomposeExpm1Pass())
         self.add_pass(DecomposeLogitPass())
         self.add_pass(DecomposeMaskedFill())
@@ -196,57 +212,67 @@ def _tosa_pipeline(
         self.add_pass(DecomposeSignPass())
         self.add_pass(DecomposeFloorDividePass())
         self.add_pass(DecomposeDivTensorModePass())
+        self.add_pass(DecomposeGeluPass())
+        self.add_pass(DecomposeAddSubAlphaPass())
+        self.add_pass(DecomposeGroupedConv())
+        self.add_pass(Conv1dUnsqueezePass())
+
+        # Scalars -> tensors, match tensor dtypes and ranks.
+
         self.add_pass(ReplaceScalarWithTensorByProfilePass())
+        self.add_pass(MatchArgDtypePass())
+        self.add_pass(UnsqueezeScalarPlaceholdersPass(exported_program))
+        # TODO: Move DecomposeNotEqualPass to before or after this block of
+        # passes. Ticket: MLETORCH-1540
+        self.add_pass(DecomposeNotEqualPass())
+        self.add_pass(MatchArgRanksPass(exported_program))
+        self.add_pass(FuseConstantArgsPass(exported_program))
+
+        # Node transformation passes (post scalar-removal)
+
         self.add_pass(DecomposeRemainderPass())
         self.add_pass(DecomposeDivTensorModePass())
         self.add_pass(DecomposeEmbeddingPass())
         self.add_pass(FuseBatchnorm2DPass(exported_program))
         self.add_pass(ConvertMmToBmmPass())
         self.add_pass(DecomposeGluPass())
-        self.add_pass(DecomposeLinearPass())
         self.add_pass(DecomposeLeakyReLUPass())
-        self.add_pass(DecomposeNotEqualPass())
         self.add_pass(DecomposeDivPass())
-        self.add_pass(DecomposeAddSubAlphaPass())
         self.add_pass(DecomposeSoftmaxPass())
-        self.add_pass(DecomposeGeluPass())
         self.add_pass(ConvertFullLikeToFullPass())
         self.add_pass(ConvertMinMaxPass())
         self.add_pass(ConvertAnyDefaultDimDimsPass())
-        self.add_pass(MatchArgDtypePass())
-        self.add_pass(UnsqueezeScalarPlaceholdersPass(exported_program))
-        self.add_pass(MatchArgRanksPass(exported_program))
         self.add_pass(DecomposeAdaptiveAvgPool2dPass())
         self.add_pass(DecomposeAvgPool2d())
         self.add_pass(
             DecorateFp32toInt32CastingPass()
         )  # Require that no new fp32->int32 is introduced after this pass
         self.add_pass(ComputeConstantOpsAOT(exported_program))
-
-        self.add_pass(DecomposeGroupedConv())
         self.add_pass(ConvertExpandCopyToRepeatPass())
         self.add_pass(UnsqueezeBeforeRepeatPass())
         self.add_pass(DecomposeCumsumPass(exported_program))
-        self.add_pass(Conv1dUnsqueezePass())
         self.add_pass(DecomposeMaxPool2DPass())
         self.add_pass(SizeAdjustInputPass())
         self.add_pass(DecomposeSelectPass())
         self.add_pass(ConvertSqueezesToViewPass())
         self.add_pass(CastToInt32Pass())
         self.add_pass(BroadcastArgsPass())
-
         self.add_pass(ConvertPermuteSingletonToViewPass())
         self.add_pass(FuseViewCopyTransform())
-        self.add_pass(FuseConstantArgsPass(exported_program))
         self.add_pass(DecomposeConv2dWithInt16ActivationPass())
-        self.add_pass(CastInt64BuffersToInt32Pass(exported_program))
+        self.add_pass(DecomposeSumPass())
         self.add_pass(InsertTableOpsPass(exported_program))
+
+        # Aten -> TOSA transformation passes
+
         self.add_pass(RewriteUpsamplePass())
         self.add_pass(RewriteConv2dPass(exported_program))
         self.add_pass(RewriteMatmulPass())
+
+        # Postprocessing/cleanup passes
+
+        self.add_pass(CastInt64BuffersToInt32Pass(exported_program))
         self.add_pass(FuseEqualPlaceholdersPass(exported_program))
-        self.add_pass(InsertRescaleInt32Pass())
-        self.add_pass(DecomposeSumPass())
         self.add_pass(ToTosaMemoryFormatPass(exported_program))
         self.add_pass(RemoveNoopPass())
         self.add_pass(InsertRescalePass())

backends/arm/_passes/decompose_linear_pass.py

@@ -12,6 +12,7 @@
     create_node,
     get_first_fake_tensor,
 )
+from executorch.backends.arm._passes.insert_rescales_pass import InsertRescaleInt32Pass
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 
@@ -26,7 +27,7 @@ class DecomposeLinearPass(ArmPass):
         output           = view(conv2d)
     """
 
-    _passes_required_after: Set[Type[ExportPass]] = set()
+    _passes_required_after: Set[Type[ExportPass]] = {InsertRescaleInt32Pass}
 
     def call(self, graph_module):
         for node in graph_module.graph.nodes:

backends/arm/_passes/match_arg_ranks_pass.py

@@ -57,6 +57,7 @@ def __init__(self, exported_program: ExportedProgram) -> None:
         exir_ops.edge.aten.lt.Tensor,
         exir_ops.edge.aten.le.Tensor,
         exir_ops.edge.aten.pow.Tensor_Tensor,
+        exir_ops.edge.aten.remainder.Tensor,
         exir_ops.edge.aten.where.self,
         exir_ops.edge.aten.bitwise_and.Tensor,
         exir_ops.edge.aten.bitwise_xor.Tensor,