[CoreAI] Rename manifest packaging to aimodel / aimodelc

backends/apple/coreai/CMakeLists.txt

@@ -463,6 +463,7 @@ if(_coreai_run_sdk_tests
   set(_coreai_stress_fixture "${CMAKE_CURRENT_BINARY_DIR}/stress_fixture")
   add_custom_command(
     OUTPUT "${_coreai_stress_fixture}/model.pte"
+           "${_coreai_stress_fixture}/model_multifunction.pte"
            "${_coreai_stress_fixture}/fixture.json"
     COMMAND
       "${PYTHON_EXECUTABLE}"
@@ -474,8 +475,10 @@ if(_coreai_run_sdk_tests
     VERBATIM
   )
   add_custom_target(
-    coreai_stress_fixture DEPENDS "${_coreai_stress_fixture}/model.pte"
-                                  "${_coreai_stress_fixture}/fixture.json"
+    coreai_stress_fixture
+    DEPENDS "${_coreai_stress_fixture}/model.pte"
+            "${_coreai_stress_fixture}/model_multifunction.pte"
+            "${_coreai_stress_fixture}/fixture.json"
   )
   add_executable(
     coreai_multithread_test runtime/test/coreai_multithread_test.cpp

backends/apple/coreai/README.md

@@ -26,14 +26,18 @@ python -m pytest backends/apple/coreai
 
 ## Runtime
 
-The delegate manifest describes `.aimodel` source bundles, which require
-`inline` packaging, and AOT-compiled `.aimodelc` bundles, which require
-`aot_compiled_inline`. Both formats require a `function` name, and ordered
+The delegate manifest's `packaging` field names the embedded format: `aimodel`
+for a portable `.aimodel` bundle, located by `path`, or `aimodelc` for
+AOT-compiled per-architecture `.aimodelc` bundles, located by `archs`. Both
+formats require exactly one of a `function` name or `enumerated_functions` (see
+[Enumerated input shapes](#enumerated-input-shapes)), and ordered
 `input_names` and `output_names` matching the converter's bindings. Dictionary
 iteration order is not a binding contract. `export_hash` identifies the exported
 program used to namespace embedded assets. The `files` object maps relative
 filenames to byte sizes; `bundle_digests` maps each bundle basename to its
-export-time SHA-256 digest. Asset paths, including bundle and directory names,
+export-time SHA-256 digest. The optional boolean `multifunction` marks a delegate
+exported as one function of a model shared with other delegates. Asset paths,
+including bundle and directory names,
 must remain distinct under case-insensitive, canonically normalized Unicode
 comparison. Accepted names are preserved, not rewritten.
 
@@ -107,6 +111,88 @@ bookmark restoration. A missing bookmark or SDK-confirmed cache miss falls back
 to source materialization and persistent specialization, then binds the requested
 function and ordered I/O on that same acquired model.
 
+### Multi-function export
+
+By default each Core AI delegate is its own model with a single `main` function.
+`CoreAIPartitioner(multifunction=True)` instead lowers every Core AI partition of
+every method into one model, with one function per partition named
+`<method>_<partition index>`:
+
+```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(),
+)
+```
+
+Here the delegates bind `forward_0` and `prefill_0`; a graph break adds
+`forward_1`, `prefill_1` and so on. Weights shared between methods are stored
+once. Every partition must carry identical compile specs (AOT config, deployment
+floor), except `input_enumerations`, which each method may set on its own (see
+below).
+
+Each delegate keeps its own manifest naming its function and ordered I/O; all of
+them reference the same embedded bundle. Function bindings are not part of the
+cache key, so the first load specializes the model and the other delegates
+restore it from the same bookmark. Each delegate still loads its own model
+instance unless the `coreai_share_multifunction_model` runtime option is set
+(see [Runtime options](#runtime-options)). Multi-function manifests carry
+`"multifunction": true`.
+
+### Enumerated input shapes
+
+`CoreAIPartitioner(input_enumerations=...)` compiles a dynamic dim to a fixed set
+of sizes instead of leaving it dynamic. Entries align with the ET model inputs and
+map a dim index to its allowed values; the dim must be exported dynamic:
+
+```python
+batch = torch.export.Dim("batch", min=1, max=64)
+ep = torch.export.export(model, (x,), dynamic_shapes={"x": {0: batch}})
+lowered = to_edge_transform_and_lower(
+    ep,
+    transform_passes=get_default_passes(),
+    partitioner=[CoreAIPartitioner(input_enumerations=[{0: [4, 16, 32]}])],
+    compile_config=get_default_compile_config(),
+)
+```
+
+Core AI replaces the delegate's dynamic function with one static function per
+combination of enumerated values, so the manifest lists them under
+`enumerated_functions` instead of naming 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 delegate
+binds every listed function at load and runs the one whose shapes equal the call's
+inputs; inputs matching none of them fail with `InvalidArgument` without running.
+All functions come from one model, so they share weights, one specialization and
+one cache entry. The function count is the product of the enumerated value counts
+across dims, and specialization and load time grow with it: on an M4 Max, a warm
+load of a two-layer MLP (256 to 1024 to 256 features) took 7 ms with 1 function, 41 ms with 16 and 248 ms
+with 64.
+
+With `multifunction=True`, 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]}])],
+}
+```
+
+Every enumerated partition expands into its own functions (`prefill_0_<key>`, and
+`prefill_1_<key>` past a graph break) inside the one shared model, while `decode_0`
+stays a single function. The other compile specs must still match across methods.
+
 ### Cache directory
 
 Backend initialization reads the string runtime spec `coreai_cache_dir`. If
@@ -128,9 +214,9 @@ even though the new complete contents are already visible. Use a local APFS
 volume; unsupported full-device flushes are reported as errors, not downgraded
 to `fsync`.
 
-Inline bundles are read through `NamedDataMap` and materialized as unchanged
-files under the cache directory. Named data can be supplied externally; inline
-packaging does not require all bytes to reside in one physical PTE. The exporter
+Embedded bundles are read through `NamedDataMap` and materialized as unchanged
+files under the cache directory. Named data can be supplied externally; embedded
+bundles do not require all bytes to reside in one physical PTE.
 computes each bundle digest from all delivered relative filenames and file bytes,
 not just the SDK's bytecode-only `main.hash`. Each AOT architecture has an
 independent digest; only the selected architecture is materialized. NDS key names
@@ -191,7 +277,7 @@ choose the SDK artifact directory. The backend does not inspect SDK-private
 files, alter their backup flags, or use an App Group cache.
 
 Each selected bundle has a `bundle_identity` derived from its bundle digest,
-name, and source/AOT mode. This identifies the delivered bundle independently of
+name, and packaging (`aimodel` or `aimodelc`). This identifies the delivered bundle independently of
 `export_hash`, which namespaces embedded assets. A readable cache key combines
 that bundle identity with
 the platform and Core AI device architecture, plus a hash of a versioned
@@ -361,6 +447,46 @@ BackendOptions<2> options;
 ET_CHECK_OK_OR_RETURN_ERROR(options.set_option("coreai_compute_unit", "cpu_only"));
 ```
 
+`coreai_expect_frequent_reshapes` (bool, default `false`) sets the SDK's
+`SpecializationOptions.expectFrequentReshapes` on top of the selected compute
+unit. Set it for models whose input shapes change between calls. It is part of
+the cache key, so each value is specialized and cached separately, and delegates
+only share a model (`coreai_share_multifunction_model`) when they agree on it.
+For AOT exports, the matching export-time setting is
+`AOTCompileConfig(expect_frequent_reshapes=True)`.
+
+```cpp
+BackendOptions<2> options;
+ET_CHECK_OK_OR_RETURN_ERROR(
+    options.set_option("coreai_expect_frequent_reshapes", true));
+```
+
+`coreai_share_multifunction_model` (bool, default `false`) makes the delegates of
+a multi-function export (`CoreAIPartitioner(multifunction=True)`) share one
+loaded Core AI model per process instead of each loading its own. Without it,
+every delegate restores a separate model with its own copy of the weights. With
+it, the first delegate loads the model and the others reuse it until the last
+one is destroyed. Delegates that start loading the same model concurrently wait
+for that single load; loads of other models proceed in parallel. Delegates sharing a model run one at a time: concurrent
+execution on one model is serialized by a per-model lock, because the SDK does
+not document it as safe and it produced wrong outputs in testing. Setting the
+option on a delegate that isn't from a multi-function export fails
+initialization with `InvalidArgument`.
+
+Sharing is process-wide, not per program, but scoped to the cache directory.
+Separate `Module`s that load the same PTE with the option set and the same
+`coreai_cache_dir` share one model, so their `execute` calls block each other.
+`Module`s with different cache directories never share: each loads its own model
+and keeps its own bookmark. A symlink to a cache directory counts as that
+directory. Different `coreai_compute_unit` values also load separate models,
+since each is a separate specialization.
+
+```cpp
+BackendOptions<2> options;
+ET_CHECK_OK_OR_RETURN_ERROR(
+    options.set_option("coreai_share_multifunction_model", true));
+```
+
 ## SDK Bridge
 
 The private Swift module `CoreAIBridge` separates error/completion support,

backends/apple/coreai/compiler/enumerated_shapes.py

@@ -166,7 +166,15 @@ def _eval_dim(dim, subs: Dict[str, int], uname: str, axis: int) -> int:
             f"substitution; enumerate those dims too, or leave this input "
             f"fully dynamic"
         )
