Skip to content

Commit 863bdaa

Browse files
Add coreml backend recipes (#13121)
Summary: Adds coreml recipes as discussed similar to xnnpack backend recipes to use them for export flow. Fixes #13100 Pull Request resolved: #13121 Test Plan: `python -m unittest backends/apple/coreml/test/test_coreml_recipes.py` Differential Revision: D79654041 Pulled By: abhinaykukkadapu
1 parent 112a09f commit 863bdaa

File tree

8 files changed

+475
-14
lines changed

8 files changed

+475
-14
lines changed

backends/apple/coreml/TARGETS

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,26 @@ runtime.python_library(
6060
],
6161
)
6262

63+
runtime.python_library(
64+
name = "recipes",
65+
srcs = glob([
66+
"recipes/*.py",
67+
]),
68+
visibility = [
69+
"@EXECUTORCH_CLIENTS",
70+
],
71+
deps = [
72+
"fbsource//third-party/pypi/coremltools:coremltools",
73+
":backend",
74+
"//caffe2:torch",
75+
"//executorch/exir:lib",
76+
"//executorch/exir/backend:compile_spec_schema",
77+
"//executorch/exir/backend:partitioner",
78+
"//executorch/exir/backend:utils",
79+
"//executorch/export:lib",
80+
],
81+
)
82+
6383
runtime.cxx_python_extension(
6484
name = "executorchcoreml",
6585
srcs = [
@@ -103,6 +123,7 @@ runtime.python_test(
103123
"fbsource//third-party/pypi/pytest:pytest",
104124
":partitioner",
105125
":quantizer",
126+
":recipes",
106127
"//caffe2:torch",
107128
"//pytorch/vision:torchvision",
108129
],
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
# pyre-strict
8+
9+
from executorch.export import recipe_registry
10+
11+
from .coreml_recipe_provider import CoreMLRecipeProvider
12+
from .coreml_recipe_types import CoreMLRecipeType
13+
14+
# Auto-register CoreML backend recipe provider
15+
recipe_registry.register_backend_recipe_provider(CoreMLRecipeProvider())
16+
17+
__all__ = [
18+
"CoreMLRecipeProvider",
19+
"CoreMLRecipeType",
20+
]
Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
# pyre-strict
8+
9+
from typing import Any, Optional, Sequence
10+
11+
from executorch.exir import EdgeCompileConfig
12+
13+
try:
14+
import coremltools as ct
15+
except ImportError:
16+
ct = None
17+
18+
from executorch.backends.apple.coreml.compiler import CoreMLBackend
19+
from executorch.backends.apple.coreml.partition.coreml_partitioner import (
20+
CoreMLPartitioner,
21+
)
22+
from executorch.backends.apple.coreml.recipes.coreml_recipe_types import (
23+
COREML_BACKEND,
24+
CoreMLRecipeType,
25+
)
26+
from executorch.export import (
27+
BackendRecipeProvider,
28+
ExportRecipe,
29+
LoweringRecipe,
30+
RecipeType,
31+
)
32+
33+
34+
class CoreMLRecipeProvider(BackendRecipeProvider):
35+
@property
36+
def backend_name(self) -> str:
37+
return COREML_BACKEND
38+
39+
def get_supported_recipes(self) -> Sequence[RecipeType]:
40+
return list(CoreMLRecipeType)
41+
42+
def create_recipe(
43+
self, recipe_type: RecipeType, **kwargs: Any
44+
) -> Optional[ExportRecipe]:
45+
"""Create CoreML recipe with precision and compute unit combinations"""
46+
47+
if recipe_type not in self.get_supported_recipes():
48+
return None
49+
50+
if ct is None:
51+
raise ImportError(
52+
"coremltools is required for CoreML recipes. "
53+
"Install it with: pip install coremltools"
54+
)
55+
56+
# Validate kwargs
57+
self._validate_recipe_kwargs(recipe_type, **kwargs)
58+
59+
# Map compute unit string to CoreML compute unit
60+
compute_unit_map = {
61+
CoreMLRecipeType.FP32_CPU: ct.ComputeUnit.CPU_ONLY,
62+
CoreMLRecipeType.FP32_GPU: ct.ComputeUnit.CPU_AND_GPU,
63+
CoreMLRecipeType.FP32_NEURAL_ENGINE: ct.ComputeUnit.CPU_AND_NE,
64+
CoreMLRecipeType.FP32_ALL: ct.ComputeUnit.ALL,
65+
CoreMLRecipeType.FP16_CPU: ct.ComputeUnit.CPU_ONLY,
66+
CoreMLRecipeType.FP16_GPU: ct.ComputeUnit.CPU_AND_GPU,
67+
CoreMLRecipeType.FP16_NEURAL_ENGINE: ct.ComputeUnit.CPU_AND_NE,
68+
CoreMLRecipeType.FP16_ALL: ct.ComputeUnit.ALL,
69+
}
70+
71+
# Parse recipe type to get precision and compute unit
72+
precision = None
73+
if recipe_type in [
74+
CoreMLRecipeType.FP32_CPU,
75+
CoreMLRecipeType.FP32_GPU,
76+
CoreMLRecipeType.FP32_NEURAL_ENGINE,
77+
CoreMLRecipeType.FP32_ALL,
78+
]:
79+
precision = ct.precision.FLOAT32
80+
elif recipe_type in [
81+
CoreMLRecipeType.FP16_CPU,
82+
CoreMLRecipeType.FP16_GPU,
83+
CoreMLRecipeType.FP16_NEURAL_ENGINE,
84+
CoreMLRecipeType.FP16_ALL,
85+
]:
86+
precision = ct.precision.FLOAT16
87+
88+
compute_unit = compute_unit_map.get(recipe_type, None)
89+
if precision is None or compute_unit is None:
90+
raise ValueError(
91+
f"Unknown precision or compute unit for recipe: {recipe_type.value}"
92+
)
93+
94+
return self._build_recipe(recipe_type, precision, compute_unit, **kwargs)
95+
96+
def _validate_recipe_kwargs(self, recipe_type: RecipeType, **kwargs: Any) -> None:
97+
if not kwargs:
98+
return
99+
expected_keys = {"minimum_deployment_target"}
100+
unexpected = set(kwargs.keys()) - expected_keys
101+
if unexpected:
102+
raise ValueError(
103+
f"CoreML Recipes only accept 'minimum_deployment_target' parameter. "
104+
f"Unexpected parameters: {list(unexpected)}"
105+
)
106+
if "minimum_deployment_target" in kwargs:
107+
minimum_deployment_target = kwargs["minimum_deployment_target"]
108+
if not isinstance(minimum_deployment_target, ct.target):
109+
raise ValueError(
110+
f"Parameter 'minimum_deployment_target' must be an enum of type ct.target, got {type(minimum_deployment_target)}"
111+
)
112+
113+
def _build_recipe(
114+
self,
115+
recipe_type: RecipeType,
116+
precision: ct.precision,
117+
compute_unit: ct.ComputeUnit,
118+
minimum_deployment_target: ct.target = ct.target.iOS15,
119+
) -> ExportRecipe:
120+
lowering_recipe = self._get_coreml_lowering_recipe(
121+
compute_precision=precision,
122+
compute_unit=compute_unit,
123+
minimum_deployment_target=minimum_deployment_target,
124+
)
125+
126+
return ExportRecipe(
127+
name=recipe_type.value,
128+
quantization_recipe=None, # TODO - add quantization recipe
129+
lowering_recipe=lowering_recipe,
130+
)
131+
132+
def _get_coreml_lowering_recipe(
133+
self,
134+
compute_unit: ct.ComputeUnit,
135+
compute_precision: ct.precision,
136+
minimum_deployment_target: ct.target,
137+
) -> LoweringRecipe:
138+
compile_specs = CoreMLBackend.generate_compile_specs(
139+
compute_unit=compute_unit,
140+
minimum_deployment_target=minimum_deployment_target,
141+
compute_precision=compute_precision,
142+
)
143+
144+
partitioner = CoreMLPartitioner(
145+
compile_specs=compile_specs,
146+
take_over_mutable_buffer=(minimum_deployment_target >= ct.target.iOS18),
147+
)
148+
149+
edge_compile_config = EdgeCompileConfig(
150+
_check_ir_validity=False,
151+
_skip_dim_order=True,
152+
)
153+
154+
return LoweringRecipe(
155+
partitioners=[partitioner], edge_compile_config=edge_compile_config
156+
)
Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
# pyre-strict
8+
9+
from executorch.export import RecipeType
10+
11+
12+
COREML_BACKEND: str = "coreml"
13+
14+
15+
class CoreMLRecipeType(RecipeType):
16+
"""CoreML-specific recipe types - combinations of precision and compute units"""
17+
18+
# FP32 precision with different compute units
19+
FP32_CPU = "coreml_fp32_cpu"
20+
FP32_GPU = "coreml_fp32_gpu"
21+
FP32_NEURAL_ENGINE = "coreml_fp32_neural_engine"
22+
FP32_ALL = "coreml_fp32_all"
23+
24+
# FP16 precision with different compute units
25+
FP16_CPU = "coreml_fp16_cpu"
26+
FP16_GPU = "coreml_fp16_gpu"
27+
FP16_NEURAL_ENGINE = "coreml_fp16_neural_engine"
28+
FP16_ALL = "coreml_fp16_all"
29+
30+
@classmethod
31+
def get_backend_name(cls) -> str:
32+
return COREML_BACKEND

0 commit comments

Comments
 (0)