[CoreAI] Add coreai_expect_frequent_reshapes runtime option by metascroy · Pull Request #24522 · pytorch/executorch

metascroy · 2026-11-03T19:15:00Z

Adds a coreai_expect_frequent_reshapes runtime option that exposes the Core AI SDK's SpecializationOptions.expectFrequentReshapes.

Stacked on #24476.

Usage

BackendOptions<2> options;
options.set_option("coreai_expect_frequent_reshapes", true);
options.set_option("coreai_compute_unit", "prefer_gpu");   // optional; combines with it
LoadBackendOptionsMap map;
map.set_options("CoreAIBackend", options.view());
module.load(map);

Bool, default false. Set it for models whose input shapes change between calls. For AOT exports, the matching export-time setting is AOTCompileConfig(expect_frequent_reshapes=True), which already existed.

Why this option

Checked against the Xcode 27 SDK by type-checking each member through import CoreAI. Of SpecializationOptions, the public surface is .default, .cpuOnly, init(preferredComputeUnitKind:), read-only allowedComputeUnitKinds / preferredComputeUnitKind, and read/write expectFrequentReshapes. Everything except expectFrequentReshapes is already covered by coreai_compute_unit. The framework exports other symbols (.interpreterOnly, init(allowedComputeUnitKinds:preferredComputeUnitKind:), enableEncodingFunctions, ...), but none of them are usable from import CoreAI.

The SDK has no doc comment for expectFrequentReshapes. The README describes it by its name only.

How it works

CoreAIBackend::init reads the option with the shared runtime_option helper. A non-boolean value fails init with InvalidArgument.
ETCoreAICreateModelLoader takes the flag next to the compute unit. The Swift loader sets expectFrequentReshapes on the SpecializationOptions chosen for that compute unit.
No cache-key change is needed. The specialization identifier already encodes expectFrequentReshapes, so each value gets its own specialization, bookmark and shared model (coreai_share_multifunction_model only shares between delegates that agree on it).
Cache eviction creates its loader with the flag off, since eviction targets the bookmark and the specialization doesn't matter.

Test plan

CoreAIDelegateTest.ExpectFrequentReshapesOptionReachesLoader: absent / false / true reach the loader factory, alongside coreai_compute_unit.
CoreAIDelegateTest.OptionsRejectWrongTypes: an integer value fails init.
SwiftBridgeTests.testComputeUnitsMapToSDKSpecializationOptions: every compute unit, with and without reshapes, maps to the expected SpecializationOptions, and all 10 combinations have distinct identifiers.
Real SDK (M4 Max, multi-function PTE with sharing on): with the option, the model specializes under its own cache key, a warm load restores it, and outputs match eager. Without it, the default key is used. That gives 2 bookmarks in one cache directory.
All 6 native CTest targets pass.

pytorch-bot · 2026-11-03T19:15:05Z

🔗 Helpful Links

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

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

Exposes the SDK's SpecializationOptions.expectFrequentReshapes as a boolean backend option (default false), applied on top of the coreai_compute_unit selection. It is the only public specialization setting the backend did not expose yet; the export-time counterpart for AOT is AOTCompileConfig(expect_frequent_reshapes=True). The specialization identifier already encodes expectFrequentReshapes, so each value is specialized, cached and shared (coreai_share_multifunction_model) separately with no cache-key change. ETCoreAICreateModelLoader takes the flag next to the compute unit; cache eviction passes NO since it targets the bookmark.

…24525) 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`: ```json "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: ```python 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.

metascroy requested a review from shoumikhin as a code owner November 3, 2026 19:15

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 temporarily deployed to cadence November 3, 2026 19:15 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 3, 2026 19:35 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 3, 2026 19:42 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 3, 2026 19:51 — with GitHub Actions Inactive

metascroy mentioned this pull request Nov 3, 2026

[CoreAI] Run enumerated-shape exports by dispatching on input shape #24525

Merged

metascroy force-pushed the coreai-expect-frequent-reshapes branch from b2fb6ae to 6e9fd39 Compare November 3, 2026 21:15

metascroy temporarily deployed to cadence November 3, 2026 21:16 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 3, 2026 21:34 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 3, 2026 21:40 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 3, 2026 21:53 — with GitHub Actions Inactive

metascroy force-pushed the coreai-multifunction branch from 7ac946a to 2b55aa0 Compare November 3, 2026 23:12

metascroy requested a review from larryliu0820 as a code owner November 3, 2026 23:12

metascroy force-pushed the coreai-expect-frequent-reshapes branch from 6e9fd39 to 4827a47 Compare November 3, 2026 23:12

metascroy temporarily deployed to cadence November 3, 2026 23:12 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 3, 2026 23:26 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 3, 2026 23:29 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 3, 2026 23:30 — with GitHub Actions Inactive

metascroy force-pushed the coreai-multifunction branch from 2b55aa0 to 216529d Compare November 4, 2026 00:11

metascroy force-pushed the coreai-expect-frequent-reshapes branch from 4827a47 to 4994aef Compare November 4, 2026 00:11

metascroy temporarily deployed to cadence November 4, 2026 00:11 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 4, 2026 00:25 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 4, 2026 00:28 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 4, 2026 00:29 — with GitHub Actions Inactive

SS-JIA approved these changes Nov 4, 2026

View reviewed changes

metascroy merged commit 211bc99 into coreai-multifunction Nov 4, 2026

metascroy deleted the coreai-expect-frequent-reshapes branch November 4, 2026 19:34

metascroy temporarily deployed to cadence November 4, 2026 19:34 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 4, 2026 19:52 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 4, 2026 19:54 — with GitHub Actions Inactive

metascroy temporarily deployed to cadence November 4, 2026 19:56 — with GitHub Actions Inactive

Navigation Menu

Repository navigation

[CoreAI] Add coreai_expect_frequent_reshapes runtime option - #24522

[CoreAI] Add coreai_expect_frequent_reshapes runtime option#24522
metascroy merged 2 commits into
coreai-multifunctionfrom
coreai-expect-frequent-reshapes

metascroy commented Nov 3, 2026

Uh oh!

pytorch-bot Bot commented Nov 3, 2026 •

edited

Loading

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

metascroy commented Nov 3, 2026

Usage

Why this option

How it works

Test plan

Uh oh!

pytorch-bot Bot commented Nov 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pytorch-bot Bot commented Nov 3, 2026 •

edited

Loading