From 884d214b16401c28a312c9b905c3226d81f3ac1b Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Mon, 2 Nov 2026 16:00:40 -0800 Subject: [PATCH 1/2] [CoreAI] Rename manifest packaging to aimodel / aimodelc The packaging values "inline" and "aot_compiled_inline" described where the bundle is stored rather than what it is. Name them after the embedded format instead: "aimodel" for the portable bundle and "aimodelc" for AOT-compiled per-architecture bundles. The runtime's kind names in logs and the cache identity ("source" / "aot") become "aimodel" / "aimodelc" too, so "source" only refers to the bundle copy materialized for the SDK. This is a format change: PTEs exported with the old names are rejected and must be re-exported, and existing cache entries miss once and are re-specialized. --- backends/apple/coreai/README.md | 15 ++++++++------- backends/apple/coreai/compiler/preprocess.py | 10 +++++----- backends/apple/coreai/runtime/coreai_cache_key.mm | 2 +- .../coreai/runtime/coreai_load_coordinator.mm | 2 +- backends/apple/coreai/runtime/coreai_manifest.mm | 8 ++++---- .../coreai/runtime/test/coreai_delegate_test.mm | 2 +- .../coreai/runtime/test/coreai_manifest_fixture.h | 4 ++-- .../runtime/test/coreai_source_integrity_test.mm | 2 +- backends/apple/coreai/test/test_preprocess.py | 14 ++++++-------- 9 files changed, 29 insertions(+), 30 deletions(-) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 36c4bb5c873..ec0fd4856ca 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -26,9 +26,10 @@ 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 a `function` name, 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 @@ -128,9 +129,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 +192,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 diff --git a/backends/apple/coreai/compiler/preprocess.py b/backends/apple/coreai/compiler/preprocess.py index 0bc064f2fbb..1d6b455bd7e 100644 --- a/backends/apple/coreai/compiler/preprocess.py +++ b/backends/apple/coreai/compiler/preprocess.py @@ -78,9 +78,9 @@ class COMPILE_SPEC_KEYS(Enum): class AssetPackaging(str, Enum): # Portable .aimodel embedded in the .pte via NamedDataStore. - INLINE = "inline" + AIMODEL = "aimodel" # AOT-compiled per-arch .aimodelc bundles embedded in the .pte. - AOT_COMPILED_INLINE = "aot_compiled_inline" + AIMODELC = "aimodelc" @dataclass(frozen=True) @@ -213,7 +213,7 @@ def _deliver( **manifest_extra, **metadata, } - if packaging == AssetPackaging.INLINE: + if packaging == AssetPackaging.AIMODEL: manifest["path"] = f"{export_hash}/{manifest['path']}" else: manifest["archs"] = { @@ -520,7 +520,7 @@ def _preprocess_portable( _save_asset(program, Path(tmp) / "model.aimodel", min_os) return _deliver( Path(tmp), - AssetPackaging.INLINE, + AssetPackaging.AIMODEL, { **bindings, # relative path the runtime resolves against its base @@ -538,7 +538,7 @@ def _preprocess_aot_compiled( out = _compile_aot(program, opts, Path(tmp)) return _deliver( out, - AssetPackaging.AOT_COMPILED_INLINE, + AssetPackaging.AIMODELC, { **bindings, "platform": opts["platform"], diff --git a/backends/apple/coreai/runtime/coreai_cache_key.mm b/backends/apple/coreai/runtime/coreai_cache_key.mm index 8f63ce7fd25..784fd803f5b 100644 --- a/backends/apple/coreai/runtime/coreai_cache_key.mm +++ b/backends/apple/coreai/runtime/coreai_cache_key.mm @@ -80,7 +80,7 @@ bool valid_device(NSString* device) { bool aot_compiled) { // Bumping the leading version makes every existing cache entry unreachable. auto encoded = encode_json( - @[ @"1", digest, aot_compiled ? @"aot" : @"source", bundle_name ]); + @[ @"1", digest, aot_compiled ? @"aimodelc" : @"aimodel", bundle_name ]); if (!encoded.ok()) { return encoded.error(); } diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm index 3cfe7b42a19..9438dc42bc9 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.mm +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -206,7 +206,7 @@ Error evict_one(id loader, NSData* bookmark) { ET_LOG( Info, "Core AI specialization started: kind=%s, key=%s", - host_bundle.bundle.aot_compiled ? "aot" : "source", + host_bundle.bundle.aot_compiled ? "aimodelc" : "aimodel", cache_key.UTF8String); [loader specializeModelAtURL:url completion:completion]; }); diff --git a/backends/apple/coreai/runtime/coreai_manifest.mm b/backends/apple/coreai/runtime/coreai_manifest.mm index af139c8f09d..007d049b07c 100644 --- a/backends/apple/coreai/runtime/coreai_manifest.mm +++ b/backends/apple/coreai/runtime/coreai_manifest.mm @@ -174,11 +174,11 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { InvalidProgram, "Missing Core AI packaging"); ET_CHECK_OR_RETURN_ERROR( - [packaging isEqualToString:@"inline"] || - [packaging isEqualToString:@"aot_compiled_inline"], + [packaging isEqualToString:@"aimodel"] || + [packaging isEqualToString:@"aimodelc"], NotSupported, - "Unsupported Core AI packaging; expected .aimodel or AOT .aimodelc"); - const bool aot_compiled = [packaging hasPrefix:@"aot_compiled_"]; + "Unsupported Core AI packaging; expected aimodel or aimodelc"); + const bool aot_compiled = [packaging isEqualToString:@"aimodelc"]; NSString* export_hash = dict[@"export_hash"]; NSString* function = dict[@"function"]; NSArray* inputs = dict[@"input_names"]; diff --git a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm index 84530ac9c92..f39c6a7996e 100644 --- a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm @@ -153,7 +153,7 @@ void expect_preflight_rejection( TEST_F(CoreAIDelegateTest, PreflightRejectsFutureDeploymentForSourceAndAot) { @autoreleasepool { for (bool aot : {false, true}) { - SCOPED_TRACE(aot ? "AOT" : "source"); + SCOPED_TRACE(aot ? "aimodelc" : "aimodel"); auto dict = aot ? aot_manifest_dict() : manifest_dict(); dict[@"min_deployment_version"] = @"999999.0"; expect_preflight_rejection( diff --git a/backends/apple/coreai/runtime/test/coreai_manifest_fixture.h b/backends/apple/coreai/runtime/test/coreai_manifest_fixture.h index bf9af55c0bb..8cd2e99a276 100644 --- a/backends/apple/coreai/runtime/test/coreai_manifest_fixture.h +++ b/backends/apple/coreai/runtime/test/coreai_manifest_fixture.h @@ -22,7 +22,7 @@ inline NSString* fixture_identity(char digit) { inline NSMutableDictionary* manifest_dict() { return [@{ - @"packaging" : @"inline", + @"packaging" : @"aimodel", @"export_hash" : @"ab", @"path" : @"ab/model.aimodel", @"function" : @"main", @@ -37,7 +37,7 @@ inline NSMutableDictionary* manifest_dict() { inline NSMutableDictionary* aot_manifest_dict() { auto dict = manifest_dict(); [dict removeObjectForKey:@"path"]; - dict[@"packaging"] = @"aot_compiled_inline"; + dict[@"packaging"] = @"aimodelc"; dict[@"platform"] = @"macOS"; dict[@"archs"] = @{ @"arch_b" : @"ab/model.arch_b.aimodelc", diff --git a/backends/apple/coreai/runtime/test/coreai_source_integrity_test.mm b/backends/apple/coreai/runtime/test/coreai_source_integrity_test.mm index 264042988e7..02596987687 100644 --- a/backends/apple/coreai/runtime/test/coreai_source_integrity_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_source_integrity_test.mm @@ -58,7 +58,7 @@ TEST_F(CoreAISourceTest, RejectsMalformedStoredBundlesWithoutReplacement) { for (bool aot : {false, true}) { - SCOPED_TRACE(aot ? "aot" : "source"); + SCOPED_TRACE(aot ? "aimodelc" : "aimodel"); for (NSString* damage in @[ @"missing-file", @"extra-file", diff --git a/backends/apple/coreai/test/test_preprocess.py b/backends/apple/coreai/test/test_preprocess.py index 2d8b0d39fd5..e81eb75ad58 100644 --- a/backends/apple/coreai/test/test_preprocess.py +++ b/backends/apple/coreai/test/test_preprocess.py @@ -142,7 +142,7 @@ class PortablePreprocessTest(unittest.TestCase): def test_inline_embeds_files_in_nds(self): result = CoreAIBackend.preprocess(_edge_program(), []) manifest = json.loads(result.processed_bytes) - self.assertEqual(manifest["packaging"], "inline") + self.assertEqual(manifest["packaging"], "aimodel") self.assertTrue(manifest["files"], "expected at least one asset file") out = result.data_store_output self.assertIsNotNone(out) @@ -215,7 +215,7 @@ def test_same_program_hash_with_different_metadata_merges(self): (bundle / "metadata.json").write_text( json.dumps({"creationDate": timestamp}) ) - result = _deliver(root, AssetPackaging.INLINE, {"path": bundle.name}) + result = _deliver(root, AssetPackaging.AIMODEL, {"path": bundle.name}) manifest = json.loads(result.processed_bytes) manifests.append(manifest) self.assertEqual( @@ -240,9 +240,7 @@ 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.AOT_COMPILED_INLINE, {"archs": archs} - ) + result = _deliver(root, AssetPackaging.AIMODELC, {"archs": archs}) manifest = json.loads(result.processed_bytes) manifests.append(manifest) for arch, name in archs.items(): @@ -378,7 +376,7 @@ def test_inline_embeds_compiled_bundles(self, _build): [_aot_spec({"platform": "iOS", "architectures": ["h15g", "h16"]})], ) manifest = json.loads(result.processed_bytes) - self.assertEqual(manifest["packaging"], "aot_compiled_inline") + self.assertEqual(manifest["packaging"], "aimodelc") self.assertEqual(manifest["platform"], "iOS") self.assertEqual(sorted(manifest["archs"]), ["h15g", "h16"]) # Compiled bundle contents are embedded (files under a *.aimodelc dir). @@ -400,7 +398,7 @@ def test_defaults_to_all_architectures(self, _build): # Empty architectures => coreai-build decides (our fake yields one). result = CoreAIBackend.preprocess(_edge_program(), [_aot_spec({})]) manifest = json.loads(result.processed_bytes) - self.assertEqual(manifest["packaging"], "aot_compiled_inline") + self.assertEqual(manifest["packaging"], "aimodelc") self.assertEqual(manifest["platform"], "macOS") # default self.assertTrue(manifest["archs"]) @@ -448,7 +446,7 @@ def test_aot_inline_embeds_aimodelc(self): CoreAIPartitioner(aot_compile_config=AOTCompileConfig(platform="macOS")) ) lbm, manifest = _lowered_manifest(lowered) - self.assertEqual(manifest["packaging"], "aot_compiled_inline") + self.assertEqual(manifest["packaging"], "aimodelc") self.assertEqual(manifest["platform"], "macOS") self.assertGreaterEqual(len(manifest["archs"]), 1) self.assertTrue(any(".aimodelc/" in f for f in manifest["files"])) From 0508853d8e43b722cd3518d5ee848ebce1a845d1 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Wed, 4 Nov 2026 11:35:21 -0800 Subject: [PATCH 2/2] [CoreAI] Add multi-function export and shared runtime loading (#24476) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an opt-in way to export several ExecuTorch methods (and graph-broken partitions) as **one Core AI model with multiple functions**, so methods that share weights ship and cache them once, plus an opt-in runtime option so those delegates also **share one loaded model in memory**. Stacked on #24474. ## Usage ```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(), ) ``` Each Core AI partition becomes one function named `_`: here `forward_0` and `prefill_0`, and a graph break adds `forward_1`, `prefill_1`, and so on. Graph breaks are allowed; partitions running on CPU stay in ExecuTorch as usual. To also share one loaded model at runtime, set the backend option when loading: ```cpp BackendOptions<2> options; options.set_option("coreai_share_multifunction_model", true); LoadBackendOptionsMap map; map.set_options("CoreAIBackend", options.view()); module.load(map); ``` ## How it works * **Export:** `CoreAIBackend.preprocess_multimethod` receives every method's partitions at once. * It stages all of them in one `TorchConverter` (one entrypoint per partition), calls `to_coreai()` once, and saves (or AOT-compiles with `coreai-build`) once. * The bundle is embedded in the `NamedDataStore` once, under a single `export_hash`. * Each delegate still gets its own manifest naming its `function`, `input_names` and `output_names`; all manifests point at the same bundle. * **Cache:** function bindings aren't part of the cache key, so all the delegates share one cache entry: the first load specializes and the others restore from the same bookmark. * **Manifest:** multi-function manifests carry `"multifunction": true` (optional boolean; absent means false). * **Runtime sharing (`coreai_share_multifunction_model`, bool, default `false`):** * Without it, each delegate restores its own model with its own copy of the weights, as before. * With it, the first delegate acquires the model and the others reuse it from a process-wide table until the last delegate is destroyed. * Loads are tracked per key. Delegates that start loading the same model at once wait for that one load, while loads of other models or cache directories proceed in parallel. If the load fails, it leaves no entry and one of the waiters retries. * Sharing requires the same cache directory (identified by the open directory's device and inode, so a symlink counts as the same directory) and the same cache key, which includes the compute-unit specialization. Different cache directories or compute units load separate models. * Delegates sharing a model run one at a time through a per-model lock. The SDK doesn't document concurrent execution on one model as safe, and without the lock `coreai_multithread_test` produced wrong outputs. * Setting the option on a delegate that isn't from a multi-function export fails init with `InvalidArgument`. * **Option off:** lowering is unchanged. `preprocess` and the multi-method path now share the conversion and delivery helpers, and the existing tests still pass. * **Errors:** export fails with a clear message if partitions have different compile specs (including the option set on only some methods), or if `input_enumerations` is combined with the option. Enumerated shapes rename functions internally and need their own design. ## Results Two-method model (8 × `Linear(1024, 1024)`, fp32; `forward` with 1 row, `prefill` with 16 rows), run on an M4 Max with a standalone C++ runner: | Variant | PTE size | Cold load | Cache entries | Outputs | |---|---|---|---|---| | `multifunction=True` | 33.6 MB | 1 specialization + 1 bookmark restore | 1 | match eager | | Default (one model per delegate) | 67.2 MB | 2 specializations | 2 | match eager | | `multifunction=True` + graph break (4 functions) | 33.6 MB | 1 specialization + 3 bookmark restores | 1 | match eager | Shared weights are stored once: in a direct converter check, adding a second function to an 8.4 MB model added 2.9 KB. Runtime sharing, measured as resident memory after one run per delegate, for K delegates of the same 33.6 MB model: | K delegates | Option off | `coreai_share_multifunction_model=true` | |---|---|---| | 1 | +91 MB | +91 MB | | 2 | +157 MB | +92 MB | | 4 | +288 MB | +92 MB | | 8 | +551 MB | +92 MB | With sharing on, 2, 4 and 8 threads × 50 concurrent runs produced exact outputs, in about the same total time as unshared on this model. ## Test plan * New `backends/apple/coreai/test/test_multifunction.py`: * Partitioner option and compile spec; rejects `input_enumerations`. * Two methods share one bundle (`export_hash`, `path`, files, data-store keys) with `forward_0`/`prefill_0`, and the combined bundle is about the size of a single method's. * A graph break yields `forward_0`, `forward_1`, `prefill_0`, `prefill_1` in one bundle. * With the option off, there's one `main` model per delegate (as today). * Mismatched compile specs, and the option set on only some methods, raise errors. * AOT with a mocked `coreai-build` (one build, shared `archs`) and with the real `coreai-build`. * Runtime: loads every method from one PTE through the Python runtime and compares to eager, with and without a graph break. * Runtime sharing (native): * `ParsesOptionalMultifunctionFlag`: flag absent, true, false, and non-boolean values. * `ShareOptionRequiresMultifunctionExport`, `OptionsRejectWrongTypes`. * `ShareOptionSharesOneModelUntilLastDelegate`, `SharedModelsAreScopedToCacheDirectory` (separate directories don't share; a symlink alias does). * `SharedModelSerializesConcurrentExecution`, `ConcurrentFirstSharedLoadsAcquireOnce`, `FailedSharedAcquisitionLeavesNoEntry`. * `SharedLoadsOfDifferentModelsRunConcurrently`, `WaiterAcquiresAfterConcurrentSharedLoadFails`. * The concurrency and directory tests each fail when the lock, per-key tracking or directory scoping they cover is removed. * `coreai_multithread_test` adds cold and warm rounds on a multi-function stress fixture with sharing on (real SDK). * All 151 Core AI Python tests pass. All 6 native CTest targets pass. * Backend suite operators: `coreai_fp32` 377 passed, `coreai_fp16` 368 passed, 0 failures. * ufmt and flake8 are clean. --- backends/apple/coreai/CMakeLists.txt | 7 +- backends/apple/coreai/README.md | 129 +++++- .../coreai/compiler/enumerated_shapes.py | 26 +- backends/apple/coreai/compiler/io_compat.py | 11 +- backends/apple/coreai/compiler/preprocess.py | 312 ++++++++++--- .../apple/coreai/partition/partitioner.py | 8 + .../apple/coreai/runtime/ETCoreAIBridge.h | 3 +- .../coreai/runtime/ETCoreAIModelLoader.swift | 28 +- .../apple/coreai/runtime/coreai_backend.mm | 144 ++++-- .../apple/coreai/runtime/coreai_bridge.mm | 7 +- .../coreai/runtime/coreai_fault_injection.h | 2 + .../coreai/runtime/coreai_load_coordinator.h | 22 + .../coreai/runtime/coreai_load_coordinator.mm | 134 +++++- .../apple/coreai/runtime/coreai_manifest.h | 26 +- .../apple/coreai/runtime/coreai_manifest.mm | 119 ++++- .../apple/coreai/runtime/coreai_tensor_io.h | 10 +- .../apple/coreai/runtime/coreai_tensor_io.mm | 13 +- .../runtime/test/SwiftBridgeTests.swift | 31 +- .../runtime/test/coreai_delegate_fixture.h | 10 + .../runtime/test/coreai_delegate_test.mm | 419 +++++++++++++++-- .../coreai/runtime/test/coreai_fake_loader.h | 13 + .../coreai/runtime/test/coreai_fake_loader.mm | 33 +- .../runtime/test/coreai_manifest_test.mm | 136 +++++- .../runtime/test/coreai_multithread_test.cpp | 33 +- .../runtime/test/export_stress_fixture.py | 33 +- backends/apple/coreai/test/test_io_compat.py | 3 +- .../apple/coreai/test/test_multifunction.py | 435 ++++++++++++++++++ backends/apple/coreai/test/test_preprocess.py | 8 +- .../apple/coreai/test/test_shape_config.py | 101 +++- 29 files changed, 2037 insertions(+), 219 deletions(-) create mode 100644 backends/apple/coreai/test/test_multifunction.py 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..3b1dbfef0ed 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -29,12 +29,15 @@ python -m pytest backends/apple/coreai 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 a `function` name, and ordered +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. @@ -108,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 +`_`: + +```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_`, and +`prefill_1_` 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 @@ -362,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, diff --git a/backends/apple/coreai/compiler/enumerated_shapes.py b/backends/apple/coreai/compiler/enumerated_shapes.py index e4ccad15d28..b014d1504ac 100644 --- a/backends/apple/coreai/compiler/enumerated_shapes.py +++ b/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 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..dd67cb8e6a9 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,145 @@ 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, + ) + bound_inputs, bound_outputs = coreai_io_names(program, entrypoint) + return { + "function": entrypoint, + "input_names": bound_inputs, + "output_names": bound_outputs, + } + + +def _enumerate_binding( + program, + edge_program: ExportedProgram, + binding: Dict[str, Any], + input_names: List[str], + compile_specs: List[CompileSpec], +) -> Dict[str, Any]: + """Apply this partition's enumerated shapes and rewrite its binding. + + ``set_static_shape_config`` erases the bound function, so a specialized + delegate lists every per-shape function with its input shapes, in + ``input_names`` order, and the runtime picks one by the inputs it gets. + """ + raw = _get_compile_spec(compile_specs, COMPILE_SPEC_KEYS.INPUT_ENUMERATIONS) + if not raw: + return binding + entrypoint = binding["function"] + shapes_config = apply_enumerated_shapes( + program, + edge_program, + json.loads(raw.decode()), + input_names=input_names, + entrypoint=entrypoint, + ) + if not shapes_config: + return binding + bound_inputs = binding["input_names"] + functions = [] + for key, shapes in shapes_config.items(): + if set(shapes) != set(bound_inputs): + raise ValueError( + f"enumerated shapes for {entrypoint!r} cover inputs " + f"{sorted(shapes)} but the delegate binds {bound_inputs}" + ) + functions.append( + { + "function": f"{entrypoint}_{key}", + "input_shapes": [list(shapes[name]) for name in bound_inputs], + } + ) + enumerated = {k: v for k, v in binding.items() if k != "function"} + enumerated["enumerated_functions"] = functions + return enumerated + + +def _check_unique_functions(bindings: List[Dict[str, Any]]) -> None: + """Reject two delegates binding the same function name. + + Enumerated functions are named ``__``, which a method + whose name looks like another method's index and key could reproduce. + """ + seen = set() + for binding in bindings: + names = ( + [binding["function"]] + if "function" in binding + else [f["function"] for f in binding["enumerated_functions"]] + ) + for name in names: + if name in seen: + raise ValueError( + f"multifunction export generates the function name {name!r} " + "twice; rename the method whose name collides with it" + ) + seen.add(name) + + +def _spec_items( + compile_specs: List[CompileSpec], ignore: Tuple[str, ...] = () +) -> List[Tuple[str, bytes]]: + return sorted( + (spec.key, bytes(spec.value)) + for spec in compile_specs + if spec.key not in ignore ) - return converter.to_coreai() # 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 +410,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,52 +582,110 @@ 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 + input_names, output_names = _boundary_names(edge_program) + program = _convert_to_aiprogram( + [(edge_program, MAIN_ENTRYPOINT, input_names, output_names)] ) + # Enumerated shapes replace ``main`` with renamed clones, so bind first. + binding = _binding( + program, edge_program, MAIN_ENTRYPOINT, input_names, output_names + ) + binding = _enumerate_binding( + program, edge_program, binding, input_names, compile_specs + ) + return CoreAIBackend._package(program, compile_specs, [binding])[0] - mutated_inputs = edge_program.graph_signature.user_inputs_to_mutate - if mutated_inputs: + @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) + + missing = sorted( + { + method + for method, _, _, specs in partitions + if _get_compile_spec(specs, COMPILE_SPEC_KEYS.MULTIFUNCTION) is None + } + ) + if missing: raise ValueError( - "Core AI delegates cannot mutate ExecuTorch-owned user inputs: " - f"{sorted(mutated_inputs.values())}" + "multifunction=True must be set on the CoreAIPartitioner of every " + f"method lowered to Core AI; methods {missing} were lowered " + "without it" ) - 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))] - program = _convert_to_aiprogram( - edge_program, input_names=input_names, output_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 + # Each partition applies its own enumerations, so methods may enumerate + # different values; every other spec shapes the shared model. + ignore = (COMPILE_SPEC_KEYS.INPUT_ENUMERATIONS.value,) + reference = _spec_items(partitions[0][3], ignore) + differing = sorted( + { + method + for method, _, _, specs in partitions + if _spec_items(specs, ignore) != reference + } ) - # 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 + if differing: + raise ValueError( + "multifunction requires identical Core AI compile specs in every " + f"partition; methods {differing} differ from {partitions[0][0]!r}" + ) + 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], ) - if raw_enum: - apply_enumerated_shapes( - program, - edge_program, - json.loads(raw_enum.decode()), - input_names=input_names, + program = _convert_to_aiprogram(entries) + # Bind every partition before specializing any, since specializing a + # partition erases its function. + bindings = [_binding(program, *entry) for entry in entries] + bindings = [ + { + **_enumerate_binding( + program, edge_program, binding, input_names, specs + ), + "multifunction": True, + } + for binding, (edge_program, _, input_names, _), (_, _, _, specs) in zip( + bindings, entries, partitions ) - bindings = { - "function": MAIN_ENTRYPOINT, - "input_names": bound_inputs, - "output_names": bound_outputs, + ] + _check_unique_functions(bindings) + 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 +696,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 +729,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..e04a73a76aa 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,9 @@ 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 _. specs = [] if min_deployment_version is not None: specs.append( @@ -344,6 +348,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__, @@ -365,6 +371,8 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: # Core AI derives op support and ops_to_not_decompose from coreai-torch's # tables; the deprecated to_edge() + to_backend() flow decomposes ops # (e.g. fused SDPA) before the partitioner runs, breaking that contract. + # It also lowers each partition on its own, so multifunction exports + # would silently become separate models without shared weights. if not self._is_to_edge_transform_and_lower(): raise RuntimeError( "CoreAIPartitioner must be used with to_edge_transform_and_lower(). " diff --git a/backends/apple/coreai/runtime/ETCoreAIBridge.h b/backends/apple/coreai/runtime/ETCoreAIBridge.h index 904a19dcc5c..ad66ecdc236 100644 --- a/backends/apple/coreai/runtime/ETCoreAIBridge.h +++ b/backends/apple/coreai/runtime/ETCoreAIBridge.h @@ -148,7 +148,8 @@ extern "C" { BOOL ETCoreAIIsAvailable(void); NSString* _Nullable ETCoreAIDeviceArchitectureName(void); id ETCoreAICreateModelLoader( - ETCoreAIComputeUnit computeUnit); + ETCoreAIComputeUnit computeUnit, + BOOL expectFrequentReshapes); #ifdef __cplusplus } #endif diff --git a/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift b/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift index 049b329d6fa..69251f231f8 100644 --- a/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift +++ b/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift @@ -238,8 +238,12 @@ public final class ETCoreAISwiftModelLoader: NSObject, ETCoreAIModelLoading { super.init() } - @objc public convenience init?(computeUnit: ETCoreAIComputeUnit) { - guard let options = Self.specializationOptions(for: computeUnit), + @objc public convenience init?( + computeUnit: ETCoreAIComputeUnit, expectFrequentReshapes: Bool + ) { + guard + let options = Self.specializationOptions( + for: computeUnit, expectFrequentReshapes: expectFrequentReshapes), let specialization = try? CoreAISpecialization(options: options) else { return nil @@ -247,17 +251,21 @@ public final class ETCoreAISwiftModelLoader: NSObject, ETCoreAIModelLoading { self.init(specialization: specialization) } - static func specializationOptions(for computeUnit: ETCoreAIComputeUnit) - -> SpecializationOptions? - { + static func specializationOptions( + for computeUnit: ETCoreAIComputeUnit, expectFrequentReshapes: Bool = false + ) -> SpecializationOptions? { + var options: SpecializationOptions switch computeUnit { - case .default: return .default - case .cpuOnly: return .cpuOnly - case .preferCPU: return SpecializationOptions(preferredComputeUnitKind: .cpu) - case .preferGPU: return SpecializationOptions(preferredComputeUnitKind: .gpu) - case .preferNeuralEngine: return SpecializationOptions(preferredComputeUnitKind: .neuralEngine) + case .default: options = .default + case .cpuOnly: options = .cpuOnly + case .preferCPU: options = SpecializationOptions(preferredComputeUnitKind: .cpu) + case .preferGPU: options = SpecializationOptions(preferredComputeUnitKind: .gpu) + case .preferNeuralEngine: + options = SpecializationOptions(preferredComputeUnitKind: .neuralEngine) @unknown default: return nil } + options.expectFrequentReshapes = expectFrequentReshapes + return options } public var specializationIdentifier: String { specialization.identifier } diff --git a/backends/apple/coreai/runtime/coreai_backend.mm b/backends/apple/coreai/runtime/coreai_backend.mm index fe4f6c8a9a3..069a08bc1ad 100644 --- a/backends/apple/coreai/runtime/coreai_backend.mm +++ b/backends/apple/coreai/runtime/coreai_backend.mm @@ -14,47 +14,57 @@ #import "coreai_tensor_io.h" #include +#include #include #include #include +#include +#include namespace executorch::backends::coreai { namespace { using namespace executorch::runtime; +struct BoundFunction final { + std::vector> input_shapes; + id session = nil; +}; + struct Handle final { size_t input_count = 0; size_t output_count = 0; - id session = nil; + ETCoreAILoadedModel* model = nil; + // Enumerated delegates pick a function per call by input shape; otherwise + // the single function accepts any shapes. + bool enumerated = false; + std::vector functions; }; +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; @@ -80,11 +90,12 @@ Result> bind_session( id model, - const SelectedBundle& bundle) { + const SelectedBundle& bundle, + NSString* function) { dispatch_semaphore_t ready = dispatch_semaphore_create(0); __block id session = nil; __block NSError* error = nil; - [model loadFunctionNamed:bundle.function + [model loadFunctionNamed:function inputNames:bundle.inputs outputNames:bundle.outputs completion:^(id result, NSError* failure) { @@ -106,6 +117,55 @@ return bound; } +bool inputs_match( + const std::vector>& shapes, + Span args) { + for (size_t index = 0; index < shapes.size(); ++index) { + const auto sizes = args[index]->toTensor().sizes(); + if (sizes.size() != shapes[index].size() || + !std::equal(sizes.begin(), sizes.end(), shapes[index].begin())) { + return false; + } + } + return true; +} + +std::string describe_shape( + executorch::aten::ArrayRef sizes) { + std::string text = "["; + for (size_t dim = 0; dim < sizes.size(); ++dim) { + text += (dim == 0 ? "" : ", ") + std::to_string(sizes[dim]); + } + return text + "]"; +} + +// Each enumerated function accepts only its own static shapes. Called after +// argument validation, so every input is a tensor. +Result> select_session( + const Handle& handle, + Span args) { + if (!handle.enumerated) { + return handle.functions.front().session; + } + for (const auto& function : handle.functions) { + if (inputs_match(function.input_shapes, args)) { + return function.session; + } + } + std::string shapes; + for (size_t index = 0; index < handle.input_count; ++index) { + shapes += (index == 0 ? "" : ", ") + + describe_shape(args[index]->toTensor().sizes()); + } + ET_LOG( + Error, + "Core AI inputs with shapes %s match none of the %zu enumerated shapes " + "this delegate was exported with", + shapes.c_str(), + handle.functions.size()); + return Error::InvalidArgument; +} + class CoreAIBackend final : public BackendInterface { public: bool is_available() const override { @@ -137,34 +197,55 @@ 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))"); + auto reshapes = runtime_option( + context, "coreai_expect_frequent_reshapes", false); + if (!reshapes.ok()) { + return reshapes.error(); + } id loader = - ETCoreAICreateModelLoader(compute_unit.get()); + ETCoreAICreateModelLoader(compute_unit.get(), reshapes.get()); ET_CHECK_OR_RETURN_ERROR( loader != nil, Internal, "Cannot create Core AI loader"); auto cache_directory = create_cache_directory(configured.get()); 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(); - } - auto session = bind_session(prepared.get(), host_bundle->bundle); - if (!session.ok()) { - return session.error(); + loader, + share.get()); + if (!loaded.ok()) { + return loaded.error(); } auto handle = std::unique_ptr(new (std::nothrow) Handle()); ET_CHECK_OR_RETURN_ERROR( handle != nullptr, MemoryAllocationFailed, "Cannot allocate Core AI handle"); + for (const auto& function : host_bundle->bundle.functions) { + auto session = bind_session( + loaded.get().model, host_bundle->bundle, function.name); + if (!session.ok()) { + return session.error(); + } + handle->functions.push_back({function.input_shapes, session.get()}); + } handle->input_count = host_bundle->bundle.inputs.count; handle->output_count = host_bundle->bundle.outputs.count; - handle->session = session.get(); + handle->model = loaded.get(); + handle->enumerated = host_bundle->bundle.enumerated; processed->Free(); return static_cast(handle.release()); } @@ -178,8 +259,11 @@ 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); + args, handle.input_count, handle.output_count, [&] { + return select_session(handle, args); + }); } } diff --git a/backends/apple/coreai/runtime/coreai_bridge.mm b/backends/apple/coreai/runtime/coreai_bridge.mm index 3defec3b174..557d1247398 100644 --- a/backends/apple/coreai/runtime/coreai_bridge.mm +++ b/backends/apple/coreai/runtime/coreai_bridge.mm @@ -25,6 +25,9 @@ BOOL ETCoreAIIsAvailable(void) { } id ETCoreAICreateModelLoader( - ETCoreAIComputeUnit computeUnit) { - return [[ETCoreAISwiftModelLoader alloc] initWithComputeUnit:computeUnit]; + ETCoreAIComputeUnit computeUnit, + BOOL expectFrequentReshapes) { + return [[ETCoreAISwiftModelLoader alloc] + initWithComputeUnit:computeUnit + expectFrequentReshapes:expectFrequentReshapes]; } diff --git a/backends/apple/coreai/runtime/coreai_fault_injection.h b/backends/apple/coreai/runtime/coreai_fault_injection.h index 28c8cd0fbbb..9091ead3541 100644 --- a/backends/apple/coreai/runtime/coreai_fault_injection.h +++ b/backends/apple/coreai/runtime/coreai_fault_injection.h @@ -29,6 +29,8 @@ enum class FaultPoint { AfterSDK, BeforeEvict, AfterEvict, + // Observation only: the result is ignored. + BeforeSharedLoadWait, }; #if defined(COREAI_FAULT_INJECTION) && COREAI_FAULT_INJECTION 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..47932c0c9ce 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,96 @@ 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; + } + static_cast(fault_point(FaultPoint::BeforeSharedLoadWait)); + load_finished.wait(guard); + } + [loading addObject:shared_key]; + } + // Runs on every exit, so a key cannot stay marked as loading and hang later + // loads of the same model. + struct FinishSharedLoad final { + NSString* key; + ETCoreAILoadedModel* loaded = nil; + ~FinishSharedLoad() { + if (key == nil) { + return; + } + { + std::lock_guard guard(mutex); + [loading removeObject:key]; + if (loaded != nil) { + [shared_models setObject:loaded forKey:key]; + } + } + load_finished.notify_all(); + } + } finish{shared_key}; + auto model = acquire_cached_model( + host_bundle, named_data, cache_directory, loader, key.get()); + if (!model.ok()) { + return model.error(); + } + finish.loaded = [[ETCoreAILoadedModel alloc] initWithModel:model.get()]; + return finish.loaded; +} Error clear_cache_entry( const CacheDirectory& cache_directory, @@ -257,7 +373,7 @@ Error clear_cache_entry( ET_CHECK_OK_OR_RETURN_ERROR(require_coreai_available()); // Eviction targets the bookmark, so the loader's specialization is // irrelevant. - loader = ETCoreAICreateModelLoader(ETCoreAIComputeUnitDefault); + loader = ETCoreAICreateModelLoader(ETCoreAIComputeUnitDefault, NO); ET_CHECK_OR_RETURN_ERROR( loader != nil, Internal, "Cannot create Core AI model loader"); } diff --git a/backends/apple/coreai/runtime/coreai_manifest.h b/backends/apple/coreai/runtime/coreai_manifest.h index 9627761c0ce..635fa2dc9b0 100644 --- a/backends/apple/coreai/runtime/coreai_manifest.h +++ b/backends/apple/coreai/runtime/coreai_manifest.h @@ -11,12 +11,25 @@ #import #include +#include +#include + namespace executorch::backends::coreai { +// A function the delegate can bind. One without input shapes accepts any input +// shapes; enumerated-shape exports list one per shape, and each accepts exactly +// its input_shapes, in input-name order. +struct ManifestFunction { + NSString* name; + std::vector> input_shapes; +}; + class Manifest final { public: NSString* const export_hash; - NSString* const function; + const std::vector functions; + // Listed under enumerated_functions; each function has static input shapes. + const bool enumerated; NSArray* const inputs; NSArray* const outputs; NSDictionary* const files; @@ -26,6 +39,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; @@ -35,7 +50,8 @@ class Manifest final { private: Manifest( NSString* export_hash, - NSString* function, + std::vector functions, + bool enumerated, NSArray* inputs, NSArray* outputs, NSDictionary* files, @@ -43,7 +59,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); }; @@ -54,7 +71,8 @@ class SelectedBundle final { NSString* const path; NSString* const bundle_name; NSString* const digest; - NSString* const function; + const std::vector functions; + const bool enumerated; NSArray* const inputs; NSArray* const outputs; NSDictionary* const files; diff --git a/backends/apple/coreai/runtime/coreai_manifest.mm b/backends/apple/coreai/runtime/coreai_manifest.mm index 007d049b07c..87e17046d9c 100644 --- a/backends/apple/coreai/runtime/coreai_manifest.mm +++ b/backends/apple/coreai/runtime/coreai_manifest.mm @@ -9,13 +9,16 @@ #import "coreai_manifest.h" #include #include +#include +#include #import "coreai_validation.h" namespace executorch::backends::coreai { Manifest::Manifest( NSString* export_hash, - NSString* function, + std::vector functions, + bool enumerated, NSArray* inputs, NSArray* outputs, NSDictionary* files, @@ -23,9 +26,11 @@ bool aot_compiled, NSString* platform, NSDictionary* architectures, - NSOperatingSystemVersion minimum_version) + NSOperatingSystemVersion minimum_version, + bool multifunction) : export_hash(export_hash), - function(function), + functions(std::move(functions)), + enumerated(enumerated), inputs(inputs), outputs(outputs), files(files), @@ -33,7 +38,8 @@ aot_compiled(aot_compiled), platform(platform), architectures(architectures), - minimum_version(minimum_version) {} + minimum_version(minimum_version), + multifunction(multifunction) {} SelectedBundle::SelectedBundle( const Manifest& manifest, @@ -43,7 +49,8 @@ path(path), bundle_name(path.lastPathComponent), digest(manifest.bundle_digests[path.lastPathComponent]), - function(manifest.function), + functions(manifest.functions), + enumerated(manifest.enumerated), inputs(manifest.inputs), outputs(manifest.outputs), files([files copy]), @@ -67,9 +74,74 @@ bool names_array(id value) { return true; } +bool is_boolean(id value) { + return value != nil && + CFGetTypeID((__bridge CFTypeRef)value) == CFBooleanGetTypeID(); +} + +bool is_dimension(id value) { + return [value isKindOfClass:NSNumber.class] && !is_boolean(value) && + !CFNumberIsFloatType((__bridge CFNumberRef)value) && + [value longLongValue] > 0 && + [value longLongValue] <= std::numeric_limits::max(); +} + +// Each entry names a distinct function and gives every input's shape; no two +// entries may accept the same shapes, so dispatch is unambiguous. +Error parse_enumerated_functions( + id value, + NSUInteger input_count, + std::vector& functions) { + ET_CHECK_OR_RETURN_ERROR( + [value isKindOfClass:NSArray.class] && [value count] > 0, + InvalidProgram, + "Core AI enumerated functions must be a nonempty array"); + NSMutableSet* names = [NSMutableSet set]; + std::set>> seen_shapes; + for (id entry in value) { + ET_CHECK_OR_RETURN_ERROR( + [entry isKindOfClass:NSDictionary.class], + InvalidProgram, + "Core AI enumerated function must be an object"); + id name = entry[@"function"]; + ET_CHECK_OR_RETURN_ERROR( + is_nonempty_string(name) && ![names containsObject:name], + InvalidProgram, + "Core AI enumerated functions need distinct nonempty names"); + [names addObject:name]; + id shapes = entry[@"input_shapes"]; + ET_CHECK_OR_RETURN_ERROR( + [shapes isKindOfClass:NSArray.class] && [shapes count] == input_count, + InvalidProgram, + "Core AI enumerated function %s needs one shape per input", + [name UTF8String]); + std::vector> input_shapes; + for (id shape in shapes) { + ET_CHECK_OR_RETURN_ERROR( + [shape isKindOfClass:NSArray.class], + InvalidProgram, + "Core AI enumerated shape must be an array"); + std::vector dims; + for (id dim in shape) { + ET_CHECK_OR_RETURN_ERROR( + is_dimension(dim), + InvalidProgram, + "Core AI enumerated shape dims must be positive integers"); + dims.push_back([dim longLongValue]); + } + input_shapes.push_back(std::move(dims)); + } + ET_CHECK_OR_RETURN_ERROR( + seen_shapes.insert(input_shapes).second, + InvalidProgram, + "Core AI enumerated functions accept the same input shapes"); + functions.push_back({name, std::move(input_shapes)}); + } + return Error::Ok; +} + 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())] != @@ -180,9 +252,15 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { "Unsupported Core AI packaging; expected aimodel or aimodelc"); const bool aot_compiled = [packaging isEqualToString:@"aimodelc"]; NSString* export_hash = dict[@"export_hash"]; - NSString* function = dict[@"function"]; + id function = dict[@"function"]; + id enumerated_functions = dict[@"enumerated_functions"]; 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( @@ -232,10 +310,25 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { [bundles addObject:@"model.aimodel"]; } ET_CHECK_OR_RETURN_ERROR( - is_nonempty_string(function) && names_array(inputs) && - names_array(outputs), + names_array(inputs) && names_array(outputs), InvalidProgram, "Core AI manifest requires ordered input/output names"); + ET_CHECK_OR_RETURN_ERROR( + (function == nil) != (enumerated_functions == nil), + InvalidProgram, + "Core AI manifest requires exactly one of function or " + "enumerated_functions"); + std::vector functions; + if (function != nil) { + ET_CHECK_OR_RETURN_ERROR( + is_nonempty_string(function), + InvalidProgram, + "Core AI manifest function must be a nonempty string"); + functions.push_back({function, {}}); + } else { + ET_CHECK_OK_OR_RETURN_ERROR(parse_enumerated_functions( + enumerated_functions, inputs.count, functions)); + } NSOperatingSystemVersion minimum_version{}; ET_CHECK_OR_RETURN_ERROR( parse_version(dict[@"min_deployment_version"], minimum_version), @@ -258,7 +351,8 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { } return Manifest( export_hash, - function, + std::move(functions), + enumerated_functions != nil, inputs, outputs, files, @@ -266,7 +360,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/coreai_tensor_io.h b/backends/apple/coreai/runtime/coreai_tensor_io.h index 1a49564face..f372065ddd7 100644 --- a/backends/apple/coreai/runtime/coreai_tensor_io.h +++ b/backends/apple/coreai/runtime/coreai_tensor_io.h @@ -8,15 +8,21 @@ #pragma once #include +#include #include + +#include #import "ETCoreAIBridge.h" namespace executorch::backends::coreai { +// Validates args, then runs the session select_session picks for them; the +// selector may assume every argument is a valid tensor. runtime::Error execute_session( - id session, runtime::Span args, size_t input_count, - size_t output_count); + size_t output_count, + const std::function>()>& + select_session); } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_tensor_io.mm b/backends/apple/coreai/runtime/coreai_tensor_io.mm index 746b5dac38f..ec27d9d5252 100644 --- a/backends/apple/coreai/runtime/coreai_tensor_io.mm +++ b/backends/apple/coreai/runtime/coreai_tensor_io.mm @@ -377,16 +377,21 @@ Error copy_staged( } // namespace Error execute_session( - id session, Span args, size_t input_count, - size_t output_count) { - ET_CHECK_OR_RETURN_ERROR( - session != nil, DelegateInvalidHandle, "Missing Core AI session"); + size_t output_count, + const std::function>()>& select_session) { auto validated = validate_arguments(args, input_count, output_count); if (!validated.ok()) { return validated.error(); } + auto selected = select_session(); + if (!selected.ok()) { + return selected.error(); + } + id session = selected.get(); + ET_CHECK_OR_RETURN_ERROR( + session != nil, DelegateInvalidHandle, "Missing Core AI session"); auto inputs = borrow_inputs(args, validated.get(), input_count); if (!inputs.ok()) { return inputs.error(); diff --git a/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift index 2ebd4ce98e8..745b07f84fa 100644 --- a/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift +++ b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift @@ -165,17 +165,32 @@ final class SwiftSpecializationTests: XCTestCase { var report: [String] = [] for (unit, options) in expected { XCTAssertEqual(ETCoreAISwiftModelLoader.specializationOptions(for: unit), options) - let loader = try XCTUnwrap(ETCoreAISwiftModelLoader(computeUnit: unit)) - XCTAssertEqual( - loader.specializationIdentifier, try CoreAISpecialization(options: options).identifier) - variants.append((options, loader.specializationIdentifier)) - report.append( - "\(unit.rawValue): \(options) -> \(loader.specializationIdentifier)") + for reshapes in [false, true] { + var expectedOptions = options + expectedOptions.expectFrequentReshapes = reshapes + XCTAssertEqual( + ETCoreAISwiftModelLoader.specializationOptions( + for: unit, expectFrequentReshapes: reshapes), + expectedOptions) + let loader = try XCTUnwrap( + ETCoreAISwiftModelLoader(computeUnit: unit, expectFrequentReshapes: reshapes)) + XCTAssertEqual( + loader.specializationIdentifier, + try CoreAISpecialization(options: expectedOptions).identifier) + variants.append((expectedOptions, loader.specializationIdentifier)) + report.append( + "\(unit.rawValue) reshapes=\(reshapes): \(expectedOptions) -> " + + loader.specializationIdentifier) + } } assertIdentifiersMatchExactlyWhenOptionsMatch(variants, report: report) XCTAssertEqual( - ETCoreAISwiftModelLoader(computeUnit: .default)?.specializationIdentifier, "default") - XCTAssertNil(ETCoreAISwiftModelLoader(computeUnit: ETCoreAIComputeUnit(rawValue: 99)!)) + ETCoreAISwiftModelLoader(computeUnit: .default, expectFrequentReshapes: false)? + .specializationIdentifier, + "default") + XCTAssertNil( + ETCoreAISwiftModelLoader( + computeUnit: ETCoreAIComputeUnit(rawValue: 99)!, expectFrequentReshapes: false)) } // The SDK limits allowed units to the hardware (a CPU-only machine makes 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..c7e12bb8da7 100644 --- a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm @@ -13,12 +13,18 @@ #include #include #include +#include +#include #include #include +#include +#include +#include #include "coreai_cache.h" #include "coreai_cache_entry_fixture.h" #import "coreai_delegate_fixture.h" +#include "coreai_fault_scope.h" #include "coreai_test_utils.h" namespace executorch::backends::coreai::testing { @@ -243,6 +249,95 @@ void expect_preflight_rejection( } } +NSMutableDictionary* enumerated_manifest(NSArray* functions) { + NSMutableDictionary* manifest = manifest_dict(); + [manifest removeObjectForKey:@"function"]; + manifest[@"enumerated_functions"] = functions; + return manifest; +} + +// A 1-D float tensor of `size` elements over `data`, shaped once at +// construction because TensorImpl caches its element count. +struct VectorTensor final { + executorch::aten::TensorImpl::SizesType size[1]; + executorch::aten::TensorImpl::DimOrderType order[1] = {0}; + executorch::aten::TensorImpl::StridesType stride[1] = {1}; + executorch::aten::TensorImpl impl; + EValue value; + VectorTensor(float* data, int size_value, bool dynamic = false) + : size{size_value}, + impl( + executorch::aten::ScalarType::Float, + 1, + size, + data, + order, + stride, + dynamic ? runtime::TensorShapeDynamism::DYNAMIC_BOUND + : runtime::TensorShapeDynamism::STATIC), + value(executorch::aten::Tensor(&impl)) {} +}; + +TEST_F(CoreAIDelegateTest, EnumeratedFunctionsDispatchByInputShape) { + @autoreleasepool { + bridge.state().functions = + [NSSet setWithArray:@[ @"main_two", @"main_one" ]]; + DelegateFixture fixture( + backend, bridge, enumerated_manifest(@[ + @{@"function" : @"main_two", + @"input_shapes" : @[ @[ @2 ] ]}, + @{@"function" : @"main_one", + @"input_shapes" : @[ @[ @1 ] ]}, + ])); + ASSERT_EQ(fixture.init(), Error::Ok); + EXPECT_EQ(bridge.state().binding_mismatches.load(), 0); + + FloatTensors two; + ASSERT_EQ(fixture.execute({two.args, 2}), Error::Ok); + EXPECT_TRUE([bridge.state().last_function isEqual:@"main_two"]); + EXPECT_EQ(two.output[0], 3); + EXPECT_EQ(two.output[1], 7); + + float one_input[1] = {5}; + float one_output[2] = {-1, -1}; + VectorTensor one_in(one_input, 1); + VectorTensor one_out(one_output, 2, /*dynamic=*/true); + EValue* one_args[] = {&one_in.value, &one_out.value}; + ASSERT_EQ(fixture.execute({one_args, 2}), Error::Ok); + EXPECT_TRUE([bridge.state().last_function isEqual:@"main_one"]); + EXPECT_EQ(one_out.impl.size(0), 1); + EXPECT_EQ(one_output[0], 5); + + float three_input[3] = {1, 2, 3}; + float three_output[3] = {-1, -1, -1}; + VectorTensor three_in(three_input, 3); + VectorTensor three_out(three_output, 3, /*dynamic=*/true); + EValue* three_args[] = {&three_in.value, &three_out.value}; + const int executions = bridge.state().executions.load(); + EXPECT_EQ(fixture.execute({three_args, 2}), Error::InvalidArgument); + EXPECT_EQ(bridge.state().executions.load(), executions); + EXPECT_EQ(three_output[0], -1); + } +} + +TEST_F(CoreAIDelegateTest, EnumeratedFunctionBindFailureFailsInit) { + @autoreleasepool { + bridge.state().functions = + [NSSet setWithArray:@[ @"main_two", @"main_one" ]]; + bridge.state().fail_function = @"main_one"; + DelegateFixture fixture( + backend, bridge, enumerated_manifest(@[ + @{@"function" : @"main_two", + @"input_shapes" : @[ @[ @2 ] ]}, + @{@"function" : @"main_one", + @"input_shapes" : @[ @[ @1 ] ]}, + ])); + EXPECT_EQ(fixture.init(), Error::InvalidProgram); + EXPECT_EQ(bridge.state().function_bindings.load(), 2); + EXPECT_EQ(fixture.freed, 0); + } +} + TEST_F(CoreAIDelegateTest, AdjacentInputAndOutputStorageCopiesDirectly) { @autoreleasepool { DelegateFixture fixture(backend, bridge); @@ -905,24 +1000,33 @@ 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}, + {"coreai_expect_frequent_reshapes", 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 +1039,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,17 +1067,285 @@ TensorImpl output( EXPECT_EQ(init_with_compute_unit(name), Error::InvalidArgument); EXPECT_EQ(bridge.state().last_compute_unit.load(), -1); } - { - SCOPED_TRACE("integer compute unit"); + } +} + +TEST_F(CoreAIDelegateTest, ExpectFrequentReshapesOptionReachesLoader) { + @autoreleasepool { + auto init_with_reshapes = [&](std::optional reshapes) { + bridge.state().last_expect_frequent_reshapes = -1; DelegateFixture fixture(backend, bridge); - BackendOption wrong_type{"coreai_compute_unit", 1}; + BackendOptions<3> options; EXPECT_EQ( - fixture.init_with_options({&wrong_type, 1}), Error::InvalidArgument); - EXPECT_EQ(fixture.freed, 0); + options.set_option("coreai_compute_unit", "cpu_only"), Error::Ok); + if (reshapes.has_value()) { + EXPECT_EQ( + options.set_option("coreai_expect_frequent_reshapes", *reshapes), + Error::Ok); + } + return fixture.init_with_storage(options); + }; + for (const auto& [reshapes, expected] : + {std::pair, int>{std::nullopt, 0}, + std::pair, int>{false, 0}, + std::pair, int>{true, 1}}) { + SCOPED_TRACE( + reshapes.has_value() ? (*reshapes ? "true" : "false") : "absent"); + ASSERT_EQ(init_with_reshapes(reshapes), Error::Ok); + EXPECT_EQ(bridge.state().last_expect_frequent_reshapes.load(), expected); + EXPECT_EQ( + bridge.state().last_compute_unit.load(), ETCoreAIComputeUnitCPUOnly); } } } +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("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); + auto second_waiting = std::make_shared>(false); + __weak BookmarkFakeLoader* loader = fake.loader; + // The first acquisition fails only once the second load waits on it. + fake.loader.onSpecialize = ^(NSURL*) { + if (++*calls == 1) { + *first_in_flight = true; + for (int i = 0; i < 1000 && !*second_waiting; ++i) { + usleep(5000); + } + 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(); + } + Error second_result = Error::Internal; + { + StorageFaultScope observe_wait(FaultPoint::BeforeSharedLoadWait); + observe_wait.armed = false; + observe_wait.observe = ^(FaultPoint point) { + if (point == FaultPoint::BeforeSharedLoadWait) { + *second_waiting = true; + } + }; + second_result = init_sharing(second, true, root); + } + worker.join(); + EXPECT_TRUE(second_waiting->load()); + 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); + } +} + TEST_F(CoreAIDelegateTest, ConfiguredSymlinkCacheDirectoryLoadsAndClears) { @autoreleasepool { DelegateFixture fixture(backend, bridge); diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.h b/backends/apple/coreai/runtime/test/coreai_fake_loader.h index 72d101b1855..cf6df68f1ad 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,21 @@ struct FakeBridgeState { bool multiple_io = false; bool retain_binding_completion = false; bool retain_execution_completion = false; + // Function names the fake model has; nil means only "main". + NSSet* functions = nil; + // Binding this function fails like a load failure; the others still bind. + NSString* fail_function = nil; + std::atomic function_bindings{0}; + useconds_t execute_delay_us = 0; NSString* last_bundle = nil; + // Written by the execution that last completed, before its completion. + NSString* last_function = 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}; @@ -42,6 +53,7 @@ struct FakeBridgeState { std::atomic prepared_models{0}; std::atomic loaders{0}; std::atomic last_compute_unit{-1}; + std::atomic last_expect_frequent_reshapes{-1}; std::atomic last_input_bytes{nullptr}; std::atomic last_input_byte_count{0}; std::atomic last_input_first_byte{0}; @@ -72,6 +84,7 @@ class ScopedFakeBridgeState final { @interface FakeSession : NSObject @property(nonatomic, strong) id preparedPin; +@property(nonatomic, copy) NSString* function; @end @interface FakePreparedModel : NSObject diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm index 4e3e69b21cb..609c42e8c2e 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm @@ -53,7 +53,9 @@ bool valid_binding( NSArray* expected_outputs = context->state.multiple_io ? @[ @"output_1", @"output_0" ] : @[ @"output_0" ]; - const bool valid = [function isEqualToString:@"main"] && + NSSet* functions = + context->state.functions ?: [NSSet setWithObject:@"main"]; + const bool valid = [functions containsObject:function] && [inputs isEqualToArray:expected_inputs] && [outputs isEqualToArray:expected_outputs]; if (!valid) { @@ -133,6 +135,20 @@ - (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; + NSString* function = self.function; const bool should_fail = context->state.fail_execute; const bool malformed = context->state.bad_output; dispatch_semaphore_t ready = context->state.input_ready; @@ -141,7 +157,11 @@ - (void)executeInputs:(NSArray*)inputs context.retainedExecutionCompletion = completion; } run_async(context, ^{ + if (delay > 0) { + usleep(delay); + } ETCoreAIInputTensor* input = inputs[0]; + context->state.last_function = function; context->state.last_input_bytes = input.bytes; context->state.last_input_byte_count = input.byteCount; if (ready != nil) { @@ -228,7 +248,9 @@ - (void)loadFunctionNamed:(NSString*)functionName CoreAIFakeBridgeContext* context = self.context; const bool valid = valid_binding(context, functionName, inputNames, outputNames); - const bool should_fail = context->state.fail_load; + const bool should_fail = context->state.fail_load || + [context->state.fail_function isEqualToString:functionName]; + ++context->state.function_bindings; if (context->state.retain_binding_completion) { context.retainedBindingCompletion = completion; } @@ -239,6 +261,7 @@ - (void)loadFunctionNamed:(NSString*)functionName } FakeSession* session = [[FakeSession alloc] initWithContext:context]; session.preparedPin = self; + session.function = functionName; completion(session, nil); }); } @@ -349,6 +372,7 @@ - (void)loadFunctionNamed:(NSString*)functionName } FakeSession* session = [[FakeSession alloc] initWithContext:context]; session.preparedPin = self; + session.function = functionName; completion(session, nil); }); } @@ -538,10 +562,13 @@ BOOL ETCoreAIIsAvailable(void) { } id ETCoreAICreateModelLoader( - ETCoreAIComputeUnit computeUnit) { + ETCoreAIComputeUnit computeUnit, + BOOL expectFrequentReshapes) { if (current_context == nil || current_context->state.fail_loader_factory) { return nil; } current_context->state.last_compute_unit = static_cast(computeUnit); + current_context->state.last_expect_frequent_reshapes = + expectFrequentReshapes ? 1 : 0; return current_context->state.bookmark_loader ?: [[FakeLoader alloc] init]; } diff --git a/backends/apple/coreai/runtime/test/coreai_manifest_test.mm b/backends/apple/coreai/runtime/test/coreai_manifest_test.mm index 2e0fa5393e7..cbef5562b03 100644 --- a/backends/apple/coreai/runtime/test/coreai_manifest_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_manifest_test.mm @@ -6,8 +6,10 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include +#include #include "coreai_manifest_fixture.h" #include "coreai_validation.h" @@ -30,6 +32,130 @@ EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::NotSupported); } +NSMutableDictionary* enumerated_manifest() { + NSMutableDictionary* dict = manifest_dict(); + [dict removeObjectForKey:@"function"]; + dict[@"enumerated_functions"] = @[ + @{@"function" : @"main_s0_4", + @"input_shapes" : @[ @[ @4, @8 ] ]}, + @{@"function" : @"main_s0_16", + @"input_shapes" : @[ @[ @16, @8 ] ]}, + ]; + return dict; +} + +TEST(CoreAIManifestTest, ParsesEnumeratedFunctions) { + auto parsed = parse_manifest(encode(enumerated_manifest())); + ASSERT_TRUE(parsed.ok()); + EXPECT_TRUE(parsed->enumerated); + ASSERT_EQ(parsed->functions.size(), 2u); + EXPECT_TRUE([parsed->functions[0].name isEqual:@"main_s0_4"]); + EXPECT_EQ( + parsed->functions[0].input_shapes, + (std::vector>{{4, 8}})); + EXPECT_TRUE([parsed->functions[1].name isEqual:@"main_s0_16"]); + EXPECT_EQ( + parsed->functions[1].input_shapes, + (std::vector>{{16, 8}})); +} + +TEST(CoreAIManifestTest, RejectsInvalidEnumeratedFunctions) { + { + SCOPED_TRACE("both function and enumerated_functions"); + auto dict = enumerated_manifest(); + dict[@"function"] = @"main"; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } + { + SCOPED_TRACE("neither"); + auto dict = enumerated_manifest(); + [dict removeObjectForKey:@"enumerated_functions"]; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } + NSArray* invalid = @[ + @[], + @{}, + @[ @"main_s0_4" ], + @[ + @{@"input_shapes" : @[ @[ @4, @8 ] ]} + ], + @[ + @{@"function" : @"", + @"input_shapes" : @[ @[ @4, @8 ] ]} + ], + @[ @{@"function" : @"f"} ], + @[ + @{@"function" : @"f", + @"input_shapes" : @[]} + ], + @[ + @{@"function" : @"f", + @"input_shapes" : @[ @[ @4 ], @[ @8 ] ]} + ], + @[ + @{@"function" : @"f", + @"input_shapes" : @[ @"4x8" ]} + ], + @[ + @{@"function" : @"f", + @"input_shapes" : @[ @[ @0, @8 ] ]} + ], + @[ + @{@"function" : @"f", + @"input_shapes" : @[ @[ @-4, @8 ] ]} + ], + @[ + @{@"function" : @"f", + @"input_shapes" : @[ @[ @4.5, @8 ] ]} + ], + @[ + @{@"function" : @"f", + @"input_shapes" : @[ @[ @YES, @8 ] ]} + ], + @[ + @{@"function" : @"f", + @"input_shapes" : @[ @[ @(1LL << 40), @8 ] ]} + ], + @[ + @{@"function" : @"f", + @"input_shapes" : @[ @[ @4, @8 ] ]}, + @{@"function" : @"f", + @"input_shapes" : @[ @[ @16, @8 ] ]} + ], + @[ + @{@"function" : @"f", + @"input_shapes" : @[ @[ @4, @8 ] ]}, + @{@"function" : @"g", + @"input_shapes" : @[ @[ @4, @8 ] ]} + ], + ]; + for (id functions in invalid) { + SCOPED_TRACE([functions description].UTF8String); + auto dict = enumerated_manifest(); + dict[@"enumerated_functions"] = functions; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + +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", @@ -171,7 +297,10 @@ ASSERT_TRUE(parsed.ok()); EXPECT_EQ(parsed->aot_compiled, GetParam()); EXPECT_TRUE([parsed->export_hash isEqual:dict[@"export_hash"]]); - EXPECT_TRUE([parsed->function isEqual:dict[@"function"]]); + ASSERT_EQ(parsed->functions.size(), 1u); + EXPECT_TRUE([parsed->functions[0].name isEqual:dict[@"function"]]); + EXPECT_TRUE(parsed->functions[0].input_shapes.empty()); + EXPECT_FALSE(parsed->enumerated); EXPECT_TRUE([parsed->inputs isEqual:dict[@"input_names"]]); EXPECT_TRUE([parsed->outputs isEqual:dict[@"output_names"]]); EXPECT_TRUE([parsed->files isEqual:dict[@"files"]]); @@ -397,7 +526,8 @@ EXPECT_EQ(selected->files.count, 2); NSString* bundle = selected->path.lastPathComponent; EXPECT_TRUE([selected->export_hash isEqual:parsed->export_hash]); - EXPECT_TRUE([selected->function isEqual:parsed->function]); + EXPECT_TRUE( + [selected->functions[0].name isEqual:parsed->functions[0].name]); EXPECT_TRUE([selected->inputs isEqual:parsed->inputs]); EXPECT_TRUE([selected->outputs isEqual:parsed->outputs]); EXPECT_EQ(selected->aot_compiled, parsed->aot_compiled); @@ -424,7 +554,7 @@ ASSERT_TRUE(selected.ok()); EXPECT_TRUE([selected->path isEqualToString:dict[@"path"]]); EXPECT_TRUE([selected->export_hash isEqual:parsed->export_hash]); - EXPECT_TRUE([selected->function isEqual:parsed->function]); + EXPECT_TRUE([selected->functions[0].name isEqual:parsed->functions[0].name]); EXPECT_TRUE([selected->inputs isEqual:parsed->inputs]); EXPECT_TRUE([selected->outputs isEqual:parsed->outputs]); EXPECT_TRUE([selected->files isEqual:parsed->files]); 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..7dd3747c4a5 --- /dev/null +++ b/backends/apple/coreai/test/test_multifunction.py @@ -0,0 +1,435 @@ +# 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 ( + _check_unique_functions, + 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 _lower_enumerated(*, graph_break=False, model=None, rows=(4, 8, 16)): + """``forward`` is static; ``prefill`` has a dynamic row count it enumerates.""" + torch.manual_seed(0) + model = model or _Stack().eval() + dim = torch.export.Dim("rows", min=1, max=64) + passes = list(get_default_passes()) + ([_BreakAtMul()] if graph_break else []) + return to_edge_transform_and_lower( + { + "forward": torch.export.export(model, (torch.randn(1, _WIDTH),)), + "prefill": torch.export.export( + model, (torch.randn(8, _WIDTH),), dynamic_shapes=({0: dim},) + ), + }, + transform_passes=passes, + partitioner={ + "forward": [CoreAIPartitioner(multifunction=True)], + "prefill": [ + CoreAIPartitioner( + multifunction=True, input_enumerations=[{0: list(rows)}] + ) + ], + }, + 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) + + +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") + ) + + +class MultifunctionEnumeratedShapesTest(unittest.TestCase): + def _shapes_by_function(self, manifest, prefix): + functions = manifest["enumerated_functions"] + for entry in functions: + self.assertTrue(entry["function"].startswith(prefix), entry["function"]) + return sorted(tuple(map(tuple, f["input_shapes"])) for f in functions) + + def test_enumerated_method_lists_per_shape_functions(self): + lowered = _lower_enumerated() + [(forward, _)] = _manifests(lowered, "forward") + [(prefill, _)] = _manifests(lowered, "prefill") + self.assertEqual(forward["function"], "forward_0") + self.assertNotIn("enumerated_functions", forward) + self.assertNotIn("function", prefill) + self.assertEqual( + self._shapes_by_function(prefill, "prefill_0_"), + [((4, _WIDTH),), ((8, _WIDTH),), ((16, _WIDTH),)], + ) + self.assertEqual(forward["export_hash"], prefill["export_hash"]) + self.assertTrue(forward["multifunction"] and prefill["multifunction"]) + + def test_graph_break_enumerates_every_partition(self): + lowered = _lower_enumerated(graph_break=True, rows=(4, 16)) + prefill = [m for m, _ in _manifests(lowered, "prefill")] + self.assertEqual(len(prefill), 2) + for index, manifest in enumerate(prefill): + self.assertEqual( + self._shapes_by_function(manifest, f"prefill_{index}_"), + [((4, _WIDTH),), ((16, _WIDTH),)], + ) + forward = [m["function"] for m, _ in _manifests(lowered, "forward")] + self.assertEqual(forward, ["forward_0", "forward_1"]) + + def test_mixed_partitioners_name_the_missing_flag(self): + with self.assertRaisesRegex( + ValueError, r"multifunction=True must be set.*\['prefill'\]" + ): + _lower( + { + "forward": [CoreAIPartitioner(multifunction=True)], + "prefill": [CoreAIPartitioner()], + } + ) + + def test_duplicate_function_names_are_rejected(self): + _check_unique_functions( + [ + {"function": "forward_0"}, + {"enumerated_functions": [{"function": "forward_0_s0_4"}]}, + ] + ) + with self.assertRaisesRegex(ValueError, "'forward_0_s0_4' twice"): + _check_unique_functions( + [ + {"enumerated_functions": [{"function": "forward_0_s0_4"}]}, + {"function": "forward_0_s0_4"}, + ] + ) + + def test_other_compile_specs_must_still_match(self): + dim = torch.export.Dim("rows", min=1, max=64) + model = _Stack().eval() + with self.assertRaisesRegex(ValueError, "identical Core AI compile specs"): + to_edge_transform_and_lower( + { + "forward": torch.export.export(model, (torch.randn(1, _WIDTH),)), + "prefill": torch.export.export( + model, (torch.randn(8, _WIDTH),), dynamic_shapes=({0: dim},) + ), + }, + transform_passes=get_default_passes(), + partitioner={ + "forward": [CoreAIPartitioner(multifunction=True)], + "prefill": [ + CoreAIPartitioner( + multifunction=True, + input_enumerations=[{0: [4, 8]}], + min_deployment_version="27.1", + ) + ], + }, + compile_config=get_default_compile_config(), + ) + + +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) + + def test_enumerated_shapes_dispatch_by_input_shape(self): + torch.manual_seed(0) + model = _Stack().eval() + lowered = _lower_enumerated(graph_break=True, model=model) + from executorch.extension.pybindings.portable_lib import ( + _load_program_from_buffer, + Verification, + ) + + buffer = lowered.to_executorch().buffer + with tempfile.TemporaryDirectory(prefix="coreai-enumerated-") as cache: + program = _load_program_from_buffer( + buffer, + program_verification=Verification.Minimal, + backend_options={"CoreAIBackend": {"coreai_cache_dir": cache}}, + ) + forward = program.load_method("forward") + prefill = program.load_method("prefill") + cases = [(forward, 1)] + [(prefill, rows) for rows in (4, 8, 16)] + with torch.no_grad(): + for method, rows in cases: + x = torch.randn(rows, _WIDTH) + torch.testing.assert_close( + method.call((x,), clone_outputs=True)[0], + model(x), + atol=1e-2, + rtol=1e-2, + msg=f"{rows} rows", + ) + with self.assertRaises(RuntimeError): + prefill.call((torch.randn(5, _WIDTH),), clone_outputs=True) + del forward, prefill, program + + +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(): diff --git a/backends/apple/coreai/test/test_shape_config.py b/backends/apple/coreai/test/test_shape_config.py index 70dd9963b30..e6dc7b68955 100644 --- a/backends/apple/coreai/test/test_shape_config.py +++ b/backends/apple/coreai/test/test_shape_config.py @@ -4,12 +4,18 @@ # 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 tempfile import unittest +from types import SimpleNamespace + +import sympy import torch import torch.nn as nn from executorch.backends.apple.coreai.compiler.enumerated_shapes import ( + _eval_dim, resolve_input_enumerations, ) from executorch.backends.apple.coreai.partition.partitioner import ( @@ -18,7 +24,10 @@ ) from executorch.exir import to_edge_transform_and_lower from executorch.exir.dialects.edge._ops import EdgeOpOverload -from executorch.exir.lowered_backend_module import executorch_call_delegate +from executorch.exir.lowered_backend_module import ( + executorch_call_delegate, + get_lowered_backend_modules, +) from executorch.exir.pass_base import PassResult @@ -250,10 +259,100 @@ def test_static_dim_fails_lowering(self): partitioner=[CoreAIPartitioner(input_enumerations=[{0: [4, 16]}])], ) + def test_manifest_lists_one_function_per_shape(self): + # The dynamic ``main`` is erased by the specialization, so the manifest + # must name the per-shape functions the runtime can actually bind. + graph = self._lower([{0: [4, 16, 32]}]).exported_program().graph_module + [lbm] = get_lowered_backend_modules(graph) + manifest = json.loads(bytes(lbm.processed_bytes)) + self.assertNotIn("function", manifest) + functions = manifest["enumerated_functions"] + self.assertEqual(len({f["function"] for f in functions}), 3) + for entry in functions: + self.assertTrue(entry["function"].startswith("main_"), entry["function"]) + self.assertEqual( + sorted(tuple(map(tuple, f["input_shapes"])) for f in functions), + [((4, 8),), ((16, 8),), ((32, 8),)], + ) + self.assertEqual(manifest["input_names"], ["input_0"]) + + def test_unenumerated_lowering_keeps_main(self): + graph = ( + to_edge_transform_and_lower( + _dynamic_ep(), partitioner=[CoreAIPartitioner()] + ) + .exported_program() + .graph_module + ) + [lbm] = get_lowered_backend_modules(graph) + manifest = json.loads(bytes(lbm.processed_bytes)) + self.assertEqual(manifest["function"], "main") + self.assertNotIn("enumerated_functions", manifest) + + +def _runtime_available(): + try: + from executorch.extension.pybindings.portable_lib import _is_available + except ImportError: + return False + return _is_available("CoreAIBackend") + + +@unittest.skipUnless(_runtime_available(), "requires the Core AI Python runtime") +class InputEnumerationsRuntimeTest(unittest.TestCase): + def test_runs_each_enumerated_shape_and_rejects_others(self): + from executorch.extension.pybindings.portable_lib import ( + _load_program_from_buffer, + Verification, + ) + + torch.manual_seed(0) + model = _model() + batch = torch.export.Dim("batch", min=1, max=64) + ep = torch.export.export( + model, (torch.randn(2, 8),), dynamic_shapes={"input": {0: batch}} + ) + buffer = ( + to_edge_transform_and_lower( + ep, + partitioner=[CoreAIPartitioner(input_enumerations=[{0: [4, 16, 32]}])], + ) + .to_executorch() + .buffer + ) + with tempfile.TemporaryDirectory(prefix="coreai-enumerated-") as cache: + program = _load_program_from_buffer( + buffer, + program_verification=Verification.Minimal, + backend_options={"CoreAIBackend": {"coreai_cache_dir": cache}}, + ) + method = program.load_method("forward") + with torch.no_grad(): + for rows in (4, 16, 32, 4): + x = torch.randn(rows, 8) + torch.testing.assert_close( + method.call((x,), clone_outputs=True)[0], + model(x), + atol=1e-3, + rtol=1e-3, + msg=f"{rows} rows", + ) + with self.assertRaises(RuntimeError): + method.call((torch.randn(8, 8),), clone_outputs=True) + del method, program + class InputEnumerationsValidationTest(unittest.TestCase): """User-supplied enumerations that cannot be honoured must say so.""" + def test_derived_dim_that_reaches_zero_is_rejected(self): + dim = SimpleNamespace(node=SimpleNamespace(expr=sympy.Symbol("s0") - 1)) + self.assertEqual(_eval_dim(dim, {"s0": 4}, "input_0", 0), 3) + with self.assertRaisesRegex( + ValueError, r"input 'input_0' dim 0 is s0 - 1, which is 0 when s0=1" + ): + _eval_dim(dim, {"s0": 1}, "input_0", 0) + def test_conflicting_values_for_one_symbol_are_rejected(self): """Two inputs tied to the same Dim cannot enumerate differently.