Add coreml_compute_plan.py: report which CoreML ops dispatch to ANE / GPU / CPU (#19252)

john-rocky · web-flow · commit 79fe3a30148d · 2026-05-26T17:25:09.000-07:00
### Summary CoreML decides at compile/load time which device each MIL operation will execute on, and coremltools 9.0+ exposes that through `MLComputePlan`. The recurring question on the issue tracker is *"why isn't my model running fully on the ANE?"* — for example: - #4091 — `llama model is not fully lowered to ANE` - #11541 — `CoreML model is crashing on iPhone GPU, but not on iPhone CPU or macOS GPU` - #8439 — `ANE compile OOMs on certain input shapes` - #8445 — `CPU Overhead After ANE Execution` Today the only way for an ExecuTorch user to answer it is to break out Swift / Xcode. This PR adds a Python wrapper around `MLComputePlan` so the answer is one shell command: ``` $ python coreml_compute_plan.py --model_path my_model.mlpackage \ --compute_units cpu_and_ne --show_non_ane === my_model.mlpackage === ANE: 412 / 480 ( 85.8%) CPU: 68 / 480 ( 14.2%) Non-ANE op types: 32 ios17.cast 18 ios17.gather 12 ios17.reshape 6 ios17.constexpr_blockwise_shift_scale ``` Inputs supported: | Input | Behavior | |---|---| | `.pte` | Extract every Core ML partition into a tempdir, then analyze each. | | `.mlpackage` | Compile to `.mlmodelc` in a tempdir, then analyze. | | `.mlmodelc` | Analyze directly. | The PTE path reuses the same JSON/named-data extraction logic that `extract_coreml_models.py` uses, and is inlined into the script so it can be run against a plain CoreML model without depending on the executorch package. ### Test plan Added `test_coreml_compute_plan.py` covering: - `_device_name(...)` for `None` and a stub `MLNeuralEngineComputeDevice`. - `_COMPUTE_UNIT_CHOICES` mapping (`cpu_and_ne` / `all`). - `analyze_one(...)` end-to-end on a tiny `relu(x @ x.T) + x.sum()` mlpackage built with `coremltools.convert(...)`: returns rows for every dispatched op, with a `main` function and the expected MIL op types (`matmul`, `relu`, `add`, `reduce_sum`). ``` $ python -m pytest examples/apple/coreml/scripts/test_coreml_compute_plan.py -v ============================== 7 passed in 3.68s =============================== ``` I also ran the script against a few hand-built `.mlpackage` and `.mlmodelc` files on macOS 26 with coremltools 9.0 and verified the output matches what `MLComputePlan` returns directly. Authored with Claude. cc @kimishpatel @YifanShenSZ @cymbalrush @metascroy
diff --git a/examples/apple/coreml/scripts/BUCK b/examples/apple/coreml/scripts/BUCK
@@ -16,6 +16,19 @@ fbcode_target(_kind = python_binary,
     ],
 )
 
+fbcode_target(_kind = python_binary,
+    name = "coreml_compute_plan",
+    srcs = [
+        "coreml_compute_plan.py",
+    ],
+    main_function = "executorch.examples.apple.coreml.scripts.coreml_compute_plan.main",
+    deps = [
+        "//executorch/backends/apple/coreml:executorchcoreml",
+        "//executorch/exir:schema",
+        "//executorch/exir/_serialize:lib",
+    ],
+)
+
 fbcode_target(_kind = python_binary,
     name = "export",
     srcs = [
diff --git a/examples/apple/coreml/scripts/coreml_compute_plan.py b/examples/apple/coreml/scripts/coreml_compute_plan.py
@@ -0,0 +1,236 @@
+# Copyright © 2026 Apple Inc. All rights reserved.
+#
+# Please refer to the license found in the LICENSE file in the root directory of the source tree.
+
+"""Report which CoreML operations would dispatch to ANE / GPU / CPU.
+
+The CoreML runtime decides at compile/load time which compute device each
+MIL operation will run on; that decision is exposed by ``MLComputePlan``
+in coremltools 9.0+.  This script wraps that API so users can answer
+"why isn't my model running on the ANE?" without writing Swift.
+
+Usage::
+
+    # Analyze a CoreML model directly (mlpackage or compiled mlmodelc).
+    python coreml_compute_plan.py --model_path path/to/model.mlpackage
+
+    # Analyze every Core ML partition embedded in an ExecuTorch .pte.
+    python coreml_compute_plan.py --model_path path/to/program.pte
+
+    # Show ops that fell off the ANE, grouped by op type.
+    python coreml_compute_plan.py --model_path model.mlpackage --show_non_ane
+
+    # Pick which devices the runtime is allowed to consider.
+    python coreml_compute_plan.py --model_path model.mlpackage \\
+        --compute_units cpu_and_ne
+"""
+
+import argparse
+import os
+import sys
+import tempfile
+from collections import Counter
+from typing import Iterable, List, Tuple
+
+import coremltools as ct
+from coremltools.models.compute_device import (
+    MLCPUComputeDevice,
+    MLGPUComputeDevice,
+    MLNeuralEngineComputeDevice,
+)
+from coremltools.models.compute_plan import MLComputePlan
+
+from executorch.examples.apple.coreml.scripts.extract_coreml_models import (
+    extract_coreml_models,
+)
+
+
+_DEVICE_NAMES: List[Tuple[type, str]] = [
+    (MLNeuralEngineComputeDevice, "ANE"),
+    (MLGPUComputeDevice, "GPU"),
+    (MLCPUComputeDevice, "CPU"),
+]
+
+_COMPUTE_UNIT_CHOICES = {
+    "all": ct.ComputeUnit.ALL,
+    "cpu_and_ne": ct.ComputeUnit.CPU_AND_NE,
+    "cpu_and_gpu": ct.ComputeUnit.CPU_AND_GPU,
+    "cpu_only": ct.ComputeUnit.CPU_ONLY,
+}
+
+
+def _device_name(device) -> str:
+    if device is None:
+        return "unknown"
+    for cls, name in _DEVICE_NAMES:
+        if isinstance(device, cls):
+            return name
+    return type(device).__name__
+
+
+def _iter_operations(block) -> Iterable:
+    for op in block.operations:
+        yield op
+        for nested in getattr(op, "blocks", None) or []:
+            yield from _iter_operations(nested)
+
+
+def _ensure_compiled(model_path: str, tmpdir: str) -> str:
+    """Return a `.mlmodelc` path; compile from `.mlpackage` if needed."""
+    if model_path.endswith(".mlmodelc"):
+        return model_path
+    if model_path.endswith(".mlpackage"):
+        dest = os.path.join(
+            tmpdir, os.path.basename(model_path).replace(".mlpackage", ".mlmodelc")
+        )
+        return str(ct.models.utils.compile_model(model_path, destination_path=dest))
+    raise ValueError(f"Expected a .mlpackage or .mlmodelc path, got: {model_path}")
+
+
+def analyze_one(
+    model_path: str, compute_units: ct.ComputeUnit
+) -> List[Tuple[str, str, str]]:
+    """Return [(function, operator_name, device)] for every op that has a plan.
+
+    coremltools 9.0's ``MLComputePlan.load_from_path`` only exposes usage for
+    the default function of a multifunction package, so a multifunction
+    .mlpackage is analyzed function-by-function by projecting each function
+    as the ``main`` of a temp single-function copy.
+    """
+    function_names = _mlpackage_function_names(model_path)
+    if len(function_names) <= 1:
+        return _analyze_compiled(model_path, compute_units)
+    rows: List[Tuple[str, str, str]] = []
+    with tempfile.TemporaryDirectory() as tmpdir:
+        for fname in function_names:
+            projected = _project_to_single(model_path, fname, tmpdir)
+            for _, op_name, device in _analyze_compiled(projected, compute_units):
+                rows.append((fname, op_name, device))
+    return rows
+
+
+def _analyze_compiled(
+    model_path: str, compute_units: ct.ComputeUnit
+) -> List[Tuple[str, str, str]]:
+    with tempfile.TemporaryDirectory() as tmpdir:
+        compiled = _ensure_compiled(model_path, tmpdir)
+        plan = MLComputePlan.load_from_path(compiled, compute_units=compute_units)
+        program = plan.model_structure.program
+        if program is None:
+            raise RuntimeError(
+                f"{model_path} is not an MLProgram model; this tool only supports "
+                "the MLProgram backend (the CoreML backend executorch produces today)."
+            )
+
+        rows: List[Tuple[str, str, str]] = []
+        for fname, fn in program.functions.items():
+            for op in _iter_operations(fn.block):
+                usage = plan.get_compute_device_usage_for_mlprogram_operation(op)
+                if usage is None:
+                    # Constants and similar non-dispatched ops don't have a plan.
+                    continue
+                rows.append(
+                    (
+                        fname,
+                        op.operator_name,
+                        _device_name(usage.preferred_compute_device),
+                    )
+                )
+        return rows
+
+
+def _mlpackage_function_names(model_path: str) -> List[str]:
+    """Names of the MLProgram functions inside an .mlpackage, or [] otherwise."""
+    if not model_path.endswith(".mlpackage"):
+        return []
+    spec = ct.models.MLModel(model_path, skip_model_load=True).get_spec()
+    if spec.WhichOneof("Type") != "mlProgram":
+        return []
+    return list(spec.mlProgram.functions.keys())
+
+
+def _project_to_single(src_mlpackage: str, function_name: str, tmpdir: str) -> str:
+    """Re-save ``src_mlpackage`` with only ``function_name`` exposed as ``main``."""
+    from coremltools.models.utils import MultiFunctionDescriptor, save_multifunction
+
+    dest = os.path.join(tmpdir, f"{function_name}.mlpackage")
+    desc = MultiFunctionDescriptor()
+    desc.add_function(
+        src_mlpackage,
+        src_function_name=function_name,
+        target_function_name="main",
+    )
+    desc.default_function_name = "main"
+    save_multifunction(desc, dest)
+    return dest
+
+
+def _print_report(
+    label: str, rows: List[Tuple[str, str, str]], show_non_ane: bool
+) -> None:
+    print(f"\n=== {label} ===")
+    if not rows:
+        print("  (no dispatched operations found)")
+        return
+    by_device = Counter(device for _, _, device in rows)
+    total = sum(by_device.values())
+    for device in ("ANE", "GPU", "CPU", "unknown"):
+        count = by_device.get(device, 0)
+        if count == 0:
+            continue
+        pct = 100.0 * count / total
+        print(f"  {device}: {count:5d} / {total} ({pct:5.1f}%)")
+
+    if show_non_ane:
+        non_ane = [(fn, op_name) for fn, op_name, dev in rows if dev != "ANE"]
+        if non_ane:
+            print("\n  Non-ANE op types:")
+            for op_name, count in Counter(op for _, op in non_ane).most_common():
+                print(f"    {count:5d}  {op_name}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument(
+        "--model_path",
+        required=True,
+        help="Path to a .pte, .mlpackage, or .mlmodelc.",
+    )
+    parser.add_argument(
+        "--compute_units",
+        default="cpu_and_ne",
+        choices=sorted(_COMPUTE_UNIT_CHOICES),
+        help="Which devices the runtime may use when planning dispatch.",
+    )
+    parser.add_argument(
+        "--show_non_ane",
+        action="store_true",
+        help="List op types that did not get assigned to the ANE.",
+    )
+    args = parser.parse_args()
+
+    compute_units = _COMPUTE_UNIT_CHOICES[args.compute_units]
+    model_path = args.model_path
+
+    if model_path.endswith(".pte"):
+        with open(model_path, "rb") as f:
+            pte_data = f.read()
+        with tempfile.TemporaryDirectory() as out_dir:
+            extracted = extract_coreml_models(pte_data, out_dir=out_dir)
+            if not extracted:
+                print(
+                    f"{model_path} does not contain any CoreML delegate partitions.",
+                    file=sys.stderr,
+                )
+                return 1
+            for path in extracted:
+                rows = analyze_one(str(path), compute_units)
+                _print_report(path.name, rows, args.show_non_ane)
+    else:
+        rows = analyze_one(model_path, compute_units)
+        _print_report(os.path.basename(model_path.rstrip("/")), rows, args.show_non_ane)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/examples/apple/coreml/scripts/extract_coreml_models.py b/examples/apple/coreml/scripts/extract_coreml_models.py
@@ -9,7 +9,7 @@
 import shutil
 from pathlib import Path
 
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from executorch.backends.apple.coreml import executorchcoreml
 from executorch.exir._serialize._program import deserialize_pte_binary
@@ -22,7 +22,12 @@
 COREML_BACKEND_ID = "CoreMLBackend"
 
 
-def extract_coreml_models(pte_data: bytes):
+def extract_coreml_models(
+    pte_data: bytes,
+    out_dir: Optional[Union[str, Path]] = None,
+) -> List[Path]:
+    out_root = Path(out_dir) if out_dir is not None else Path("extracted_coreml_models")
+
     pte_file = deserialize_pte_binary(pte_data)
     program = pte_file.program
 
@@ -44,6 +49,7 @@ def extract_coreml_models(pte_data: bytes):
     ]
 
     # Track extracted models to avoid duplicates (multifunction models share partitions)
+    extracted_paths: List[Path] = []
     extracted_keys: set = set()
     model_index: int = 1
 
@@ -95,7 +101,7 @@ def extract_coreml_models(pte_data: bytes):
         if model_name is None:
             model_name = f"model_{model_index}"
 
-        model_path: Path = Path() / "extracted_coreml_models" / model_name
+        model_path: Path = out_root / model_name
         if model_path.exists():
             shutil.rmtree(model_path.absolute())
         os.makedirs(model_path.absolute())
@@ -104,11 +110,14 @@ def extract_coreml_models(pte_data: bytes):
             coreml_processed_bytes, str(model_path.absolute())
         ):
             print(f"Core ML models are extracted and saved to path = {model_path}")
+            extracted_paths.append(model_path)
         model_index += 1
 
     if len(coreml_delegates) == 0:
         print("The model isn't delegated to Core ML.")
 
+    return extracted_paths
+
 
 def main() -> None:
     """
diff --git a/examples/apple/coreml/scripts/test_coreml_compute_plan.py b/examples/apple/coreml/scripts/test_coreml_compute_plan.py
