Arm backend: Convert ExportPasses to ArmPasses

All passes in backends/arm/_passes have been set to inherit from
`ArmPass` instead of `ExportPass`.

Furthermore, the attribute `exported_program` in `ArmPass` has been
removed to instead let each subclasses set it. The reason is that it was
type annotated as an `Optional[ExportedProgram]]` in `ArmPass` and Mypy
complained about that when a subclass tried to access it without
checking for `None`. By moving the attribute down to each subclass that
needs it, the confusion around whether the value is `None` or not is
elimininated.

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

Author: Martin Lindström

backends/arm/_passes/_debug_passes.py

@@ -4,12 +4,13 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
+from executorch.backends.arm._passes import ArmPass
 from executorch.devtools.visualization.visualization_utils import visualize_graph
 from executorch.exir import ExportedProgram
-from executorch.exir.pass_base import ExportPass, PassResult
+from executorch.exir.pass_base import PassResult
 
 
-class VisualizePass(ExportPass):
+class VisualizePass(ArmPass):
     """
     This pass visualizes the graph at the point of insertion in the pass manager
     """

backends/arm/_passes/add_bias_pass.py

@@ -10,6 +10,7 @@
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.arm.tosa.mapping import TosaSpecialDtype
 from executorch.backends.transforms.utils import create_constant_placeholder
+from executorch.exir import ExportedProgram
 
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
@@ -26,6 +27,10 @@ class AddBiasPass(ArmPass):
 
     targeted_ops = (exir_ops.edge.aten.convolution.default,)
 
+    def __init__(self, exported_program: ExportedProgram) -> None:
+        super().__init__()
+        self.exported_program = exported_program
+
     def call(self, graph_module):
         modified = False
         for node in graph_module.graph.nodes:

backends/arm/_passes/annotate_decomposed_matmul.py

@@ -10,6 +10,7 @@
 from typing import cast, List, Set, Type
 
 import torch
+from executorch.backends.arm._passes.arm_pass import ArmPass
 from executorch.backends.arm._passes.arm_pass_utils import create_node
 from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
     FoldAndAnnotateQParamsPass,
@@ -23,7 +24,7 @@
 from torch.fx.passes.utils.source_matcher_utils import get_source_partitions
 
 
-class AnnotateDecomposedMatmulPass(ExportPass):
+class AnnotateDecomposedMatmulPass(ArmPass):
     """
     torch.matmul and it's equivalent operator @ can be decomposed in many ways, for instance:
     dq -> matmul -> q can become

backends/arm/_passes/arm_pass.py

@@ -9,17 +9,12 @@
 from abc import abstractmethod
 from typing import List, Optional, Set, Type
 
-import torch
 from executorch.exir.pass_base import ExportPass, NodeMetadata
 
 
 class ArmPass(ExportPass):
     """Base class for Arm passes"""
 
-    def __init__(self, exported_program: Optional[torch.export.ExportedProgram] = None):
-        super(ArmPass, self).__init__()
-        self.exported_program = exported_program
-
     @property
     @abstractmethod
     def _passes_required_after(self) -> Set[Type[ExportPass]]:

backends/arm/_passes/arm_pass_manager.py

@@ -207,7 +207,7 @@ def _tosa_INT_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         # needs to happen before AddBiasPass, but after the table ops are inserted
         # to be able to validate that conv2d has right dtype arguments.
         self.add_pass(DecomposeConv2dWithInt16ActivationPass())
-        self.add_pass(RewriteUpsamplePass(exported_program))
+        self.add_pass(RewriteUpsamplePass())
         self.add_pass(AddBiasPass(exported_program))
 
         self.add_pass(FuseEqualPlaceholdersPass(exported_program))
@@ -292,7 +292,7 @@ def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(FuseViewCopyTransform())
         self.add_pass(FuseConstantArgsPass(exported_program))
         self.add_pass(CastInt64BuffersToInt32Pass(exported_program))
-        self.add_pass(RewriteUpsamplePass(exported_program))
+        self.add_pass(RewriteUpsamplePass())
         self.add_pass(AddBiasPass(exported_program))
         self.add_pass(InsertTableOpsPass(exported_program))
         self.add_pass(FuseEqualPlaceholdersPass(exported_program))

backends/arm/_passes/cast_bool_to_int8_pass.py

@@ -10,11 +10,12 @@
 
 import torch
 
+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
 
 
-class CastBoolToInt8Pass(ExportPass):
+class CastBoolToInt8Pass(ArmPass):
     """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()

backends/arm/_passes/cast_int64_pass.py

@@ -9,21 +9,23 @@
 from typing import Set, Type
 
 import torch
+from executorch.backends.arm._passes.arm_pass import ArmPass
 from executorch.exir.pass_base import ExportPass, PassResult
 from torch._export.utils import is_buffer
+from torch.export import ExportedProgram
 
 logger = logging.getLogger(__name__)
 
 
-class CastInt64BuffersToInt32Pass(ExportPass):
+class CastInt64BuffersToInt32Pass(ArmPass):
     """
     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__()
+    def __init__(self, exported_program: ExportedProgram):
+        super().__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

@@ -7,11 +7,12 @@
 
 import torch
 
+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
 
 
-class CastToInt32Pass(ExportPass):
+class CastToInt32Pass(ArmPass):
     """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()

backends/arm/_passes/conv1d_unsqueeze_pass.py

@@ -8,14 +8,16 @@
 
 from typing import Set, Type
 
+from executorch.backends.arm._passes import ArmPass
+
 from executorch.backends.arm._passes.add_bias_pass import AddBiasPass
 from executorch.backends.arm._passes.size_adjust_input_pass import SizeAdjustInputPass
 
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
 
-class Conv1dUnsqueezePass(ExportPass):
+class Conv1dUnsqueezePass(ArmPass):
     """
     This pass is used to change conv1d ops into conv2d since TOSA only
     supports 2d and 3d convolution. This is done by modifying the graph to do the
@@ -38,7 +40,11 @@ def call_operator(self, op, args, kwargs, meta):
         x = args[0]
         x_unsqueezed_shape = list(x.data.shape) + [1]
         x = super().call_operator(
-            exir_ops.edge.aten.view_copy.default, (x, x_unsqueezed_shape), {}, meta
+            exir_ops.edge.aten.view_copy.default,
+            (x, x_unsqueezed_shape),
+            {},
+            meta,
+            updated=True,
         )
 
         w_meta = meta.copy()
@@ -48,7 +54,11 @@ def call_operator(self, op, args, kwargs, meta):
         w = args[1]
         w_unsqueezed_shape = list(w.data.shape) + [1]
         w = super().call_operator(
-            exir_ops.edge.aten.view_copy.default, (w, w_unsqueezed_shape), {}, w_meta
+            exir_ops.edge.aten.view_copy.default,
+            (w, w_unsqueezed_shape),
+            {},
+            w_meta,
+            updated=True,
         )
 
         new_args = (
@@ -63,12 +73,16 @@ def call_operator(self, op, args, kwargs, meta):
             args[8],
         )
         x = super().call_operator(
-            exir_ops.edge.aten.convolution.default, new_args, kwargs, meta
+            exir_ops.edge.aten.convolution.default, new_args, kwargs, meta, updated=True
         )
 
         x_squeezed_shape = list(x.data.shape)[:-1]
         x = super().call_operator(
-            exir_ops.edge.aten.view_copy.default, (x, x_squeezed_shape), {}, meta
+            exir_ops.edge.aten.view_copy.default,
+            (x, x_squeezed_shape),
+            {},
+            meta,
+            updated=True,
         )
 
         return x

backends/arm/_passes/convert_any_default_dim_dims_pass.py

@@ -6,6 +6,7 @@
 from typing import Set, Type
 
 import torch
+from executorch.backends.arm._passes import ArmPass
 from executorch.backends.arm._passes.convert_squeezes_to_view import (
     ConvertSqueezesToViewPass,
 )
@@ -18,7 +19,7 @@
 )
 
 
-class ConvertAnyDefaultDimDimsPass(ExportPass):
+class ConvertAnyDefaultDimDimsPass(ArmPass):
     """
     Converts any.default, any.dim and any.dims to a sequence of any.dim by unrolling multi-dimensional reduction.
     Please refer to KeepDimsFalseToSqueezePass for an explanation of this coversion.