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