Add PostToBackend stage to Recipes

After a discussion in #14588, it was decided to
create an additional recipe stage to run passes after partitioning.
This is meant for backends that convert ops directly instead of
partitioning.

Signed-off-by: Erik Lundell <[email protected]>
Change-Id: Iddf74accb739d4dff16fa46c6fad88ffccfe2f3b

Author: Erik-Lundell

export/export.py

@@ -182,6 +182,7 @@ def _get_default_pipeline(self) -> List[StageType]:
             StageType.QUANTIZE,  # Optional stage, returns original model if quant recipe is invalid
             StageType.TORCH_EXPORT,
             StageType.TO_EDGE_TRANSFORM_AND_LOWER,
+            StageType.POST_TO_BACKEND,
             StageType.TO_EXECUTORCH,
         ]
 

export/recipe.py

@@ -17,6 +17,7 @@
 
 from executorch.exir.backend.partitioner import Partitioner
 from executorch.exir.capture import EdgeCompileConfig, ExecutorchBackendConfig
+from executorch.exir.pass_base import ExportPass
 from executorch.exir.pass_manager import PassType
 from torchao.core.config import AOBaseConfig
 from torchao.quantization.pt2e.quantizer import Quantizer
@@ -122,6 +123,7 @@ class LoweringRecipe:
         edge_transform_passes: Optional list of callables that take (method_name: str, exported_program: ExportedProgram) as arguments
                                and return a list of passes (PassType) to be executed during lowering stages.
         edge_compile_config: Optional edge compilation configuration
+        post_to_backend_passes: Optional list of passes to run after all partitioners have ran.
     """
 
     partitioners: Optional[List[Partitioner]] = None
@@ -130,6 +132,7 @@ class LoweringRecipe:
     ) = None
     # pyre-ignore[11]: Type not defined
     edge_compile_config: Optional[EdgeCompileConfig] = None
+    post_to_backend_passes: list[PassType | type[ExportPass]] | None = None
 
 
 @experimental(

export/stages.py

@@ -9,12 +9,14 @@
 import logging
 from abc import ABC, abstractmethod
 from collections import defaultdict
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, cast, Dict, List, Optional
 
 import torch
 from executorch.devtools.backend_debug import get_delegation_info
-from executorch.exir import EdgeCompileConfig, ExportedProgram
+from executorch.exir import EdgeCompileConfig, EdgeProgramManager, ExportedProgram
 from executorch.exir.backend.backend_api import validation_disabled
+from executorch.exir.pass_base import ExportPass, RequireExportedProgram
+from executorch.exir.pass_manager import PassManager
 from executorch.exir.program import to_edge, to_edge_transform_and_lower
 from executorch.export.recipe import LoweringRecipe, QuantizationRecipe
 from executorch.export.types import StageType
@@ -118,7 +120,7 @@ def __init__(
         self.strict = strict
 
     @property
-    def stage_type(self) -> str:
+    def stage_type(self) -> StageType:
         return StageType.TORCH_EXPORT
 
     @property
@@ -197,7 +199,7 @@ def from_recipe(
         )
 
     @property
-    def stage_type(self) -> str:
+    def stage_type(self) -> StageType:
         return StageType.TO_EDGE_TRANSFORM_AND_LOWER
 
     @property
@@ -266,7 +268,7 @@ def __init__(self, backend_config: Any) -> None:
         self._backend_config = backend_config
 
     @property
-    def stage_type(self) -> str:
+    def stage_type(self) -> StageType:
         return StageType.TO_EXECUTORCH
 
     @property
@@ -304,7 +306,7 @@ def __init__(self, quantization_recipe: Optional[QuantizationRecipe]) -> None:
         self._transformed_models: Dict[str, nn.Module] = {}
 
     @property
-    def stage_type(self) -> str:
+    def stage_type(self) -> StageType:
         return StageType.SOURCE_TRANSFORM
 
     @property
@@ -358,7 +360,7 @@ def __init__(self, quantization_recipe: Optional[QuantizationRecipe]) -> None:
         self._quantization_recipe = quantization_recipe
 
     @property
-    def stage_type(self) -> str:
+    def stage_type(self) -> StageType:
         return StageType.QUANTIZE
 
     @property
@@ -459,7 +461,7 @@ def from_recipe(cls, lowering_recipe: Optional["LoweringRecipe"]) -> "ToEdgeStag
         )
 
     @property
-    def stage_type(self) -> str:
+    def stage_type(self) -> StageType:
         return StageType.TO_EDGE
 
     @property
@@ -520,7 +522,7 @@ def from_recipe(
         )
 
     @property
-    def stage_type(self) -> str:
+    def stage_type(self) -> StageType:
         return StageType.TO_BACKEND
 
     @property
@@ -583,3 +585,102 @@ def delegation_info(self) -> Any:
         Returns the delegation info.
         """
         return self._artifact.get_context("delegation_info")
