Fix issues in the export recipe framework

export/export.py

@@ -297,8 +297,8 @@ def _build_stages(self, stages: List[StageType]) -> Dict[StageType, Stage]:
         """Build the stage registry from the given stages."""
         stage_registry: Dict[StageType, Stage] = {}
 
-        stage = None
         for stage_type in stages or self._get_default_pipeline():
+            stage = None
             if stage_type == StageType.SOURCE_TRANSFORM:
                 stage = SourceTransformStage(
                     self._quant_recipe,
@@ -329,10 +329,10 @@ def _build_stages(self, stages: List[StageType]) -> Dict[StageType, Stage]:
                 stage = ExecutorchStage(self._export_recipe.executorch_backend_config)
             else:
                 logging.info(
-                    f"{stage_type} is unknown, you have to register it before executing export()"
+                    f"{stage_type} is unknown, register it with session.register_stage()"
                 )
 
-            if stage:
+            if stage is not None:
                 stage_registry[stage_type] = stage
         return stage_registry
 
@@ -422,7 +422,7 @@ def _validate_pipeline_sequence(
             stage_instance = self._stage_registry.get(current_stage)
             if stage_instance is None:
                 raise ValueError(
-                    f"Stage {current_stage} not found in registry, , register it using session.register_stage()"
+                    f"Stage {current_stage} not found in registry, register it using session.register_stage()"
                 )
 
             valid_predecessors = stage_instance.valid_predecessor_stages
@@ -440,6 +440,10 @@ def _run_pipeline(self) -> None:
             stages=self._pipeline_stages,
         )
 
+        # After validation: a rejected pipeline must not destroy the previous
+        # run. In place, since get_stage_artifacts() hands out this dict.
+        self._stage_to_artifacts.clear()
+
         current_artifact = PipelineArtifact(data=self._model, context=self._run_context)
 
         # Execute stages from registry in the order specified by pipeline_stages

export/recipe.py

@@ -5,6 +5,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 import copy
+import dataclasses
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass
 from enum import Enum, EnumMeta
@@ -32,6 +33,46 @@
 """
 
 
+# Operator lists say which ops to keep, not in what order.
+_UNORDERED_EDGE_CONFIG_FIELDS = ("preserve_ops", "_core_aten_ops_exception_list")
+
+
+def _edge_config_key(config: EdgeCompileConfig) -> tuple:
+    return tuple(
+        (
+            frozenset(getattr(config, f.name) or [])
+            if f.name in _UNORDERED_EDGE_CONFIG_FIELDS
+            else getattr(config, f.name)
+        )
+        for f in dataclasses.fields(config)
+    )
+
+
+def _edge_compile_configs_agree(
+    left: EdgeCompileConfig, right: EdgeCompileConfig
+) -> bool:
+    return left is right or _edge_config_key(left) == _edge_config_key(right)
+
+
+def _edge_compile_config_conflict(configs: List[tuple[str, EdgeCompileConfig]]) -> str:
+    """Report only the fields the configs actually disagree on."""
+
+    def render(config: EdgeCompileConfig, name: str) -> str:
+        value = getattr(config, name)
+        if name in _UNORDERED_EDGE_CONFIG_FIELDS:
+            return str(sorted(str(op) for op in value or []))
+        return str(value)
+
+    keys = [_edge_config_key(config) for _, config in configs]
+    return "; ".join(
+        f"{field.name} ("
+        + ", ".join(f"{name}={render(config, field.name)}" for name, config in configs)
+        + ")"
+        for i, field in enumerate(dataclasses.fields(configs[0][1]))
+        if any(key[i] != keys[0][i] for key in keys[1:])
+    )
+
+
 class RecipeTypeMeta(EnumMeta, ABCMeta):
     """Metaclass that combines EnumMeta and ABCMeta"""
 
@@ -249,6 +290,17 @@ def _combine_recipes(  # noqa: C901
         Returns:
             Combined ExportRecipe for multi-backend deployment
         """
+        overriding = [
+            r.name or f"recipes[{i}]"
+            for i, r in enumerate(backend_recipes)
+            if r.pipeline_stages
+        ]
+        if overriding:
+            raise ValueError(
+                "Cannot combine recipes that override pipeline_stages, there is no "
+                f"correct way to merge the orderings: {overriding}"
+            )
+
         # Extract components from individual recipes
         all_partitioners = []
         all_quantizers = []
@@ -300,18 +352,30 @@ def _combine_recipes(  # noqa: C901
                 ),
             )
 
-        # Create combined lowering recipe
-        combined_lowering_recipe = None
-        if all_partitioners or all_transform_passes:
-            edge_compile_config = None
-            for recipe in backend_recipes:
-                if (
-                    recipe.lowering_recipe
-                    and recipe.lowering_recipe.edge_compile_config
-                ):
-                    edge_compile_config = recipe.lowering_recipe.edge_compile_config
-                    break
+        # By value, not identity: every provider builds a fresh config object,
+        # so asking for the same thing twice is not a conflict.
+        distinct: List[tuple[str, EdgeCompileConfig]] = []
+        for i, recipe in enumerate(backend_recipes):
+            config = (
+                recipe.lowering_recipe.edge_compile_config
+                if recipe.lowering_recipe
+                else None
+            )
+            if config is None or any(
+                _edge_compile_configs_agree(config, seen) for _, seen in distinct
+            ):
+                continue
+            distinct.append((recipe.name or f"recipes[{i}]", config))
 
+        if len(distinct) > 1:
+            raise ValueError(
+                "Cannot combine recipes whose edge_compile_configs disagree on "
+                + _edge_compile_config_conflict(distinct)
+            )
+        edge_compile_config = copy.deepcopy(distinct[0][1]) if distinct else None
+
+        combined_lowering_recipe = None
+        if all_partitioners or all_transform_passes or edge_compile_config:
             combined_lowering_recipe = LoweringRecipe(
                 partitioners=all_partitioners if all_partitioners else None,
                 edge_transform_passes=(

export/stages.py

@@ -30,6 +30,12 @@
 from torchao.utils import unwrap_tensor_subclass
 
 
+def _drop_empty(
+    passes_by_method: Dict[str, List[PassType]]
+) -> Dict[str, List[PassType]]:
+    return {method: p for method, p in passes_by_method.items() if p}
+
+
 class PipelineArtifact:
     def __init__(
         self,
@@ -242,8 +248,9 @@ def run(self, artifact: PipelineArtifact) -> None:
             if pass_manager:
                 break
 
-        # Use PassManager directly if found, otherwise use dict
-        final_passes = pass_manager if pass_manager else transform_passes
+        # An empty dict is not no passes: EdgeProgramManager deep-copies every
+        # method the dict does not name, so it would copy to apply nothing.
+        final_passes = pass_manager or _drop_empty(transform_passes) or None
 
         with validation_disabled():
             edge_program_manager = to_edge_transform_and_lower(
@@ -619,11 +626,10 @@ def run(self, artifact: PipelineArtifact) -> None:
             if pass_manager:
                 break
 
-        # Use PassManager directly if found, otherwise use dict
-        final_passes = pass_manager if pass_manager else transform_passes
-
-        # Apply edge transform passes
-        edge_program_manager = edge_program_manager.transform(final_passes)
+        # 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: