From 00bfbe65884cd3f896eaf209e11f3bfbb0a9d318 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 3 Nov 2026 12:03:36 -0800 Subject: [PATCH 1/3] [CoreAI] Run enumerated-shape exports by dispatching on input shape Enumerated-shape exports could not load: set_static_shape_config erases the delegate's dynamic function and replaces it with one static function per shape, but the manifest still named the erased function, so every load failed with "Model has no function named 'main'". Export now lists the per-shape functions under enumerated_functions, each with its input shapes in input_names order, instead of a single function. The runtime binds every listed function at init and, per call, runs the one whose shapes equal the inputs; inputs matching none fail with InvalidArgument before anything runs. All functions come from one model, so they share weights, one specialization and one cache entry. Enumerations also work with multifunction export now. Each method's partitioner may enumerate its own values (or none), and each enumerated partition expands into __ functions in the shared model. Every other compile spec must still match across methods. --- backends/apple/coreai/README.md | 55 +++++++- .../coreai/compiler/enumerated_shapes.py | 16 ++- backends/apple/coreai/compiler/preprocess.py | 89 ++++++++++--- .../apple/coreai/partition/partitioner.py | 2 - .../apple/coreai/runtime/coreai_backend.mm | 96 ++++++++++++-- .../apple/coreai/runtime/coreai_manifest.h | 17 ++- .../apple/coreai/runtime/coreai_manifest.mm | 93 ++++++++++++- .../runtime/test/coreai_delegate_test.mm | 89 +++++++++++++ .../coreai/runtime/test/coreai_fake_loader.h | 8 ++ .../coreai/runtime/test/coreai_fake_loader.mm | 12 +- .../runtime/test/coreai_manifest_test.mm | 115 ++++++++++++++++- .../apple/coreai/test/test_multifunction.py | 122 +++++++++++++++++- .../apple/coreai/test/test_shape_config.py | 89 ++++++++++++- 13 files changed, 748 insertions(+), 55 deletions(-) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index f8ca55f632a..fee1bf8f66a 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -29,7 +29,8 @@ python -m pytest backends/apple/coreai The delegate manifest's `packaging` field names the embedded format: `aimodel` for a portable `.aimodel` bundle, located by `path`, or `aimodelc` for AOT-compiled per-architecture `.aimodelc` bundles, located by `archs`. Both -formats require a `function` name, and ordered +formats require exactly one of a `function` name or `enumerated_functions` (see +[Enumerated input shapes](#enumerated-input-shapes)), and ordered `input_names` and `output_names` matching the converter's bindings. Dictionary iteration order is not a binding contract. `export_hash` identifies the exported program used to namespace embedded assets. The `files` object maps relative @@ -129,7 +130,8 @@ lowered = to_edge_transform_and_lower( Here the delegates bind `forward_0` and `prefill_0`; a graph break adds `forward_1`, `prefill_1` and so on. Weights shared between methods are stored once. Every partition must carry identical compile specs (AOT config, deployment -floor), and `input_enumerations` is not supported with this option yet. +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 @@ -139,6 +141,55 @@ 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. + +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 diff --git a/backends/apple/coreai/compiler/enumerated_shapes.py b/backends/apple/coreai/compiler/enumerated_shapes.py index e4ccad15d28..c1bb8d06801 100644 --- a/backends/apple/coreai/compiler/enumerated_shapes.py +++ b/backends/apple/coreai/compiler/enumerated_shapes.py @@ -175,7 +175,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 +185,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 +222,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 +248,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/preprocess.py b/backends/apple/coreai/compiler/preprocess.py index 39e5c3e9fa3..1067df79ce9 100644 --- a/backends/apple/coreai/compiler/preprocess.py +++ b/backends/apple/coreai/compiler/preprocess.py @@ -316,8 +316,59 @@ def _binding( } -def _spec_items(compile_specs: List[CompileSpec]) -> List[Tuple[str, bytes]]: - return sorted((spec.key, bytes(spec.value)) for spec in compile_specs) +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 _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 + ) # Asset embedding helpers. @@ -517,16 +568,9 @@ def preprocess( binding = _binding( program, edge_program, MAIN_ENTRYPOINT, input_names, output_names ) - raw_enum = _get_compile_spec( - compile_specs, COMPILE_SPEC_KEYS.INPUT_ENUMERATIONS + binding = _enumerate_binding( + program, edge_program, binding, input_names, compile_specs ) - if raw_enum: - apply_enumerated_shapes( - program, - edge_program, - json.loads(raw_enum.decode()), - input_names=input_names, - ) return CoreAIBackend._package(program, compile_specs, [binding])[0] @classmethod @@ -546,12 +590,15 @@ def preprocess_multimethod( ): return super().preprocess_multimethod(edge_programs, compile_specs) - reference = _spec_items(partitions[0][3]) + # 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) != reference + if _spec_items(specs, ignore) != reference } ) if differing: @@ -559,9 +606,6 @@ def preprocess_multimethod( "multifunction requires identical Core AI compile specs in every " f"partition; methods {differing} differ from {partitions[0][0]!r}" ) - if _get_compile_spec(partitions[0][3], COMPILE_SPEC_KEYS.INPUT_ENUMERATIONS): - raise ValueError("multifunction does not support input_enumerations yet") - entries = [ (program, f"{method}_{index}", *_boundary_names(program)) for method, index, program, _ in partitions @@ -572,8 +616,19 @@ def preprocess_multimethod( [entrypoint for _, entrypoint, _, _ in entries], ) 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 = [ - {**_binding(program, *entry), "multifunction": True} for entry in entries + { + **_enumerate_binding( + program, edge_program, binding, input_names, specs + ), + "multifunction": True, + } + for binding, (edge_program, _, input_names, _), (_, _, _, specs) in zip( + bindings, entries, partitions + ) ] results = CoreAIBackend._package(program, partitions[0][3], bindings) diff --git a/backends/apple/coreai/partition/partitioner.py b/backends/apple/coreai/partition/partitioner.py index 300298fc2be..efd84be7282 100644 --- a/backends/apple/coreai/partition/partitioner.py +++ b/backends/apple/coreai/partition/partitioner.py @@ -333,8 +333,6 @@ def __init__( # # multifunction lowers every Core AI partition of every method into one # Core AI model, one function per partition named _. - if multifunction and input_enumerations is not None: - raise ValueError("multifunction does not support input_enumerations yet") specs = [] if min_deployment_version is not None: specs.append( diff --git a/backends/apple/coreai/runtime/coreai_backend.mm b/backends/apple/coreai/runtime/coreai_backend.mm index 2f8172817dc..ba568fcb485 100644 --- a/backends/apple/coreai/runtime/coreai_backend.mm +++ b/backends/apple/coreai/runtime/coreai_backend.mm @@ -14,19 +14,27 @@ #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; ETCoreAILoadedModel* model = nil; - id session = nil; + std::vector functions; }; template @@ -79,11 +87,12 @@ Result> bind_session( id model, - const SelectedBundle& bundle) { + const SelectedBundle& bundle, + NSString* function) { dispatch_semaphore_t ready = dispatch_semaphore_create(0); __block id session = nil; __block NSError* error = nil; - [model loadFunctionNamed:bundle.function + [model loadFunctionNamed:function inputNames:bundle.inputs outputNames:bundle.outputs completion:^(id result, NSError* failure) { @@ -105,6 +114,68 @@ 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 + "]"; +} + +// Enumerated-shape delegates bind one function per shape, and each accepts +// only its own static shapes, so pick the one these inputs match. +Result> select_session( + const Handle& handle, + Span args) { + if (handle.functions.size() == 1 && + handle.functions.front().input_shapes.empty()) { + return handle.functions.front().session; + } + ET_CHECK_OR_RETURN_ERROR( + args.size() >= handle.input_count, + InvalidArgument, + "Core AI delegate expected %zu inputs", + handle.input_count); + for (size_t index = 0; index < handle.input_count; ++index) { + ET_CHECK_OR_RETURN_ERROR( + args[index] != nullptr && args[index]->isTensor(), + InvalidArgument, + "Core AI input %zu must be a tensor", + index); + } + 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 { @@ -168,19 +239,22 @@ bool is_available() const override { if (!loaded.ok()) { return loaded.error(); } - auto session = bind_session(loaded.get().model, host_bundle->bundle); - if (!session.ok()) { - return session.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->model = loaded.get(); - handle->session = session.get(); processed->Free(); return static_cast(handle.release()); } @@ -194,9 +268,13 @@ Error execute( ET_CHECK_OR_RETURN_ERROR( opaque != nullptr, DelegateInvalidHandle, "Null Core AI handle"); auto& handle = *static_cast(opaque); + auto session = select_session(handle, args); + if (!session.ok()) { + return session.error(); + } std::lock_guard guard([handle.model executionMutex]); return execute_session( - handle.session, args, handle.input_count, handle.output_count); + session.get(), args, handle.input_count, handle.output_count); } } diff --git a/backends/apple/coreai/runtime/coreai_manifest.h b/backends/apple/coreai/runtime/coreai_manifest.h index ef6e3158706..ac4202848cb 100644 --- a/backends/apple/coreai/runtime/coreai_manifest.h +++ b/backends/apple/coreai/runtime/coreai_manifest.h @@ -11,12 +11,23 @@ #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; NSArray* const inputs; NSArray* const outputs; NSDictionary* const files; @@ -37,7 +48,7 @@ class Manifest final { private: Manifest( NSString* export_hash, - NSString* function, + std::vector functions, NSArray* inputs, NSArray* outputs, NSDictionary* files, @@ -57,7 +68,7 @@ class SelectedBundle final { NSString* const path; NSString* const bundle_name; NSString* const digest; - NSString* const function; + const std::vector functions; 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 24da4da3523..2bce1e23227 100644 --- a/backends/apple/coreai/runtime/coreai_manifest.mm +++ b/backends/apple/coreai/runtime/coreai_manifest.mm @@ -9,13 +9,15 @@ #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, NSArray* inputs, NSArray* outputs, NSDictionary* files, @@ -26,7 +28,7 @@ NSOperatingSystemVersion minimum_version, bool multifunction) : export_hash(export_hash), - function(function), + functions(std::move(functions)), inputs(inputs), outputs(outputs), files(files), @@ -45,7 +47,7 @@ path(path), bundle_name(path.lastPathComponent), digest(manifest.bundle_digests[path.lastPathComponent]), - function(manifest.function), + functions(manifest.functions), inputs(manifest.inputs), outputs(manifest.outputs), files([files copy]), @@ -74,6 +76,67 @@ bool is_boolean(id value) { 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] && !is_boolean(value) && !CFNumberIsFloatType((__bridge CFNumberRef)value) && @@ -186,7 +249,8 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { "Unsupported Core AI packaging; expected aimodel or aimodelc"); const bool aot_compiled = [packaging isEqualToString:@"aimodelc"]; NSString* export_hash = dict[@"export_hash"]; - NSString* function = dict[@"function"]; + id function = dict[@"function"]; + id enumerated_functions = dict[@"enumerated_functions"]; NSArray* inputs = dict[@"input_names"]; NSArray* outputs = dict[@"output_names"]; id multifunction = dict[@"multifunction"]; @@ -243,10 +307,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), @@ -269,7 +348,7 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { } return Manifest( export_hash, - function, + std::move(functions), inputs, outputs, files, diff --git a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm index 8fa3f831228..5af523bbe98 100644 --- a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm @@ -248,6 +248,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); diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.h b/backends/apple/coreai/runtime/test/coreai_fake_loader.h index 1b59715e6b7..cf6df68f1ad 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.h +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.h @@ -31,8 +31,15 @@ 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}; @@ -77,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 91ba7af57d9..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) { @@ -146,6 +148,7 @@ - (void)executeInputs:(NSArray*)inputs 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; @@ -158,6 +161,7 @@ - (void)executeInputs:(NSArray*)inputs 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) { @@ -244,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; } @@ -255,6 +261,7 @@ - (void)loadFunctionNamed:(NSString*)functionName } FakeSession* session = [[FakeSession alloc] initWithContext:context]; session.preparedPin = self; + session.function = functionName; completion(session, nil); }); } @@ -365,6 +372,7 @@ - (void)loadFunctionNamed:(NSString*)functionName } FakeSession* session = [[FakeSession alloc] initWithContext:context]; session.preparedPin = self; + session.function = functionName; completion(session, nil); }); } diff --git a/backends/apple/coreai/runtime/test/coreai_manifest_test.mm b/backends/apple/coreai/runtime/test/coreai_manifest_test.mm index 259dc91ba53..dd52cb0174b 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,110 @@ 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()); + 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()); @@ -190,7 +296,9 @@ 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_TRUE([parsed->inputs isEqual:dict[@"input_names"]]); EXPECT_TRUE([parsed->outputs isEqual:dict[@"output_names"]]); EXPECT_TRUE([parsed->files isEqual:dict[@"files"]]); @@ -416,7 +524,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); @@ -443,7 +552,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/test/test_multifunction.py b/backends/apple/coreai/test/test_multifunction.py index 28f27dcc997..a9a990ecd46 100644 --- a/backends/apple/coreai/test/test_multifunction.py +++ b/backends/apple/coreai/test/test_multifunction.py @@ -77,6 +77,32 @@ def _lower(partitioner, *, graph_break=False, model=None, inputs=None): ) +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 [ @@ -103,10 +129,6 @@ def test_option_adds_compile_spec(self): ] self.assertNotIn(COMPILE_SPEC_KEYS.MULTIFUNCTION.value, default_keys) - def test_rejects_input_enumerations(self): - with self.assertRaisesRegex(ValueError, "input_enumerations"): - CoreAIPartitioner(multifunction=True, input_enumerations=[{0: [1, 8]}]) - class MultifunctionPreprocessTest(unittest.TestCase): def test_methods_share_one_bundle(self): @@ -210,6 +232,65 @@ def fake_build(aimodel_path, out_dir, opts): ) +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_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 @@ -287,6 +368,39 @@ def test_each_method_runs_its_function(self): 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_shape_config.py b/backends/apple/coreai/test/test_shape_config.py index 70dd9963b30..ed0bcdf3bc2 100644 --- a/backends/apple/coreai/test/test_shape_config.py +++ b/backends/apple/coreai/test/test_shape_config.py @@ -4,6 +4,8 @@ # 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 import torch @@ -18,7 +20,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,6 +255,88 @@ 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.""" From 215f2aa9580e8f5be886f01669cfba92e19aa530 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 3 Nov 2026 13:15:13 -0800 Subject: [PATCH 2/3] [CoreAI] Clearer export errors for enumerated multi-function exports - Mixing multifunction and plain CoreAIPartitioners across methods now names the methods missing multifunction=True, instead of reporting that compile specs differ. - Multi-function export rejects two delegates generating the same function name, which a method named like another method's _ could cause. - A derived boundary dim that an enumerated value makes non-positive (s0 - 1 at s0=1) is rejected naming the input, dim and value, instead of failing in set_static_shape_config without them. - README notes that enumerated function counts multiply across dims, with measured load times. --- backends/apple/coreai/README.md | 5 ++- .../coreai/compiler/enumerated_shapes.py | 10 +++++- backends/apple/coreai/compiler/preprocess.py | 36 +++++++++++++++++++ .../apple/coreai/test/test_multifunction.py | 31 +++++++++++++++- .../apple/coreai/test/test_shape_config.py | 12 +++++++ 5 files changed, 91 insertions(+), 3 deletions(-) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index fee1bf8f66a..665632b2977 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -174,7 +174,10 @@ combination of enumerated values, so the manifest lists them under 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. +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 256→1024→256 MLP 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: diff --git a/backends/apple/coreai/compiler/enumerated_shapes.py b/backends/apple/coreai/compiler/enumerated_shapes.py index c1bb8d06801..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( diff --git a/backends/apple/coreai/compiler/preprocess.py b/backends/apple/coreai/compiler/preprocess.py index 1067df79ce9..dd67cb8e6a9 100644 --- a/backends/apple/coreai/compiler/preprocess.py +++ b/backends/apple/coreai/compiler/preprocess.py @@ -361,6 +361,28 @@ def _enumerate_binding( 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]]: @@ -590,6 +612,19 @@ def preprocess_multimethod( ): 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( + "multifunction=True must be set on the CoreAIPartitioner of every " + f"method lowered to Core AI; methods {missing} were lowered " + "without it" + ) # 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,) @@ -630,6 +665,7 @@ def preprocess_multimethod( bindings, entries, partitions ) ] + _check_unique_functions(bindings) results = CoreAIBackend._package(program, partitions[0][3], bindings) by_method: Dict[str, List[PreprocessResult]] = { diff --git a/backends/apple/coreai/test/test_multifunction.py b/backends/apple/coreai/test/test_multifunction.py index a9a990ecd46..7dd3747c4a5 100644 --- a/backends/apple/coreai/test/test_multifunction.py +++ b/backends/apple/coreai/test/test_multifunction.py @@ -17,7 +17,10 @@ get_default_compile_config, get_default_passes, ) -from executorch.backends.apple.coreai.compiler.preprocess import COMPILE_SPEC_KEYS +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, @@ -265,6 +268,32 @@ def test_graph_break_enumerates_every_partition(self): 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() diff --git a/backends/apple/coreai/test/test_shape_config.py b/backends/apple/coreai/test/test_shape_config.py index ed0bcdf3bc2..e6dc7b68955 100644 --- a/backends/apple/coreai/test/test_shape_config.py +++ b/backends/apple/coreai/test/test_shape_config.py @@ -7,11 +7,15 @@ 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 ( @@ -341,6 +345,14 @@ def test_runs_each_enumerated_shape_and_rejects_others(self): 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. From 6eab9896eba4b1d3efa13ae7b3021eaebbf58dc5 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 3 Nov 2026 14:49:51 -0800 Subject: [PATCH 3/3] [CoreAI] Address review feedback on enumerated dispatch and shared loads - Manifests mark enumerated delegates with an explicit flag that the handle carries, instead of select_session inferring it from empty input shapes. - execute_session validates arguments once, then asks the delegate to pick a session for them, so select_session no longer repeats the tensor checks. - A scope guard clears a shared load's in-flight key and wakes its waiters on every exit, so a key cannot stay marked as loading. - WaiterAcquiresAfterConcurrentSharedLoadFails observes the waiter reaching the wait (new BeforeSharedLoadWait fault point) before failing the first load, instead of relying on a sleep. - Partitioner comment notes that to_backend() would also split multifunction exports into separate models. - README drops arrows from prose. --- backends/apple/coreai/README.md | 2 +- .../apple/coreai/partition/partitioner.py | 2 ++ .../apple/coreai/runtime/coreai_backend.mm | 31 ++++++---------- .../coreai/runtime/coreai_fault_injection.h | 2 ++ .../coreai/runtime/coreai_load_coordinator.mm | 36 +++++++++++-------- .../apple/coreai/runtime/coreai_manifest.h | 4 +++ .../apple/coreai/runtime/coreai_manifest.mm | 4 +++ .../apple/coreai/runtime/coreai_tensor_io.h | 10 ++++-- .../apple/coreai/runtime/coreai_tensor_io.mm | 13 ++++--- .../runtime/test/coreai_delegate_test.mm | 21 +++++++++-- .../runtime/test/coreai_manifest_test.mm | 2 ++ 11 files changed, 82 insertions(+), 45 deletions(-) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 665632b2977..3b1dbfef0ed 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -176,7 +176,7 @@ 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 256→1024→256 MLP took 7 ms with 1 function, 41 ms with 16 and 248 ms +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 diff --git a/backends/apple/coreai/partition/partitioner.py b/backends/apple/coreai/partition/partitioner.py index efd84be7282..e04a73a76aa 100644 --- a/backends/apple/coreai/partition/partitioner.py +++ b/backends/apple/coreai/partition/partitioner.py @@ -371,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/coreai_backend.mm b/backends/apple/coreai/runtime/coreai_backend.mm index ba568fcb485..069a08bc1ad 100644 --- a/backends/apple/coreai/runtime/coreai_backend.mm +++ b/backends/apple/coreai/runtime/coreai_backend.mm @@ -34,6 +34,9 @@ size_t input_count = 0; size_t output_count = 0; 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; }; @@ -136,27 +139,14 @@ bool inputs_match( return text + "]"; } -// Enumerated-shape delegates bind one function per shape, and each accepts -// only its own static shapes, so pick the one these inputs match. +// 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.functions.size() == 1 && - handle.functions.front().input_shapes.empty()) { + if (!handle.enumerated) { return handle.functions.front().session; } - ET_CHECK_OR_RETURN_ERROR( - args.size() >= handle.input_count, - InvalidArgument, - "Core AI delegate expected %zu inputs", - handle.input_count); - for (size_t index = 0; index < handle.input_count; ++index) { - ET_CHECK_OR_RETURN_ERROR( - args[index] != nullptr && args[index]->isTensor(), - InvalidArgument, - "Core AI input %zu must be a tensor", - index); - } for (const auto& function : handle.functions) { if (inputs_match(function.input_shapes, args)) { return function.session; @@ -255,6 +245,7 @@ bool is_available() const override { handle->input_count = host_bundle->bundle.inputs.count; handle->output_count = host_bundle->bundle.outputs.count; handle->model = loaded.get(); + handle->enumerated = host_bundle->bundle.enumerated; processed->Free(); return static_cast(handle.release()); } @@ -268,13 +259,11 @@ Error execute( ET_CHECK_OR_RETURN_ERROR( opaque != nullptr, DelegateInvalidHandle, "Null Core AI handle"); auto& handle = *static_cast(opaque); - auto session = select_session(handle, args); - if (!session.ok()) { - return session.error(); - } std::lock_guard guard([handle.model executionMutex]); return execute_session( - session.get(), 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_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.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm index fa378ef611f..47932c0c9ce 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.mm +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -310,29 +310,37 @@ Error evict_one(id loader, NSData* bookmark) { if (![loading containsObject:shared_key]) { break; } + static_cast(fault_point(FaultPoint::BeforeSharedLoadWait)); load_finished.wait(guard); } [loading addObject:shared_key]; } - auto model = acquire_cached_model( - host_bundle, named_data, cache_directory, loader, key.get()); - ETCoreAILoadedModel* loaded = model.ok() - ? [[ETCoreAILoadedModel alloc] initWithModel:model.get()] - : nil; - if (share) { - { - std::lock_guard guard(mutex); - [loading removeObject:shared_key]; - if (loaded != nil) { - [shared_models setObject:loaded forKey:shared_key]; + // 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(); } - 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(); } - return loaded; + finish.loaded = [[ETCoreAILoadedModel alloc] initWithModel:model.get()]; + return finish.loaded; } Error clear_cache_entry( diff --git a/backends/apple/coreai/runtime/coreai_manifest.h b/backends/apple/coreai/runtime/coreai_manifest.h index ac4202848cb..635fa2dc9b0 100644 --- a/backends/apple/coreai/runtime/coreai_manifest.h +++ b/backends/apple/coreai/runtime/coreai_manifest.h @@ -28,6 +28,8 @@ class Manifest final { public: NSString* const export_hash; 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; @@ -49,6 +51,7 @@ class Manifest final { Manifest( NSString* export_hash, std::vector functions, + bool enumerated, NSArray* inputs, NSArray* outputs, NSDictionary* files, @@ -69,6 +72,7 @@ class SelectedBundle final { NSString* const bundle_name; NSString* const digest; 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 2bce1e23227..87e17046d9c 100644 --- a/backends/apple/coreai/runtime/coreai_manifest.mm +++ b/backends/apple/coreai/runtime/coreai_manifest.mm @@ -18,6 +18,7 @@ Manifest::Manifest( NSString* export_hash, std::vector functions, + bool enumerated, NSArray* inputs, NSArray* outputs, NSDictionary* files, @@ -29,6 +30,7 @@ bool multifunction) : export_hash(export_hash), functions(std::move(functions)), + enumerated(enumerated), inputs(inputs), outputs(outputs), files(files), @@ -48,6 +50,7 @@ bundle_name(path.lastPathComponent), digest(manifest.bundle_digests[path.lastPathComponent]), functions(manifest.functions), + enumerated(manifest.enumerated), inputs(manifest.inputs), outputs(manifest.outputs), files([files copy]), @@ -349,6 +352,7 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { return Manifest( export_hash, std::move(functions), + enumerated_functions != nil, inputs, outputs, files, 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/coreai_delegate_test.mm b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm index 5af523bbe98..c7e12bb8da7 100644 --- a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm @@ -24,6 +24,7 @@ #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 { @@ -1283,12 +1284,15 @@ Error init_sharing( 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 after the second load has queued behind it. + // The first acquisition fails only once the second load waits on it. fake.loader.onSpecialize = ^(NSURL*) { if (++*calls == 1) { *first_in_flight = true; - usleep(100000); + for (int i = 0; i < 1000 && !*second_waiting; ++i) { + usleep(5000); + } loader.specializeError = [NSError errorWithDomain:NSPOSIXErrorDomain code:EIO userInfo:nil]; @@ -1304,8 +1308,19 @@ Error init_sharing( while (!*first_in_flight) { std::this_thread::yield(); } - const Error second_result = init_sharing(second, true, root); + 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); diff --git a/backends/apple/coreai/runtime/test/coreai_manifest_test.mm b/backends/apple/coreai/runtime/test/coreai_manifest_test.mm index dd52cb0174b..cbef5562b03 100644 --- a/backends/apple/coreai/runtime/test/coreai_manifest_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_manifest_test.mm @@ -47,6 +47,7 @@ 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( @@ -299,6 +300,7 @@ 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"]]);