+
+
+class PostToBackendStage(Stage):
+    """
+    Stage: Run passes after all partitioners have done their partitioning.
+    """
+
+    def __init__(
+        self,
+        pass_list_or_manager: (
+            list[PassType | type[ExportPass]] | PassManager | None
+        ) = None,
+        edge_compile_config: EdgeCompileConfig | None = None,
+    ) -> None:
+        super().__init__()
+        if pass_list_or_manager is None:
+            pass_list_or_manager = []
+
+        self._pass_list_or_manager = pass_list_or_manager
+        self._edge_compile_config = edge_compile_config
+
+    @classmethod
+    def from_recipe(
+        cls, lowering_recipe: Optional["LoweringRecipe"]
+    ) -> "PostToBackendStage":
+        if lowering_recipe is None:
+            return cls()
+
+        return cls(
+            pass_list=lowering_recipe.post_to_backend_passes,
+            edge_compile_config=lowering_recipe.edge_compile_config,
+        )
+
+    @property
+    def stage_type(self) -> StageType:
+        return StageType.POST_TO_BACKEND
+
+    @property
+    def valid_predecessor_stages(self) -> List["StageType"]:
+        return [StageType.TO_BACKEND, StageType.TO_EDGE_TRANSFORM_AND_LOWER]
+
+    @property
+    def can_start_pipeline(self) -> bool:
+        return False
+
+    def run(self, artifact: PipelineArtifact) -> None:
+        """
+        Run list of passes using edge_program_manager.transform().
+
+        Args:
+            artifact: PipelineArtifact which's data field is expected to contain an edge_program_manager.
+        """
+
+        if self._pass_list_or_manager:
+            edge_program_manager = cast(EdgeProgramManager, artifact.data)
+
+            if isinstance(self._pass_list_or_manager, PassManager):
+                edge_program_manager = edge_program_manager.transform(
+                    self._pass_list_or_manager, self._edge_compile_config
+                )
+            else:
+                exported_program = edge_program_manager.exported_program()
+                pass_instances: list[PassType] = []
+                for _pass in self._pass_list_or_manager:
+                    if isinstance(_pass, type):
+                        if not issubclass(_pass, ExportPass):
+                            raise RuntimeError(
+                                f"Pass {_pass} was not subclass of ExportPass."
+                            )
+                        if issubclass(_pass, RequireExportedProgram):
+                            pass_instance = _pass(
+                                exported_program=exported_program  #  type: ignore
+                            )
+                        else:
+                            pass_instance = _pass()
+                        pass_instances.append(pass_instance)
+                    else:
+                        pass_instances.append(_pass)
+
+                edge_program_manager = edge_program_manager.transform(
+                    pass_instances, self._edge_compile_config
+                )
+            # Get delegation info
+            delegation_info = get_delegation_info(
+                edge_program_manager.exported_program().graph_module
+            )
+
+            self._artifact = artifact.copy_with_new_data(edge_program_manager)
+            self._artifact.add_context("delegation_info", delegation_info)
+        else:
+            # If pass_list_or_manager is None or empty list, do nothing.
+            self._artifact = artifact
+
+    @property
+    def delegation_info(self) -> Any:
+        """
+        Returns the delegation info.
+        """
+        return self._artifact.get_context("delegation_info")

