From 1a12548f831d76d0cfc20946eccc2ccc199a957c Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 30 Jul 2026 16:57:13 -0700 Subject: [PATCH] [ExecuTorch][Core AI] Core backend: partitioner, preprocess, IO compat, copy-op pass --- backends/apple/coreai/README.md | 7 + backends/apple/coreai/__init__.py | 48 ++ backends/apple/coreai/compiler/__init__.py | 5 + backends/apple/coreai/compiler/constants.py | 12 + backends/apple/coreai/compiler/io_compat.py | 168 ++++++ backends/apple/coreai/compiler/preprocess.py | 532 ++++++++++++++++++ backends/apple/coreai/partition/__init__.py | 5 + .../apple/coreai/partition/partitioner.py | 303 ++++++++++ .../apple/coreai/passes/replace_copy_ops.py | 66 +++ .../passes/test/test_replace_copy_ops.py | 74 +++ backends/apple/coreai/test/__init__.py | 5 + backends/apple/coreai/test/test_io_compat.py | 264 +++++++++ .../apple/coreai/test/test_partitioner.py | 285 ++++++++++ backends/apple/coreai/test/test_preprocess.py | 305 ++++++++++ 14 files changed, 2079 insertions(+) create mode 100644 backends/apple/coreai/README.md create mode 100644 backends/apple/coreai/__init__.py create mode 100644 backends/apple/coreai/compiler/__init__.py create mode 100644 backends/apple/coreai/compiler/constants.py create mode 100644 backends/apple/coreai/compiler/io_compat.py create mode 100644 backends/apple/coreai/compiler/preprocess.py create mode 100644 backends/apple/coreai/partition/__init__.py create mode 100644 backends/apple/coreai/partition/partitioner.py create mode 100644 backends/apple/coreai/passes/replace_copy_ops.py create mode 100644 backends/apple/coreai/passes/test/test_replace_copy_ops.py create mode 100644 backends/apple/coreai/test/__init__.py create mode 100644 backends/apple/coreai/test/test_io_compat.py create mode 100644 backends/apple/coreai/test/test_partitioner.py create mode 100644 backends/apple/coreai/test/test_preprocess.py diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md new file mode 100644 index 00000000000..0c9ad10a1fc --- /dev/null +++ b/backends/apple/coreai/README.md @@ -0,0 +1,7 @@ +# Core AI backend (ExecuTorch) + +> ⚠️ **Under construction — not for use.** +> This backend is a work in progress. It is **not** ready for production or +> general use, and provides **no backward- or forward-compatibility guarantees** +> (no BC/FC). APIs, compile specs, serialized manifests, and file layouts may +> change or break at any time without notice. Do not depend on it. diff --git a/backends/apple/coreai/__init__.py b/backends/apple/coreai/__init__.py new file mode 100644 index 00000000000..1bfa23416e5 --- /dev/null +++ b/backends/apple/coreai/__init__.py @@ -0,0 +1,48 @@ +# 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 typing import List + +from executorch.backends.apple.coreai.compiler.preprocess import ( + AOTCompileConfig, + coreai_sidecar_dir, + CoreAIBackend, +) +from executorch.backends.apple.coreai.partition.partitioner import CoreAIPartitioner +from executorch.exir import EdgeCompileConfig +from executorch.exir.pass_base import PassType + + +def get_default_passes() -> List[PassType]: + """Default edge transform passes for Core AI lowering. + + ``NarrowToCoreAIDtypesPass`` casts int64/float64 graph inputs to 32-bit at + the boundary (preserving the model's external I/O dtype) so index-style + inputs can be delegated. Core AI supports only up-to-32-bit dtypes. + """ + from executorch.backends.apple.coreai.passes import NarrowToCoreAIDtypesPass + + return [NarrowToCoreAIDtypesPass()] + + +def get_default_compile_config() -> EdgeCompileConfig: + """Default ``EdgeCompileConfig`` for Core AI lowering. + + ``_skip_dim_order=True`` keeps ExecuTorch on ``aten._to_copy`` instead of + emitting ``dim_order_ops._to_dim_order_copy``, which coreai-torch cannot + lower (its validator requires dim-order ops be decomposed). + """ + return EdgeCompileConfig(_skip_dim_order=True) + + +__all__ = [ + "AOTCompileConfig", + "CoreAIBackend", + "CoreAIPartitioner", + "coreai_sidecar_dir", + "get_default_compile_config", + "get_default_passes", +] diff --git a/backends/apple/coreai/compiler/__init__.py b/backends/apple/coreai/compiler/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/apple/coreai/compiler/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/apple/coreai/compiler/constants.py b/backends/apple/coreai/compiler/constants.py new file mode 100644 index 00000000000..46e3e275ab0 --- /dev/null +++ b/backends/apple/coreai/compiler/constants.py @@ -0,0 +1,12 @@ +# 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. + +"""Shared constants for the Core AI backend compiler.""" + +# The single entrypoint / graph name of a per-delegate coreai AIProgram. Each +# ExecuTorch delegate is converted to its own coreai program whose only graph is +# named "main" (``save_asset`` emits ``main.hash`` / ``main.mlirb`` accordingly). +MAIN_ENTRYPOINT = "main" diff --git a/backends/apple/coreai/compiler/io_compat.py b/backends/apple/coreai/compiler/io_compat.py new file mode 100644 index 00000000000..c99a05c61f7 --- /dev/null +++ b/backends/apple/coreai/compiler/io_compat.py @@ -0,0 +1,168 @@ +# 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. + +"""Delegate-boundary compatibility checks for the Core AI backend. + +ExecuTorch calls a delegate with plain tensors whose dtypes/shapes come from the +edge subgraph boundary. The converted ``.aimodel`` must declare I/O that matches +what ExecuTorch will feed/read, or the runtime will mismatch. :func:` +assert_io_compatible` compares the two *before* compilation and fails loudly. +""" + +import logging +from typing import Any, List, Sequence, Tuple + +import torch +from executorch.backends.apple.coreai.compiler.constants import MAIN_ENTRYPOINT +from torch.export.exported_program import ExportedProgram + +logger = logging.getLogger(__name__) + +# coreai's element-type spelling for each torch dtype (from coreai_torch's +# _type_mapping: signed ints -> siN, unsigned -> uiN, bool -> i1, floats -> fN). +# These are the un-narrowed spellings, asserted exactly. coreai narrows +# int64->si32 and float64->f32 in get_tensor_type, so those two dtypes surface +# here as mismatches by design (that is the signal). +_TORCH_TO_COREAI_ELT = { + torch.float16: "f16", + torch.float32: "f32", + torch.float64: "f64", + torch.bfloat16: "bf16", + torch.bool: "i1", + torch.int8: "si8", + torch.uint8: "ui8", + torch.int16: "si16", + torch.uint16: "ui16", + torch.int32: "si32", + torch.uint32: "ui32", + torch.int64: "si64", +} + + +def _coreai_io(program, entrypoint: str = MAIN_ENTRYPOINT): + """Core AI graph I/O as lists of (element_type_str, shape_or_None). + + Ranked tensors -> (element_type_str, shape_tuple). Any non-tensor coreai + type is represented as (str(type), None) so the comparison can flag it, + rather than assuming a tensor and crashing on ``.element_type`` / ``.shape``. + """ + from coreai._compiler.ir import RankedTensorType + + func_type = program.get_graph(entrypoint).function_type.value + + def _describe(t): + if RankedTensorType.isinstance(t): + return (str(t.element_type), tuple(t.shape)) + return (str(t), None) + + inputs = [_describe(t) for t in func_type.inputs] + outputs = [_describe(t) for t in func_type.results] + return inputs, outputs + + +def _edge_io(edge_program: ExportedProgram): + """Edge subgraph I/O as lists of (dtype_or_typename, shape_or_None). + + Tensors -> (torch_dtype, shape_tuple). Non-tensor I/O (symint / const + int/float) -> (type_name_str, None), kept as entries so both sides stay + positionally aligned with :func:`_coreai_io` (which also emits non-tensor + entries). + """ + placeholders = { + n.name: n for n in edge_program.graph.nodes if n.op == "placeholder" + } + inputs = [] + for name in edge_program.graph_signature.user_inputs: + # A non-tensor input appears in user_inputs as its literal value, not a + # placeholder name. + node = placeholders.get(name) + val = node.meta.get("val") if node is not None else name + if hasattr(val, "dtype"): + inputs.append((val.dtype, tuple(val.shape))) + else: + inputs.append((type(val).__name__, None)) + outputs = [] + for arg in edge_program.graph.output_node().args[0]: + val = arg.meta.get("val") if hasattr(arg, "meta") else arg + if hasattr(val, "dtype"): + outputs.append((val.dtype, tuple(val.shape))) + else: + outputs.append((type(val).__name__, None)) + return inputs, outputs + + +def io_mismatches( + coreai_io: Sequence[Tuple[str, Tuple[Any, ...]]], + edge_io: Sequence[Tuple[Any, Tuple[Any, ...]]], + kind: str, +) -> List[str]: + """Return human-readable mismatch messages between coreai and edge I/O. + + Compares count, per-tensor dtype (all mapped types: floats and ints), rank, + and static (concrete) dims. Symbolic (dynamic) dims are skipped: coreai is + specialized to a concrete size while the edge dim stays symbolic. + """ + if len(coreai_io) != len(edge_io): + return [ + f"{kind} count mismatch: .aimodel has {len(coreai_io)}, " + f"ExecuTorch feeds {len(edge_io)}" + ] + errors: List[str] = [] + for i, ((elt, cshape), (dtype, eshape)) in enumerate(zip(coreai_io, edge_io)): + coreai_nontensor = cshape is None + edge_nontensor = eshape is None + if coreai_nontensor or edge_nontensor: + # At least one side is non-tensor; only a class mismatch matters. + if coreai_nontensor != edge_nontensor: + errors.append( + f"{kind} {i}: tensor/non-tensor mismatch " + f"(.aimodel={elt}, ExecuTorch feeds {dtype})" + ) + continue + expected = _TORCH_TO_COREAI_ELT.get(dtype) + if expected is not None and elt != expected: + errors.append( + f"{kind} {i}: dtype mismatch (.aimodel={elt}, ExecuTorch feeds " + f"{dtype}, expected {expected})" + ) + if len(cshape) != len(eshape): + errors.append( + f"{kind} {i}: rank mismatch (.aimodel={cshape}, " + f"ExecuTorch={eshape})" + ) + continue + for d, (cdim, edim) in enumerate(zip(cshape, eshape)): + if isinstance(edim, int) and cdim != edim: + errors.append( + f"{kind} {i} dim {d}: static shape mismatch " + f"(.aimodel={cdim}, ExecuTorch={edim})" + ) + return errors + + +def assert_io_compatible(program, edge_program: ExportedProgram) -> None: + """Raise if the ``.aimodel`` boundary I/O is incompatible with ExecuTorch.""" + coreai_in, coreai_out = _coreai_io(program) + edge_in, edge_out = _edge_io(edge_program) + logger.info( + "Core AI delegate boundary I/O:\n" + " inputs: .aimodel=%s\n" + " ExecuTorch=%s\n" + " outputs: .aimodel=%s\n" + " ExecuTorch=%s", + coreai_in, + edge_in, + coreai_out, + edge_out, + ) + errors = io_mismatches(coreai_in, edge_in, "input") + io_mismatches( + coreai_out, edge_out, "output" + ) + if errors: + raise ValueError( + "Core AI delegate boundary is incompatible with ExecuTorch:\n " + + "\n ".join(errors) + ) diff --git a/backends/apple/coreai/compiler/preprocess.py b/backends/apple/coreai/compiler/preprocess.py new file mode 100644 index 00000000000..ef890368a9d --- /dev/null +++ b/backends/apple/coreai/compiler/preprocess.py @@ -0,0 +1,532 @@ +# 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. + +# Core AI backend for delegating an EdgeProgram to Apple's Core AI framework +# via the ``coreai-torch`` converter. + +import contextlib +import copy +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 executorch.backends.apple.coreai.compiler.constants import MAIN_ENTRYPOINT +from executorch.backends.apple.coreai.compiler.io_compat import assert_io_compatible +from executorch.backends.apple.coreai.passes.replace_copy_ops import ( + ReplaceCopyOpsWithFunctionalPass, +) +from executorch.exir._serialize._named_data_store import NamedDataStore +from executorch.exir.backend.backend_details import ( + BackendDetails, + ExportedProgram, + PreprocessResult, +) +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.dialects.edge._ops import EdgeOpOverload + +logger = logging.getLogger(__name__) + +# Alignment for the embedded asset files so the large blobs (e.g. ``main.mlirb``) +# 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 / deployment version. This is a save-time property of the + # .aimodel (save_asset(minimum_os=...)) that ALSO feeds coreai-build, so it + # applies to every delivery path, AOT or not. Numeric string, e.g. "27.0". + MIN_DEPLOYMENT_VERSION = "coreai_min_deployment_version" + + # Ahead-of-time compilation (xcrun coreai-build). + # Presence implies compiling the .aimodel to per-architecture .aimodelc + # bundles at build time (requires macOS + the Metal Toolchain). The value is + # a JSON object of build-only coreai-build options; an empty {} means all + # defaults: + # {"platform": "iOS|macOS|watchOS|visionOS|tvOS", + # "preferred_compute": "gpu|neural-engine|none", + # "architectures": [...], # empty/absent => all supported + # "expect_frequent_reshapes": bool} + # (min_deployment_version is a separate general spec that also applies to + # the portable .aimodel, so it is not part of this AOT-only blob.) + AOT_COMPILE_CONFIG = "coreai_aot_compile_config" + + +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) +class AOTCompileConfig: + """Build-time ``coreai-build`` options for AOT compilation. + + Constructible in Python and (de)serializable to/from JSON so it can ride as + the ``AOT_COMPILE_CONFIG`` compile spec. ``min_deployment_version`` is not + part of this object; it is a general spec (it also sets the portable + .aimodel's OS floor), so it is passed separately to the partitioner. + """ + + platform: Optional[str] = None # iOS / macOS / watchOS / visionOS / tvOS + preferred_compute: Optional[str] = None # gpu / neural-engine / none + architectures: Optional[List[str]] = None # None / empty => all supported + expect_frequent_reshapes: bool = False + + # Fields allowed in the JSON form (kept in sync with the dataclass fields). + _ALLOWED = ( + "platform", + "preferred_compute", + "architectures", + "expect_frequent_reshapes", + ) + + def to_dict(self) -> Dict[str, Any]: + """Minimal JSON-able dict (omits None/defaults).""" + d: Dict[str, Any] = {} + if self.platform is not None: + d["platform"] = self.platform + if self.preferred_compute is not None: + d["preferred_compute"] = self.preferred_compute + if self.architectures: + d["architectures"] = list(self.architectures) + if self.expect_frequent_reshapes: + d["expect_frequent_reshapes"] = True + return d + + def to_json(self) -> str: + return json.dumps(self.to_dict()) + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "AOTCompileConfig": + unexpected = set(data) - set(cls._ALLOWED) + if unexpected: + raise ValueError( + f"unexpected AOTCompileConfig field(s): {sorted(unexpected)}; " + f"allowed: {sorted(cls._ALLOWED)}" + ) + arch = data.get("architectures") + return cls( + platform=data.get("platform"), + preferred_compute=data.get("preferred_compute"), + architectures=list(arch) if arch else None, + expect_frequent_reshapes=bool(data.get("expect_frequent_reshapes", False)), + ) + + @classmethod + def from_json(cls, text: str) -> "AOTCompileConfig": + return cls.from_dict(json.loads(text)) + + +def _get_compile_spec( + compile_specs: List[CompileSpec], key: COMPILE_SPEC_KEYS +) -> Optional[bytes]: + for spec in compile_specs: + if spec.key == key.value: + return spec.value + return None + + +@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. + + Example:: + + with coreai_sidecar_dir("build/model"): + lowered = to_edge_transform_and_lower( + ep, partitioner=[CoreAIPartitioner(uses_sidecar=True)] + ) + """ + 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 _prepare_program_for_conversion(edge_program: ExportedProgram) -> ExportedProgram: + """Rewrite an edge-dialect program into the ATen form ``coreai-torch`` expects. + + ``coreai-torch`` keys its lowering table on plain ATen overload names + (``"addmm.default"``), but ExecuTorch hands the backend edge ops whose + ``__name__`` is prefixed (``"aten.addmm.default"``) and whose view ops are + functionalized ``*_copy`` variants. So we: + + 1. Remap edge ``*_copy`` ops to their functional ATen forms + (:class:`ReplaceCopyOpsWithFunctionalPass`), then + 2. Unwrap any remaining ``EdgeOpOverload`` to its underlying ATen overload. + """ + ep = copy.deepcopy(edge_program) + ReplaceCopyOpsWithFunctionalPass()(ep.graph_module) + for node in ep.graph.nodes: + if node.op == "call_function" and isinstance(node.target, EdgeOpOverload): + node.target = node.target._op + ep.graph_module.recompile() + return ep + + +def _convert_to_aiprogram(edge_program: ExportedProgram): + from coreai_torch import TorchConverter + + aten_program = _prepare_program_for_conversion(edge_program) + converter = TorchConverter() + converter.add_exported_program(aten_program) + return converter.to_coreai() + + +# Asset embedding helpers. +def _embed_dir_inline( + root_dir: Path, model_hash: str, manifest_extra: Dict[str, Any] +) -> PreprocessResult: + """Flatten every file under ``root_dir`` into the NamedDataStore. + + Keys are ``coreai/{hash}/{relpath}``; the manifest lists the relpaths plus + any ``manifest_extra`` (packaging, archs, ...). + """ + store = NamedDataStore() + files: List[str] = [] + for path in sorted(root_dir.rglob("*")): + if path.is_file(): + rel = path.relative_to(root_dir).as_posix() + files.append(rel) + store.add_named_data( + _nds_key(model_hash, rel), + path.read_bytes(), + alignment=_ASSET_ALIGNMENT, + ) + manifest = {"hash": model_hash, "files": files, **manifest_extra} + return PreprocessResult( + processed_bytes=json.dumps(manifest).encode("utf-8"), + data_store_output=store.get_named_data_store_output(), + ) + + +# AOT (coreai-build) helpers. +def _aot_compile_options(compile_specs: List[CompileSpec]) -> Dict[str, Any]: + """Normalized coreai-build opts from AOT_COMPILE_CONFIG + general specs. + + ``min_deployment_version`` is pulled from its own (general) spec, since it + also applies to the portable .aimodel. + """ + raw = _get_compile_spec(compile_specs, COMPILE_SPEC_KEYS.AOT_COMPILE_CONFIG) + config = AOTCompileConfig.from_json(raw.decode()) if raw else AOTCompileConfig() + min_dep = _get_compile_spec(compile_specs, COMPILE_SPEC_KEYS.MIN_DEPLOYMENT_VERSION) + return { + "platform": config.platform or "macOS", + # None => let coreai-build use its own default; matches save_asset default. + "min_deployment_version": min_dep.decode() if min_dep else None, + "preferred_compute": config.preferred_compute or "none", + # empty list => compile for all supported architectures + "architectures": list(config.architectures or []), + "expect_frequent_reshapes": config.expect_frequent_reshapes, + } + + +def _min_os_version(compile_specs: List[CompileSpec]): + """Map the MIN_DEPLOYMENT_VERSION spec to a coreai ``OSVersion`` (or None). + + Applied to every delivery path via ``save_asset(minimum_os=...)``. Accepts a + numeric string ("27.0" / "27") or an ``OSVersion`` name ("v27"). Returns + None when unset so ``save_asset`` uses its own default. + """ + raw = _get_compile_spec(compile_specs, COMPILE_SPEC_KEYS.MIN_DEPLOYMENT_VERSION) + if raw is None: + return None + from coreai.authoring import OSVersion + + text = raw.decode() + name = text if text.startswith("v") else f"v{text.split('.')[0]}" + try: + return OSVersion(name) + except ValueError as e: + available = [m.value for m in OSVersion] + raise ValueError( + f"unsupported min_deployment_version {text!r} (maps to {name!r}); " + f"available coreai OSVersion values: {available}" + ) from e + + +def _save_asset(program, path: Path, min_os) -> None: + """``program.save_asset`` honoring an optional ``minimum_os``.""" + if min_os is None: + program.save_asset(path) + else: + program.save_asset(path, minimum_os=min_os) + + +def _save_and_hash(program, path: Path, min_os=None) -> str: + """Save the ``.aimodel`` and return its content hash (``main.hash``).""" + _save_asset(program, path, min_os) + return (path / f"{MAIN_ENTRYPOINT}.hash").read_bytes().hex() + + +def _run_coreai_build(aimodel_path: Path, out_dir: Path, opts: Dict[str, Any]) -> None: + """Invoke ``xcrun coreai-build compile`` to produce per-arch .aimodelc.""" + cmd = [ + "xcrun", + "coreai-build", + "compile", + str(aimodel_path), + "--output", + str(out_dir), + "--platform", + opts["platform"], + ] + if opts["min_deployment_version"]: + cmd += ["--min-deployment-version", opts["min_deployment_version"]] + if opts["preferred_compute"] and opts["preferred_compute"] != "none": + cmd += ["--preferred-compute", opts["preferred_compute"]] + for arch in opts["architectures"]: + cmd += ["--architecture", arch] + if opts["expect_frequent_reshapes"]: + cmd += ["--expect-frequent-reshapes"] + + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=300) + except FileNotFoundError as e: + raise RuntimeError( + "`xcrun coreai-build` not found. AOT compilation requires macOS with " + "the Metal Toolchain installed. Disable aot_compile or install the " + "toolchain (xcodebuild -downloadComponent MetalToolchain)." + ) from e + except subprocess.TimeoutExpired as e: + raise RuntimeError( + f"`coreai-build compile` timed out after {e.timeout}s" + ) from e + if result.returncode != 0: + raise RuntimeError( + f"`coreai-build compile` failed (rc={result.returncode}):\n{result.stderr}" + ) + + +def _compiled_arch_bundles(out_dir: Path) -> List[Tuple[str, Path]]: + """Return [(arch, bundle_dir), ...] for each ``model..aimodelc``.""" + bundles = [] + for path in sorted(out_dir.iterdir()): + if path.is_dir() and path.name.endswith(".aimodelc"): + name = path.name + # ``model..aimodelc`` -> ```` + arch = name[len("model.") : -len(".aimodelc")] + bundles.append((arch, path)) + if not bundles: + raise RuntimeError(f"coreai-build produced no .aimodelc bundles in {out_dir}") + return bundles + + +def _compile_aot(program, opts: Dict[str, Any], tmp_dir: Path) -> Tuple[str, Path]: + """Save the intermediate ``.aimodel`` and AOT-compile it to ``.aimodelc``. + + Returns ``(model_hash, compiled_out_dir)``. The intermediate ``.aimodel`` + keeps the default OS floor; ``min_deployment_version`` targets the compiled + ``.aimodelc`` via ``coreai-build``, so it is not double-specified. + """ + model_hash = _save_and_hash(program, tmp_dir / "model.aimodel") + out = tmp_dir / "compiled" + out.mkdir() + _run_coreai_build(tmp_dir / "model.aimodel", out, opts) + return model_hash, out + + +@final +class CoreAIBackend(BackendDetails): + """AOT lowering of an edge program to a Core AI asset. + + Delivery is chosen by compile specs, along two orthogonal axes: + + * **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). + Runtime execution is not wired up yet. + """ + + @staticmethod + def preprocess( + edge_program: ExportedProgram, + compile_specs: List[CompileSpec], + ) -> PreprocessResult: + aot_compiled = ( + _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)" + ) + + program = _convert_to_aiprogram(edge_program) + # Fail fast if the .aimodel boundary I/O won't match what ET feeds/reads. + assert_io_compatible(program, edge_program) + + # min-deployment-version is a single knob applied to whichever artifact + # ships: for aot-compiled delivery the .aimodelc (via coreai-build + # --min-deployment-version), for portable delivery the .aimodel's floor + # (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 aot_compiled: + opts = _aot_compile_options(compile_specs) + if uses_sidecar: + return CoreAIBackend._preprocess_aot_compiled_sidecar( + program, opts, sidecar_dir + ) + return CoreAIBackend._preprocess_aot_compiled_inline(program, opts) + + min_os = _min_os_version(compile_specs) + if uses_sidecar: + return CoreAIBackend._preprocess_sidecar(program, sidecar_dir, min_os) + _maybe_warn_sidecar_env_ignored() + return CoreAIBackend._preprocess_inline(program, min_os) + + # Portable .aimodel delivery. + @staticmethod + def _preprocess_inline(program, min_os) -> PreprocessResult: + with TemporaryDirectory() as tmp: + asset_dir = Path(tmp) / "model.aimodel" + model_hash = _save_and_hash(program, asset_dir, min_os) + return _embed_dir_inline( + asset_dir, model_hash, {"packaging": AssetPackaging.INLINE.value} + ) + + @staticmethod + def _preprocess_sidecar(program, sidecar_dir: str, min_os) -> PreprocessResult: + out_dir = Path(sidecar_dir) + out_dir.mkdir(parents=True, exist_ok=True) + with TemporaryDirectory() as tmp: + staged = Path(tmp) / "model.aimodel" + model_hash = _save_and_hash(program, staged, min_os) + bundle_name = f"{model_hash}.aimodel" + final = out_dir / bundle_name + if not final.exists(): + shutil.move(str(staged), str(final)) + + manifest = { + "packaging": AssetPackaging.SIDECAR.value, + "hash": model_hash, + "path": bundle_name, + } + return PreprocessResult(processed_bytes=json.dumps(manifest).encode("utf-8")) + + # AOT-compiled .aimodelc delivery (per architecture). + @staticmethod + def _preprocess_aot_compiled_inline( + program, opts: Dict[str, Any] + ) -> PreprocessResult: + with TemporaryDirectory() as tmp: + model_hash, out = _compile_aot(program, opts, Path(tmp)) + archs = [arch for arch, _ in _compiled_arch_bundles(out)] + return _embed_dir_inline( + out, + model_hash, + { + "packaging": AssetPackaging.AOT_COMPILED_INLINE.value, + "platform": opts["platform"], + "min_deployment_version": opts["min_deployment_version"], + "archs": archs, + }, + ) + + @staticmethod + def _preprocess_aot_compiled_sidecar( + program, opts: Dict[str, Any], sidecar_dir: str + ) -> PreprocessResult: + out_base = Path(sidecar_dir) + out_base.mkdir(parents=True, exist_ok=True) + with TemporaryDirectory() as tmp: + model_hash, out = _compile_aot(program, opts, Path(tmp)) + dest = out_base / model_hash + archs: Dict[str, str] = {} + for arch, bundle in _compiled_arch_bundles(out): + target = dest / bundle.name + if not target.exists(): + target.parent.mkdir(parents=True, exist_ok=True) + shutil.move(str(bundle), str(target)) + # relative path the runtime resolves against its base dir + archs[arch] = f"{model_hash}/{bundle.name}" + + manifest = { + "packaging": AssetPackaging.AOT_COMPILED_SIDECAR.value, + "hash": model_hash, + "platform": opts["platform"], + "min_deployment_version": opts["min_deployment_version"], + "archs": archs, + } + return PreprocessResult(processed_bytes=json.dumps(manifest).encode("utf-8")) diff --git a/backends/apple/coreai/partition/__init__.py b/backends/apple/coreai/partition/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/apple/coreai/partition/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/apple/coreai/partition/partitioner.py b/backends/apple/coreai/partition/partitioner.py new file mode 100644 index 00000000000..af70c22a6c8 --- /dev/null +++ b/backends/apple/coreai/partition/partitioner.py @@ -0,0 +1,303 @@ +# 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 inspect +import logging +from typing import Callable, List, Optional, Tuple + +import torch + +from executorch.backends.apple.coreai.compiler.preprocess import ( + AOTCompileConfig, + COMPILE_SPEC_KEYS, + CoreAIBackend, +) +from executorch.backends.apple.coreai.passes.replace_copy_ops import functional_aten_op +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.backend.partitioner import ( + DelegationSpec, + Partitioner, + PartitionResult, +) +from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch.export.exported_program import ExportedProgram +from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner +from torch.fx.passes.operator_support import OperatorSupportBase + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + +# Nodes carrying this meta key are never claimed by the delegate, even if Core AI +# could otherwise lower them. Set it (e.g. via :func:`do_not_delegate`) on any +# node you want to keep running outside Core AI. +DO_NOT_DELEGATE_TAG = "coreai_do_not_delegate" + + +def do_not_delegate(node: torch.fx.Node) -> None: + """Mark a node so the CoreAIPartitioner leaves it out of the delegate.""" + node.meta[DO_NOT_DELEGATE_TAG] = True + + +def _resolvers() -> Tuple[dict, dict]: + """Return ``coreai-torch``'s (aten, higher-order) lowering tables. + + These dicts are the source of truth for which ops the Core AI converter + can lower; the partitioner mirrors them so a tagged subgraph is guaranteed + to be convertible. Imported lazily so this module is importable in + environments without ``coreai-torch`` installed. + """ + from coreai_torch._aten_to_core import ( + _aten_to_core_resolver, + _higher_order_resolver, + ) + + return _aten_to_core_resolver, _higher_order_resolver + + +def _underlying_target(target): + """Unwrap an ExecuTorch ``EdgeOpOverload`` to its plain ATen overload. + + ExecuTorch edge ops wrap the ATen overload in an ``EdgeOpOverload`` whose + ``__name__`` is prefixed (``"aten.view.default"``); the overload at + ``target._op`` has the bare name (``"view.default"``) that the converter + keys on. Plain ``OpOverload``s also expose a ``_op`` attribute, but its + ``__name__`` is empty, so we must only unwrap genuine edge ops. + """ + if isinstance(target, EdgeOpOverload): + return target._op + return target + + +def _coreai_op_name(target) -> Optional[str]: + """Return the ``coreai-torch`` resolver key for an fx call_function target.""" + return getattr(_underlying_target(target), "__name__", None) + + +def _coreai_namespace(target) -> Optional[str]: + return getattr(_underlying_target(target), "namespace", None) + + +def is_coreai_supported_target(target) -> bool: + """Whether ``coreai-torch`` has a lowering for this fx target.""" + name = _coreai_op_name(target) + if name is None: + return False + + aten_resolver, higher_order_resolver = _resolvers() + namespace = _coreai_namespace(target) + + if namespace in ("coreai", "coreaix"): + return True + if namespace == "higher_order": + return name in higher_order_resolver + # ATen ops (namespace "aten") and namespace-less targets such as + # operator.getitem / sym_* are all keyed in the aten resolver. + if name in aten_resolver: + return True + # Edge ``*_copy`` variants (e.g. permute_copy) are not in the resolver, but + # their functional forms (permute) are. Claim them here so they land in the + # delegate; CoreAIBackend.preprocess remaps them via + # ReplaceCopyOpsWithFunctionalPass before conversion. + func_op = functional_aten_op(target) + if func_op is not None: + return getattr(func_op, "__name__", None) in aten_resolver + return False + + +def _node_tensor_vals(node: torch.fx.Node): + """Yield the meta ``val``s for a node's inputs and its own result. + + Flattens tuple/list results so multi-output ops are covered. + """ + for arg in node.all_input_nodes: + yield arg.meta.get("val") + out = node.meta.get("val") + if isinstance(out, (tuple, list)): + yield from out + else: + yield out + + +class _OperatorsSupportedForCoreAIBackend(OperatorSupportBase): + def __init__(self, log: bool = False) -> None: + super().__init__() + self._log = log + self._logged_msgs = set() + + def log_once(self, msg: str) -> None: + if self._log and msg not in self._logged_msgs: + logger.info(msg) + self._logged_msgs.add(msg) + + def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: + # get_attr nodes (e.g. subgraphs referenced by higher-order ops) can + # always ride along with the delegate. + if node.op == "get_attr": + return True + if node.op != "call_function": + return False + + # Respect an explicit user opt-out on the node. + if node.meta.get(DO_NOT_DELEGATE_TAG, False): + self.log_once( + f"Node {node.name} tagged {DO_NOT_DELEGATE_TAG}; leaving out of delegate" + ) + return False + + name = _coreai_op_name(node.target) or "" + if not is_coreai_supported_target(node.target): + self.log_once(f"Core AI cannot lower op, leaving out of delegate: {name}") + return False + + # Core AI has no scalar-symint graph input, and it narrows i64/f64 to + # 32-bit, so it can't faithfully carry an i64/f64 tensor across the + # delegate boundary. Reject any node with a SymInt/SymFloat/SymBool + # operand or an i64/f64 tensor operand/result; the default + # NarrowToCoreAIDtypesPass casts i64/f64 to 32-bit at the EP boundary so + # only the boundary cast (which keeps an i64/f64 tensor) is left out. + for val in _node_tensor_vals(node): + if isinstance(val, (torch.SymInt, torch.SymFloat, torch.SymBool)): + self.log_once( + f"Node {node.name} has a symbolic scalar operand; " + "leaving out of delegate" + ) + return False + if isinstance(val, torch.Tensor) and val.dtype in ( + torch.int64, + torch.float64, + ): + self.log_once( + f"Node {node.name} has a {val.dtype} tensor; Core AI narrows " + "64-bit dtypes and can't carry them across the delegate " + "boundary. Run NarrowToCoreAIDtypesPass (included in " + "coreai.get_default_passes()) to cast int64/float64 inputs to " + "32-bit at the boundary." + ) + return False + return True + + +class CoreAIPartitioner(Partitioner): + """Partition the largest subgraphs that Core AI can lower. + + Op support is derived directly from ``coreai-torch``'s lowering tables, and + :meth:`ops_to_not_decompose` is derived from ``coreai-torch``'s own + decomposition table, so ExecuTorch preserves exactly the ops the converter + expects to see (e.g. fused SDPA) rather than decomposing them. + """ + + def __init__( + self, + *, + uses_sidecar: bool = False, + aot_compile_config: Optional[AOTCompileConfig] = None, + min_deployment_version: Optional[str] = None, + 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 + # single JSON AOT_COMPILE_CONFIG spec (presence implies AOT). + # min_deployment_version is a general spec (it also sets the portable + # .aimodel's OS floor), so it is emitted separately. + 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( + COMPILE_SPEC_KEYS.MIN_DEPLOYMENT_VERSION.value, + str(min_deployment_version).encode(), + ) + ) + if aot_compile_config is not None: + specs.append( + CompileSpec( + COMPILE_SPEC_KEYS.AOT_COMPILE_CONFIG.value, + aot_compile_config.to_json().encode(), + ) + ) + + self.delegation_spec = DelegationSpec( + backend_id=CoreAIBackend.__name__, + compile_specs=specs, + ) + self.take_over_constant_data = take_over_constant_data + self.take_over_mutable_buffer = take_over_mutable_buffer + + @staticmethod + def _is_to_edge_transform_and_lower() -> bool: + """Whether we are being called from ``to_edge_transform_and_lower``.""" + return any( + frame.function == "to_edge_transform_and_lower" for frame in inspect.stack() + ) + + def partition(self, exported_program: ExportedProgram) -> PartitionResult: + # Core AI derives op support and ops_to_not_decompose from coreai-torch's + # tables; the deprecated to_edge() + to_backend() flow decomposes ops + # (e.g. fused SDPA) before the partitioner runs, breaking that contract. + if not self._is_to_edge_transform_and_lower(): + raise RuntimeError( + "CoreAIPartitioner must be used with to_edge_transform_and_lower(). " + "The to_edge() + to_backend() workflow is not supported because it " + "decomposes ops that Core AI lowers directly. Please use:\n" + " exir.to_edge_transform_and_lower(\n" + ' {"forward": exported_program},\n' + " partitioner=[CoreAIPartitioner()],\n" + " compile_config=get_default_compile_config(),\n" + " )" + ) + logger.info("CoreAIPartitioner::partition") + partition_tags = {} + delegation_spec = self.delegation_spec + + capability_partitioner = CapabilityBasedPartitioner( + exported_program.graph_module, + _OperatorsSupportedForCoreAIBackend(log=True), + allows_single_node_partition=True, + ) + partition_list = capability_partitioner.propose_partitions() + for partition in partition_list: + for node in partition.nodes: + tag = f"tag{partition.id}" + node.meta["delegation_tag"] = tag + partition_tags[tag] = delegation_spec + + if self.take_over_constant_data: + tag_constant_data(exported_program) + if self.take_over_mutable_buffer: + tag_mutated_buffer(exported_program) + + return PartitionResult( + tagged_exported_program=exported_program, + partition_tags=partition_tags, + ) + + def ops_to_not_decompose( + self, ep: ExportedProgram + ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]: + # Preserve exactly the ops that Core AI removes from the default + # decomposition table (e.g. scaled_dot_product_attention, silu) so they + # reach the converter in the fused form it lowers optimally. + from coreai_torch import get_decomp_table + + default_table = torch.export.default_decompositions() + coreai_table = get_decomp_table() + do_not_decompose = [ + op + for op in default_table + if op not in coreai_table and isinstance(op, torch._ops.OpOverload) + ] + return do_not_decompose, None diff --git a/backends/apple/coreai/passes/replace_copy_ops.py b/backends/apple/coreai/passes/replace_copy_ops.py new file mode 100644 index 00000000000..b819c0b2097 --- /dev/null +++ b/backends/apple/coreai/passes/replace_copy_ops.py @@ -0,0 +1,66 @@ +# 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 torch + +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_base import PassResult +from torch.fx import GraphModule + + +def functional_aten_op(target) -> object: + """Return the functional ATen overload for an edge ``*_copy`` op, or ``None``. + + ExecuTorch's edge dialect functionalizes view ops into ``*_copy`` variants + (``permute_copy``, ``view_copy``, ``slice_copy``, ...). ``coreai-torch`` + only registers the non-copy forms (``permute``, ``view``, ``slice``), so we + map ``aten::permute_copy.default`` -> ``torch.ops.aten.permute.default``. + + Returns ``None`` for anything that is not an edge ``*_copy`` op or whose + functional counterpart cannot be resolved. + """ + if not isinstance(target, EdgeOpOverload): + return None + name = target._op.__name__ # e.g. "permute_copy.default" + op_name, _, overload = name.partition(".") + if not op_name.endswith("_copy"): + return None + base = op_name[: -len("_copy")] + packet = getattr(torch.ops.aten, base, None) + if packet is None: + return None + return getattr(packet, overload or "default", None) + + +class ReplaceCopyOpsWithFunctionalPass: + """Preprocess-only pass: retarget edge ``*_copy`` ops to their functional + ATen forms so ``coreai-torch`` can lower them. + + Only rewrites when ``coreai-torch`` actually supports the functional form; + unsupported ops are left untouched so conversion still raises an informative + error rather than silently dropping the op. This is safe to run only inside + the backend's ``preprocess`` because view/copy semantics are equivalent in + Core AI's value-based IR. It must not be applied to the shared ExecuTorch + edge graph, which relies on the functionalized ``*_copy`` forms. + """ + + def __call__(self, graph_module: GraphModule) -> PassResult: + from executorch.backends.apple.coreai.partition.partitioner import ( + is_coreai_supported_target, + ) + + modified = False + for node in graph_module.graph.nodes: + if node.op != "call_function": + continue + func_op = functional_aten_op(node.target) + if func_op is not None and is_coreai_supported_target(func_op): + node.target = func_op + modified = True + + if modified: + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/apple/coreai/passes/test/test_replace_copy_ops.py b/backends/apple/coreai/passes/test/test_replace_copy_ops.py new file mode 100644 index 00000000000..25fe60864a6 --- /dev/null +++ b/backends/apple/coreai/passes/test/test_replace_copy_ops.py @@ -0,0 +1,74 @@ +# 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. + +"""Unit tests for the edge ``*_copy`` -> functional ATen rewrite pass.""" + +import unittest + +import torch +import torch.nn as nn + +from executorch.backends.apple.coreai.passes.replace_copy_ops import ( + functional_aten_op, + ReplaceCopyOpsWithFunctionalPass, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload + + +def _edge_graph_module(): + model = nn.Linear(8, 8).eval() + ep = torch.export.export(model, (torch.randn(2, 8),)) + return to_edge(ep).exported_program().graph_module + + +class FunctionalAtenOpTest(unittest.TestCase): + def test_maps_copy_variants_to_functional(self): + self.assertIs( + functional_aten_op(exir_ops.edge.aten.permute_copy.default), + torch.ops.aten.permute.default, + ) + self.assertIs( + functional_aten_op(exir_ops.edge.aten.slice_copy.Tensor), + torch.ops.aten.slice.Tensor, + ) + + def test_returns_none_for_non_copy_and_non_edge(self): + self.assertIsNone(functional_aten_op(exir_ops.edge.aten.addmm.default)) + self.assertIsNone(functional_aten_op(torch.ops.aten.add.Tensor)) + + +class ReplaceCopyOpsWithFunctionalPassTest(unittest.TestCase): + def test_rewrites_copy_ops_in_place(self): + gm = _edge_graph_module() + self.assertTrue( + any( + n.op == "call_function" + and isinstance(n.target, EdgeOpOverload) + and n.target._op.__name__ == "permute_copy.default" + for n in gm.graph.nodes + ), + "precondition: edge graph should contain permute_copy", + ) + + result = ReplaceCopyOpsWithFunctionalPass()(gm) + self.assertTrue(result.modified) + + targets = [n.target for n in gm.graph.nodes if n.op == "call_function"] + self.assertIn(torch.ops.aten.permute.default, targets) + self.assertFalse( + any( + isinstance(t, EdgeOpOverload) + and t._op.__name__.split(".")[0].endswith("_copy") + for t in targets + ), + "no supported *_copy edge op should remain after the pass", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/apple/coreai/test/__init__.py b/backends/apple/coreai/test/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/apple/coreai/test/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/apple/coreai/test/test_io_compat.py b/backends/apple/coreai/test/test_io_compat.py new file mode 100644 index 00000000000..9303ed86a8c --- /dev/null +++ b/backends/apple/coreai/test/test_io_compat.py @@ -0,0 +1,264 @@ +# 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. + +"""Tests for the Core AI delegate-boundary compatibility check. + +* :class:`IoMismatchesTest`: unit tests for the ``io_mismatches`` utility. +* :class:`BoundaryLoweringTest`: e2e lowering of models with various input and + output types through the full flow (``to_edge_transform_and_lower`` -> + ``to_executorch``), asserting the expected Core AI delegation and no leftover + ops. The boundary check runs inside ``preprocess``. +""" + +import unittest + +import torch +import torch.nn as nn + +from executorch.backends.apple.coreai import ( + get_default_compile_config, + get_default_passes, +) +from executorch.backends.apple.coreai.compiler.io_compat import io_mismatches +from executorch.backends.apple.coreai.partition.partitioner import CoreAIPartitioner +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.lowered_backend_module import executorch_call_delegate +from torch.export import Dim + + +class _Sym: + """Stand-in for a symbolic (dynamic) edge dim.""" + + +class IoMismatchesTest(unittest.TestCase): + def test_io_mismatches(self): + # (name, coreai_io, edge_io, expected substring or None if compatible) + cases = [ + ("compatible", [("f32", (2, 8))], [(torch.float32, (2, 8))], None), + ("count", [("f32", (2, 8))], [], "count mismatch"), + ( + "float dtype", + [("f16", (2, 8))], + [(torch.float32, (2, 8))], + "dtype mismatch", + ), + ("int exact", [("si8", (4,))], [(torch.int8, (4,))], None), + ( + "int64 narrowed", + [("si32", (4,))], + [(torch.int64, (4,))], + "dtype mismatch", + ), + ("rank", [("f32", (2, 8))], [(torch.float32, (2,))], "rank mismatch"), + ( + "static shape", + [("f32", (2, 8))], + [(torch.float32, (2, 16))], + "static shape mismatch", + ), + ( + "symbolic dim skipped", + [("f32", (2, 8))], + [(torch.float32, (_Sym(), 8))], + None, + ), + ("rank-0 compatible", [("f32", ())], [(torch.float32, ())], None), + ( + "rank-0 vs rank-1", + [("f32", ())], + [(torch.float32, (1,))], + "rank mismatch", + ), + ("both non-tensor", [("int", None)], [("int", None)], None), + ( + "tensor vs non-tensor", + [("!scalar", None)], + [(torch.float32, (2, 8))], + "non-tensor", + ), + ] + for name, coreai_io, edge_io, expected in cases: + with self.subTest(name): + errs = io_mismatches(coreai_io, edge_io, "input") + if expected is None: + self.assertEqual(errs, [], f"{name}: {errs}") + else: + self.assertTrue(any(expected in e for e in errs), f"{name}: {errs}") + + +def _lower( + module, + example_inputs, + dynamic_shapes=None, + *, + expect_delegates=1, + allow_leftover=False, +): + """Full AOT flow through Core AI; returns the lowered program. + + ``_skip_dim_order`` keeps ExecuTorch on ``aten._to_copy`` (which coreai + supports) instead of emitting ``dim_order_ops._to_dim_order_copy``. + + Args: + expect_delegates: assert this exact number of Core AI delegates; pass + ``None`` to skip the count check. + allow_leftover: if ``False`` (default), assert no ops remain outside the + delegate (ignoring getitem and boundary ``_to_copy`` casts). Set + ``True`` for models Core AI only partially delegates (symint / i64 + producing ops). + """ + ep = torch.export.export( + module.eval(), example_inputs, dynamic_shapes=dynamic_shapes + ) + lowered = to_edge_transform_and_lower( + ep, + partitioner=[CoreAIPartitioner()], + transform_passes=get_default_passes(), + compile_config=get_default_compile_config(), + ) + gm = lowered.exported_program().graph_module + delegates = [ + n + for n in gm.graph.nodes + if n.op == "call_function" and n.target is executorch_call_delegate + ] + leftover = [ + str(n.target) + for n in gm.graph.nodes + if n.op == "call_function" + and n.target is not executorch_call_delegate + and "getitem" not in str(n.target) + and "_to_copy" not in str(n.target) # boundary dtype-narrowing casts + ] + if expect_delegates is not None: + assert ( + len(delegates) == expect_delegates + ), f"expected {expect_delegates} Core AI delegate(s), got {len(delegates)}" + if not allow_leftover: + assert not leftover, f"not fully delegated; leftover ops: {leftover}" + lowered.to_executorch() + return lowered + + +class _Add(nn.Module): + def forward(self, x): + return x + x + + +class _And(nn.Module): + def forward(self, x): + return torch.logical_and(x, x) + + +class _AddScalar(nn.Module): + def forward(self, x, n): + return x + n + + +class _ConstIntOut(nn.Module): + def forward(self, x): + return x + x, 2 + + +class _ConstFloatOut(nn.Module): + def forward(self, x): + return x + x, 4.0 + + +class _Embedding(nn.Module): + def __init__(self): + super().__init__() + self.emb = nn.Embedding(16, 4) + + def forward(self, tok): + return self.emb(tok) + + +class _Rank0Out(nn.Module): + def forward(self, x): + return (x + x).sum() # rank-0 (0-d) tensor output + + +class _ProducedInt64Out(nn.Module): + def forward(self, x): + # produced (non-const-folded) int64 output alongside a delegated tensor + return x + x, (x > 0).to(torch.int64) + + +class BoundaryLoweringTest(unittest.TestCase): + def test_supported_input_dtypes(self): + # dtypes Core AI carries directly (no 64-bit narrowing); all fully + # delegate. int64/float64 (narrowed) are covered separately below. + cases = { + "float32": (_Add(), torch.randn(2, 8, dtype=torch.float32)), + "float16": (_Add(), torch.randn(2, 8).to(torch.float16)), + "bfloat16": (_Add(), torch.randn(2, 8).to(torch.bfloat16)), + "int8": (_Add(), torch.randint(-8, 8, (2, 8), dtype=torch.int8)), + "int16": (_Add(), torch.randint(-8, 8, (2, 8), dtype=torch.int16)), + "int32": (_Add(), torch.randint(-8, 8, (2, 8), dtype=torch.int32)), + "uint8": (_Add(), torch.randint(0, 8, (2, 8), dtype=torch.uint8)), + "bool": (_And(), torch.randint(0, 2, (2, 8)).bool()), + } + for name, (module, x) in cases.items(): + with self.subTest(dtype=name): + _lower(module, (x,)) + + # 64-bit narrowing (input cast + output widen via the default pass). + def test_float64(self): + _lower(_Add(), (torch.randn(2, 8, dtype=torch.float64),)) + + def test_int64(self): + _lower(_Add(), (torch.randint(-8, 8, (2, 8), dtype=torch.int64),)) + + def test_embedding_int64_index(self): + # Real int64 case: token indices stay int64 at the model boundary; the + # default narrow pass casts them to int32 for the delegated embedding. + _lower(_Embedding(), (torch.randint(0, 16, (2, 8), dtype=torch.int64),)) + + def test_symint_input(self): + # A genuine dynamic scalar (SymInt) can't be a Core AI graph input; the + # symint-consuming op is left outside the delegate, so the model lowers + # (partial delegation) without crashing. + _lower( + _AddScalar(), + (torch.randn(2, 8), 3), + dynamic_shapes={"x": None, "n": Dim.DYNAMIC}, + allow_leftover=True, + ) + + # Non-tensor (const) outputs. + def test_const_int_output(self): + _lower(_ConstIntOut(), (torch.randn(2, 8),)) + + def test_const_float_output(self): + _lower(_ConstFloatOut(), (torch.randn(2, 8),)) + + # Rank-0 (0-d) output. + def test_rank0_output(self): + _lower(_Rank0Out(), (torch.randn(2, 8),)) + + # Produced (non-const-folded) int64 output. + def test_produced_int64_output(self): + # coreai narrows i64, so the i64-producing op stays outside the delegate + # (guard) and runs portable; the model still lowers and the int64 output + # dtype is preserved. + lowered = _lower( + _ProducedInt64Out(), + (torch.randn(2, 8),), + expect_delegates=None, # partial: only the f32 part is delegated + allow_leftover=True, + ) + out_dtypes = [ + a.meta["val"].dtype + for a in lowered.exported_program().graph_module.graph.output_node().args[0] + if isinstance(a, torch.fx.Node) + and isinstance(a.meta.get("val"), torch.Tensor) + ] + self.assertIn(torch.int64, out_dtypes) # int64 output preserved + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/apple/coreai/test/test_partitioner.py b/backends/apple/coreai/test/test_partitioner.py new file mode 100644 index 00000000000..0c3aa3a7679 --- /dev/null +++ b/backends/apple/coreai/test/test_partitioner.py @@ -0,0 +1,285 @@ +# 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 import get_default_compile_config +from executorch.backends.apple.coreai.compiler.preprocess import COMPILE_SPEC_KEYS +from executorch.backends.apple.coreai.partition.partitioner import ( + _OperatorsSupportedForCoreAIBackend, + CoreAIPartitioner, + do_not_delegate, + DO_NOT_DELEGATE_TAG, + is_coreai_supported_target, +) +from executorch.exir import to_edge, to_edge_transform_and_lower +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.lowered_backend_module import executorch_call_delegate +from executorch.exir.pass_base import PassResult + + +def _make_node(op, target): + node = torch.fx.Graph().create_node(op, target, name="n") + return node + + +class IsSupportedTargetTest(unittest.TestCase): + def test_supported_aten_ops(self): + for op in ( + torch.ops.aten.add.Tensor, + torch.ops.aten.mm.default, + torch.ops.aten.addmm.default, + torch.ops.aten.view.default, + torch.ops.aten.permute.default, + ): + with self.subTest(op=str(op)): + self.assertTrue(is_coreai_supported_target(op)) + + def test_unsupported_aten_op(self): + self.assertFalse( + is_coreai_supported_target(torch.ops.aten.linalg_solve.default) + ) + + def test_supported_edge_op_unwrapped(self): + # An edge op whose non-copy ATen form Core AI supports. + self.assertTrue(is_coreai_supported_target(exir_ops.edge.aten.view.default)) + self.assertTrue(is_coreai_supported_target(exir_ops.edge.aten.addmm.default)) + + def test_edge_copy_variants_supported_via_functional_form(self): + # The ExecuTorch edge dialect emits *_copy variants which Core AI's + # resolver does not list directly, but their functional forms (view, + # permute) are supported. The partitioner claims them so they land in + # the delegate; CoreAIBackend.preprocess remaps them before conversion. + self.assertTrue( + is_coreai_supported_target(exir_ops.edge.aten.view_copy.default) + ) + self.assertTrue( + is_coreai_supported_target(exir_ops.edge.aten.permute_copy.default) + ) + + def test_edge_copy_variant_without_functional_support_is_unsupported(self): + # transpose_copy -> transpose, which Core AI does not lower (only + # permute is in the resolver), so it stays unsupported. + self.assertFalse( + is_coreai_supported_target(exir_ops.edge.aten.transpose_copy.int) + ) + + +class OperatorSupportTest(unittest.TestCase): + def setUp(self): + self.support = _OperatorsSupportedForCoreAIBackend() + + def test_rejects_placeholder_and_output(self): + for op in ("placeholder", "output", "call_module", "call_method"): + with self.subTest(op=op): + self.assertFalse( + self.support.is_node_supported({}, _make_node(op, "x")) + ) + + def test_accepts_get_attr(self): + self.assertTrue(self.support.is_node_supported({}, _make_node("get_attr", "w"))) + + def test_call_function_follows_resolver(self): + self.assertTrue( + self.support.is_node_supported( + {}, _make_node("call_function", torch.ops.aten.add.Tensor) + ) + ) + self.assertFalse( + self.support.is_node_supported( + {}, _make_node("call_function", torch.ops.aten.linalg_solve.default) + ) + ) + + def test_do_not_delegate_tag_overrides_support(self): + node = _make_node("call_function", torch.ops.aten.add.Tensor) + self.assertTrue(self.support.is_node_supported({}, node)) + do_not_delegate(node) + self.assertFalse(self.support.is_node_supported({}, node)) + + +class OpsToNotDecomposeTest(unittest.TestCase): + def test_preserves_coreai_composite_ops(self): + ep = torch.export.export(nn.Linear(4, 4), (torch.randn(1, 4),)) + ops, filt = CoreAIPartitioner().ops_to_not_decompose(ep) + self.assertIsNone(filt) + # Core AI keeps SDPA fused rather than decomposing it. + self.assertIn(torch.ops.aten.scaled_dot_product_attention.default, ops) + + +class LinearE2ETest(unittest.TestCase): + def _delegate_and_leftover_ops(self, edge_program): + graph = edge_program.graph + delegate_calls = [ + n + for n in graph.nodes + if n.op == "call_function" and n.target is executorch_call_delegate + ] + leftover = [ + n + for n in graph.nodes + if n.op == "call_function" + and n.target is not executorch_call_delegate + and "getitem" not in str(n.target) + ] + return delegate_calls, leftover + + def test_linear_lowers_to_coreai(self): + model = nn.Linear(8, 8).eval() + example_inputs = (torch.randn(2, 8),) + ep = torch.export.export(model, example_inputs) + + lowered = to_edge_transform_and_lower(ep, partitioner=[CoreAIPartitioner()]) + edge_program = lowered.exported_program() + + delegate_calls, leftover = self._delegate_and_leftover_ops(edge_program) + self.assertGreater( + len(delegate_calls), + 0, + "Expected Core AI to lower at least part of a Linear model", + ) + # With the copy-op remap, permute_copy is delegated too, so a plain + # Linear should lower fully with nothing left outside the delegate. + self.assertEqual( + leftover, + [], + f"Unexpected ops left outside the delegate: {[str(n.target) for n in leftover]}", + ) + + def test_graph_break_second_linear_not_delegated(self): + # linear -> relu -> linear -> relu, with the SECOND linear tagged + # do-not-delegate. Tagging happens at the edge stage via transform_passes + # so the meta survives into partition(). Expect a graph break: >= 2 + # separate Core AI delegates with the tagged linear running outside. + def _addmm_nodes(gm): + return [ + n + for n in gm.graph.nodes + if n.op == "call_function" + and isinstance(n.target, EdgeOpOverload) + and n.target._op.__name__ == "addmm.default" + ] + + class _TagSecondLinear: + def __call__(self, gm): + addmms = _addmm_nodes(gm) + if len(addmms) >= 2: + second = addmms[1] + do_not_delegate(second) + # Exclude its weight-transpose too, so the whole 2nd linear + # sits outside the delegate rather than leaving a dangling + # single-node permute delegate. + for inp in second.all_input_nodes: + if isinstance( + inp.target, EdgeOpOverload + ) and inp.target._op.__name__.endswith("_copy.default"): + do_not_delegate(inp) + return PassResult(gm, True) + + model = nn.Sequential( + nn.Linear(8, 8), nn.ReLU(), nn.Linear(8, 8), nn.ReLU() + ).eval() + ep = torch.export.export(model, (torch.randn(2, 8),)) + lowered = to_edge_transform_and_lower( + ep, + partitioner=[CoreAIPartitioner()], + transform_passes=[_TagSecondLinear()], + ) + graph_module = lowered.exported_program().graph_module + + delegate_calls = [ + n + for n in graph_module.graph.nodes + if n.op == "call_function" and n.target is executorch_call_delegate + ] + # Graph break: the tagged linear splits delegation into >= 2 delegates. + self.assertGreaterEqual(len(delegate_calls), 2) + # The 2nd linear's addmm runs OUTSIDE any delegate (visible at top level). + self.assertEqual(len(_addmm_nodes(graph_module)), 1) + + def test_do_not_delegate_tag_excludes_node_from_partition(self): + class _TwoOps(nn.Module): + def forward(self, x): + return x + x, x * 2.0 + + ep = torch.export.export(_TwoOps(), (torch.randn(4),)) + edge_ep = to_edge(ep).exported_program() + + # Test the opt-out via the operator-support check (the mechanism the + # partitioner uses); partition() itself only runs inside + # to_edge_transform_and_lower. + support = _OperatorsSupportedForCoreAIBackend() + node = next( + n + for n in edge_ep.graph.nodes + if n.op == "call_function" and is_coreai_supported_target(n.target) + ) + self.assertTrue(support.is_node_supported({}, node)) + do_not_delegate(node) + self.assertIn(DO_NOT_DELEGATE_TAG, node.meta) + self.assertFalse(support.is_node_supported({}, node)) + + def test_symint_operand_not_supported(self): + # add.Tensor is in coreai's resolver, but an operand that is a scalar + # SymInt would become a delegate boundary input coreai can't type, so + # the node must be left out of the delegate. + from torch.export import Dim + + class _AddScalar(nn.Module): + def forward(self, x, n): + return x + n + + ep = torch.export.export( + _AddScalar().eval(), + (torch.randn(2, 8), 3), + dynamic_shapes={"x": None, "n": Dim.DYNAMIC}, + ) + add = next(n for n in ep.graph.nodes if n.op == "call_function") + support = _OperatorsSupportedForCoreAIBackend() + # op is supported by name... + self.assertTrue(is_coreai_supported_target(add.target)) + # ...but the symint operand blocks delegation. + self.assertFalse(support.is_node_supported({}, add)) + # Core AI claims edge *_copy variants (whose functional form it supports), + # so a fully supported model lowers with no ops left outside the delegate. + model = nn.Linear(8, 8).eval() + ep = torch.export.export(model, (torch.randn(2, 8),)) + lowered = to_edge_transform_and_lower( + ep, + partitioner=[CoreAIPartitioner()], + compile_config=get_default_compile_config(), + ) + gm = lowered.exported_program().graph_module + untagged = [ + str(n.target) + for n in gm.graph.nodes + if n.op == "call_function" + and n.target is not executorch_call_delegate + and "getitem" not in str(n.target) + ] + self.assertEqual( + untagged, [], f"Expected all ops delegated, but these were not: {untagged}" + ) + + +class PartitionerCompileSpecTest(unittest.TestCase): + def test_uses_sidecar_spec_is_embedded(self): + # The delivery mode is serialized (the runtime needs it); the build dir + # is not (it is an env var); only the mode flag rides along. + specs = CoreAIPartitioner(uses_sidecar=True).delegation_spec.compile_specs + self.assertEqual([s.key for s in specs], [COMPILE_SPEC_KEYS.USES_SIDECAR.value]) + self.assertNotIn(b"/", specs[0].value) + + def test_inline_partitioner_embeds_no_specs(self): + self.assertEqual(CoreAIPartitioner().delegation_spec.compile_specs, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/apple/coreai/test/test_preprocess.py b/backends/apple/coreai/test/test_preprocess.py new file mode 100644 index 00000000000..fbfac65f77c --- /dev/null +++ b/backends/apple/coreai/test/test_preprocess.py @@ -0,0 +1,305 @@ +# 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. + +"""Tests for ``CoreAIBackend.preprocess`` and its AOT-compile config. + +Covers all four asset-delivery combinations: portable ``.aimodel`` vs +AOT-compiled ``.aimodelc``, each either inline (embedded in the .pte) or as a +sidecar, plus ``AOTCompileConfig`` parsing. The compiled-delivery cases mock +``coreai-build`` so they run without the Metal Toolchain; the real-toolchain +integration is in :class:`CoreAIAOTCompileTest` (gated on ``coreai-build``). +""" + +import json +import os +import subprocess +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +import torch +import torch.nn as nn + +from executorch.backends.apple.coreai.compiler.preprocess import ( + _aot_compile_options, + AOTCompileConfig, + COMPILE_SPEC_KEYS, + coreai_sidecar_dir, + CoreAIBackend, +) +from executorch.backends.apple.coreai.partition.partitioner import CoreAIPartitioner +from executorch.exir import to_edge, to_edge_transform_and_lower +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.lowered_backend_module import ( + executorch_call_delegate, + get_lowered_backend_modules, +) + + +def _coreai_build_available() -> bool: + try: + return ( + subprocess.run( + ["xcrun", "--find", "coreai-build"], capture_output=True + ).returncode + == 0 + ) + except FileNotFoundError: + return False + + +class _Elementwise(nn.Module): + # Parameter-free so the raw edge program converts directly, keeping these + # tests focused on preprocess delivery/packaging (weight handling is covered + # by the delegation e2e tests). + def forward(self, x): + return x + x * 2.0 + + +def _edge_program(): + ep = torch.export.export(_Elementwise().eval(), (torch.randn(2, 8),)) + return to_edge(ep).exported_program() + + +def _fake_run_coreai_build(aimodel_path, out_dir, opts): + """Stand-in for ``xcrun coreai-build``: drop fake per-arch .aimodelc dirs.""" + for arch in opts["architectures"] or ["h15g"]: + bundle = Path(out_dir) / f"model.{arch}.aimodelc" + bundle.mkdir(parents=True, exist_ok=True) + (bundle / "model.mil").write_bytes(b"fake-compiled") + + +def _aot_spec(config: dict) -> CompileSpec: + return CompileSpec( + COMPILE_SPEC_KEYS.AOT_COMPILE_CONFIG.value, json.dumps(config).encode() + ) + + +_SIDECAR_SPEC = CompileSpec(COMPILE_SPEC_KEYS.USES_SIDECAR.value, b"1") +_MOCK_BUILD = "executorch.backends.apple.coreai.compiler.preprocess._run_coreai_build" + + +class AOTCompileConfigTest(unittest.TestCase): + """``AOTCompileConfig`` / ``_aot_compile_options`` parsing (no coreai-build).""" + + def test_parses_known_fields(self): + opts = _aot_compile_options( + [_aot_spec({"platform": "iOS", "architectures": ["h17p"]})] + ) + self.assertEqual(opts["platform"], "iOS") + self.assertEqual(opts["architectures"], ["h17p"]) + + def test_defaults_when_empty(self): + opts = _aot_compile_options([_aot_spec({})]) + self.assertEqual(opts["platform"], "macOS") + self.assertEqual(opts["architectures"], []) + self.assertEqual(opts["preferred_compute"], "none") + self.assertFalse(opts["expect_frequent_reshapes"]) + + def test_rejects_unexpected_field(self): + # A typo like "platfrom" must fail loudly, not silently default. + with self.assertRaises(ValueError): + _aot_compile_options([_aot_spec({"platfrom": "iOS"})]) + + def test_config_json_roundtrip(self): + cfg = AOTCompileConfig( + platform="iOS", + preferred_compute="neural-engine", + architectures=["h17p"], + expect_frequent_reshapes=True, + ) + self.assertEqual(AOTCompileConfig.from_json(cfg.to_json()), cfg) + + def test_config_from_dict_rejects_unexpected(self): + with self.assertRaises(ValueError): + AOTCompileConfig.from_dict({"platfrom": "iOS"}) + + def test_empty_config_roundtrip(self): + self.assertEqual( + AOTCompileConfig.from_json(AOTCompileConfig().to_json()), + AOTCompileConfig(), + ) + + +class PortablePreprocessTest(unittest.TestCase): + """Portable ``.aimodel`` delivery (no coreai-build).""" + + def test_inline_embeds_files_in_nds(self): + result = CoreAIBackend.preprocess(_edge_program(), []) + manifest = json.loads(result.processed_bytes) + self.assertEqual(manifest["packaging"], "inline") + self.assertTrue(manifest["files"], "expected at least one asset file") + out = result.data_store_output + self.assertIsNotNone(out) + for rel in manifest["files"]: + self.assertIn(f"coreai/{manifest['hash']}/{rel}", out.pte_data) + self.assertEqual(out.external_data, {}) # nothing external for inline + + def test_sidecar_writes_bundle_and_does_not_embed(self): + with tempfile.TemporaryDirectory() as d: + with coreai_sidecar_dir(d): + result = CoreAIBackend.preprocess(_edge_program(), [_SIDECAR_SPEC]) + manifest = json.loads(result.processed_bytes) + self.assertEqual(manifest["packaging"], "sidecar") + # Only the relative, hash-named bundle is referenced (no build path). + self.assertEqual(manifest["path"], f"{manifest['hash']}.aimodel") + self.assertNotIn(d, result.processed_bytes.decode()) + bundle = Path(d) / manifest["path"] + self.assertTrue(bundle.is_dir()) + self.assertTrue((bundle / "main.mlirb").exists()) + self.assertIsNone(result.data_store_output) # nothing embedded + + def test_sidecar_without_env_var_raises(self): + # uses_sidecar set but COREAI_SIDECAR_DIR unset -> fail fast. + env = {k: v for k, v in os.environ.items() if k != "COREAI_SIDECAR_DIR"} + with mock.patch.dict(os.environ, env, clear=True): + with self.assertRaises(ValueError): + CoreAIBackend.preprocess(_edge_program(), [_SIDECAR_SPEC]) + + def test_inline_with_env_var_warns(self): + # Env var set but delegate is inline -> soft warning, not an error. + import executorch.backends.apple.coreai.compiler.preprocess as cp + + with tempfile.TemporaryDirectory() as d: + with coreai_sidecar_dir(d): + cp._WARNED_SIDECAR_ENV_IGNORED = False # allow the once-warning + with self.assertLogs(cp.logger, level="WARNING") as cm: + result = CoreAIBackend.preprocess(_edge_program(), []) + self.assertEqual(json.loads(result.processed_bytes)["packaging"], "inline") + self.assertTrue(any("uses_sidecar=True" in m for m in cm.output)) + + def test_coreai_sidecar_dir_sets_and_restores_env(self): + self.assertNotIn("COREAI_SIDECAR_DIR", os.environ) + with coreai_sidecar_dir("/tmp/scoped"): + self.assertEqual(os.environ["COREAI_SIDECAR_DIR"], "/tmp/scoped") + self.assertNotIn("COREAI_SIDECAR_DIR", os.environ) # restored on exit + + def test_end_to_end_sidecar(self): + with tempfile.TemporaryDirectory() as d: + ep = torch.export.export(nn.Linear(8, 8).eval(), (torch.randn(2, 8),)) + with coreai_sidecar_dir(d): + to_edge_transform_and_lower( + ep, partitioner=[CoreAIPartitioner(uses_sidecar=True)] + ) + bundles = list(Path(d).glob("*.aimodel")) + self.assertEqual(len(bundles), 1) + self.assertTrue((bundles[0] / "main.mlirb").exists()) + + +class CompiledPreprocessTest(unittest.TestCase): + """AOT-compiled ``.aimodelc`` delivery, with coreai-build mocked (ungated).""" + + @mock.patch(_MOCK_BUILD, side_effect=_fake_run_coreai_build) + def test_inline_embeds_compiled_bundles(self, _build): + result = CoreAIBackend.preprocess( + _edge_program(), + [_aot_spec({"platform": "iOS", "architectures": ["h15g", "h16"]})], + ) + manifest = json.loads(result.processed_bytes) + self.assertEqual(manifest["packaging"], "aot_compiled_inline") + self.assertEqual(manifest["platform"], "iOS") + self.assertEqual(sorted(manifest["archs"]), ["h15g", "h16"]) + # Compiled bundle contents are embedded (files under a *.aimodelc dir). + self.assertTrue(any(".aimodelc/" in f for f in manifest["files"])) + self.assertIsNotNone(result.data_store_output) + + @mock.patch(_MOCK_BUILD, side_effect=_fake_run_coreai_build) + def test_sidecar_writes_compiled_bundles(self, _build): + with tempfile.TemporaryDirectory() as d: + with coreai_sidecar_dir(d): + result = CoreAIBackend.preprocess( + _edge_program(), + [_aot_spec({"architectures": ["h15g"]}), _SIDECAR_SPEC], + ) + manifest = json.loads(result.processed_bytes) + self.assertEqual(manifest["packaging"], "aot_compiled_sidecar") + self.assertEqual(list(manifest["archs"]), ["h15g"]) + self.assertIsNone(result.data_store_output) # bundles on disk only + self.assertNotIn(d, result.processed_bytes.decode()) + for rel in manifest["archs"].values(): + bundle = Path(d) / rel + self.assertTrue(bundle.is_dir()) + self.assertTrue(str(bundle).endswith(".aimodelc")) + + @mock.patch(_MOCK_BUILD, side_effect=_fake_run_coreai_build) + def test_defaults_to_all_architectures(self, _build): + # Empty architectures => coreai-build decides (our fake yields one). + result = CoreAIBackend.preprocess(_edge_program(), [_aot_spec({})]) + manifest = json.loads(result.processed_bytes) + self.assertEqual(manifest["packaging"], "aot_compiled_inline") + self.assertEqual(manifest["platform"], "macOS") # default + self.assertTrue(manifest["archs"]) + + +def _lower_linear(partitioner): + model = nn.Sequential(nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 32)).eval() + ep = torch.export.export(model, (torch.randn(2, 32),)) + return to_edge_transform_and_lower(ep, partitioner=[partitioner]) + + +def _lowered_manifest(lowered): + lbms = get_lowered_backend_modules(lowered.exported_program().graph_module) + assert len(lbms) == 1, lbms + return lbms[0], json.loads(bytes(lbms[0].processed_bytes)) + + +@unittest.skipUnless( + _coreai_build_available(), + "requires macOS with the Metal Toolchain (xcrun coreai-build)", +) +class CoreAIAOTCompileTest(unittest.TestCase): + """Real coreai-build integration (only runs when the toolchain is present).""" + + def test_aot_inline_embeds_aimodelc(self): + lowered = _lower_linear( + CoreAIPartitioner(aot_compile_config=AOTCompileConfig(platform="macOS")) + ) + lbm, manifest = _lowered_manifest(lowered) + self.assertEqual(manifest["packaging"], "aot_compiled_inline") + self.assertEqual(manifest["platform"], "macOS") + self.assertGreaterEqual(len(manifest["archs"]), 1) + self.assertTrue(any(".aimodelc/" in f for f in manifest["files"])) + self.assertIsNotNone(lbm.named_data_store_output) + self.assertGreater(len(bytes(lowered.to_executorch().buffer)), 0) + + def test_aot_sidecar_writes_aimodelc_and_keeps_pte_small(self): + with tempfile.TemporaryDirectory() as sidecar: + with coreai_sidecar_dir(sidecar): + lowered = _lower_linear( + CoreAIPartitioner( + aot_compile_config=AOTCompileConfig(platform="macOS"), + uses_sidecar=True, + ) + ) + lbm, manifest = _lowered_manifest(lowered) + pte = bytes(lowered.to_executorch().buffer) + + self.assertEqual(manifest["packaging"], "aot_compiled_sidecar") + self.assertGreaterEqual(len(manifest["archs"]), 1) + # Sidecar: compiled bundle *contents* must not be embedded in the .pte. + self.assertIsNone(lbm.named_data_store_output) + self.assertNotIn(b"mpsExecutable", pte) + for rel in manifest["archs"].values(): + bundle = os.path.join(sidecar, rel) + self.assertTrue(os.path.isdir(bundle), bundle) + self.assertTrue(bundle.endswith(".aimodelc")) + + def test_delegates_present(self): + lowered = _lower_linear( + CoreAIPartitioner(aot_compile_config=AOTCompileConfig(platform="macOS")) + ) + gm = lowered.exported_program().graph_module + self.assertTrue( + any( + n.op == "call_function" and n.target is executorch_call_delegate + for n in gm.graph.nodes + ) + ) + + +if __name__ == "__main__": + unittest.main()