Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion backends/apple/coreai/.llms/skills/code-writing.md
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ repository root, in an env with ExecuTorch installed in editable mode and
python -m pytest backends/apple/coreai
```

Standard pytest options apply, e.g. `-v` or `-k sidecar`.
Standard pytest options apply, e.g. `-v` or `-k manifest`.

The real-toolchain AOT test (`CoreAIAOTCompileTest`) is gated on
`xcrun coreai-build`; it runs on macOS with the Metal Toolchain and skips
Expand Down
2 changes: 0 additions & 2 deletions backends/apple/coreai/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@

from executorch.backends.apple.coreai.compiler.preprocess import (
AOTCompileConfig,
coreai_sidecar_dir,
CoreAIBackend,
)
from executorch.backends.apple.coreai.partition.partitioner import CoreAIPartitioner
Expand Down Expand Up @@ -48,7 +47,6 @@ def get_default_compile_config() -> EdgeCompileConfig:
"AOTCompileConfig",
"CoreAIBackend",
"CoreAIPartitioner",
"coreai_sidecar_dir",
"get_default_compile_config",
"get_default_passes",
]
193 changes: 18 additions & 175 deletions backends/apple/coreai/compiler/preprocess.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,19 +7,16 @@
# Core AI backend for delegating an EdgeProgram to Apple's Core AI framework
# via the ``coreai-torch`` converter.

import contextlib
import copy
import hashlib
import json
import logging
import os
import shutil
import subprocess
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Any, Dict, final, Iterator, List, Optional, Tuple
from typing import Any, Dict, final, List, Optional, Tuple

from executorch.backends.apple.coreai.compiler.constants import MAIN_ENTRYPOINT
from executorch.backends.apple.coreai.compiler.enumerated_shapes import (
Expand Down Expand Up @@ -48,18 +45,9 @@
# are mmap-friendly on device. Applied to every file flattened into the NDS.
_ASSET_ALIGNMENT = 16

# Build-time only: directory where sidecar bundles are written. Read from the
# environment (set it via :func:`coreai_sidecar_dir`) so no build-machine path
# is ever serialized into the .pte. The runtime load directory is a separate,
# runtime-provided concern (default: the .pte's dir).
SIDECAR_DIR_ENV = "COREAI_SIDECAR_DIR"


class COMPILE_SPEC_KEYS(Enum):
# Whether this delegate's asset is delivered as a sidecar bundle (vs
# embedded in the .pte via NamedDataStore). Serialized: the runtime needs
# to know how to load the asset. Carries only the delivery mode, no path.
USES_SIDECAR = "coreai_uses_sidecar"
# Minimum OS

# Minimum OS / deployment version. This is a save-time property of the
# .aimodel (save_asset(minimum_os=...)) that ALSO feeds coreai-build, so it
Expand Down Expand Up @@ -91,12 +79,8 @@ class COMPILE_SPEC_KEYS(Enum):
class AssetPackaging(str, Enum):
# Portable .aimodel embedded in the .pte via NamedDataStore.
INLINE = "inline"
# Portable .aimodel written as a sidecar next to the .pte.
SIDECAR = "sidecar"
# AOT-compiled per-arch .aimodelc bundles embedded in the .pte.
AOT_COMPILED_INLINE = "aot_compiled_inline"
# AOT-compiled per-arch .aimodelc bundles written as a sidecar.
AOT_COMPILED_SIDECAR = "aot_compiled_sidecar"


@dataclass(frozen=True)
Expand Down Expand Up @@ -173,96 +157,12 @@ def _get_compile_spec(
return None


def _reject_existing_asset_dir(dest: Path) -> None:
"""Refuse to write over an asset directory a previous build left behind.

Each delegate owns one ``<hash>/`` directory, and the hash covers the
``.aimodel`` only. AOT options are applied after it is computed, so a
rebuild with a different platform or OS floor lands on the same path: left
alone, the manifest would advertise the new options while the bundles on
disk stayed from the earlier build. Two delegates in one build always
differ here, so this only fires across builds.
"""
if dest.exists():
raise RuntimeError(
f"sidecar asset directory already exists: {dest}\n"
f"A previous build wrote it, and its contents may have been "
f"compiled with different options. Remove it, or point "
f"COREAI_SIDECAR_DIR at a clean directory."
)


@contextlib.contextmanager
def coreai_sidecar_dir(path: str) -> Iterator[None]:
"""Set ``COREAI_SIDECAR_DIR`` for the duration of a ``with`` block.

Sidecar delivery (``uses_sidecar=True``, including AOT+sidecar) requires this
env var to name the build-time output directory. The prior value is
restored on exit.

The directory must not already hold asset directories: they are build
output, and stale ones are indistinguishable from this build's. Loose
files beside them (``.DS_Store``, a ``.pte``) are left alone. Setting the
env var directly skips this check, but each asset directory is still
guarded individually as it is written.

Example::

with coreai_sidecar_dir("build/model"):
lowered = to_edge_transform_and_lower(
ep, partitioner=[CoreAIPartitioner(uses_sidecar=True)]
)
"""
existing = Path(path)
if existing.exists() and not existing.is_dir():
raise RuntimeError(f"sidecar output path is not a directory: {path}")
if existing.is_dir() and any(p.is_dir() for p in existing.iterdir()):
raise RuntimeError(
f"sidecar output directory already holds assets: {path}\n"
f"Assets from an earlier build cannot be told apart from this "
f"one's. Remove the directory or choose another."
)
prev = os.environ.get(SIDECAR_DIR_ENV)
os.environ[SIDECAR_DIR_ENV] = path
try:
yield
finally:
if prev is None:
os.environ.pop(SIDECAR_DIR_ENV, None)
else:
os.environ[SIDECAR_DIR_ENV] = prev


def _nds_key(model_hash: str, relative_path: str) -> str:
# Namespaced by content hash so identical assets dedup to one entry and the
# hash doubles as an on-device cache key.
return f"coreai/{model_hash}/{relative_path}"


# Warn at most once per process that a sidecar dir was set but ignored.
_WARNED_SIDECAR_ENV_IGNORED = False


def _maybe_warn_sidecar_env_ignored() -> None:
"""Soft guardrail: env var set but this delegate uses inline delivery.

Likely a misconfiguration (user set the sidecar dir but forgot
``uses_sidecar=True``). Not an error: the env var is process-wide and may
legitimately be set for a different delegate, so we only warn, once.
"""
global _WARNED_SIDECAR_ENV_IGNORED
if _WARNED_SIDECAR_ENV_IGNORED:
return
if os.environ.get(SIDECAR_DIR_ENV):
_WARNED_SIDECAR_ENV_IGNORED = True
logger.warning(
"%s is set but this Core AI delegate uses inline delivery, "
"so the sidecar directory is ignored. Did you mean to pass "
"uses_sidecar=True to CoreAIPartitioner?",
SIDECAR_DIR_ENV,
)


def _asset_metadata(staging: Path) -> Dict[str, Any]:
"""Per-file sizes and a SHA-256 digest per bundle under ``staging``.

Expand Down Expand Up @@ -297,44 +197,20 @@ def _deliver(
model_hash: str,
packaging: AssetPackaging,
manifest_extra: Dict[str, Any],
sidecar_dir: Optional[str],
) -> PreprocessResult:
"""Ship a staging directory of bundles, embedded or on disk.
"""Embed a staging directory of bundles in the .pte with its manifest.

``staging`` holds the ``.aimodel`` / ``.aimodelc`` bundles, never their
contents, so both deliveries land on the same ``<hash>/<bundle>/`` layout:
inline under the ``coreai/`` NamedDataStore prefix, sidecar under the build
output directory. Keeping one implementation is what stops the two from
drifting apart again.

The sidecar write is staged and renamed into place, so a failure partway
cannot leave a half-populated ``<hash>/`` behind for
:func:`_reject_existing_asset_dir` to trip over on the next build.
contents, so assets land under ``coreai/<hash>/<bundle>/`` in the
NamedDataStore.
"""
manifest = {
"packaging": packaging.value,
"hash": model_hash,
**manifest_extra,
**_asset_metadata(staging),
}
if sidecar_dir is None:
return _embed_dir_inline(staging, model_hash, manifest)

dest = Path(sidecar_dir) / model_hash
dest.parent.mkdir(parents=True, exist_ok=True)
_reject_existing_asset_dir(dest)
pending = dest.with_name(f".{model_hash}.partial")
if pending.exists():
shutil.rmtree(pending)
try:
pending.mkdir(parents=True)
for bundle in sorted(staging.iterdir()):
shutil.move(str(bundle), str(pending / bundle.name))
pending.rename(dest)
except BaseException:
shutil.rmtree(pending, ignore_errors=True)
raise
return PreprocessResult(processed_bytes=json.dumps(manifest).encode("utf-8"))
return _embed_dir_inline(staging, model_hash, manifest)


def _prepare_program_for_conversion(edge_program: ExportedProgram) -> ExportedProgram:
Expand Down Expand Up @@ -384,8 +260,7 @@ def _embed_dir_inline(
"""Flatten every file under ``root_dir`` into the NamedDataStore.

``root_dir`` is the directory *containing* the ``.aimodel`` / ``.aimodelc``
bundles, never a bundle itself, so the bundle name survives into the keys
and the asset reconstructs to the layout the sidecar routes write on disk.
bundles, never a bundle itself, so the bundle name survives into the keys.

Keys are ``coreai/{hash}/{relpath}``; ``manifest_extra`` becomes the
manifest (packaging, files, archs, ...).
Expand Down Expand Up @@ -566,16 +441,12 @@ def _compile_aot(program, opts: Dict[str, Any], tmp_dir: Path) -> Tuple[str, Pat
class CoreAIBackend(BackendDetails):
"""AOT lowering of an edge program to a Core AI asset.

Delivery is chosen by compile specs, along two orthogonal axes:
Delivery is chosen by compile specs: the portable ``.aimodel`` (default),
or AOT-compiled per-architecture ``.aimodelc`` bundles (``aot_compile``, via
``xcrun coreai-build``; architecture selection is one / a list / all). The
bundles are embedded in the ``.pte`` via NamedDataStore.

* **Format**: the portable ``.aimodel`` (default), or AOT-compiled
per-architecture ``.aimodelc`` bundles (``aot_compile``, via
``xcrun coreai-build``; architecture selection is one / a list / all).
* **Location**: embedded in the ``.pte`` via NamedDataStore (default), or a
``sidecar`` written to ``$COREAI_SIDECAR_DIR`` (``uses_sidecar``).

``processed_bytes`` is always a small JSON manifest naming what/where; the
bytes live in the NamedDataStore (inline) or on disk (sidecar).
``processed_bytes`` is always a small JSON manifest naming what/where.
Runtime execution is not wired up yet.
"""

Expand All @@ -588,20 +459,6 @@ def preprocess(
_get_compile_spec(compile_specs, COMPILE_SPEC_KEYS.AOT_COMPILE_CONFIG)
is not None
)
uses_sidecar = (
_get_compile_spec(compile_specs, COMPILE_SPEC_KEYS.USES_SIDECAR) is not None
)

# Sidecar delivery (portable or AOT) needs a build-time output dir.
sidecar_dir = None
if uses_sidecar:
sidecar_dir = os.environ.get(SIDECAR_DIR_ENV)
if not sidecar_dir:
raise ValueError(
"sidecar asset delivery requires the "
f"{SIDECAR_DIR_ENV} environment variable to name the "
"build-time output directory (set it via coreai_sidecar_dir)"
)

mutated_inputs = edge_program.graph_signature.user_inputs_to_mutate
if mutated_inputs:
Expand Down Expand Up @@ -643,57 +500,44 @@ def preprocess(
# (via save_asset(minimum_os=...)). In the aot-compiled path the temp
# .aimodel is discarded, so it keeps the default floor (no
# double-specification).
if not uses_sidecar:
_maybe_warn_sidecar_env_ignored()
if aot_compiled:
return CoreAIBackend._preprocess_aot_compiled(
program, _aot_compile_options(compile_specs), sidecar_dir, bindings
program, _aot_compile_options(compile_specs), bindings
)
return CoreAIBackend._preprocess_portable(
program,
_min_os_version(compile_specs),
sidecar_dir,
bindings,
program, _min_os_version(compile_specs), bindings
)

# Portable .aimodel delivery.
@staticmethod
def _preprocess_portable(
program, min_os, sidecar_dir: Optional[str], bindings: Dict[str, Any]
program, min_os, bindings: Dict[str, Any]
) -> PreprocessResult:
with TemporaryDirectory() as tmp:
model_hash = _save_and_hash(program, Path(tmp) / "model.aimodel", min_os)
return _deliver(
Path(tmp),
model_hash,
AssetPackaging.SIDECAR if sidecar_dir else AssetPackaging.INLINE,
AssetPackaging.INLINE,
{
**bindings,
# relative path the runtime resolves against its base
"path": f"{model_hash}/model.aimodel",
"min_deployment_version": _os_version_text(min_os),
},
sidecar_dir,
)

# AOT-compiled .aimodelc delivery (per architecture).
@staticmethod
def _preprocess_aot_compiled(
program,
opts: Dict[str, Any],
sidecar_dir: Optional[str],
bindings: Dict[str, Any],
program, opts: Dict[str, Any], bindings: Dict[str, Any]
) -> PreprocessResult:
with TemporaryDirectory() as tmp:
model_hash, out = _compile_aot(program, opts, Path(tmp))
return _deliver(
out,
model_hash,
(
AssetPackaging.AOT_COMPILED_SIDECAR
if sidecar_dir
else AssetPackaging.AOT_COMPILED_INLINE
),
AssetPackaging.AOT_COMPILED_INLINE,
{
**bindings,
"platform": opts["platform"],
Expand All @@ -703,5 +547,4 @@ def _preprocess_aot_compiled(
for arch, bundle in _compiled_arch_bundles(out)
},
},
sidecar_dir,
)
10 changes: 0 additions & 10 deletions backends/apple/coreai/partition/partitioner.py
Original file line number Diff line number Diff line change
Expand Up @@ -197,7 +197,6 @@ class CoreAIPartitioner(Partitioner):
def __init__(
self,
*,
uses_sidecar: bool = False,
aot_compile_config: Optional[AOTCompileConfig] = None,
min_deployment_version: Optional[str] = None,
input_enumerations: Optional[
Expand All @@ -206,13 +205,6 @@ def __init__(
take_over_constant_data: bool = True,
take_over_mutable_buffer: bool = True,
) -> None:
# uses_sidecar selects sidecar delivery (vs inline). It is embedded as a
# compile spec because the runtime needs to know how to load the asset,
# but it carries only the mode, no path. The build-time output directory
# comes from the COREAI_SIDECAR_DIR env var (see preprocess.py /
# coreai_sidecar_dir), never a compile spec, so no build-machine path is
# serialized.
#
# aot_compile_config requests ahead-of-time ``xcrun coreai-build
# compile`` in preprocess, emitting per-architecture ``.aimodelc``
# bundles instead of the portable ``.aimodel``. It is serialized as a
Expand All @@ -225,8 +217,6 @@ def __init__(
# symbols in :meth:`partition` and propagated to each subgraph boundary
# in preprocess, so it is not built into a compile spec here.
specs = []
if uses_sidecar:
specs.append(CompileSpec(COMPILE_SPEC_KEYS.USES_SIDECAR.value, b"1"))
if min_deployment_version is not None:
specs.append(
CompileSpec(
Expand Down
Loading
Loading