Repository navigation
[CoreAI] Rename manifest packaging to aimodel / aimodelc - #24474
Merged
Merged
Conversation
🔗 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. |
metascroy
force-pushed
the
coreai-expose-compute-unit
branch
from
November 3, 2026 18:59
d05c245 to
d703edf
Compare
metascroy
force-pushed
the
coreai-rename-packaging
branch
from
November 3, 2026 18:59
cfdb000 to
73ded2d
Compare
SS-JIA
approved these changes
Nov 3, 2026
metascroy
force-pushed
the
coreai-rename-packaging
branch
from
November 3, 2026 23:12
73ded2d to
5d03055
Compare
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.
metascroy
force-pushed
the
coreai-rename-packaging
branch
from
November 4, 2026 00:11
5d03055 to
884d214
Compare
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.
This branch was previously deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Renames the Core AI delegate manifest's
packagingvalues 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.aimodelversus AOT-compiled.aimodelcbundles. 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
packaging"inline"/"aot_compiled_inline""aimodel"/"aimodelc"AssetPackaging(compiler/preprocess.py)INLINE/AOT_COMPILED_INLINEAIMODEL/AIMODELCruntime/coreai_manifest.mm)aot_compiled_prefix== "aimodelc"coreai_load_coordinator.mm)kind=source/kind=aotkind=aimodel/kind=aimodelccoreai_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
Unsupported Core AI packaging; expected aimodel or aimodelcand must be re-exported.Test plan
test_preprocess.py: 25/25 passed, including the realcoreai-buildAOT test.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 carrypackaging: "aimodel".kind=aimodel, outputs match the reference.kind=aimodelcwith exact outputs; second load restores from the cache bookmark.