export/tests/test_export_stages.py

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2025 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -10,12 +11,14 @@
 from unittest.mock import Mock, patch, PropertyMock
 
 import torch
+from executorch.exir.pass_base import ExportPass, RequireExportedProgram
 from executorch.exir.program import EdgeProgramManager, ExecutorchProgramManager
 from executorch.export import AOQuantizationConfig, QuantizationRecipe, StageType
 from executorch.export.stages import (
     EdgeTransformAndLowerStage,
     ExecutorchStage,
     PipelineArtifact,
+    PostToBackendStage,
     QuantizeStage,
     SourceTransformStage,
     ToBackendStage,
@@ -35,6 +38,18 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.linear(x)
 
 
+class DummyExportPassWithProgram(RequireExportedProgram, ExportPass):
+    def __init__(self, exported_program: ExportedProgram) -> None:
+        super().__init__()
+        self.received_exported_program = exported_program
+
+
+class DummyExportPassWithoutProgram(ExportPass):
+    def __init__(self) -> None:
+        super().__init__()
+        self.initialized = True
+
+
 class TestPipelineArtifact(unittest.TestCase):
 
     def test_copy_with_new_data(self) -> None:
@@ -587,3 +602,51 @@ def test_run_edge_manager_none(self) -> None:
         with self.assertRaises(RuntimeError) as cm:
             stage.run(artifact)
         self.assertIn("Edge program manager is not set", str(cm.exception))
+
+
+class TestPostToBackendStage(unittest.TestCase):
+    @patch("executorch.export.stages.get_delegation_info")
+    def test_run_with_mixed_pass_types(self, mock_get_delegation_info: Mock) -> None:
+        mock_get_delegation_info.return_value = {"delegation": "info"}
+
+        exported_program = Mock(spec=ExportedProgram)
+        edge_program_manager = Mock(spec=EdgeProgramManager)
+        transformed_manager = Mock(spec=EdgeProgramManager)
+        transformed_exported_program = Mock(spec=ExportedProgram)
+        transformed_graph_module = Mock()
+        transformed_exported_program.graph_module = transformed_graph_module
+
+        edge_program_manager.exported_program.return_value = exported_program
+        edge_program_manager.transform.return_value = transformed_manager
+        transformed_manager.exported_program.return_value = transformed_exported_program
+
+        passthrough_pass = Mock()
+        non_program_pass_instance = DummyExportPassWithoutProgram()
+
+        stage = PostToBackendStage(
+            pass_list_or_manager=[
+                passthrough_pass,
+                DummyExportPassWithProgram,
+                DummyExportPassWithoutProgram,
+                non_program_pass_instance,
+            ]
+        )
+
+        artifact = PipelineArtifact(data=edge_program_manager, context={})
+        stage.run(artifact)
+
+        edge_program_manager.transform.assert_called_once()
+        pass_instances, compile_config = edge_program_manager.transform.call_args[0]
+
+        self.assertEqual(len(pass_instances), 4)
+        self.assertIs(pass_instances[0], passthrough_pass)
+        self.assertIsInstance(pass_instances[1], DummyExportPassWithProgram)
+        self.assertIsInstance(pass_instances[2], DummyExportPassWithoutProgram)
+        self.assertIs(pass_instances[1].received_exported_program, exported_program)
+        self.assertIs(pass_instances[3], non_program_pass_instance)
+        self.assertIsNone(compile_config)
+
+        result_artifact = stage.get_artifacts()
+        self.assertIs(result_artifact.data, transformed_manager)
+        self.assertEqual(stage.delegation_info, {"delegation": "info"})
+        mock_get_delegation_info.assert_called_once_with(transformed_graph_module)

export/types.py

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2025 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -18,4 +19,5 @@ class StageType(str, Enum):
     TO_EDGE_TRANSFORM_AND_LOWER = "to_edge_transform_and_lower"
     TO_EDGE = "to_edge"
     TO_BACKEND = "to_backend"
+    POST_TO_BACKEND = "post_to_backend"
     TO_EXECUTORCH = "to_executorch"