[ExecuTorch][Core AI][4/x] Export recipes (fp32/fp16)

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",
+]

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,
+        )

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

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()