diff --git a/backends/apple/coreai/recipes/__init__.py b/backends/apple/coreai/recipes/__init__.py new file mode 100644 index 00000000000..3c9c93287b1 --- /dev/null +++ b/backends/apple/coreai/recipes/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.export import recipe_registry + +from .recipe_provider import CoreAIRecipeProvider +from .recipe_types import CoreAIRecipeType + +# Auto-register the Core AI backend recipe provider on import. +recipe_registry.register_backend_recipe_provider(CoreAIRecipeProvider()) + +__all__ = [ + "CoreAIRecipeProvider", + "CoreAIRecipeType", +] diff --git a/backends/apple/coreai/recipes/recipe_provider.py b/backends/apple/coreai/recipes/recipe_provider.py new file mode 100644 index 00000000000..9c3602db19b --- /dev/null +++ b/backends/apple/coreai/recipes/recipe_provider.py @@ -0,0 +1,93 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +from typing import Any, Optional, Sequence + +from executorch.backends.apple.coreai.partition.partitioner import CoreAIPartitioner +from executorch.backends.apple.coreai.recipes.recipe_types import ( + COREAI_BACKEND, + CoreAIRecipeType, +) +from executorch.export import ( + BackendRecipeProvider, + ExportRecipe, + LoweringRecipe, + RecipeType, +) + +logger = logging.getLogger(__name__) + + +def _cast_fp32_to_fp16_pass(_method_name, exported_program): + """aten_transform pass: cast the whole program to FP16 before edge lowering. + + coreai-torch does not auto-downcast (unlike coremltools), so we run + coreai_opt's ``cast_fp32_to_fp16`` on the ExportedProgram pre-``to_edge``. + Because this runs on the whole model before partitioning, the graph is + uniformly FP16 and no delegate boundary mismatch occurs (the model's own + I/O becomes FP16). + """ + from coreai_opt.casting import cast_fp32_to_fp16 + + return cast_fp32_to_fp16(exported_program) + + +def _default_edge_transform_passes(_method_name, _exported_program): + """Factory the recipe framework calls as ``(method_name, ep) -> [passes]``. + + LoweringRecipe.edge_transform_passes entries are pass *factories* (see + executorch/export/stages.py), unlike ``to_edge_transform_and_lower``'s + ``transform_passes`` which take the passes directly. Both ultimately apply + the same GraphModule passes from ``get_default_passes``. + """ + from executorch.backends.apple.coreai import get_default_passes + + return get_default_passes() + + +class CoreAIRecipeProvider(BackendRecipeProvider): + """Provides Core AI export recipes (FP32 / FP16).""" + + @property + def backend_name(self) -> str: + return COREAI_BACKEND + + def get_supported_recipes(self) -> Sequence[RecipeType]: + return list(CoreAIRecipeType) + + def create_recipe( + self, recipe_type: RecipeType, **kwargs: Any + ) -> Optional[ExportRecipe]: + if recipe_type not in self.get_supported_recipes(): + return None + + if kwargs: + logger.warning( + "Core AI recipe '%s' ignoring unexpected parameters: %s", + recipe_type.value, + list(kwargs.keys()), + ) + + if recipe_type == CoreAIRecipeType.FP32: + return self._build_recipe(recipe_type, fp16=False) + if recipe_type == CoreAIRecipeType.FP16: + return self._build_recipe(recipe_type, fp16=True) + return None + + def _build_recipe(self, recipe_type: RecipeType, fp16: bool) -> ExportRecipe: + from executorch.backends.apple.coreai import get_default_compile_config + + lowering_recipe = LoweringRecipe( + partitioners=[CoreAIPartitioner()], + edge_transform_passes=[_default_edge_transform_passes], + edge_compile_config=get_default_compile_config(), + ) + return ExportRecipe( + name=recipe_type.value, + aten_transform_passes=[_cast_fp32_to_fp16_pass] if fp16 else None, + lowering_recipe=lowering_recipe, + ) diff --git a/backends/apple/coreai/recipes/recipe_types.py b/backends/apple/coreai/recipes/recipe_types.py new file mode 100644 index 00000000000..205067d0a64 --- /dev/null +++ b/backends/apple/coreai/recipes/recipe_types.py @@ -0,0 +1,26 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.export import RecipeType + + +COREAI_BACKEND: str = "coreai" + + +class CoreAIRecipeType(RecipeType): + """Core AI-specific recipe types.""" + + # FP32 precision: lower as-is to Core AI. + FP32 = "coreai_fp32" + + # FP16 precision: cast the exported program to FP16 (via coreai_opt's + # cast_fp32_to_fp16) before edge lowering, then lower to Core AI. Note this + # makes the model's external I/O FP16. + FP16 = "coreai_fp16" + + @classmethod + def get_backend_name(cls) -> str: + return COREAI_BACKEND diff --git a/backends/apple/coreai/test/test_recipes.py b/backends/apple/coreai/test/test_recipes.py new file mode 100644 index 00000000000..35a1f821517 --- /dev/null +++ b/backends/apple/coreai/test/test_recipes.py @@ -0,0 +1,86 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +import torch.nn as nn + +from executorch.backends.apple.coreai.recipes import ( # noqa: F401 (registers provider) + CoreAIRecipeType, +) +from executorch.backends.apple.coreai.recipes.recipe_provider import ( + CoreAIRecipeProvider, +) +from executorch.exir.lowered_backend_module import executorch_call_delegate +from executorch.export import export, ExportRecipe + + +def _model(): + return nn.Sequential(nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 32)).eval() + + +def _delegate_count(edge_manager): + gm = edge_manager.exported_program().graph_module + return sum( + 1 + for n in gm.graph.nodes + if n.op == "call_function" and n.target is executorch_call_delegate + ) + + +def _input_dtypes(edge_manager): + gm = edge_manager.exported_program().graph_module + return [ + n.meta["val"].dtype + for n in gm.graph.nodes + if n.op == "placeholder" and hasattr(n.meta.get("val"), "dtype") + ] + + +class CoreAIRecipeProviderTest(unittest.TestCase): + def test_fp32_recipe_has_no_cast_pass(self): + recipe = CoreAIRecipeProvider().create_recipe(CoreAIRecipeType.FP32) + self.assertIsNotNone(recipe) + self.assertIsNone(recipe.aten_transform_passes) + self.assertEqual(len(recipe.lowering_recipe.partitioners), 1) + + def test_fp16_recipe_has_cast_pass(self): + recipe = CoreAIRecipeProvider().create_recipe(CoreAIRecipeType.FP16) + self.assertIsNotNone(recipe) + self.assertEqual(len(recipe.aten_transform_passes), 1) + self.assertEqual(len(recipe.lowering_recipe.partitioners), 1) + + def test_extra_kwargs_are_ignored(self): + # Unexpected kwargs are warned about, not fatal. + recipe = CoreAIRecipeProvider().create_recipe(CoreAIRecipeType.FP32, foo=1) + self.assertIsNotNone(recipe) + + +class CoreAIRecipeLoweringTest(unittest.TestCase): + def setUp(self): + self.example_inputs = [(torch.randn(2, 32),)] + + def test_fp32_recipe_lowers_to_coreai(self): + recipe = ExportRecipe.get_recipe(CoreAIRecipeType.FP32) + session = export(_model(), self.example_inputs, recipe) + edge = session.get_edge_program_manager() + self.assertGreaterEqual(_delegate_count(edge), 1) + self.assertTrue(all(d == torch.float32 for d in _input_dtypes(edge))) + self.assertGreater(len(session.get_pte_buffer()), 0) + + def test_fp16_recipe_casts_and_lowers(self): + recipe = ExportRecipe.get_recipe(CoreAIRecipeType.FP16) + session = export(_model(), self.example_inputs, recipe) + edge = session.get_edge_program_manager() + self.assertGreaterEqual(_delegate_count(edge), 1) + # FP16 recipe casts the whole model -> edge inputs are fp16. + self.assertTrue(all(d == torch.float16 for d in _input_dtypes(edge))) + self.assertGreater(len(session.get_pte_buffer()), 0) + + +if __name__ == "__main__": + unittest.main()