-    return int(resolved)
+    value = int(resolved)
+    if value <= 0:
+        values = ", ".join(sorted(f"{s}={v}" for s, v in mapping.items()))
+        raise ValueError(
+            f"input '{uname}' dim {axis} is {expr}, which is {value} when "
+            f"{values}; Core AI needs every dim to be positive, so drop that "
+            "enumerated value"
+        )
+    return value
 
 
 def apply_enumerated_shapes(
@@ -175,7 +183,8 @@ def apply_enumerated_shapes(
     enumerations: Dict[str, Sequence[int]],
     *,
     input_names: Sequence[str],
-) -> None:
+    entrypoint: str = MAIN_ENTRYPOINT,
+) -> Dict[str, Dict[str, tuple]]:
     """Attach enumerated shapes to the coreai program (all delivery modes).
 
     ``enumerations`` is ``{symbol_name: [value, ...]}`` (from
@@ -184,15 +193,19 @@ def apply_enumerated_shapes(
     combination of the enumerated symbol values, then attach the resulting shapes
     via ``set_static_shape_config``. ``input_names`` are the names the converter
     was given for the user inputs, in boundary order.
+
+    Returns ``{key: {input_name: shape}}`` for the attached specializations, or
+    an empty dict when this subgraph has nothing enumerated. ``entrypoint`` is
+    erased and replaced by one function named ``f"{entrypoint}_{key}"`` per key.
     """
     if not enumerations:
-        return
+        return {}
 
     placeholders = {
         n.name: n for n in edge_program.graph.nodes if n.op == "placeholder"
     }
     user_inputs = list(edge_program.graph_signature.user_inputs)
-    coreai_names = set(graph_input_names(program))
+    coreai_names = set(graph_input_names(program, entrypoint))
 
     # Every entry has to name every boundary input, since
     # set_static_shape_config treats a key as a whole-graph shape: an input left
@@ -217,7 +230,7 @@ def apply_enumerated_shapes(
         # Nothing here is enumerated, so there is nothing to specialize. Routine
         # under multi-partition lowering, and legitimate when the enumerated
         # input is not delegated at all.
-        return
+        return {}
 
     # One axis per enumerated symbol, shared across every input that uses it.
     # _eval_dim rejects any boundary dim left symbolic, so a specialization
@@ -243,4 +256,5 @@ def apply_enumerated_shapes(
             len(shapes_config),
             active,
         )
-        program.set_static_shape_config(MAIN_ENTRYPOINT, shapes_config)
+        program.set_static_shape_config(entrypoint, shapes_config)
+    return shapes_config

backends/apple/coreai/compiler/io_compat.py

@@ -185,13 +185,15 @@ def io_mismatches(
     return errors
 
 
-def coreai_io_names(program) -> Tuple[List[str], List[str]]:
+def coreai_io_names(
+    program, entrypoint: str = MAIN_ENTRYPOINT
+) -> Tuple[List[str], List[str]]:
     """Boundary I/O names in graph argument/result order, never dictionary order.
 
     State arguments (marked ``MutableBuffers.buffer_mutation``) and the results
     they name are skipped: they are not ExecuTorch boundary I/O.
     """
-    graph = program._get_graph(MAIN_ENTRYPOINT)
+    graph = program._get_graph(entrypoint)
     arg_attrs = list(graph.arg_attrs)
     res_attrs = list(graph.res_attrs)
     for kind, attrs, types in (
@@ -227,9 +229,10 @@ def assert_io_compatible(
     *,
     input_names: Optional[Sequence[str]] = None,
     output_names: Optional[Sequence[str]] = None,
+    entrypoint: str = MAIN_ENTRYPOINT,
 ) -> None:
     """Raise if the ``.aimodel`` boundary I/O is incompatible with ExecuTorch."""
-    coreai_in, coreai_out = _coreai_io(program)
+    coreai_in, coreai_out = _coreai_io(program, entrypoint)
     edge_in, edge_out = _edge_io(edge_program)
     logger.info(
         "Core AI delegate boundary I/O:\n"
@@ -246,7 +249,7 @@ def assert_io_compatible(
         coreai_out, edge_out, "output"
     )
     if input_names is not None or output_names is not None:
-        actual_inputs, actual_outputs = coreai_io_names(program)
+        actual_inputs, actual_outputs = coreai_io_names(program, entrypoint)
         for kind, actual, expected in (
             ("input", actual_inputs, input_names),
             ("output", actual_outputs, output_names),