Skip to content

[CoreAI] Rename manifest packaging to aimodel / aimodelc - #24474

Merged
metascroy merged 2 commits into
coreai-expose-compute-unitfrom
coreai-rename-packaging
Nov 4, 2026
Merged

metascroy merged 2 commits into
coreai-expose-compute-unitfrom
coreai-rename-packaging

Conversation

@metascroy

Copy link
Copy Markdown
Contributor

Renames the Core AI delegate manifest's packaging values so they name the embedded format instead of where it's stored.

Stacked on #24375.

Why

"inline" and "aot_compiled_inline" describe storage (embedded in the .pte), which is the same for both formats. They don't say what is actually different: a portable .aimodel versus AOT-compiled .aimodelc bundles. The runtime also used a third vocabulary, "source" / "aot", in logs and the cache identity. "source" already meant something else there: the bundle copy the runtime writes to the cache for the SDK, which happens for AOT bundles too.

Changes

Where Before After
Manifest packaging "inline" / "aot_compiled_inline" "aimodel" / "aimodelc"
Python AssetPackaging (compiler/preprocess.py) INLINE / AOT_COMPILED_INLINE AIMODEL / AIMODELC
Runtime manifest check (runtime/coreai_manifest.mm) accepts the old values; AOT detected by an aot_compiled_ prefix accepts only the new values; AOT means == "aimodelc"
Specialization log (coreai_load_coordinator.mm) kind=source / kind=aot kind=aimodel / kind=aimodelc
Cache identity (coreai_cache_key.mm) "source" / "aot" "aimodel" / "aimodelc"

Test fixtures, test labels and the README follow the new names. In the README, "source" now only means the copy written for the SDK.

Compatibility

  • Old PTEs: this is a manifest format change. PTEs exported with the old names fail to load with Unsupported Core AI packaging; expected aimodel or aimodelc and must be re-exported.
  • Existing caches: the cache identity changes, so existing bookmarks aren't found and each model specializes again on its first load. Old entries stay on disk unused.

Test plan

  • test_preprocess.py: 25/25 passed, including the real coreai-build AOT test.
  • Native: all 6 CTest targets pass (coreai_host_test, coreai_pte_export_test, coreai_swift_bridge_test, coreai_sdk_source_lifetime_test, coreai_runtime_smoke, coreai_multithread_test). The native fixtures are regenerated and carry packaging: "aimodel".
  • Manual end-to-end with a standalone runner on an M4 Max:
    • Portable PTE: kind=aimodel, outputs match the reference.
    • AOT PTE: first load kind=aimodelc with exact outputs; second load restores from the cache bookmark.
    • A PTE exported before this change is rejected with the error above.
  • ufmt, flake8 and clang-format are clean.

@metascroy
metascroy requested a review from shoumikhin as a code owner November 3, 2026 00:00
@pytorch-bot

pytorch-bot Bot commented Nov 3, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24474

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 3, 2026
@metascroy
metascroy force-pushed the coreai-expose-compute-unit branch from d05c245 to d703edf Compare November 3, 2026 18:59
@metascroy
metascroy force-pushed the coreai-rename-packaging branch from cfdb000 to 73ded2d Compare November 3, 2026 18:59
The packaging values "inline" and "aot_compiled_inline" described where the
bundle is stored rather than what it is. Name them after the embedded format
instead: "aimodel" for the portable bundle and "aimodelc" for AOT-compiled
per-architecture bundles.

The runtime's kind names in logs and the cache identity ("source" / "aot")
become "aimodel" / "aimodelc" too, so "source" only refers to the bundle copy
materialized for the SDK. This is a format change: PTEs exported with the old
names are rejected and must be re-exported, and existing cache entries miss
once and are re-specialized.
Adds an opt-in way to export several ExecuTorch methods (and
graph-broken partitions) as **one Core AI model with multiple
functions**, so methods that share weights ship and cache them once,
plus an opt-in runtime option so those delegates also **share one loaded
model in memory**.

Stacked on #24474.

## Usage

```python
lowered = to_edge_transform_and_lower(
    {"forward": forward_ep, "prefill": prefill_ep},
    transform_passes=get_default_passes(),
    partitioner=[CoreAIPartitioner(multifunction=True)],
    compile_config=get_default_compile_config(),
)
```

Each Core AI partition becomes one function named `<method>_<partition
index>`: here `forward_0` and `prefill_0`, and a graph break adds
`forward_1`, `prefill_1`, and so on. Graph breaks are allowed;
partitions running on CPU stay in ExecuTorch as usual.

To also share one loaded model at runtime, set the backend option when
loading:

```cpp
BackendOptions<2> options;
options.set_option("coreai_share_multifunction_model", true);
LoadBackendOptionsMap map;
map.set_options("CoreAIBackend", options.view());
module.load(map);
```

## How it works

