Arm backend: Add pre-decomposition partitioner pipeline (pytorch#22514)

Route the EXIR pre-decomposition hook through an Arm pass pipeline.
Store compile specs on concrete partitioners so the pipeline can be
configured consistently for TOSA, Ethos-U, and VGF.

Keep the pipeline a no-op until backend-specific passes are registered,
and update generated partitioner docs and public API manifests.

Assisted by Codex.

Change-Id: I1963bad577714907e5ba22eed765333d9816308c
Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>

Author: YufengShi-dudu

backends/arm/_passes/arm_pass_manager.py

@@ -458,6 +458,12 @@ def _tosa_context(self, graph_module: GraphModule) -> TosaLoweringContext:
         shape_env = _get_shape_env_from_gm(graph_module)
         return TosaLoweringContext(self.tosa_spec, shape_env)
 
+    def transform_for_pre_decomposition_pipeline(
+        self, exported_program: ExportedProgram
+    ) -> ExportedProgram:
+        """Apply Arm passes before default ATen decompositions."""
+        return exported_program
+
     def _transform_graph_module(self, graph_module: GraphModule):
         # TFA and control-flow submodule paths operate on bare GraphModules
         # without a standalone ExportedProgram to keep in sync.

backends/arm/ethosu/partitioner.py

@@ -34,6 +34,7 @@ def __init__(
         self.delegation_spec = DelegationSpec(
             EthosUBackend.__name__, compile_spec._to_list()
         )
+        self.compile_spec = compile_spec
         self.additional_checks = additional_checks
         self.tosa_spec = compile_spec.tosa_spec
         self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec)

backends/arm/public_api_manifests/api_manifest_running.toml

@@ -60,6 +60,10 @@ signature = "EthosUPartitioner.partition(self, exported_program: torch.export.ex
 kind = "function"
 signature = "EthosUPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None"
 
+[python.EthosUPartitioner.transform_for_pre_decomposition]
+kind = "function"
+signature = "EthosUPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram"
+
 [python.EthosUQuantizer]
 kind = "class"
 signature = "EthosUQuantizer(compile_spec: 'EthosUCompileSpec', use_composable_quantizer: 'bool' = True) -> 'None'"
@@ -180,6 +184,10 @@ signature = "VgfPartitioner.partition(self, exported_program: torch.export.expor
 kind = "function"
 signature = "VgfPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None"
 
+[python.VgfPartitioner.transform_for_pre_decomposition]
+kind = "function"
+signature = "VgfPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram"
+
 [python.VgfQuantizer]
 kind = "class"
 signature = "VgfQuantizer(compile_spec: 'VgfCompileSpec', use_composable_quantizer: 'bool' = True) -> 'None'"

backends/arm/tosa/partitioner.py

@@ -21,6 +21,7 @@
 from typing import Callable, cast, List, Mapping, Optional, Sequence, Tuple
 
 import torch
+from executorch.backends.arm._passes.arm_pass_manager import ArmPassManager
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.arm._passes.convert_expand_copy_to_repeat import (
     calculate_multiples,
@@ -32,6 +33,7 @@
     is_exact_tosa_boundary_bilinear_downscale,
 )
 
+from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec
 from executorch.backends.arm.common.type import ensure_type
 from executorch.backends.arm.constants import DQ_OPS, Q_OPS
 from executorch.backends.arm.operator_support.tosa_supported_operators import (
@@ -361,6 +363,8 @@ class TOSAPartitioner(Partitioner):
 
     """
 
+    compile_spec: ArmCompileSpec
+
     def __init__(
         self,
         compile_spec: TosaCompileSpec,
@@ -381,12 +385,34 @@ def __init__(
         self.delegation_spec = DelegationSpec(
             TOSABackend.__name__, compile_spec._to_list()
         )
+        self.compile_spec = compile_spec
         self.tosa_spec = compile_spec.tosa_spec
         self.additional_checks = additional_checks
         self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec)
         self._custom_partition_ops: set[torch._ops.OpOverload] = set()
         self.intermediate_path = compile_spec._get_intermediate_path()
 
+    def transform_for_pre_decomposition(
+        self, exported_program: ExportedProgram
+    ) -> ExportedProgram:
+        """Apply required Arm passes before default ATen decompositions.
+
+        EXIR invokes this backend extension hook automatically through
+        ``to_edge_transform_and_lower``. Model export users should not call it
+        directly.
+
+        Args:
+            exported_program (ExportedProgram): The ATen-dialect program to
+                transform.
+
+        Returns:
+            ExportedProgram: The transformed ATen-dialect program.
+
+        """
+        return ArmPassManager(
+            self.compile_spec
+        ).transform_for_pre_decomposition_pipeline(exported_program)
+
     def register_custom_partition_op(self, op: torch._ops.OpOverload) -> None:
         """Register a custom op to be considered supported."""
         self._custom_partition_ops.add(op)

backends/arm/vgf/partitioner.py

@@ -35,6 +35,7 @@ def __init__(
         self.delegation_spec = DelegationSpec(
             VgfBackend.__name__, compile_spec._to_list()
         )
+        self.compile_spec = compile_spec
         self.additional_checks = additional_checks
         self.tosa_spec = compile_spec.tosa_spec
         self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec)

docs/source/backends/arm-ethos-u/arm-ethos-u-partitioner.md

@@ -50,3 +50,19 @@ Returns:
 def EthosUPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None:
 ```
 Register a custom op to be considered supported.
+
+```python
+def EthosUPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram:
+```
+Apply required Arm passes before default ATen decompositions.
+
+EXIR invokes this backend extension hook automatically through
+``to_edge_transform_and_lower``. Model export users should not call it
+directly.
+
+Args:
+- **exported_program (ExportedProgram)**: The ATen-dialect program to
+        transform.
+
+Returns:
+- **ExportedProgram**: The transformed ATen-dialect program.

docs/source/backends/arm-vgf/arm-vgf-partitioner.md

@@ -50,3 +50,19 @@ Returns:
 def VgfPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None:
 ```
 Register a custom op to be considered supported.
+
+```python
+def VgfPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram:
+```
+Apply required Arm passes before default ATen decompositions.
+
+EXIR invokes this backend extension hook automatically through
+``to_edge_transform_and_lower``. Model export users should not call it
+directly.
+
+Args:
+- **exported_program (ExportedProgram)**: The ATen-dialect program to
+        transform.
+
+Returns:
+- **ExportedProgram**: The transformed ATen-dialect program.