[metascroy]

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

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:

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.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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.