* **Export:** `CoreAIBackend.preprocess_multimethod` receives every
method's partitions at once.
* It stages all of them in one `TorchConverter` (one entrypoint per
partition), calls `to_coreai()` once, and saves (or AOT-compiles with
`coreai-build`) once.
* The bundle is embedded in the `NamedDataStore` once, under a single
`export_hash`.
* Each delegate still gets its own manifest naming its `function`,
`input_names` and `output_names`; all manifests point at the same
bundle.
* **Cache:** function bindings aren't part of the cache key, so all the
delegates share one cache entry: the first load specializes and the
others restore from the same bookmark.
* **Manifest:** multi-function manifests carry `"multifunction": true`
(optional boolean; absent means false).
* **Runtime sharing (`coreai_share_multifunction_model`, bool, default
`false`):**
* Without it, each delegate restores its own model with its own copy of
the weights, as before.
* With it, the first delegate acquires the model and the others reuse it
from a process-wide table until the last delegate is destroyed.
* Loads are tracked per key. Delegates that start loading the same model
at once wait for that one load, while loads of other models or cache
directories proceed in parallel. If the load fails, it leaves no entry
and one of the waiters retries.
* Sharing requires the same cache directory (identified by the open
directory's device and inode, so a symlink counts as the same directory)
and the same cache key, which includes the compute-unit specialization.
Different cache directories or compute units load separate models.
* Delegates sharing a model run one at a time through a per-model lock.
The SDK doesn't document concurrent execution on one model as safe, and
without the lock `coreai_multithread_test` produced wrong outputs.
* Setting the option on a delegate that isn't from a multi-function
export fails init with `InvalidArgument`.
* **Option off:** lowering is unchanged. `preprocess` and the
multi-method path now share the conversion and delivery helpers, and the
existing tests still pass.
* **Errors:** export fails with a clear message if partitions have
different compile specs (including the option set on only some methods),
or if `input_enumerations` is combined with the option. Enumerated
shapes rename functions internally and need their own design.

## Results

Two-method model (8 × `Linear(1024, 1024)`, fp32; `forward` with 1 row,
`prefill` with 16 rows), run on an M4 Max with a standalone C++ runner:

| Variant | PTE size | Cold load | Cache entries | Outputs |
|---|---|---|---|---|
| `multifunction=True` | 33.6 MB | 1 specialization + 1 bookmark restore
| 1 | match eager |
| Default (one model per delegate) | 67.2 MB | 2 specializations | 2 |
match eager |
| `multifunction=True` + graph break (4 functions) | 33.6 MB | 1
specialization + 3 bookmark restores | 1 | match eager |

Shared weights are stored once: in a direct converter check, adding a
second function to an 8.4 MB model added 2.9 KB.

Runtime sharing, measured as resident memory after one run per delegate,
for K delegates of the same 33.6 MB model:

| K delegates | Option off | `coreai_share_multifunction_model=true` |
|---|---|---|
| 1 | +91 MB | +91 MB |
| 2 | +157 MB | +92 MB |
| 4 | +288 MB | +92 MB |
| 8 | +551 MB | +92 MB |

With sharing on, 2, 4 and 8 threads × 50 concurrent runs produced exact
outputs, in about the same total time as unshared on this model.

## Test plan

* New `backends/apple/coreai/test/test_multifunction.py`:
  * Partitioner option and compile spec; rejects `input_enumerations`.
* Two methods share one bundle (`export_hash`, `path`, files, data-store
keys) with `forward_0`/`prefill_0`, and the combined bundle is about the
size of a single method's.
* A graph break yields `forward_0`, `forward_1`, `prefill_0`,
`prefill_1` in one bundle.
* With the option off, there's one `main` model per delegate (as today).
* Mismatched compile specs, and the option set on only some methods,
raise errors.
* AOT with a mocked `coreai-build` (one build, shared `archs`) and with
the real `coreai-build`.
* Runtime: loads every method from one PTE through the Python runtime
and compares to eager, with and without a graph break.
* Runtime sharing (native):
* `ParsesOptionalMultifunctionFlag`: flag absent, true, false, and
non-boolean values.
  * `ShareOptionRequiresMultifunctionExport`, `OptionsRejectWrongTypes`.
* `ShareOptionSharesOneModelUntilLastDelegate`,
`SharedModelsAreScopedToCacheDirectory` (separate directories don't
share; a symlink alias does).
* `SharedModelSerializesConcurrentExecution`,
`ConcurrentFirstSharedLoadsAcquireOnce`,
`FailedSharedAcquisitionLeavesNoEntry`.
* `SharedLoadsOfDifferentModelsRunConcurrently`,
`WaiterAcquiresAfterConcurrentSharedLoadFails`.
* The concurrency and directory tests each fail when the lock, per-key
tracking or directory scoping they cover is removed.
* `coreai_multithread_test` adds cold and warm rounds on a
multi-function stress fixture with sharing on (real SDK).
* All 151 Core AI Python tests pass. All 6 native CTest targets pass.
* Backend suite operators: `coreai_fp32` 377 passed, `coreai_fp16` 368
passed, 0 failures.
* ufmt and flake8 are clean.
@metascroy
metascroy merged commit 4566bed into coreai-expose-compute-unit Nov 4, 2026
239 checks passed
@metascroy
metascroy deleted the coreai-rename-packaging branch November 4, 2026 19:36

This branch was previously deployed

1 inactive deployment
cadence — 0508853d Deployed Nov 4, 2026 by metascroy via hifi-op-test / hifi4 #36616
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants