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 b4912248fa8..3b1dbfef0ed 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -26,14 +26,18 @@ python -m pytest backends/apple/coreai ## Runtime -The delegate manifest describes `.aimodel` source bundles, which require -`inline` packaging, and AOT-compiled `.aimodelc` bundles, which require -`aot_compiled_inline`. Both formats require a `function` name, and ordered +The delegate manifest's `packaging` field names the embedded format: `aimodel` +for a portable `.aimodel` bundle, located by `path`, or `aimodelc` for +AOT-compiled per-architecture `.aimodelc` bundles, located by `archs`. Both +formats require exactly one of a `function` name or `enumerated_functions` (see +[Enumerated input shapes](#enumerated-input-shapes)), and ordered `input_names` and `output_names` matching the converter's bindings. Dictionary iteration order is not a binding contract. `export_hash` identifies the exported program used to namespace embedded assets. The `files` object maps relative filenames to byte sizes; `bundle_digests` maps each bundle basename to its -export-time SHA-256 digest. Asset paths, including bundle and directory names, +export-time SHA-256 digest. The optional boolean `multifunction` marks a delegate +exported as one function of a model shared with other delegates. Asset paths, +including bundle and directory names, must remain distinct under case-insensitive, canonically normalized Unicode comparison. Accepted names are preserved, not rewritten. @@ -107,6 +111,88 @@ bookmark restoration. A missing bookmark or SDK-confirmed cache miss falls back to source materialization and persistent specialization, then binds the requested function and ordered I/O on that same acquired model. +### Multi-function export + +By default each Core AI delegate is its own model with a single `main` function. +`CoreAIPartitioner(multifunction=True)` instead lowers every Core AI partition of +every method into one model, with one function per partition named +`_`: + +```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 @@ -128,9 +214,9 @@ even though the new complete contents are already visible. Use a local APFS volume; unsupported full-device flushes are reported as errors, not downgraded to `fsync`. -Inline bundles are read through `NamedDataMap` and materialized as unchanged -files under the cache directory. Named data can be supplied externally; inline -packaging does not require all bytes to reside in one physical PTE. The exporter +Embedded bundles are read through `NamedDataMap` and materialized as unchanged +files under the cache directory. Named data can be supplied externally; embedded +bundles do not require all bytes to reside in one physical PTE. computes each bundle digest from all delivered relative filenames and file bytes, not just the SDK's bytecode-only `main.hash`. Each AOT architecture has an independent digest; only the selected architecture is materialized. NDS key names @@ -191,7 +277,7 @@ choose the SDK artifact directory. The backend does not inspect SDK-private files, alter their backup flags, or use an App Group cache. Each selected bundle has a `bundle_identity` derived from its bundle digest, -name, and source/AOT mode. This identifies the delivered bundle independently of +name, and packaging (`aimodel` or `aimodelc`). This identifies the delivered bundle independently of `export_hash`, which namespaces embedded assets. A readable cache key combines that bundle identity with the platform and Core AI device architecture, plus a hash of a versioned @@ -338,6 +424,69 @@ Use separate option maps to choose different cache directories for different mod The path names the cache directory, not a bundle. ExecuTorch option strings have a 255-byte limit. +`coreai_compute_unit` selects the Core AI specialization options. Each value is +one configuration the SDK provides: + +| Value | Core AI `SpecializationOptions` | Allowed compute units | +|---|---|---| +| `default` (or absent) | `.default` | CPU, GPU, Neural Engine | +| `cpu_only` | `.cpuOnly` | CPU | +| `prefer_cpu` | `init(preferredComputeUnitKind: .cpu)` | CPU, GPU, Neural Engine | +| `prefer_gpu` | `init(preferredComputeUnitKind: .gpu)` | CPU, GPU, Neural Engine | +| `prefer_neural_engine` | `init(preferredComputeUnitKind: .neuralEngine)` | CPU, GPU, Neural Engine | + +`cpu_only` is the only restriction: a preference maximizes use of that unit but +falls back to any allowed unit for operations it cannot run. The SDK has no +other combinations, such as CPU and Neural Engine only. Unknown values fail +initialization with `InvalidArgument`. The selection is part of the cache key, +so each value is specialized and cached separately; clearing a PTE removes every +variant. + +```cpp +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 0bc064f2fbb..dd67cb8e6a9 100644 --- a/backends/apple/coreai/compiler/preprocess.py +++ b/backends/apple/coreai/compiler/preprocess.py @@ -75,12 +75,17 @@ 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. - 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) @@ -195,32 +200,44 @@ 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, - } - if packaging == AssetPackaging.INLINE: + manifest = {**manifest_extra, **metadata} + if packaging == AssetPackaging.AIMODEL: manifest["path"] = f"{export_hash}/{manifest['path']}" else: manifest["archs"] = { 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.INLINE, + 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.AOT_COMPILED_INLINE, + 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 a8fb3d2cda5..ad66ecdc236 100644 --- a/backends/apple/coreai/runtime/ETCoreAIBridge.h +++ b/backends/apple/coreai/runtime/ETCoreAIBridge.h @@ -21,6 +21,17 @@ typedef NS_ENUM(NSInteger, ETCoreAIScalarType) { ETCoreAIScalarTypeBool = 3, }; +// One value per Core AI SpecializationOptions configuration: the default (all +// compute units), cpuOnly, or a preferred compute unit kind. Preferences still +// allow every unit for operations the preferred unit cannot run. +typedef NS_ENUM(NSInteger, ETCoreAIComputeUnit) { + ETCoreAIComputeUnitDefault = 0, + ETCoreAIComputeUnitCPUOnly = 1, + ETCoreAIComputeUnitPreferCPU = 2, + ETCoreAIComputeUnitPreferGPU = 3, + ETCoreAIComputeUnitPreferNeuralEngine = 4, +}; + typedef NS_ENUM(NSInteger, ETCoreAIErrorCode) { // Invalid model descriptors, functions or bindings. ETCoreAIErrorInvalidModel = 1, @@ -136,7 +147,9 @@ extern "C" { #endif BOOL ETCoreAIIsAvailable(void); NSString* _Nullable ETCoreAIDeviceArchitectureName(void); -id ETCoreAICreateModelLoader(void); +id ETCoreAICreateModelLoader( + ETCoreAIComputeUnit computeUnit, + BOOL expectFrequentReshapes); #ifdef __cplusplus } #endif diff --git a/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift b/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift index adf80fbc91f..69251f231f8 100644 --- a/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift +++ b/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift @@ -238,6 +238,36 @@ public final class ETCoreAISwiftModelLoader: NSObject, ETCoreAIModelLoading { super.init() } + @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 + } + self.init(specialization: specialization) + } + + static func specializationOptions( + for computeUnit: ETCoreAIComputeUnit, expectFrequentReshapes: Bool = false + ) -> SpecializationOptions? { + var options: SpecializationOptions + switch computeUnit { + 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 } @objc public static var deviceArchitectureName: String { AIModel.deviceArchitectureName } diff --git a/backends/apple/coreai/runtime/coreai_backend.mm b/backends/apple/coreai/runtime/coreai_backend.mm index 7b52edbd767..069a08bc1ad 100644 --- a/backends/apple/coreai/runtime/coreai_backend.mm +++ b/backends/apple/coreai/runtime/coreai_backend.mm @@ -14,41 +14,88 @@ #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 = + runtime_option(context, "coreai_compute_unit", "default"); + if (!option.ok()) { + return option.error(); + } + static constexpr struct { + const char* name; + ETCoreAIComputeUnit value; + } kComputeUnits[] = { + {"default", ETCoreAIComputeUnitDefault}, + {"cpu_only", ETCoreAIComputeUnitCPUOnly}, + {"prefer_cpu", ETCoreAIComputeUnitPreferCPU}, + {"prefer_gpu", ETCoreAIComputeUnitPreferGPU}, + {"prefer_neural_engine", ETCoreAIComputeUnitPreferNeuralEngine}, + }; + for (const auto& unit : kComputeUnits) { + if (std::strcmp(option.get(), unit.name) == 0) { + return unit.value; + } + } + ET_LOG( + Error, + "Unknown Core AI compute unit '%s'; expected default, cpu_only, " + "prefer_cpu, prefer_gpu or prefer_neural_engine", + option.get()); + return Error::InvalidArgument; +} + 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) { @@ -70,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 { @@ -97,33 +193,59 @@ bool is_available() const override { if (!configured.ok()) { return configured.error(); } - id loader = ETCoreAICreateModelLoader(); + auto compute_unit = compute_unit_option(context); + 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(), 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()); } @@ -137,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 539bdb408cb..557d1247398 100644 --- a/backends/apple/coreai/runtime/coreai_bridge.mm +++ b/backends/apple/coreai/runtime/coreai_bridge.mm @@ -24,6 +24,10 @@ BOOL ETCoreAIIsAvailable(void) { return nil; } -id ETCoreAICreateModelLoader(void) { - return [[ETCoreAISwiftModelLoader alloc] init]; +id ETCoreAICreateModelLoader( + ETCoreAIComputeUnit computeUnit, + BOOL expectFrequentReshapes) { + return [[ETCoreAISwiftModelLoader alloc] + initWithComputeUnit:computeUnit + expectFrequentReshapes:expectFrequentReshapes]; } 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_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 bd4f45a7de1..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()) { @@ -206,7 +232,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]; }); @@ -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, @@ -255,7 +371,9 @@ Error clear_cache_entry( } else { if (loader == nil) { ET_CHECK_OK_OR_RETURN_ERROR(require_coreai_available()); - loader = ETCoreAICreateModelLoader(); + // Eviction targets the bookmark, so the loader's specialization is + // irrelevant. + 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 af139c8f09d..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())] != @@ -174,15 +246,21 @@ 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"]; + 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 fbbe6274ef7..745b07f84fa 100644 --- a/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift +++ b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift @@ -153,6 +153,65 @@ final class SwiftSpecializationTests: XCTestCase { loader.specializationIdentifier, try CoreAISpecialization(options: options).identifier) } + func testComputeUnitsMapToSDKSpecializationOptions() throws { + let expected: [(ETCoreAIComputeUnit, SpecializationOptions)] = [ + (.default, .default), + (.cpuOnly, .cpuOnly), + (.preferCPU, SpecializationOptions(preferredComputeUnitKind: .cpu)), + (.preferGPU, SpecializationOptions(preferredComputeUnitKind: .gpu)), + (.preferNeuralEngine, SpecializationOptions(preferredComputeUnitKind: .neuralEngine)), + ] + var variants: [(options: SpecializationOptions, identifier: String)] = [] + var report: [String] = [] + for (unit, options) in expected { + XCTAssertEqual(ETCoreAISwiftModelLoader.specializationOptions(for: unit), options) + 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, 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 + // cpuOnly equal default), so identifiers must match exactly when options do. + private func assertIdentifiersMatchExactlyWhenOptionsMatch( + _ variants: [(options: SpecializationOptions, identifier: String)], + report: [String], + file: StaticString = #filePath, + line: UInt = #line + ) { + for a in variants.indices { + for b in variants.indices where a < b { + XCTAssertEqual( + variants[a].identifier == variants[b].identifier, + variants[a].options == variants[b].options, + "variants \(a) and \(b):\n" + report.joined(separator: "\n"), + file: file, line: line) + } + } + } + func testDescriptorIsVersionedAndUsesStableComputeUnitNames() throws { var options = SpecializationOptions(preferredComputeUnitKind: .gpu) options.expectFrequentReshapes = !SpecializationOptions.default.expectFrequentReshapes 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 c637ebd3374..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 { @@ -153,7 +159,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( @@ -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,25 +1000,349 @@ TensorImpl output( } } -TEST_F(CoreAIDelegateTest, StorageOptionsRejectWrongTypeAndRelativePaths) { +TEST_F(CoreAIDelegateTest, OptionsRejectWrongTypes) { + @autoreleasepool { + 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({&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); + } +} + +TEST_F(CoreAIDelegateTest, ComputeUnitOptionSelectsSpecialization) { + @autoreleasepool { + // A successful init frees the processed manifest, so each case gets a + // fresh fixture. + auto init_with_compute_unit = [&](const char* name) { + bridge.state().last_compute_unit = -1; + DelegateFixture fixture(backend, bridge); + BackendOptions<2> options; + if (name != nullptr) { + EXPECT_EQ(options.set_option("coreai_compute_unit", name), Error::Ok); + } + return fixture.init_with_storage(options); + }; { - SCOPED_TRACE("integer cache directory"); - BackendOption wrong_type{"coreai_cache_dir", 1}; + SCOPED_TRACE("absent"); + ASSERT_EQ(init_with_compute_unit(nullptr), Error::Ok); EXPECT_EQ( - fixture.init_with_options({&wrong_type, 1}), Error::InvalidArgument); + bridge.state().last_compute_unit.load(), ETCoreAIComputeUnitDefault); } - { - SCOPED_TRACE("relative cache directory"); - BackendOptions<1> relative; - ASSERT_EQ(relative.set_option("coreai_cache_dir", "relative"), Error::Ok); + const std::pair valid[] = { + {"default", ETCoreAIComputeUnitDefault}, + {"cpu_only", ETCoreAIComputeUnitCPUOnly}, + {"prefer_cpu", ETCoreAIComputeUnitPreferCPU}, + {"prefer_gpu", ETCoreAIComputeUnitPreferGPU}, + {"prefer_neural_engine", ETCoreAIComputeUnitPreferNeuralEngine}, + }; + for (const auto& [name, unit] : valid) { + SCOPED_TRACE(name); + ASSERT_EQ(init_with_compute_unit(name), Error::Ok); + EXPECT_EQ(bridge.state().last_compute_unit.load(), unit); + } + for (const char* name : {"gpu", "neural_engine", "CPU_ONLY", ""}) { + SCOPED_TRACE(name); + EXPECT_EQ(init_with_compute_unit(name), Error::InvalidArgument); + EXPECT_EQ(bridge.state().last_compute_unit.load(), -1); + } + } +} + +TEST_F(CoreAIDelegateTest, ExpectFrequentReshapesOptionReachesLoader) { + @autoreleasepool { + auto init_with_reshapes = [&](std::optional reshapes) { + bridge.state().last_expect_frequent_reshapes = -1; + DelegateFixture fixture(backend, bridge); + BackendOptions<3> options; EXPECT_EQ( - fixture.init_with_options( - {relative.view().data(), relative.view().size()}), - Error::InvalidArgument); + 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); } } diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.h b/backends/apple/coreai/runtime/test/coreai_fake_loader.h index 5fb287ed79a..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,17 +31,29 @@ 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}; std::atomic sessions{0}; 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}; @@ -71,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 5eefda12f9d..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); }); } @@ -537,9 +561,14 @@ BOOL ETCoreAIIsAvailable(void) { return current_context->state.device_architecture; } -id ETCoreAICreateModelLoader(void) { +id ETCoreAICreateModelLoader( + 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_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_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/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/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 2d8b0d39fd5..32fb21ba229 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,9 @@ 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}, [{}] + )[0] manifest = json.loads(result.processed_bytes) manifests.append(manifest) self.assertEqual( @@ -241,8 +243,8 @@ def test_aot_namespace_covers_every_architecture(self): for payload in (b"original", b"modified"): (root / archs["h16"] / "model.mil").write_bytes(payload) result = _deliver( - root, AssetPackaging.AOT_COMPILED_INLINE, {"archs": archs} - ) + root, AssetPackaging.AIMODELC, {"archs": archs}, [{}] + )[0] manifest = json.loads(result.processed_bytes) manifests.append(manifest) for arch, name in archs.items(): @@ -378,7 +380,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 +402,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 +450,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"])) 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.