[metascroy]

Makes Core AI exports that use input_enumerations load and run, including with multi-function export.

Stacked on #24522.

Problem

Every enumerated-shape PTE failed at load_method:

Core AI execution error: Model has no function named 'main'   (0x23)

AIProgram.set_static_shape_config erases the delegate's dynamic main and replaces it with one static clone per shape (main_s31_4, main_s31_16, ...). The manifest still named main, and the runtime binds exactly that name. On the real SDK, a specialized model exposes only the clones, each with a static input descriptor.

Export

An enumerated delegate's manifest lists its per-shape functions instead of one function:

"enumerated_functions": [
  {"function": "main_s31_4",  "input_shapes": [[4, 8]]},
  {"function": "main_s31_16", "input_shapes": [[16, 8]]},
  {"function": "main_s31_32", "input_shapes": [[32, 8]]}
]

input_shapes gives every input's shape in input_names order, the same order ExecuTorch passes arguments in. Delegates without enumerations keep function: "main".

If an enumerated value makes a derived boundary dim zero or negative (s0 - 1 at s0=1), export fails with an error naming the input, the dim and the value.

Runtime

• Manifest parsing: a manifest must have exactly one of function and enumerated_functions. Enumerated entries must have distinct names, one shape per input, positive integer dims, and no two entries with the same shapes. Enumerated delegates carry an explicit enumerated flag.
• Init: binds one session per listed function. Any bind failure fails init.
• Execute: validates the arguments once, then picks the function whose shapes equal the inputs. This happens inside the per-model execution lock. Inputs that match none fail with InvalidArgument before anything runs, and the log names the shapes received.
• Sharing: all functions come from one model, so they share weights, one specialization and one cache entry.
• Cost: the number of functions is the product of the enumerated value counts. On an M4 Max, a warm load took 7 ms with 1 function, 41 ms with 16 and 248 ms with 64 (documented in the README).

Multi-function export

CoreAIPartitioner(multifunction=True) accepts input_enumerations. Give each method its own partitioner to enumerate only some methods, or different values per method:

partitioner={
    "decode":  [CoreAIPartitioner(multifunction=True)],
    "prefill": [CoreAIPartitioner(multifunction=True, input_enumerations=[{0: [64, 128]}])],
}

• Naming: each enumerated partition expands into its own functions in the one shared model, prefill_0_<key>, plus prefill_1_<key> past a graph break. Unenumerated partitions stay single functions (decode_0).
• Binding order: every partition is bound before any is specialized, because specializing a partition erases its function.
• Compile specs: input_enumerations may differ between methods. Every other compile spec (AOT config, deployment floor) must match.
• Errors: export fails with a clear error when some methods omit multifunction=True, naming those methods, and when two delegates would generate the same function name.
• Runtime sharing: works with coreai_share_multifunction_model; every delegate and per-shape function shares one loaded model.

Shared-load robustness

This PR also tightens #24476's per-key shared loading:

• A scope guard clears a load's in-flight key, publishes the model and wakes waiters on every exit, so a key can't stay marked as loading.
• WaiterAcquiresAfterConcurrentSharedLoadFails confirms the second load is actually waiting before the first one fails. It uses a test-only BeforeSharedLoadWait fault point.

Test plan

Python. All 171 CoreAI tests pass.

• test_shape_config.py: test_manifest_lists_one_function_per_shape, test_unenumerated_lowering_keeps_main, test_derived_dim_that_reaches_zero_is_rejected. The real-runtime test test_runs_each_enumerated_shape_and_rejects_others runs each shape and checks rejection.
• test_multifunction.py:
  • test_enumerated_method_lists_per_shape_functions, test_graph_break_enumerates_every_partition.
  • test_other_compile_specs_must_still_match, test_mixed_partitioners_name_the_missing_flag, test_duplicate_function_names_are_rejected.
  • The real-runtime test test_enumerated_shapes_dispatch_by_input_shape covers multi-method with a graph break.

Native. All 6 CTest targets pass.

• ParsesEnumeratedFunctions, RejectsInvalidEnumeratedFunctions (16 malformed cases, plus both or neither field present).
• EnumeratedFunctionsDispatchByInputShape: checks which function ran, the resized output, and that an unmatched shape runs nothing.
• EnumeratedFunctionBindFailureFailsInit.
• The sharing tests, including WaiterAcquiresAfterConcurrentSharedLoadFails, passed in 5 of 5 runs.

Real SDK, M4 Max:

• Single method: rows 4/16/32 match eager (max error ≤1.8e-7) on cold and warm loads. Rows 8 is rejected with InvalidArgument.
• Multi-method, C++ runner: a static forward plus an enumerated prefill, with a graph break. That's 4 delegates and 8 functions in one model. With sharing off and on, cold and warm, 100 of 100 calls pass in each configuration (worst max error 1.4e-6), and an unmatched shape is rejected with 0x12.
• AOT (aimodelc, h16c): the enumerated functions compile and run (max error ≤1.8e-7).

Compatibility

PTEs exported with input_enumerations before this change never loaded, so nothing that worked before breaks. Every other export keeps the function manifest field.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Planned maintenance affecting GB300 CI jobs

❌ 2 New Failures, 1 Unrelated Failure

As of commit 6eab989 with merge base ee4b85b ():

NEW FAILURES - The following jobs have failed:

• pull / android / run-emulator (gh)
  Warning: An error occurred while preparing SDK package Android Emulator: Error on ZipFile unknown archive.
• pull / test-qnn-wheel-packages-linux (3.13) / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-arm-backend-vkml (test_ootb_tests_vgf) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

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