Revert pytorch#23005: Prevent edge_transform_passes from running twice (pytorch#23119)

Reverts pytorch#23005.

The change in pytorch#23005 moved edge_transform_passes out of
EdgeProgramManagerTransformStage, but this causes recipe-based lowering
to lose the stage behavior expected by some existing callers. Restore
the stage wiring and its regression test.

Test plan:
- `git diff --check upstream/main...HEAD`
- Not run: `python -m pytest export/tests/test_export_session.py`
(pytest is not installed in the current environment).

Author: JakeStevens

export/recipe.py

@@ -207,12 +207,9 @@ class LoweringRecipe:
                       per-method partitioner lists. Use the dict form when
                       backends need per-method compile specs.
         edge_transform_passes: Optional list of callables that take (method_name: str, exported_program: ExportedProgram)
-                               and return either List[PassType] or PassManager. Applied during
-                               TO_EDGE_TRANSFORM_AND_LOWER as per-ExportedProgram graph-module passes.
+                               and return either List[PassType] or PassManager to be applied during edge lowering.
         edge_manager_transform_passes: Optional list of callables that take EdgeProgramManager as argument
-                                        and return passes to be applied. Applied sequentially by
-                                        EDGE_PROGRAM_MANAGER_TRANSFORM, which runs after
-                                        TO_EDGE_TRANSFORM_AND_LOWER in the default pipeline.
+                                        and return passes to be applied. Applied sequentially after TO_EDGE stage.
         edge_compile_config: Optional edge compilation configuration
         pre_partitioning_callback: Optional callable invoked just before partitioning with
                                    `(partitioners, programs)` arguments.

export/stages.py

@@ -734,6 +734,14 @@ class EdgeProgramManagerTransformStage(Stage):
 
     def __init__(
         self,
+        edge_transform_passes: (
+            None
+            | List[
+                Callable[
+                    [str, ExportedProgram], List[PassType] | GraphModulePassManager
+                ]
+            ]
+        ) = None,
         edge_manager_transform_passes: (
             None
             | List[
@@ -751,6 +759,7 @@ def __init__(
                                            backends to control pass ordering and dependencies.
         """
         super().__init__()
+        self._edge_transform_passes = edge_transform_passes or []
         self._edge_manager_transform_passes = edge_manager_transform_passes or []
 
     @classmethod
@@ -761,6 +770,7 @@ def from_recipe(
             return cls()
 
         return cls(
+            edge_transform_passes=lowering_recipe.edge_transform_passes,
             edge_manager_transform_passes=lowering_recipe.edge_manager_transform_passes,
         )
 
@@ -793,10 +803,35 @@ def run(self, artifact: PipelineArtifact) -> None:
                 f"Expected EdgeProgramManager but got {type(edge_program_manager)}"
             )
 
-        if not self._edge_manager_transform_passes:
+        if not self._edge_transform_passes and not self._edge_manager_transform_passes:
             self._artifact = artifact
             return
 
+        # Detect if any callable returns PassManager
+        pass_manager = None
+        transform_passes = defaultdict(list)
+        for method_name in edge_program_manager.methods:
+            # Resolve transform passes if it's a callable
+            ep = edge_program_manager.exported_program(method_name)
+            for pass_callable in self._edge_transform_passes or []:
+                if not callable(pass_callable):
+                    raise ValueError(
+                        "Transform passes must be a callable that resolves to passes"
+                    )
+                passes = pass_callable(method_name, ep)
+                if isinstance(passes, GraphModulePassManager):
+                    pass_manager = passes
+                    break
+                else:
+                    transform_passes[method_name].extend(passes)
+            if pass_manager:
+                break
+
+        # See EdgeTransformAndLowerStage.run.
+        final_passes = pass_manager or _drop_empty(transform_passes) or None
+        if final_passes is not None:
+            edge_program_manager = edge_program_manager.transform(final_passes)
+
         # Run edge manager transform passes
         for pass_callable in self._edge_manager_transform_passes:
             passes = pass_callable(edge_program_manager)

export/tests/test_export_session.py

@@ -757,30 +757,6 @@ def test_no_transform_passes_means_no_stage(self) -> None:
             StageType.EDGE_PROGRAM_MANAGER_TRANSFORM, session._get_default_pipeline()
         )
 
-    def test_edge_transform_passes_not_duplicated_in_default_pipeline(self) -> None:
-        # Before the fix, from_recipe() gave edge_transform_passes to
-        # EDGE_PROGRAM_MANAGER_TRANSFORM as well as TO_EDGE_TRANSFORM_AND_LOWER,
-        # so they ran twice. Verify the stage no longer holds them at all.
-        from executorch.export.stages import EdgeProgramManagerTransformStage
-
-        edge_pass = Mock()
-        epm_pass = Mock()
-
-        stage = EdgeProgramManagerTransformStage.from_recipe(
-            LoweringRecipe(
-                edge_transform_passes=[edge_pass],
-                edge_manager_transform_passes=[epm_pass],
-            )
-        )
-
-        # The stage must only know about edge_manager_transform_passes.
-        self.assertEqual(stage._edge_manager_transform_passes, [epm_pass])
-        self.assertFalse(
-            hasattr(stage, "_edge_transform_passes"),
-            "EdgeProgramManagerTransformStage must not hold edge_transform_passes "
-            "because TO_EDGE_TRANSFORM_AND_LOWER already applies them.",
-        )
-
     def test_example_inputs_required_for_nn_module(self) -> None:
         """Test that example_inputs are required for nn.Module."""
         with self.assertRaises(ValueError) as cm: