Arm backend: Add pass order validation to ArmPassManager (#14148)

Introduce a mechanism to enforce required ordering of passes in
ArmPassManager. Each ArmPass must now declare which passes are required
to run after it, ensuring ordering constraints are always upheld.

This prevents accidental breakage when modifying pass ordering in the
manager.

Ordering constraints are verified by the new method
ArmPass.validate_constraints_mandatory. We considered reusing
torch.fx.passes.infra.pass_manager.PassManager.validate_constraints, but
that utility only checks pairwise ordering and cannot enforce that a
pass is actually run, which did not meet our needs.

This patch only implements the mechanism and tests for it. Defining the
actual pass orderings are done in a later patch.

### Test plan
The change comes with added unit tests in
backends/arm/test/misc/test_pass_required_order.py


Signed-off-by: Adrian Lundell <[email protected]>
Signed-off-by: Martin Lindström <[email protected]>
Co-authored-by: Adrian Lundell <[email protected]>
Co-authored-by: Martin Lindström <[email protected]>

Author: 3 people

backends/arm/_passes/add_bias_pass.py

@@ -3,13 +3,15 @@
 # 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
+
 import torch
 from executorch.backends.arm._passes import ArmPass
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.transforms.utils import create_constant_placeholder
 
 from executorch.exir.dialects._ops import ops as exir_ops
-from executorch.exir.pass_base import PassResult
+from executorch.exir.pass_base import ExportPass, PassResult
 from torch.export.graph_signature import InputKind
 
 
@@ -19,6 +21,8 @@ class AddBiasPass(ArmPass):
     The bias is set to zero.
     """
 
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
     targeted_ops = (exir_ops.edge.aten.convolution.default,)
 
     def call(self, graph_module):

backends/arm/_passes/annotate_decomposed_matmul.py

@@ -7,7 +7,7 @@
 
 import itertools
 import operator
-from typing import cast, List
+from typing import cast, List, Set, Type
 
 import torch
 from executorch.backends.arm._passes.arm_pass_utils import create_node
@@ -29,6 +29,8 @@ class AnnotateDecomposedMatmulPass(ExportPass):
     matmul-op (can be mm or bmm).
     """
 
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
     def _match_partition_to_node(
         self, node: torch.fx.Node, partitioned_inputs: List[torch.fx.Node]
     ) -> torch.fx.Node:

backends/arm/_passes/annotate_output_dim_order_pass.py

@@ -3,9 +3,12 @@
 # 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 import ArmPass
 from executorch.backends.arm._passes.arm_pass_utils import get_output_dim_orders
-from executorch.exir.pass_base import PassResult
+from executorch.exir.pass_base import ExportPass, PassResult
 
 
 class AnnotateOutputDimOrderPass(ArmPass):
@@ -14,6 +17,8 @@ class AnnotateOutputDimOrderPass(ArmPass):
     for verifying that the dim order does not change unexpectedly in later passes.
     """
 
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
     def call(self, graph_module):
         output_node = graph_module.graph.output_node()
         output_node.meta["original_dim_orders"] = get_output_dim_orders(graph_module)

backends/arm/_passes/arm_pass.py

@@ -6,7 +6,8 @@
 # pyre-unsafe
 
 import traceback
-from typing import Optional
+from abc import abstractmethod
+from typing import List, Optional, Set, Type
 
 import torch
 from executorch.exir.pass_base import ExportPass, NodeMetadata
@@ -19,6 +20,36 @@ def __init__(self, exported_program: Optional[torch.export.ExportedProgram] = No
         super(ArmPass, self).__init__()
         self.exported_program = exported_program
 
+    @property
+    @abstractmethod
+    def _passes_required_after(self) -> Set[Type[ExportPass]]:
+        """The subclass defines passes that must run after it"""
+        pass
+
+    @staticmethod
+    def get_required_passes(pass_) -> List[str]:
+        """
+        Returns the list of passes that must be run after this pass, sorted by name.
+        """
+        if hasattr(pass_, "_passes_required_after"):
+            return sorted([ArmPass.get_name(p) for p in pass_._passes_required_after])
+        else:
+            return []
+
+    @staticmethod
+    def get_name(pass_) -> str:
+        """
+        Returns the name of the pass.
+        """
+        if isinstance(pass_, ExportPass):
+            return pass_.__class__.__name__
+        elif hasattr(pass_, "__name__"):
+            return pass_.__name__
+        else:
+            raise ValueError(
+                f"Cannot get name for pass: {pass_}. It must be an instance of ExportPass or have a __name__ attribute."
+            )
+
     def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False):
         if not updated:
             return super().call_operator(op, args, kwargs, meta)

backends/arm/_passes/arm_pass_manager.py

@@ -7,6 +7,9 @@
 
 # pyre-unsafe
 
+
+from collections import defaultdict
+
 import executorch.backends.arm.tosa.dialect  # noqa: unused
 from executorch.backends.arm._passes import (
     AddBiasPass,
@@ -94,6 +97,7 @@
     UnsqueezeScalarPlaceholdersPass,
 )
 
+from executorch.backends.arm._passes.arm_pass import ArmPass
 from executorch.backends.arm.tosa.specification import (
     TosaLoweringContext,
     TosaSpecification,
@@ -115,6 +119,32 @@ def __init__(self, tosa_spec: TosaSpecification) -> None:
         self.tosa_spec = tosa_spec
         super().__init__()
 
+    def validate_constraints_mandatory(self):
+        """
+        Validates that necessary passes have run before transforming to backend.
+
+        Note that this differs from the original validate_constraints function, which
+        only checks the order of passes.
+        """
+        passes_to_run = defaultdict(list)
+
+        for current_pass in self.passes:
+            current_pass_name = ArmPass.get_name(current_pass)
+            for required_pass_name in ArmPass.get_required_passes(current_pass):
+                passes_to_run[required_pass_name].append(current_pass_name)
+
+            passes_to_run.pop(current_pass_name, None)
+
+        if len(passes_to_run) > 0:
+            error_msg = "The following constraints for passes are not met:\n"
+            for required_pass, requiring_passes in passes_to_run.items():
+                for requiring_pass in requiring_passes:
+                    error_msg += (
+                        f"  - {required_pass} must run after {requiring_pass}\n"
+                    )
+
+            raise RuntimeError(error_msg)
+
     def _transform(self, graph_module: GraphModule):
         with TosaLoweringContext(self.tosa_spec):
             return self(graph_module).graph_module
@@ -125,7 +155,6 @@ def _tosa_INT_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(RemoveGetItemPass())
         self.add_pass(ConvertSplitToSlicePass())
         self.add_pass(ConvertMmToBmmPass())
-        self.add_pass(DecomposeLinearVectorNormPass())
         self.add_pass(
             DecomposeMeanDimPass(exported_program.graph_module, self.tosa_spec)
         )
@@ -175,6 +204,7 @@ def _tosa_INT_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(RemoveNoopPass())
         self.add_pass(InsertRescalePass())
 
+        self.validate_constraints_mandatory()
         return self._transform(exported_program.graph_module)
 
     def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
@@ -258,6 +288,7 @@ def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(RemoveNoopPass())
         self.add_pass(InsertRescalePass())
 
+        self.validate_constraints_mandatory()
         return self._transform(exported_program.graph_module)
 
     def transform_to_backend_pipeline(self, exported_program: ExportedProgram):

backends/arm/_passes/broadcast_args_pass.py

@@ -3,6 +3,8 @@
 # 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 import ArmPass
 
 from executorch.backends.arm._passes.arm_pass_utils import (
@@ -12,7 +14,7 @@
 
 from executorch.exir.dialects._ops import ops as exir_ops
 
-from executorch.exir.pass_base import PassResult
+from executorch.exir.pass_base import ExportPass, PassResult
 from torch.fx import GraphModule, Node
 
 
@@ -22,6 +24,8 @@ class BroadcastArgsPass(ArmPass):
     This is done when more than one arg needs broadcasting.
     """
 
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
     targeted_ops = {
         exir_ops.edge.aten.add.Tensor,
         exir_ops.edge.aten.sub.Tensor,

backends/arm/_passes/cast_bool_to_int8_pass.py

@@ -6,6 +6,8 @@
 # The TOSA BITWISE_AND, BITWISE_OR, and BITWISE_XOR don't handle bool as input
 # If input/output is bool lest add a cast/conversion pass before/after to/from int8.
 
+from typing import Set, Type
+
 import torch
 
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -15,6 +17,8 @@
 class CastBoolToInt8Pass(ExportPass):
     """Casts the input to int8 if it is not already and casts back the output to the original input dtype."""
 
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
     targeted_ops = {
         exir_ops.edge.aten.bitwise_and.Tensor,
         exir_ops.edge.aten.bitwise_or.Tensor,

backends/arm/_passes/cast_int64_pass.py

@@ -6,6 +6,7 @@
 # pyre-unsafe
 
 import logging
+from typing import Set, Type
 
 import torch
 from executorch.exir.pass_base import ExportPass, PassResult
@@ -19,6 +20,8 @@ class CastInt64BuffersToInt32Pass(ExportPass):
     Cast int64 buffers to int32 if the int64 data is in int32 range.
     """
 
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
     def __init__(self, exported_program: torch.export.ExportedProgram):
         super(CastInt64BuffersToInt32Pass, self).__init__()
         self.exported_program = exported_program

backends/arm/_passes/cast_to_int32_pass.py

@@ -3,6 +3,8 @@
 # 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
+
 import torch
 
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -12,6 +14,8 @@
 class CastToInt32Pass(ExportPass):
     """Casts the input to int32 if it is not already and casts back the output to the original input dtype."""
 
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
     targeted_ops = {
         exir_ops.edge.aten.bitwise_left_shift.Tensor,
         exir_ops.edge.aten.bitwise_right_shift.Tensor,

backends/arm/_passes/conv1d_unsqueeze_pass.py

@@ -6,6 +6,8 @@
 # LICENSE file in the root directory of this source tree.
 
 
+from typing import Set, Type
+
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
@@ -21,6 +23,8 @@ class Conv1dUnsqueezePass(ExportPass):
     3) squeeze the output back down to 3d.
     """
 
+    _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)