diff --git a/export/export.py b/export/export.py index 09bcacbb919..68f80f46eae 100644 --- a/export/export.py +++ b/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 diff --git a/export/recipe.py b/export/recipe.py index c205491f80e..2c5a4c17899 100644 --- a/export/recipe.py +++ b/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=( diff --git a/export/stages.py b/export/stages.py index a68ad408493..03dd2368bfa 100644 --- a/export/stages.py +++ b/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: diff --git a/export/tests/test_export_recipe.py b/export/tests/test_export_recipe.py index d22442371e2..5fc4701a5f6 100644 --- a/export/tests/test_export_recipe.py +++ b/export/tests/test_export_recipe.py @@ -8,6 +8,9 @@ import unittest from typing import Any, Dict, Optional, Sequence +from unittest.mock import Mock + +import torch from executorch.export.recipe import ExportRecipe, RecipeType from executorch.export.recipe_provider import BackendRecipeProvider @@ -129,3 +132,205 @@ def test_get_recipe_with_kwargs_verification(self) -> None: # Verify that the kwargs were passed to the backend provider's create_recipe method self.assertIsNotNone(self.provider.last_kwargs) self.assertEqual(self.provider.last_kwargs, kwargs) + + +class TestExportRecipeCombine(unittest.TestCase): + def _lowering(self, **kwargs): + from executorch.export.recipe import LoweringRecipe + + return LoweringRecipe(**kwargs) + + def test_combine_keeps_edge_compile_config_without_partitioners(self) -> None: + # A recipe that lowers without delegating still needs its to_edge + # config; it used to be dropped along with the whole LoweringRecipe. + from executorch.exir import EdgeCompileConfig + + config = EdgeCompileConfig(_check_ir_validity=False) + combined = ExportRecipe.combine( + [ + ExportRecipe( + name="a", lowering_recipe=self._lowering(edge_compile_config=config) + ), + ExportRecipe(name="b"), + ] + ) + assert combined.lowering_recipe is not None + self.assertFalse( + combined.lowering_recipe.edge_compile_config._check_ir_validity + ) + # Carried by copy, so mutating the combination cannot reach back into + # the provider's own config -- including through the mutable lists, + # which a shallow copy would still share. + combined.lowering_recipe.edge_compile_config.preserve_ops.append( + torch.ops.aten.linear.default + ) + self.assertEqual(config.preserve_ops, []) + + def test_combine_rejects_conflicting_edge_compile_configs(self) -> None: + # Picking one by position would decide the emitted graph by argument + # order: preserve_ops that one backend requires would silently vanish. + from executorch.exir import EdgeCompileConfig + + recipes = [ + ExportRecipe( + name="a", + lowering_recipe=self._lowering( + edge_compile_config=EdgeCompileConfig( + preserve_ops=[torch.ops.aten.linear.default] + ) + ), + ), + ExportRecipe( + name="b", + lowering_recipe=self._lowering( + edge_compile_config=EdgeCompileConfig(preserve_ops=[]) + ), + ), + ] + for ordering in (recipes, list(reversed(recipes))): + with self.assertRaisesRegex(ValueError, "edge_compile_configs") as cm: + ExportRecipe.combine(ordering) + # The message has to say which recipes disagree, and how. + self.assertIn("preserve_ops", str(cm.exception)) + self.assertIn("a=", str(cm.exception)) + self.assertIn("b=", str(cm.exception)) + self.assertIn("aten.linear.default", str(cm.exception)) + + def test_combine_names_the_field_that_conflicts(self) -> None: + # A summary of hand-picked fields prints identically for two configs + # that differ elsewhere, so the error names no cause at all. + from executorch.exir import EdgeCompileConfig + + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine( + [ + ExportRecipe( + name=name, + lowering_recipe=self._lowering( + edge_compile_config=EdgeCompileConfig(_skip_dim_order=skip) + ), + ) + for name, skip in (("a", True), ("b", False)) + ] + ) + self.assertIn("_skip_dim_order (a=True, b=False)", str(cm.exception)) + # Fields they agree on are noise that hides the one that matters. + self.assertNotIn("preserve_ops", str(cm.exception)) + + def test_combine_accepts_distinct_but_equal_configs(self) -> None: + # Every provider builds a fresh config object, so this -- not the + # shared-object case -- is what a real multi-backend combination looks + # like. Comparing by identity would reject all of them. + from executorch.exir import EdgeCompileConfig + + combined = ExportRecipe.combine( + [ + ExportRecipe( + name=n, + lowering_recipe=self._lowering( + edge_compile_config=EdgeCompileConfig(_check_ir_validity=False) + ), + ) + for n in ("a", "b") + ] + ) + assert combined.lowering_recipe is not None + self.assertFalse( + combined.lowering_recipe.edge_compile_config._check_ir_validity + ) + + def test_combine_ignores_preserve_ops_ordering(self) -> None: + # preserve_ops says which ops to keep, not in what order. + from executorch.exir import EdgeCompileConfig + + ops = [torch.ops.aten.linear.default, torch.ops.aten.silu.default] + combined = ExportRecipe.combine( + [ + ExportRecipe( + name="a", + lowering_recipe=self._lowering( + edge_compile_config=EdgeCompileConfig(preserve_ops=ops) + ), + ), + ExportRecipe( + name="b", + lowering_recipe=self._lowering( + edge_compile_config=EdgeCompileConfig( + preserve_ops=list(reversed(ops)) + ) + ), + ), + ] + ) + assert combined.lowering_recipe is not None + self.assertCountEqual( + combined.lowering_recipe.edge_compile_config.preserve_ops, ops + ) + + def test_combine_without_lowering_inputs_produces_none(self) -> None: + combined = ExportRecipe.combine( + [ExportRecipe(name="a"), ExportRecipe(name="b")] + ) + self.assertIsNone(combined.lowering_recipe) + + def test_combine_keeps_partitioners(self) -> None: + from executorch.exir.backend.partitioner import Partitioner + + a, b = Mock(spec=Partitioner), Mock(spec=Partitioner) + combined = ExportRecipe.combine( + [ + ExportRecipe( + name="a", lowering_recipe=self._lowering(partitioners=[a]) + ), + ExportRecipe( + name="b", lowering_recipe=self._lowering(partitioners=[b]) + ), + ] + ) + assert combined.lowering_recipe is not None + self.assertEqual(combined.lowering_recipe.partitioners, [a, b]) + + def test_combine_rejects_pipeline_stages(self) -> None: + from executorch.export.types import StageType + + # Unnamed recipes fall back to their position; named ones are named. + with self.assertRaisesRegex(ValueError, r"pipeline_stages.*recipes\[0\]"): + ExportRecipe.combine( + [ + ExportRecipe( + pipeline_stages=[StageType.TO_EDGE_TRANSFORM_AND_LOWER] + ), + ExportRecipe(name="b"), + ] + ) + with self.assertRaisesRegex(ValueError, r"pipeline_stages.*'stagey'"): + ExportRecipe.combine( + [ + ExportRecipe( + name="stagey", + pipeline_stages=[StageType.TO_EDGE_TRANSFORM_AND_LOWER], + ), + ExportRecipe(name="b"), + ] + ) + + def test_combine_keeps_edge_transform_passes(self) -> None: + # QNN supplies these and is combined with XNNPACK by target_recipes. + first = lambda name, ep: [] # noqa: E731 + second = lambda name, ep: [] # noqa: E731 + combined = ExportRecipe.combine( + [ + ExportRecipe( + name="a", + lowering_recipe=self._lowering(edge_transform_passes=[first]), + ), + ExportRecipe( + name="b", + lowering_recipe=self._lowering(edge_transform_passes=[second]), + ), + ] + ) + assert combined.lowering_recipe is not None + self.assertEqual( + combined.lowering_recipe.edge_transform_passes, [first, second] + ) diff --git a/export/tests/test_export_session.py b/export/tests/test_export_session.py index d3a3d68d42b..c658bf45ed9 100644 --- a/export/tests/test_export_session.py +++ b/export/tests/test_export_session.py @@ -17,7 +17,7 @@ LoweringRecipe, QuantizationRecipe, ) -from executorch.export.stages import PipelineArtifact +from executorch.export.stages import PipelineArtifact, Stage from executorch.export.types import StageType @@ -971,3 +971,49 @@ def test_get_edge_program_manager_before_edge_stage_fails(self) -> None: with self.assertRaises(RuntimeError) as cm: session.get_edge_program_manager() self.assertIn("Edge program manager is not available", str(cm.exception)) + + +class TestStageArtifactsAreScopedToOneRun(unittest.TestCase): + def test_rerunning_export_discards_the_previous_run(self) -> None: + # A stage that raises must not leave the previous run's artifact behind + # for the accessors to hand back as if it were current. + model = SimpleTestModel() + inputs = [(torch.randn(1, 10),)] + session = ExportSession( + model=model, example_inputs=inputs, export_recipe=ExportRecipe(name="t") + ) + session.export() + first = session.get_stage_artifacts()[StageType.TO_EXECUTORCH] + + artifacts = session.get_stage_artifacts() + + failing = Mock(spec=Stage) + failing.run.side_effect = RuntimeError("boom") + failing.stage_type = StageType.TO_EXECUTORCH + failing.valid_predecessor_stages = [StageType.TO_EDGE_TRANSFORM_AND_LOWER] + failing.can_start_pipeline = False + session.register_stage(StageType.TO_EXECUTORCH, failing) + + with self.assertRaises(RuntimeError): + session.export() + self.assertNotIn(StageType.TO_EXECUTORCH, session.get_stage_artifacts()) + # Cleared in place, so a dict the caller captured before the re-run + # reflects the clear too. + self.assertNotIn(StageType.TO_EXECUTORCH, artifacts) + self.assertIsNotNone(first) + + def test_a_rejected_pipeline_leaves_the_previous_run_intact(self) -> None: + # Validation runs before anything is cleared: a pipeline that never + # executes must not destroy results the caller still has. + session = ExportSession( + model=SimpleTestModel(), + example_inputs=[(torch.randn(1, 10),)], + export_recipe=ExportRecipe(name="t"), + ) + session.export() + before = session.get_executorch_program_manager() + + session._pipeline_stages = [StageType.TO_EXECUTORCH] + with self.assertRaises(ValueError): + session.export() + self.assertIs(session.get_executorch_program_manager(), before) diff --git a/export/tests/test_export_stages.py b/export/tests/test_export_stages.py index e9acee0ea26..2f124bdb7b7 100644 --- a/export/tests/test_export_stages.py +++ b/export/tests/test_export_stages.py @@ -13,6 +13,7 @@ from executorch.exir.program import EdgeProgramManager, ExecutorchProgramManager from executorch.export import AOQuantizationConfig, QuantizationRecipe, StageType from executorch.export.stages import ( + EdgeProgramManagerTransformStage, EdgeTransformAndLowerStage, ExecutorchStage, PipelineArtifact, @@ -551,3 +552,70 @@ 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 TestEmptyPassDictIsNotApplied(unittest.TestCase): + """`EdgeProgramManager.transform` deep-copies the graph and weights of every + method the pass dict does not name, so handing it an empty dict copies + methods 2..n in order to apply nothing.""" + + def _manager(self) -> Mock: + manager = Mock(spec=EdgeProgramManager) + manager.methods = {"forward", "decode"} + manager.transform.return_value = Mock(spec=EdgeProgramManager) + manager.exported_program.return_value = Mock() + return manager + + def test_edge_program_manager_stage_skips_empty_transform(self) -> None: + manager = self._manager() + stage = EdgeProgramManagerTransformStage( + edge_manager_transform_passes=[lambda epm: []] + ) + stage.run(PipelineArtifact(data=manager, context={})) + + manager.transform.assert_not_called() + self.assertIs(stage.get_artifacts().data, manager) + + def test_edge_program_manager_stage_still_applies_real_passes(self) -> None: + manager = self._manager() + pass_ = Mock() + stage = EdgeProgramManagerTransformStage( + edge_manager_transform_passes=[lambda epm: [pass_]] + ) + stage.run(PipelineArtifact(data=manager, context={})) + + manager.transform.assert_called_once_with([pass_]) + self.assertIs(stage.get_artifacts().data, manager.transform.return_value) + + @patch("executorch.export.stages.get_delegation_info") + @patch("executorch.export.stages.to_edge_transform_and_lower") + def test_lower_stage_passes_none_not_empty_dict( + self, mock_lower: Mock, mock_delegation_info: Mock + ) -> None: + mock_lower.return_value = Mock(spec=EdgeProgramManager) + mock_delegation_info.return_value = {} + stage = EdgeTransformAndLowerStage() + stage.run( + PipelineArtifact(data={"forward": Mock(spec=ExportedProgram)}, context={}) + ) + + self.assertIsNone(mock_lower.call_args.kwargs["transform_passes"]) + + +class TestUnknownStageIsNotRegistered(unittest.TestCase): + def test_unknown_stage_type_gets_no_stage(self) -> None: + # The loop used to hold the previous iteration's instance, so an + # unrecognised stage type silently registered the stage before it and + # the "register it first" guard could never fire. + from executorch.export import ExportRecipe + from executorch.export.export import ExportSession + + session = ExportSession( + model=SimpleTestModel(), + example_inputs=[(torch.randn(1, 10),)], + export_recipe=ExportRecipe(name="t"), + ) + registry = session._build_stages( + [StageType.TORCH_EXPORT, "not_a_stage", StageType.TO_EXECUTORCH] + ) + self.assertNotIn("not_a_stage", registry) diff --git a/pytest-windows.ini b/pytest-windows.ini index 6d5b5c2881e..39f238cdc00 100644 --- a/pytest-windows.ini +++ b/pytest-windows.ini @@ -33,7 +33,6 @@ addopts = exir/tests/ # executorch/export export/tests - --ignore=export/tests/test_export_stages.py # kernels/ kernels/prim_ops/test kernels/quantized diff --git a/pytest.ini b/pytest.ini index 3e6138451a9..576ed2b78ce 100644 --- a/pytest.ini +++ b/pytest.ini @@ -53,7 +53,6 @@ addopts = --ignore=exir/tests/test_common.py --ignore=exir/tests/test_op_convert.py - --ignore=export/tests/test_export_stages.py # Ignore tests with missing dependencies or build issues --ignore=extension/flat_tensor/test/test_serialize.py