[CoreAI] Add multi-function export and shared runtime loading

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

@@ -34,7 +34,9 @@ formats require a `function` name, and ordered
 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.
 
@@ -108,6 +110,35 @@ 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), and `input_enumerations` is not supported with this option yet.
+
+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`.
+
 ### Cache directory
 
 Backend initialization reads the string runtime spec `coreai_cache_dir`. If
@@ -362,6 +393,32 @@ BackendOptions<2> options;
 ET_CHECK_OK_OR_RETURN_ERROR(options.set_option("coreai_compute_unit", "cpu_only"));
 ```
 
+`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/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),