[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.

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.

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.
• Runtime: unchanged, and so is the manifest format. 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.
• 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.

Process memory after loading and running is about the same either way (~184 MB), because each delegate still loads its own model instance. In-process sharing of one loaded model could be a follow-up if it turns out to matter.

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.
• 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.

❌ 3 New Failures, 1 Unclassified Failure

As of commit 92c2805 with merge base 822cb60 ():

NEW FAILURES - The following jobs have failed:

• pull / android / run-emulator (gh)
  The process '/usr/bin/sh' failed with exit code 1
• pull / nxp-mcuxpresso-test / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
• pull / test-models-linux-basic (mv3, xnnpack-quantization-delegation, cmake, mt-l-arm64g4-16-62, executo... / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• Core AI / test-coreai / test-coreai (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 8

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