diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index b5c1bfae8a3..f82569c963d 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/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 diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index ec0fd4856ca..8d242505aaa 100644 --- a/backends/apple/coreai/README.md +++ b/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 +`_`: + +```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, diff --git a/backends/apple/coreai/compiler/io_compat.py b/backends/apple/coreai/compiler/io_compat.py index 3968f044368..4f3ea578fc0 100644 --- a/backends/apple/coreai/compiler/io_compat.py +++ b/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), diff --git a/backends/apple/coreai/compiler/preprocess.py b/backends/apple/coreai/compiler/preprocess.py index 1d6b455bd7e..39e5c3e9fa3 100644 --- a/backends/apple/coreai/compiler/preprocess.py +++ b/backends/apple/coreai/compiler/preprocess.py @@ -75,6 +75,11 @@ class COMPILE_SPEC_KEYS(Enum): # the portable .aimodel, so it is not part of this AOT-only blob.) AOT_COMPILE_CONFIG = "coreai_aot_compile_config" + # Lower every Core AI partition of every method into one Core AI model, one + # function per partition named _. Value: b"1". + # Every partition must carry identical compile specs. + MULTIFUNCTION = "coreai_multifunction" + class AssetPackaging(str, Enum): # Portable .aimodel embedded in the .pte via NamedDataStore. @@ -195,24 +200,22 @@ def _deliver( staging: Path, packaging: AssetPackaging, manifest_extra: Dict[str, Any], -) -> PreprocessResult: - """Embed a staging directory of bundles in the .pte with its manifest. + bindings: List[Dict[str, Any]], +) -> List[PreprocessResult]: + """Embed a staging directory of bundles in the .pte, one manifest per binding. ``staging`` holds the ``.aimodel`` / ``.aimodelc`` bundles, never their contents, so assets land under ``coreai///`` in the NamedDataStore. ``manifest_extra`` paths are staging-relative bundle names. + Each binding (``function``, ``input_names``, ``output_names``) gets its own + manifest; all of them reference the same embedded bundles. """ metadata = _asset_metadata(staging) # The SDK program hash excludes metadata such as creation timestamps. export_hash = hashlib.sha256( json.dumps(metadata["bundle_digests"], sort_keys=True).encode("utf-8") ).hexdigest() - manifest = { - "packaging": packaging.value, - "export_hash": export_hash, - **manifest_extra, - **metadata, - } + manifest = {**manifest_extra, **metadata} if packaging == AssetPackaging.AIMODEL: manifest["path"] = f"{export_hash}/{manifest['path']}" else: @@ -220,7 +223,21 @@ def _deliver( arch: f"{export_hash}/{bundle}" for arch, bundle in manifest["archs"].items() } - return _embed_dir_inline(staging, export_hash, manifest) + store_output = _embed_dir(staging, export_hash) + return [ + PreprocessResult( + processed_bytes=json.dumps( + { + "packaging": packaging.value, + "export_hash": export_hash, + **binding, + **manifest, + } + ).encode("utf-8"), + data_store_output=store_output, + ) + for binding in bindings + ] def _prepare_program_for_conversion(edge_program: ExportedProgram) -> ExportedProgram: @@ -245,35 +262,72 @@ def _prepare_program_for_conversion(edge_program: ExportedProgram) -> ExportedPr def _convert_to_aiprogram( - edge_program: ExportedProgram, - *, - input_names: Optional[List[str]] = None, - output_names: Optional[List[str]] = None, + entries: List[Tuple[ExportedProgram, str, List[str], List[str]]], ): + """Convert ``(edge_program, entrypoint, input_names, output_names)`` entries + into one Core AI program with one function per entry.""" from coreai_torch import TorchConverter - aten_program = _prepare_program_for_conversion(edge_program) converter = TorchConverter() - converter.add_exported_program( - aten_program, + for edge_program, entrypoint, input_names, output_names in entries: + converter.add_exported_program( + _prepare_program_for_conversion(edge_program), + input_names=input_names, + output_names=output_names, + entrypoint_name=entrypoint, + ) + return converter.to_coreai() + + +def _boundary_names(edge_program: ExportedProgram) -> Tuple[List[str], List[str]]: + mutated_inputs = edge_program.graph_signature.user_inputs_to_mutate + if mutated_inputs: + raise ValueError( + "Core AI delegates cannot mutate ExecuTorch-owned user inputs: " + f"{sorted(mutated_inputs.values())}" + ) + edge_inputs, edge_outputs = _edge_io(edge_program) + return ( + [f"input_{i}" for i in range(len(edge_inputs))], + [f"output_{i}" for i in range(len(edge_outputs))], + ) + + +def _binding( + program, + edge_program: ExportedProgram, + entrypoint: str, + input_names: List[str], + output_names: List[str], +) -> Dict[str, Any]: + # Fail fast if the .aimodel boundary I/O won't match what ET feeds/reads. + assert_io_compatible( + program, + edge_program, input_names=input_names, output_names=output_names, - entrypoint_name=MAIN_ENTRYPOINT, + entrypoint=entrypoint, ) - return converter.to_coreai() + bound_inputs, bound_outputs = coreai_io_names(program, entrypoint) + return { + "function": entrypoint, + "input_names": bound_inputs, + "output_names": bound_outputs, + } + + +def _spec_items(compile_specs: List[CompileSpec]) -> List[Tuple[str, bytes]]: + return sorted((spec.key, bytes(spec.value)) for spec in compile_specs) # Asset embedding helpers. -def _embed_dir_inline( - root_dir: Path, export_hash: str, manifest_extra: Dict[str, Any] -) -> PreprocessResult: - """Flatten every file under ``root_dir`` into the NamedDataStore. +def _embed_dir(root_dir: Path, export_hash: str): + """Flatten every file under ``root_dir`` into a NamedDataStore output. ``root_dir`` is the directory *containing* the ``.aimodel`` / ``.aimodelc`` bundles, never a bundle itself, so the bundle name survives into the keys. - Keys are ``coreai/{export_hash}/{relpath}``; ``manifest_extra`` becomes the - manifest (packaging, files, archs, ...). + Keys are ``coreai/{export_hash}/{relpath}``. """ store = NamedDataStore() for path in sorted(root_dir.rglob("*")): @@ -283,10 +337,7 @@ def _embed_dir_inline( path.read_bytes(), alignment=_ASSET_ALIGNMENT, ) - return PreprocessResult( - processed_bytes=json.dumps(manifest_extra).encode("utf-8"), - data_store_output=store.get_named_data_store_output(), - ) + return store.get_named_data_store_output() def _os_version_text(min_os) -> Optional[str]: @@ -458,29 +509,14 @@ def preprocess( edge_program: ExportedProgram, compile_specs: List[CompileSpec], ) -> PreprocessResult: - aot_compiled = ( - _get_compile_spec(compile_specs, COMPILE_SPEC_KEYS.AOT_COMPILE_CONFIG) - is not None - ) - - mutated_inputs = edge_program.graph_signature.user_inputs_to_mutate - if mutated_inputs: - raise ValueError( - "Core AI delegates cannot mutate ExecuTorch-owned user inputs: " - f"{sorted(mutated_inputs.values())}" - ) - edge_inputs, edge_outputs = _edge_io(edge_program) - input_names = [f"input_{i}" for i in range(len(edge_inputs))] - output_names = [f"output_{i}" for i in range(len(edge_outputs))] + input_names, output_names = _boundary_names(edge_program) program = _convert_to_aiprogram( - edge_program, input_names=input_names, output_names=output_names + [(edge_program, MAIN_ENTRYPOINT, input_names, output_names)] ) - # Fail fast if the .aimodel boundary I/O won't match what ET feeds/reads. - assert_io_compatible( - program, edge_program, input_names=input_names, output_names=output_names + # Enumerated shapes replace ``main`` with renamed clones, so bind first. + binding = _binding( + program, edge_program, MAIN_ENTRYPOINT, input_names, output_names ) - # Enumerated shapes replace ``main`` with renamed clones, so read first. - bound_inputs, bound_outputs = coreai_io_names(program) raw_enum = _get_compile_spec( compile_specs, COMPILE_SPEC_KEYS.INPUT_ENUMERATIONS ) @@ -491,19 +527,74 @@ def preprocess( json.loads(raw_enum.decode()), input_names=input_names, ) - bindings = { - "function": MAIN_ENTRYPOINT, - "input_names": bound_inputs, - "output_names": bound_outputs, + return CoreAIBackend._package(program, compile_specs, [binding])[0] + + @classmethod + def preprocess_multimethod( + cls, + edge_programs: Dict[str, List[ExportedProgram]], + compile_specs: Dict[str, List[List[CompileSpec]]], + ) -> Dict[str, List[PreprocessResult]]: + partitions = [ + (method, index, program, compile_specs[method][index]) + for method, programs in edge_programs.items() + for index, program in enumerate(programs) + ] + if not any( + _get_compile_spec(specs, COMPILE_SPEC_KEYS.MULTIFUNCTION) is not None + for _, _, _, specs in partitions + ): + return super().preprocess_multimethod(edge_programs, compile_specs) + + reference = _spec_items(partitions[0][3]) + differing = sorted( + { + method + for method, _, _, specs in partitions + if _spec_items(specs) != reference + } + ) + if differing: + raise ValueError( + "multifunction requires identical Core AI compile specs in every " + f"partition; methods {differing} differ from {partitions[0][0]!r}" + ) + if _get_compile_spec(partitions[0][3], COMPILE_SPEC_KEYS.INPUT_ENUMERATIONS): + raise ValueError("multifunction does not support input_enumerations yet") + + entries = [ + (program, f"{method}_{index}", *_boundary_names(program)) + for method, index, program, _ in partitions + ] + logger.info( + "Lowering %d Core AI partitions into one model: %s", + len(entries), + [entrypoint for _, entrypoint, _, _ in entries], + ) + program = _convert_to_aiprogram(entries) + bindings = [ + {**_binding(program, *entry), "multifunction": True} for entry in entries + ] + results = CoreAIBackend._package(program, partitions[0][3], bindings) + + by_method: Dict[str, List[PreprocessResult]] = { + method: [] for method in edge_programs } + for (method, _, _, _), result in zip(partitions, results): + by_method[method].append(result) + return by_method + @staticmethod + def _package( + program, compile_specs: List[CompileSpec], bindings: List[Dict[str, Any]] + ) -> List[PreprocessResult]: # min-deployment-version is a single knob applied to whichever artifact # ships: for aot-compiled delivery the .aimodelc (via coreai-build # --min-deployment-version), for portable delivery the .aimodel's floor # (via save_asset(minimum_os=...)). In the aot-compiled path the temp # .aimodel is discarded, so it keeps the default floor (no # double-specification). - if aot_compiled: + if _get_compile_spec(compile_specs, COMPILE_SPEC_KEYS.AOT_COMPILE_CONFIG): return CoreAIBackend._preprocess_aot_compiled( program, _aot_compile_options(compile_specs), bindings ) @@ -514,33 +605,32 @@ def preprocess( # Portable .aimodel delivery. @staticmethod def _preprocess_portable( - program, min_os, bindings: Dict[str, Any] - ) -> PreprocessResult: + program, min_os, bindings: List[Dict[str, Any]] + ) -> List[PreprocessResult]: with TemporaryDirectory() as tmp: _save_asset(program, Path(tmp) / "model.aimodel", min_os) return _deliver( Path(tmp), AssetPackaging.AIMODEL, { - **bindings, # relative path the runtime resolves against its base "path": "model.aimodel", "min_deployment_version": _os_version_text(min_os), }, + bindings, ) # AOT-compiled .aimodelc delivery (per architecture). @staticmethod def _preprocess_aot_compiled( - program, opts: Dict[str, Any], bindings: Dict[str, Any] - ) -> PreprocessResult: + program, opts: Dict[str, Any], bindings: List[Dict[str, Any]] + ) -> List[PreprocessResult]: with TemporaryDirectory() as tmp: out = _compile_aot(program, opts, Path(tmp)) return _deliver( out, AssetPackaging.AIMODELC, { - **bindings, "platform": opts["platform"], "min_deployment_version": opts["min_deployment_version"], "archs": { @@ -548,4 +638,5 @@ def _preprocess_aot_compiled( for arch, bundle in _compiled_arch_bundles(out) }, }, + bindings, ) diff --git a/backends/apple/coreai/partition/partitioner.py b/backends/apple/coreai/partition/partitioner.py index 4b4b482bd80..300298fc2be 100644 --- a/backends/apple/coreai/partition/partitioner.py +++ b/backends/apple/coreai/partition/partitioner.py @@ -317,6 +317,7 @@ def __init__( ] = None, take_over_constant_data: bool = True, take_over_mutable_buffer: bool = True, + multifunction: bool = False, ) -> None: # aot_compile_config requests ahead-of-time ``xcrun coreai-build # compile`` in preprocess, emitting per-architecture ``.aimodelc`` @@ -329,6 +330,11 @@ def __init__( # ({dim_index: [values]}, aligned to inputs); it is resolved to export # symbols in :meth:`partition` and propagated to each subgraph boundary # in preprocess, so it is not built into a compile spec here. + # + # multifunction lowers every Core AI partition of every method into one + # Core AI model, one function per partition named _. + if multifunction and input_enumerations is not None: + raise ValueError("multifunction does not support input_enumerations yet") specs = [] if min_deployment_version is not None: specs.append( @@ -344,6 +350,8 @@ def __init__( aot_compile_config.to_json().encode(), ) ) + if multifunction: + specs.append(CompileSpec(COMPILE_SPEC_KEYS.MULTIFUNCTION.value, b"1")) self.delegation_spec = DelegationSpec( backend_id=CoreAIBackend.__name__, diff --git a/backends/apple/coreai/runtime/coreai_backend.mm b/backends/apple/coreai/runtime/coreai_backend.mm index fe4f6c8a9a3..bcaadded533 100644 --- a/backends/apple/coreai/runtime/coreai_backend.mm +++ b/backends/apple/coreai/runtime/coreai_backend.mm @@ -25,36 +25,35 @@ struct Handle final { size_t input_count = 0; size_t output_count = 0; + ETCoreAILoadedModel* model = nil; id session = nil; }; +template +Result +runtime_option(BackendInitContext& context, const char* key, T fallback) { + auto option = context.get_runtime_spec(key); + if (!option.ok() && option.error() == Error::NotFound) { + return fallback; + } + return option; +} + Result cache_directory_option(BackendInitContext& context) { - auto option = context.get_runtime_spec("coreai_cache_dir"); + auto option = + runtime_option(context, "coreai_cache_dir", nullptr); if (!option.ok()) { - if (option.error() == Error::NotFound) { - return cache_directory_path(nullptr); - } return option.error(); } - ET_CHECK_OR_RETURN_ERROR( - option.get() != nullptr, - InvalidArgument, - "Core AI cache directory option must not be null"); return cache_directory_path(option.get()); } Result compute_unit_option(BackendInitContext& context) { - auto option = context.get_runtime_spec("coreai_compute_unit"); + auto option = + runtime_option(context, "coreai_compute_unit", "default"); if (!option.ok()) { - if (option.error() == Error::NotFound) { - return ETCoreAIComputeUnitDefault; - } return option.error(); } - ET_CHECK_OR_RETURN_ERROR( - option.get() != nullptr, - InvalidArgument, - "Core AI compute unit option must not be null"); static constexpr struct { const char* name; ETCoreAIComputeUnit value; @@ -137,6 +136,16 @@ bool is_available() const override { if (!compute_unit.ok()) { return compute_unit.error(); } + auto share = runtime_option( + context, "coreai_share_multifunction_model", false); + if (!share.ok()) { + return share.error(); + } + ET_CHECK_OR_RETURN_ERROR( + !share.get() || manifest->multifunction, + InvalidArgument, + "coreai_share_multifunction_model requires a multi-function export " + "(CoreAIPartitioner(multifunction=True))"); id loader = ETCoreAICreateModelLoader(compute_unit.get()); ET_CHECK_OR_RETURN_ERROR( @@ -145,15 +154,16 @@ bool is_available() const override { if (!cache_directory.ok()) { return cache_directory.error(); } - auto prepared = acquire_model( + auto loaded = acquire_loaded_model( host_bundle.get(), context.get_named_data_map(), cache_directory.get(), - loader); - if (!prepared.ok()) { - return prepared.error(); + loader, + share.get()); + if (!loaded.ok()) { + return loaded.error(); } - auto session = bind_session(prepared.get(), host_bundle->bundle); + auto session = bind_session(loaded.get().model, host_bundle->bundle); if (!session.ok()) { return session.error(); } @@ -164,6 +174,7 @@ bool is_available() const override { "Cannot allocate Core AI handle"); handle->input_count = host_bundle->bundle.inputs.count; handle->output_count = host_bundle->bundle.outputs.count; + handle->model = loaded.get(); handle->session = session.get(); processed->Free(); return static_cast(handle.release()); @@ -178,6 +189,7 @@ Error execute( ET_CHECK_OR_RETURN_ERROR( opaque != nullptr, DelegateInvalidHandle, "Null Core AI handle"); auto& handle = *static_cast(opaque); + std::lock_guard guard([handle.model executionMutex]); return execute_session( handle.session, args, handle.input_count, handle.output_count); } diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.h b/backends/apple/coreai/runtime/coreai_load_coordinator.h index 0cff9087e7d..1290287f738 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.h +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.h @@ -12,6 +12,17 @@ #import "coreai_cache_root.h" #import "coreai_manifest.h" +#include + +// A loaded model and the lock that serializes execution on it; delegates that +// share the model share this object. The SDK does not document concurrent +// execution on one model as safe, and it produced wrong outputs in testing. +@interface ETCoreAILoadedModel : NSObject +@property(nonatomic, readonly) id model; +- (instancetype)init NS_UNAVAILABLE; +- (std::mutex&)executionMutex; +@end + namespace executorch::backends::coreai { struct SelectedHostBundle final { @@ -31,6 +42,17 @@ runtime::Result> acquire_model( const CacheDirectory& cache_directory, id loader); +// With share, returns the model another delegate in this process already holds +// for the same cache directory and cache key, otherwise acquires one and +// records it. Distinct cache directories never share, so each keeps its own +// bookmark. Entries last only while some caller retains the result. +runtime::Result acquire_loaded_model( + const SelectedHostBundle& host_bundle, + const runtime::NamedDataMap* named_data, + const CacheDirectory& cache_directory, + id loader, + bool share); + // Callers must keep affected loads and sessions quiescent across processes. runtime::Error clear_cache_entry( const CacheDirectory& cache_directory, diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm index 9438dc42bc9..e1922126c9c 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.mm +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -14,8 +14,31 @@ #import "coreai_source_bundle.h" #include +#include +#include +#include #include +@interface ETCoreAILoadedModel () +- (instancetype)initWithModel:(id)model; +@end + +@implementation ETCoreAILoadedModel { + std::mutex _executionMutex; +} + +- (instancetype)initWithModel:(id)model { + if ((self = [super init])) { + _model = model; + } + return self; +} + +- (std::mutex&)executionMutex { + return _executionMutex; +} +@end + namespace executorch::backends::coreai { namespace { using runtime::Error; @@ -145,23 +168,26 @@ Error evict_one(id loader, NSData* bookmark) { std::move(selected.get()), [platform copy], [architecture copy]}; } -Result> acquire_model( +namespace { +Result model_cache_key( const SelectedHostBundle& host_bundle, - const runtime::NamedDataMap* named_data, - const CacheDirectory& cache_directory, id loader) { ET_CHECK_OR_RETURN_ERROR( loader != nil, InvalidArgument, "Missing Core AI model loader"); NSString* specialization = [loader.specializationIdentifier copy]; - auto key = cache_variant_key( + return cache_variant_key( host_bundle.bundle, host_bundle.platform, host_bundle.architecture, specialization); - if (!key.ok()) { - return key.error(); - } - NSString* cache_key = key.get(); +} + +Result> acquire_cached_model( + const SelectedHostBundle& host_bundle, + const runtime::NamedDataMap* named_data, + const CacheDirectory& cache_directory, + id loader, + NSString* cache_key) { ET_LOG(Info, "Core AI cache lookup: key=%s", cache_key.UTF8String); auto locked = LockedCacheEntry::acquire(cache_directory, cache_key); if (!locked.ok()) { @@ -226,6 +252,88 @@ Error evict_one(id loader, NSData* bookmark) { cleanup_source(*locked.get()); return model; } +} // namespace + +Result> acquire_model( + const SelectedHostBundle& host_bundle, + const runtime::NamedDataMap* named_data, + const CacheDirectory& cache_directory, + id loader) { + auto key = model_cache_key(host_bundle, loader); + if (!key.ok()) { + return key.error(); + } + return acquire_cached_model( + host_bundle, named_data, cache_directory, loader, key.get()); +} + +Result acquire_loaded_model( + const SelectedHostBundle& host_bundle, + const runtime::NamedDataMap* named_data, + const CacheDirectory& cache_directory, + id loader, + bool share) { + auto key = model_cache_key(host_bundle, loader); + if (!key.ok()) { + return key.error(); + } + static std::mutex mutex; + static std::condition_variable load_finished; + static NSMapTable* shared_models = + [NSMapTable strongToWeakObjectsMapTable]; + static NSMutableSet* loading = [NSMutableSet set]; + NSString* shared_key = nil; + if (share) { + // The open directory's identity, not its path, so aliases of one + // directory share while distinct directories stay isolated. + struct stat directory; + ET_CHECK_OR_RETURN_ERROR( + fstat(cache_directory.descriptor.get(), &directory) == 0, + AccessFailed, + "Cannot identify Core AI cache directory"); + shared_key = [NSString + stringWithFormat:@"%llu:%llu/%@", + static_cast(directory.st_dev), + static_cast(directory.st_ino), + key.get()]; + // Concurrent first loads of one key wait for a single acquisition, which + // may specialize for a long time, while other keys load in parallel. A + // failed acquisition leaves no entry, so a waiter then acquires itself. + std::unique_lock guard(mutex); + while (true) { + ETCoreAILoadedModel* existing = [shared_models objectForKey:shared_key]; + if (existing != nil) { + ET_LOG( + Info, "Core AI shared model reused: key=%s", key.get().UTF8String); + return existing; + } + if (![loading containsObject:shared_key]) { + break; + } + load_finished.wait(guard); + } + [loading addObject:shared_key]; + } + auto model = acquire_cached_model( + host_bundle, named_data, cache_directory, loader, key.get()); + ETCoreAILoadedModel* loaded = model.ok() + ? [[ETCoreAILoadedModel alloc] initWithModel:model.get()] + : nil; + if (share) { + { + std::lock_guard guard(mutex); + [loading removeObject:shared_key]; + if (loaded != nil) { + [shared_models setObject:loaded forKey:shared_key]; + } + } + load_finished.notify_all(); + } + if (!model.ok()) { + return model.error(); + } + return loaded; +} Error clear_cache_entry( const CacheDirectory& cache_directory, diff --git a/backends/apple/coreai/runtime/coreai_manifest.h b/backends/apple/coreai/runtime/coreai_manifest.h index 9627761c0ce..ef6e3158706 100644 --- a/backends/apple/coreai/runtime/coreai_manifest.h +++ b/backends/apple/coreai/runtime/coreai_manifest.h @@ -26,6 +26,8 @@ class Manifest final { NSString* const platform; NSDictionary* const architectures; const NSOperatingSystemVersion minimum_version; + // Exported with other delegates as functions of one Core AI model. + const bool multifunction; Manifest(const Manifest&) = default; Manifest(Manifest&&) = default; @@ -43,7 +45,8 @@ class Manifest final { bool aot_compiled, NSString* platform, NSDictionary* architectures, - NSOperatingSystemVersion minimum_version); + NSOperatingSystemVersion minimum_version, + bool multifunction); friend runtime::Result parse_manifest(NSData* data); }; diff --git a/backends/apple/coreai/runtime/coreai_manifest.mm b/backends/apple/coreai/runtime/coreai_manifest.mm index 007d049b07c..24da4da3523 100644 --- a/backends/apple/coreai/runtime/coreai_manifest.mm +++ b/backends/apple/coreai/runtime/coreai_manifest.mm @@ -23,7 +23,8 @@ bool aot_compiled, NSString* platform, NSDictionary* architectures, - NSOperatingSystemVersion minimum_version) + NSOperatingSystemVersion minimum_version, + bool multifunction) : export_hash(export_hash), function(function), inputs(inputs), @@ -33,7 +34,8 @@ aot_compiled(aot_compiled), platform(platform), architectures(architectures), - minimum_version(minimum_version) {} + minimum_version(minimum_version), + multifunction(multifunction) {} SelectedBundle::SelectedBundle( const Manifest& manifest, @@ -67,9 +69,13 @@ bool names_array(id value) { return true; } +bool is_boolean(id value) { + return value != nil && + CFGetTypeID((__bridge CFTypeRef)value) == CFBooleanGetTypeID(); +} + bool file_size(id value) { - return [value isKindOfClass:NSNumber.class] && - CFGetTypeID((__bridge CFTypeRef)value) != CFBooleanGetTypeID() && + return [value isKindOfClass:NSNumber.class] && !is_boolean(value) && !CFNumberIsFloatType((__bridge CFNumberRef)value) && [value compare:@0] != NSOrderedAscending && [value compare:@(std::numeric_limits::max())] != @@ -183,6 +189,11 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { NSString* function = dict[@"function"]; NSArray* inputs = dict[@"input_names"]; NSArray* outputs = dict[@"output_names"]; + id multifunction = dict[@"multifunction"]; + ET_CHECK_OR_RETURN_ERROR( + multifunction == nil || is_boolean(multifunction), + InvalidProgram, + "Core AI multifunction flag must be a boolean"); NSString* platform = nil; NSDictionary* architectures = nil; ET_CHECK_OR_RETURN_ERROR( @@ -266,7 +277,8 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { aot_compiled, platform, architectures, - minimum_version); + minimum_version, + [multifunction boolValue]); } Result parse_manifest(const void* data, size_t size) { diff --git a/backends/apple/coreai/runtime/test/coreai_delegate_fixture.h b/backends/apple/coreai/runtime/test/coreai_delegate_fixture.h index 7d9a9b64fa1..67f937f4b98 100644 --- a/backends/apple/coreai/runtime/test/coreai_delegate_fixture.h +++ b/backends/apple/coreai/runtime/test/coreai_delegate_fixture.h @@ -31,6 +31,16 @@ class DelegateFixture final { runtime::Error init_with_options( runtime::Span options, runtime::MemoryAllocator* allocator = nullptr); + // Adds coreai_cache_dir for this fixture's storage to the other options. + template + runtime::Error init_with_storage(runtime::BackendOptions& options) { + const runtime::Error error = + options.set_option("coreai_cache_dir", storage.url.path.UTF8String); + if (error != runtime::Error::Ok) { + return error; + } + return init_with_options({options.view().data(), options.view().size()}); + } void reset(); runtime::Error execute(runtime::Span args); diff --git a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm index f39c6a7996e..ef345dd75a1 100644 --- a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm @@ -13,8 +13,13 @@ #include #include #include +#include +#include #include #include +#include +#include +#include #include "coreai_cache.h" #include "coreai_cache_entry_fixture.h" @@ -905,24 +910,32 @@ TensorImpl output( } } -TEST_F(CoreAIDelegateTest, StorageOptionsRejectWrongTypeAndRelativePaths) { +TEST_F(CoreAIDelegateTest, OptionsRejectWrongTypes) { @autoreleasepool { - DelegateFixture fixture(backend, bridge); - { - SCOPED_TRACE("integer cache directory"); - BackendOption wrong_type{"coreai_cache_dir", 1}; - EXPECT_EQ( - fixture.init_with_options({&wrong_type, 1}), Error::InvalidArgument); - } - { - SCOPED_TRACE("relative cache directory"); - BackendOptions<1> relative; - ASSERT_EQ(relative.set_option("coreai_cache_dir", "relative"), Error::Ok); + const BackendOption wrong_types[] = { + {"coreai_cache_dir", 1}, + {"coreai_compute_unit", 1}, + {"coreai_share_multifunction_model", 1}, + }; + for (const auto& option : wrong_types) { + SCOPED_TRACE(option.key); + DelegateFixture fixture(backend, bridge); EXPECT_EQ( - fixture.init_with_options( - {relative.view().data(), relative.view().size()}), - Error::InvalidArgument); + fixture.init_with_options({&option, 1}), Error::InvalidArgument); + EXPECT_EQ(fixture.freed, 0); } + } +} + +TEST_F(CoreAIDelegateTest, StorageOptionsRejectRelativePaths) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + BackendOptions<1> relative; + ASSERT_EQ(relative.set_option("coreai_cache_dir", "relative"), Error::Ok); + EXPECT_EQ( + fixture.init_with_options( + {relative.view().data(), relative.view().size()}), + Error::InvalidArgument); EXPECT_EQ(fixture.freed, 0); } } @@ -935,15 +948,10 @@ TensorImpl output( bridge.state().last_compute_unit = -1; DelegateFixture fixture(backend, bridge); BackendOptions<2> options; - EXPECT_EQ( - options.set_option( - "coreai_cache_dir", fixture.storage.url.path.UTF8String), - Error::Ok); if (name != nullptr) { EXPECT_EQ(options.set_option("coreai_compute_unit", name), Error::Ok); } - return fixture.init_with_options( - {options.view().data(), options.view().size()}); + return fixture.init_with_storage(options); }; { SCOPED_TRACE("absent"); @@ -968,14 +976,239 @@ TensorImpl output( EXPECT_EQ(init_with_compute_unit(name), Error::InvalidArgument); EXPECT_EQ(bridge.state().last_compute_unit.load(), -1); } + } +} + +NSMutableDictionary* multifunction_manifest() { + NSMutableDictionary* manifest = manifest_dict(); + manifest[@"multifunction"] = @YES; + return manifest; +} + +// Uses the fixture's own cache directory unless root names another; sharing is +// scoped to the cache directory. +Error init_sharing( + DelegateFixture& fixture, + std::optional share, + NSString* root = nil) { + BackendOptions<2> options; + if (share.has_value()) { + const Error error = + options.set_option("coreai_share_multifunction_model", *share); + if (error != Error::Ok) { + return error; + } + } + if (root == nil) { + return fixture.init_with_storage(options); + } + const Error error = options.set_option("coreai_cache_dir", root.UTF8String); + if (error != Error::Ok) { + return error; + } + return fixture.init_with_options( + {options.view().data(), options.view().size()}); +} + +TEST_F(CoreAIDelegateTest, ShareOptionRequiresMultifunctionExport) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + EXPECT_EQ(init_sharing(fixture, true), Error::InvalidArgument); + EXPECT_EQ(bridge.state().prepared_models.load(), 0); + EXPECT_EQ(fixture.freed, 0); + DelegateFixture allowed(backend, bridge); + EXPECT_EQ(init_sharing(allowed, false), Error::Ok); + } +} + +TEST_F(CoreAIDelegateTest, ShareOptionSharesOneModelUntilLastDelegate) { + @autoreleasepool { { - SCOPED_TRACE("integer compute unit"); - DelegateFixture fixture(backend, bridge); - BackendOption wrong_type{"coreai_compute_unit", 1}; - EXPECT_EQ( - fixture.init_with_options({&wrong_type, 1}), Error::InvalidArgument); - EXPECT_EQ(fixture.freed, 0); + SCOPED_TRACE("off loads one model per delegate"); + DelegateFixture first(backend, bridge, multifunction_manifest()); + DelegateFixture second(backend, bridge, multifunction_manifest()); + ASSERT_EQ(init_sharing(first, std::nullopt), Error::Ok); + ASSERT_EQ(init_sharing(second, false, first.storage.url.path), Error::Ok); + EXPECT_EQ(bridge.state().prepared_models.load(), 2); } + EXPECT_EQ(bridge.state().prepared_models.load(), 0); + { + SCOPED_TRACE("on shares one model"); + DelegateFixture first(backend, bridge, multifunction_manifest()); + DelegateFixture second(backend, bridge, multifunction_manifest()); + ASSERT_EQ(init_sharing(first, true), Error::Ok); + ASSERT_EQ(init_sharing(second, true, first.storage.url.path), Error::Ok); + EXPECT_EQ(bridge.state().prepared_models.load(), 1); + EXPECT_EQ(bridge.state().sessions.load(), 2); + for (auto* fixture : {&first, &second}) { + FloatTensors tensors; + ASSERT_EQ(fixture->execute({tensors.args, 2}), Error::Ok); + EXPECT_EQ(tensors.output[0], 3); + EXPECT_EQ(tensors.output[1], 7); + } + first.reset(); + EXPECT_EQ(bridge.state().prepared_models.load(), 1); + second.reset(); + EXPECT_EQ(bridge.state().prepared_models.load(), 0); + } + } +} + +TEST_F(CoreAIDelegateTest, SharedModelsAreScopedToCacheDirectory) { + @autoreleasepool { + DelegateFixture first(backend, bridge, multifunction_manifest()); + DelegateFixture other_directory(backend, bridge, multifunction_manifest()); + DelegateFixture alias(backend, bridge, multifunction_manifest()); + ASSERT_EQ(init_sharing(first, true), Error::Ok); + ASSERT_EQ(init_sharing(other_directory, true), Error::Ok); + EXPECT_EQ(bridge.state().prepared_models.load(), 2); + NSURL* link = [alias.storage.url URLByAppendingPathComponent:@"alias"]; + ASSERT_TRUE([NSFileManager.defaultManager + createSymbolicLinkAtURL:link + withDestinationURL:first.storage.url + error:nil]); + ASSERT_EQ(init_sharing(alias, true, link.path), Error::Ok); + EXPECT_EQ(bridge.state().prepared_models.load(), 2); + } +} + +TEST_F(CoreAIDelegateTest, SharedModelSerializesConcurrentExecution) { + @autoreleasepool { + bridge.state().execute_delay_us = 1000; + // The unshared run shows the workload does overlap when nothing serializes + // it, so the shared assertion can fail. + for (const bool share : {false, true}) { + SCOPED_TRACE(share ? "shared" : "unshared"); + bridge.state().max_active_executions = 0; + DelegateFixture first(backend, bridge, multifunction_manifest()); + DelegateFixture second(backend, bridge, multifunction_manifest()); + ASSERT_EQ(init_sharing(first, share), Error::Ok); + ASSERT_EQ(init_sharing(second, share, first.storage.url.path), Error::Ok); + std::atomic failures{0}; + auto run = [&](DelegateFixture* fixture) { + for (int iteration = 0; iteration < 20; ++iteration) { + FloatTensors tensors; + if (fixture->execute({tensors.args, 2}) != Error::Ok || + tensors.output[0] != 3 || tensors.output[1] != 7) { + ++failures; + } + } + }; + std::thread worker(run, &first); + run(&second); + worker.join(); + EXPECT_EQ(failures.load(), 0); + if (share) { + EXPECT_EQ(bridge.state().max_active_executions.load(), 1); + } else { + EXPECT_GT(bridge.state().max_active_executions.load(), 1); + } + } + } +} + +TEST_F(CoreAIDelegateTest, ConcurrentFirstSharedLoadsAcquireOnce) { + @autoreleasepool { + BookmarkLoaderScope fake; + // Widens the window in which a second load could miss the registry. + fake.loader.onSpecialize = ^(NSURL*) { + usleep(50000); + }; + DelegateFixture first(backend, bridge, multifunction_manifest()); + DelegateFixture second(backend, bridge, multifunction_manifest()); + Error second_result = Error::Internal; + NSString* root = first.storage.url.path; + std::thread worker( + [&] { second_result = init_sharing(second, true, root); }); + const Error first_result = init_sharing(first, true); + worker.join(); + EXPECT_EQ(first_result, Error::Ok); + EXPECT_EQ(second_result, Error::Ok); + EXPECT_EQ( + fake.loader->specializations.load() + fake.loader->restores.load(), 1); + } +} + +TEST_F(CoreAIDelegateTest, SharedLoadsOfDifferentModelsRunConcurrently) { + @autoreleasepool { + BookmarkLoaderScope fake; + // Each specialization waits until both are in flight; if one shared load + // blocked the other, neither would ever see two in flight. + auto in_flight = std::make_shared>(0); + auto max_in_flight = std::make_shared>(0); + fake.loader.onSpecialize = ^(NSURL*) { + const int now = ++*in_flight; + int seen = max_in_flight->load(); + while (now > seen && !max_in_flight->compare_exchange_weak(seen, now)) { + } + for (int i = 0; i < 400 && max_in_flight->load() < 2; ++i) { + usleep(5000); + } + --*in_flight; + }; + DelegateFixture first(backend, bridge, multifunction_manifest()); + DelegateFixture second(backend, bridge, multifunction_manifest()); + Error second_result = Error::Internal; + std::thread worker([&] { second_result = init_sharing(second, true); }); + const Error first_result = init_sharing(first, true); + worker.join(); + EXPECT_EQ(first_result, Error::Ok); + EXPECT_EQ(second_result, Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 2); + EXPECT_EQ(max_in_flight->load(), 2); + } +} + +TEST_F(CoreAIDelegateTest, WaiterAcquiresAfterConcurrentSharedLoadFails) { + @autoreleasepool { + BookmarkLoaderScope fake; + auto calls = std::make_shared>(0); + auto first_in_flight = std::make_shared>(false); + __weak BookmarkFakeLoader* loader = fake.loader; + // The first acquisition fails after the second load has queued behind it. + fake.loader.onSpecialize = ^(NSURL*) { + if (++*calls == 1) { + *first_in_flight = true; + usleep(100000); + loader.specializeError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EIO + userInfo:nil]; + } else { + loader.specializeError = nil; + } + }; + DelegateFixture first(backend, bridge, multifunction_manifest()); + DelegateFixture second(backend, bridge, multifunction_manifest()); + NSString* root = first.storage.url.path; + Error first_result = Error::Ok; + std::thread worker([&] { first_result = init_sharing(first, true); }); + while (!*first_in_flight) { + std::this_thread::yield(); + } + const Error second_result = init_sharing(second, true, root); + worker.join(); + EXPECT_NE(first_result, Error::Ok); + EXPECT_EQ(second_result, Error::Ok); + EXPECT_EQ(calls->load(), 2); + } +} + +TEST_F(CoreAIDelegateTest, FailedSharedAcquisitionLeavesNoEntry) { + @autoreleasepool { + BookmarkLoaderScope fake; + fake.loader.specializeError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EIO + userInfo:nil]; + DelegateFixture failed(backend, bridge, multifunction_manifest()); + EXPECT_NE(init_sharing(failed, true), Error::Ok); + fake.loader.specializeError = nil; + DelegateFixture first(backend, bridge, multifunction_manifest()); + DelegateFixture second(backend, bridge, multifunction_manifest()); + NSString* root = failed.storage.url.path; + ASSERT_EQ(init_sharing(first, true, root), Error::Ok); + ASSERT_EQ(init_sharing(second, true, root), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 2); + EXPECT_EQ(fake.loader->restores.load(), 0); } } diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.h b/backends/apple/coreai/runtime/test/coreai_fake_loader.h index 72d101b1855..b7e3ddff83b 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.h +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include #import "ETCoreAIBridge.h" @@ -30,11 +31,14 @@ struct FakeBridgeState { bool multiple_io = false; bool retain_binding_completion = false; bool retain_execution_completion = false; + useconds_t execute_delay_us = 0; NSString* last_bundle = nil; dispatch_semaphore_t input_ready = nil; dispatch_semaphore_t allow_input_read = nil; std::atomic architecture_queries{0}; std::atomic executions{0}; + std::atomic active_executions{0}; + std::atomic max_active_executions{0}; std::atomic output_copies{0}; std::atomic output_owners{0}; std::atomic last_output_destination{nullptr}; diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm index 4e3e69b21cb..7ee37b238f3 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm @@ -133,6 +133,19 @@ - (void)executeInputs:(NSArray*)inputs (void (^)(NSArray*, NSError*))completion { CoreAIFakeBridgeContext* context = self.context; ++context->state.executions; + const int active = ++context->state.active_executions; + int peak = context->state.max_active_executions.load(); + while (active > peak && + !context->state.max_active_executions.compare_exchange_weak( + peak, active)) { + } + // Ends the overlap window before the caller can observe completion. + void (^reply)(NSArray*, NSError*) = completion; + completion = ^(NSArray* outputs, NSError* error) { + --context->state.active_executions; + reply(outputs, error); + }; + const useconds_t delay = context->state.execute_delay_us; const bool should_fail = context->state.fail_execute; const bool malformed = context->state.bad_output; dispatch_semaphore_t ready = context->state.input_ready; @@ -141,6 +154,9 @@ - (void)executeInputs:(NSArray*)inputs context.retainedExecutionCompletion = completion; } run_async(context, ^{ + if (delay > 0) { + usleep(delay); + } ETCoreAIInputTensor* input = inputs[0]; context->state.last_input_bytes = input.bytes; context->state.last_input_byte_count = input.byteCount; diff --git a/backends/apple/coreai/runtime/test/coreai_manifest_test.mm b/backends/apple/coreai/runtime/test/coreai_manifest_test.mm index 2e0fa5393e7..259dc91ba53 100644 --- a/backends/apple/coreai/runtime/test/coreai_manifest_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_manifest_test.mm @@ -30,6 +30,25 @@ EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::NotSupported); } +TEST(CoreAIManifestTest, ParsesOptionalMultifunctionFlag) { + auto absent = parse_manifest(encode(manifest_dict())); + ASSERT_TRUE(absent.ok()); + EXPECT_FALSE(absent->multifunction); + for (NSNumber* value in @[ @YES, @NO ]) { + auto dict = manifest_dict(); + dict[@"multifunction"] = value; + auto parsed = parse_manifest(encode(dict)); + ASSERT_TRUE(parsed.ok()); + EXPECT_EQ(parsed->multifunction, value.boolValue); + } + for (id value in @[ @1, @"true", NSNull.null ]) { + SCOPED_TRACE([value description].UTF8String); + auto dict = manifest_dict(); + dict[@"multifunction"] = value; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + TEST(CoreAIManifestTest, RejectsUnsafeSourcePaths) { for (NSString* path in @[ @"../model.aimodel", diff --git a/backends/apple/coreai/runtime/test/coreai_multithread_test.cpp b/backends/apple/coreai/runtime/test/coreai_multithread_test.cpp index b08f5dd5310..204087d814f 100644 --- a/backends/apple/coreai/runtime/test/coreai_multithread_test.cpp +++ b/backends/apple/coreai/runtime/test/coreai_multithread_test.cpp @@ -29,6 +29,7 @@ #include #include #include +#include #include namespace et = executorch::runtime; @@ -119,6 +120,7 @@ struct JoinWorkers { void round( const std::string& pte, const std::string& cache, + bool share, const std::vector& sample, const std::vector& expected, size_t threads, @@ -142,11 +144,14 @@ void round( return; } Input input(sample); - et::BackendOptions<1> options; + et::BackendOptions<2> options; et::LoadBackendOptionsMap option_map; check( options.set_option("coreai_cache_dir", cache.c_str()), "set option"); + check( + options.set_option("coreai_share_multifunction_model", share), + "set share option"); check( option_map.set_options("CoreAIBackend", options.view()), "set backend options"); @@ -234,15 +239,23 @@ int main(int argc, char** argv) { std::cout << threads << " threads x " << iterations << " iterations per round\n" << "Cache: " << cache << std::endl; - for (const char* phase : {"cold", "warm"}) { - std::cout << phase << " round" << std::endl; - round( - (fixture / "model.pte").string(), - cache, - sample, - expected, - threads, - iterations); + // The shared variant puts every thread's delegates on one Core AI model, + // so its executions must be serialized to stay correct. + for (const auto& [file, share] : + {std::pair{"model.pte", false}, + std::pair{"model_multifunction.pte", true}}) { + for (const char* phase : {"cold", "warm"}) { + std::cout << file << (share ? " (shared model) " : " ") << phase + << " round" << std::endl; + round( + (fixture / file).string(), + cache, + share, + sample, + expected, + threads, + iterations); + } } check( executorch::backends::coreai::clear_cache(cache.c_str()), diff --git a/backends/apple/coreai/runtime/test/export_stress_fixture.py b/backends/apple/coreai/runtime/test/export_stress_fixture.py index 441ac9663e8..1583223f5b1 100644 --- a/backends/apple/coreai/runtime/test/export_stress_fixture.py +++ b/backends/apple/coreai/runtime/test/export_stress_fixture.py @@ -60,19 +60,11 @@ def __call__(self, graph_module): raise RuntimeError("Missing residual boundary") -def export_fixture(output_dir): - # Distinct weights keep SDK artifacts separate from the playground model. - torch.manual_seed(20261029) - model = ResidualMLP().eval().half().requires_grad_(False) - sample = torch.randn(1, 32, dtype=torch.float16) - with torch.no_grad(): - expected = model(sample) - assert expected.dtype == torch.float16 and torch.isfinite(expected).all() - exported = torch.export.export(model, (sample,)) +def lower(exported, multifunction): lowered = to_edge_transform_and_lower( exported, transform_passes=[*get_default_passes(), TagResidualBoundary()], - partitioner=[CoreAIPartitioner()], + partitioner=[CoreAIPartitioner(multifunction=multifunction)], compile_config=get_default_compile_config(), ) graph = lowered.exported_program().graph_module @@ -88,9 +80,26 @@ def export_fixture(output_dir): and (node.op != "call_function" or node.target not in allowed) ] assert len(fallback) == 1 and fallback[0].target == exir_ops.edge.aten.relu.default + return lowered + + +def export_fixture(output_dir): + # Distinct weights keep SDK artifacts separate from the playground model. + torch.manual_seed(20261029) + model = ResidualMLP().eval().half().requires_grad_(False) + sample = torch.randn(1, 32, dtype=torch.float16) + with torch.no_grad(): + expected = model(sample) + assert expected.dtype == torch.float16 and torch.isfinite(expected).all() + exported = torch.export.export(model, (sample,)) output_dir.mkdir(parents=True, exist_ok=True) - with (output_dir / "model.pte").open("wb") as stream: - lowered.to_executorch().write_to_file(stream) + # The two partitions become two functions of one Core AI model. + for name, multifunction in ( + ("model.pte", False), + ("model_multifunction.pte", True), + ): + with (output_dir / name).open("wb") as stream: + lower(exported, multifunction).to_executorch().write_to_file(stream) (output_dir / "fixture.json").write_text( json.dumps( { diff --git a/backends/apple/coreai/test/test_io_compat.py b/backends/apple/coreai/test/test_io_compat.py index e22e43094b5..058a79712ef 100644 --- a/backends/apple/coreai/test/test_io_compat.py +++ b/backends/apple/coreai/test/test_io_compat.py @@ -22,6 +22,7 @@ get_default_compile_config, get_default_passes, ) +from executorch.backends.apple.coreai.compiler.constants import MAIN_ENTRYPOINT from executorch.backends.apple.coreai.compiler.io_compat import ( assert_io_compatible, io_mismatches, @@ -54,7 +55,7 @@ def forward(self, z, a): inputs = ["input_0", "input_1"] outputs = ["output_0", "output_1"] program = _convert_to_aiprogram( - edge, input_names=inputs[::-1], output_names=outputs + [(edge, MAIN_ENTRYPOINT, inputs[::-1], outputs)] ) with self.assertRaisesRegex(ValueError, "input names mismatch"): assert_io_compatible( diff --git a/backends/apple/coreai/test/test_multifunction.py b/backends/apple/coreai/test/test_multifunction.py new file mode 100644 index 00000000000..28f27dcc997 --- /dev/null +++ b/backends/apple/coreai/test/test_multifunction.py @@ -0,0 +1,292 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import json +import subprocess +import tempfile +import unittest +from unittest import mock + +import torch +import torch.nn as nn +from executorch.backends.apple.coreai import ( + AOTCompileConfig, + get_default_compile_config, + get_default_passes, +) +from executorch.backends.apple.coreai.compiler.preprocess import COMPILE_SPEC_KEYS +from executorch.backends.apple.coreai.partition.partitioner import ( + CoreAIPartitioner, + do_not_delegate, +) +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.lowered_backend_module import get_lowered_backend_modules +from executorch.exir.pass_base import PassResult + +_MOCK_BUILD = "executorch.backends.apple.coreai.compiler.preprocess._run_coreai_build" +_WIDTH = 256 + + +def _coreai_build_available() -> bool: + try: + return ( + subprocess.run( + ["xcrun", "--find", "coreai-build"], capture_output=True + ).returncode + == 0 + ) + except FileNotFoundError: + return False + + +class _Stack(nn.Module): + def __init__(self): + super().__init__() + self.up = nn.Linear(_WIDTH, _WIDTH) + self.down = nn.Linear(_WIDTH, _WIDTH) + + def forward(self, x): + return self.down(torch.relu(self.up(x)) * 2.0) + + +class _BreakAtMul: + def __call__(self, gm): + for node in gm.graph.nodes: + if node.op == "call_function" and "mul" in str(node.target): + do_not_delegate(node) + return PassResult(gm, True) + + +def _inputs(): + return {"forward": (torch.randn(1, _WIDTH),), "prefill": (torch.randn(8, _WIDTH),)} + + +def _lower(partitioner, *, graph_break=False, model=None, inputs=None): + torch.manual_seed(0) + model = model or _Stack().eval() + inputs = inputs or _inputs() + passes = list(get_default_passes()) + ([_BreakAtMul()] if graph_break else []) + return to_edge_transform_and_lower( + {name: torch.export.export(model, args) for name, args in inputs.items()}, + transform_passes=passes, + partitioner=partitioner, + compile_config=get_default_compile_config(), + ) + + +def _manifests(lowered, method): + graph = lowered.exported_program(method).graph_module + return [ + (json.loads(bytes(lbm.processed_bytes)), lbm) + for lbm in get_lowered_backend_modules(graph) + ] + + +def _bundle_bytes(lbm): + return sum(len(blob) for blob in lbm.named_data_store_output.buffers) + + +class MultifunctionPartitionerTest(unittest.TestCase): + def test_option_adds_compile_spec(self): + keys = [ + spec.key + for spec in CoreAIPartitioner( + multifunction=True + ).delegation_spec.compile_specs + ] + self.assertIn(COMPILE_SPEC_KEYS.MULTIFUNCTION.value, keys) + default_keys = [ + spec.key for spec in CoreAIPartitioner().delegation_spec.compile_specs + ] + self.assertNotIn(COMPILE_SPEC_KEYS.MULTIFUNCTION.value, default_keys) + + def test_rejects_input_enumerations(self): + with self.assertRaisesRegex(ValueError, "input_enumerations"): + CoreAIPartitioner(multifunction=True, input_enumerations=[{0: [1, 8]}]) + + +class MultifunctionPreprocessTest(unittest.TestCase): + def test_methods_share_one_bundle(self): + lowered = _lower([CoreAIPartitioner(multifunction=True)]) + ((forward, forward_lbm),) = _manifests(lowered, "forward") + ((prefill, prefill_lbm),) = _manifests(lowered, "prefill") + + self.assertEqual(forward["function"], "forward_0") + self.assertEqual(prefill["function"], "prefill_0") + self.assertEqual(forward["packaging"], "aimodel") + self.assertIs(forward["multifunction"], True) + self.assertIs(prefill["multifunction"], True) + for key in ("export_hash", "path", "files", "bundle_digests"): + self.assertEqual(forward[key], prefill[key], key) + self.assertEqual(forward["input_names"], ["input_0"]) + self.assertEqual(prefill["output_names"], ["output_0"]) + self.assertEqual( + set(forward_lbm.named_data_store_output.pte_data), + set(prefill_lbm.named_data_store_output.pte_data), + ) + + # Shared weights are stored once: two functions cost about one bundle. + ((single, single_lbm),) = _manifests( + _lower( + [CoreAIPartitioner()], + inputs={"forward": _inputs()["forward"]}, + ), + "forward", + ) + self.assertEqual(single["function"], "main") + self.assertLess(_bundle_bytes(forward_lbm), 1.1 * _bundle_bytes(single_lbm)) + self.assertGreater(_bundle_bytes(single_lbm), 2 * _WIDTH * _WIDTH * 4) + + def test_graph_break_partitions_become_functions(self): + lowered = _lower([CoreAIPartitioner(multifunction=True)], graph_break=True) + functions = { + method: [m["function"] for m, _ in _manifests(lowered, method)] + for method in ("forward", "prefill") + } + self.assertEqual( + functions, + { + "forward": ["forward_0", "forward_1"], + "prefill": ["prefill_0", "prefill_1"], + }, + ) + hashes = { + m["export_hash"] + for method in ("forward", "prefill") + for m, _ in _manifests(lowered, method) + } + self.assertEqual(len(hashes), 1) + + def test_disabled_keeps_one_model_per_delegate(self): + lowered = _lower([CoreAIPartitioner()]) + ((forward, _),) = _manifests(lowered, "forward") + ((prefill, _),) = _manifests(lowered, "prefill") + self.assertEqual((forward["function"], prefill["function"]), ("main", "main")) + self.assertNotIn("multifunction", forward) + self.assertNotEqual(forward["export_hash"], prefill["export_hash"]) + + def test_mismatched_compile_specs_raise(self): + partitioners = { + "forward": [CoreAIPartitioner(multifunction=True)], + "prefill": [ + CoreAIPartitioner(multifunction=True, min_deployment_version="27.0") + ], + } + with self.assertRaisesRegex(ValueError, "identical Core AI compile specs"): + _lower(partitioners) + + def test_option_on_only_some_methods_raises(self): + partitioners = { + "forward": [CoreAIPartitioner(multifunction=True)], + "prefill": [CoreAIPartitioner()], + } + with self.assertRaisesRegex(ValueError, r"\['prefill'\]"): + _lower(partitioners) + + @mock.patch(_MOCK_BUILD) + def test_aot_compiled_methods_share_archs(self, build): + def fake_build(aimodel_path, out_dir, opts): + for arch in opts["architectures"]: + bundle = out_dir / f"model.{arch}.aimodelc" + bundle.mkdir(parents=True) + (bundle / "model.mil").write_bytes(b"fake-compiled") + + build.side_effect = fake_build + aot = AOTCompileConfig(platform="iOS", architectures=["h15g", "h16"]) + lowered = _lower( + [CoreAIPartitioner(multifunction=True, aot_compile_config=aot)] + ) + ((forward, _),) = _manifests(lowered, "forward") + ((prefill, _),) = _manifests(lowered, "prefill") + self.assertEqual(build.call_count, 1) + self.assertEqual(forward["packaging"], "aimodelc") + self.assertEqual(sorted(forward["archs"]), ["h15g", "h16"]) + self.assertEqual(forward["archs"], prefill["archs"]) + self.assertEqual( + (forward["function"], prefill["function"]), ("forward_0", "prefill_0") + ) + + +def _runtime_available(): + try: + from executorch.extension.pybindings.portable_lib import _is_available + except ImportError: + return False + return _is_available("CoreAIBackend") + + +@unittest.skipUnless( + _coreai_build_available(), + "requires macOS with the Metal Toolchain (xcrun coreai-build)", +) +class MultifunctionAOTCompileTest(unittest.TestCase): + def test_real_build_shares_compiled_bundles(self): + aot = AOTCompileConfig(platform="macOS", architectures=["h16c"]) + lowered = _lower( + [CoreAIPartitioner(multifunction=True, aot_compile_config=aot)], + graph_break=True, + ) + manifests = [ + m + for method in ("forward", "prefill") + for m, _ in _manifests(lowered, method) + ] + self.assertEqual( + [m["function"] for m in manifests], + ["forward_0", "forward_1", "prefill_0", "prefill_1"], + ) + self.assertEqual(len({json.dumps(m["archs"]) for m in manifests}), 1) + self.assertEqual(list(manifests[0]["archs"]), ["h16c"]) + + +@unittest.skipUnless(_runtime_available(), "requires the Core AI Python runtime") +class MultifunctionRuntimeTest(unittest.TestCase): + def _run_all(self, lowered, inputs): + from executorch.extension.pybindings.portable_lib import ( + _load_program_from_buffer, + Verification, + ) + + buffer = lowered.to_executorch().buffer + outputs = {} + with tempfile.TemporaryDirectory(prefix="coreai-multifunction-") as cache: + program = _load_program_from_buffer( + buffer, + program_verification=Verification.Minimal, + backend_options={"CoreAIBackend": {"coreai_cache_dir": cache}}, + ) + methods = {name: program.load_method(name) for name in inputs} + for name, args in inputs.items(): + outputs[name] = methods[name].call(args, clone_outputs=True)[0] + del methods, program + return outputs + + def _check(self, graph_break): + torch.manual_seed(0) + model = _Stack().eval() + inputs = _inputs() + lowered = _lower( + [CoreAIPartitioner(multifunction=True)], + graph_break=graph_break, + model=model, + inputs=inputs, + ) + outputs = self._run_all(lowered, inputs) + with torch.no_grad(): + for name, args in inputs.items(): + torch.testing.assert_close( + outputs[name], model(*args), atol=1e-2, rtol=1e-2, msg=name + ) + + def test_each_method_runs_its_function(self): + self._check(graph_break=False) + + def test_graph_break_methods_run(self): + self._check(graph_break=True) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/apple/coreai/test/test_preprocess.py b/backends/apple/coreai/test/test_preprocess.py index e81eb75ad58..32fb21ba229 100644 --- a/backends/apple/coreai/test/test_preprocess.py +++ b/backends/apple/coreai/test/test_preprocess.py @@ -215,7 +215,9 @@ def test_same_program_hash_with_different_metadata_merges(self): (bundle / "metadata.json").write_text( json.dumps({"creationDate": timestamp}) ) - result = _deliver(root, AssetPackaging.AIMODEL, {"path": bundle.name}) + result = _deliver( + root, AssetPackaging.AIMODEL, {"path": bundle.name}, [{}] + )[0] manifest = json.loads(result.processed_bytes) manifests.append(manifest) self.assertEqual( @@ -240,7 +242,9 @@ def test_aot_namespace_covers_every_architecture(self): (bundle / "model.mil").write_bytes(b"original") for payload in (b"original", b"modified"): (root / archs["h16"] / "model.mil").write_bytes(payload) - result = _deliver(root, AssetPackaging.AIMODELC, {"archs": archs}) + result = _deliver( + root, AssetPackaging.AIMODELC, {"archs": archs}, [{}] + )[0] manifest = json.loads(result.processed_bytes) manifests.append(manifest) for arch, name in archs.items():