Add ability to run passes after edge in LoweringRecipe

This is neccessary for passes that want to handle nodes not
handled by a backend. For example, the cortex_m ReplaceQuantNodesPass
wants to only replace quant nodes that have not been delegated.

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

Author: Erik-Lundell

export/recipe.py

@@ -129,6 +129,7 @@ class LoweringRecipe:
         None | List[Callable[[str, ExportedProgram], List[PassType]]]
     ) = None
     # pyre-ignore[11]: Type not defined
+    post_edge_passes: list[PassType] | None = None
     edge_compile_config: Optional[EdgeCompileConfig] = None
 
 

export/stages.py

@@ -177,10 +177,12 @@ def __init__(
         transform_passes: (
             None | List[Callable[[str, ExportedProgram], List[PassType]]]
         ) = None,
+        post_edge_passes: list[PassType] | None = None,
         compile_config: Optional[Any] = None,
     ) -> None:
         self._partitioners = partitioners
         self._transform_passes = transform_passes
+        self._post_edge_passes = post_edge_passes
         self._compile_config = compile_config
 
     @classmethod
@@ -194,6 +196,7 @@ def from_recipe(
             partitioners=lowering_recipe.partitioners,
             transform_passes=lowering_recipe.edge_transform_passes,
             compile_config=lowering_recipe.edge_compile_config,
+            post_edge_passes=lowering_recipe.post_edge_passes,
         )
 
     @property
@@ -242,6 +245,10 @@ def run(self, artifact: PipelineArtifact) -> None:
                 compile_config=self._compile_config,
                 generate_etrecord=generate_etrecord,
             )
+            if self._post_edge_passes:
+                edge_program_manager = edge_program_manager.transform(
+                    self._post_edge_passes
+                )
 
         delegation_info = get_delegation_info(
             edge_program_manager.exported_program().graph_module
@@ -502,10 +509,12 @@ def __init__(
         transform_passes: (
             None | List[Callable[[str, ExportedProgram], List[PassType]]]
         ) = None,
+        post_edge_passes: list[PassType] | None = None,
     ) -> None:
         super().__init__()
         self._partitioners = partitioners
         self._transform_passes = transform_passes
+        self._post_edge_passes = post_edge_passes
 
     @classmethod
     def from_recipe(
@@ -517,6 +526,7 @@ def from_recipe(
         return cls(
             partitioners=lowering_recipe.partitioners,
             transform_passes=lowering_recipe.edge_transform_passes,
+            post_edge_passes=lowering_recipe.post_edge_passes,
         )
 
     @property
@@ -569,6 +579,9 @@ def run(self, artifact: PipelineArtifact) -> None:
                 for partitioner in self._partitioners:
                     edge_program_manager = edge_program_manager.to_backend(partitioner)
 
+                if self._post_edge_passes:
+                    edge_program_manager.transform(self._post_edge_passes)
+
         # Get delegation info
         delegation_info = get_delegation_info(
             edge_program_manager.exported_program().graph_module