From d703edf3f60713c30f18f71baf7ce11fa0a28aee Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Fri, 30 Oct 2026 17:27:46 -0700 Subject: [PATCH 1/3] [CoreAI] Add a coreai_compute_unit backend option Expose the Core AI SDK's SpecializationOptions configurations as a runtime backend option: default, cpu_only, prefer_cpu, prefer_gpu and prefer_neural_engine. The selection feeds the loader's specialization identity, so each value is specialized and cached under its own key. --- backends/apple/coreai/README.md | 23 ++++++++ .../apple/coreai/runtime/ETCoreAIBridge.h | 14 ++++- .../coreai/runtime/ETCoreAIModelLoader.swift | 22 ++++++++ .../apple/coreai/runtime/coreai_backend.mm | 43 ++++++++++++++- .../apple/coreai/runtime/coreai_bridge.mm | 5 +- .../coreai/runtime/coreai_load_coordinator.mm | 4 +- .../runtime/test/SwiftBridgeTests.swift | 22 ++++++++ .../runtime/test/coreai_delegate_test.mm | 52 +++++++++++++++++++ .../coreai/runtime/test/coreai_fake_loader.h | 1 + .../coreai/runtime/test/coreai_fake_loader.mm | 4 +- 10 files changed, 184 insertions(+), 6 deletions(-) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index b4912248fa8..36c4bb5c873 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -338,6 +338,29 @@ 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")); +``` + ## SDK Bridge The private Swift module `CoreAIBridge` separates error/completion support, diff --git a/backends/apple/coreai/runtime/ETCoreAIBridge.h b/backends/apple/coreai/runtime/ETCoreAIBridge.h index a8fb3d2cda5..904a19dcc5c 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,8 @@ extern "C" { #endif BOOL ETCoreAIIsAvailable(void); NSString* _Nullable ETCoreAIDeviceArchitectureName(void); -id ETCoreAICreateModelLoader(void); +id ETCoreAICreateModelLoader( + ETCoreAIComputeUnit computeUnit); #ifdef __cplusplus } #endif diff --git a/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift b/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift index adf80fbc91f..049b329d6fa 100644 --- a/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift +++ b/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift @@ -238,6 +238,28 @@ public final class ETCoreAISwiftModelLoader: NSObject, ETCoreAIModelLoading { super.init() } + @objc public convenience init?(computeUnit: ETCoreAIComputeUnit) { + guard let options = Self.specializationOptions(for: computeUnit), + let specialization = try? CoreAISpecialization(options: options) + else { + return nil + } + self.init(specialization: specialization) + } + + static func specializationOptions(for computeUnit: ETCoreAIComputeUnit) + -> SpecializationOptions? + { + switch computeUnit { + case .default: return .default + case .cpuOnly: return .cpuOnly + case .preferCPU: return SpecializationOptions(preferredComputeUnitKind: .cpu) + case .preferGPU: return SpecializationOptions(preferredComputeUnitKind: .gpu) + case .preferNeuralEngine: return SpecializationOptions(preferredComputeUnitKind: .neuralEngine) + @unknown default: return nil + } + } + 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..fe4f6c8a9a3 100644 --- a/backends/apple/coreai/runtime/coreai_backend.mm +++ b/backends/apple/coreai/runtime/coreai_backend.mm @@ -14,6 +14,7 @@ #import "coreai_tensor_io.h" #include +#include #include #include @@ -42,6 +43,41 @@ return cache_directory_path(option.get()); } +Result compute_unit_option(BackendInitContext& context) { + auto option = context.get_runtime_spec("coreai_compute_unit"); + if (!option.ok()) { + if (option.error() == Error::NotFound) { + return ETCoreAIComputeUnitDefault; + } + return option.error(); + } + ET_CHECK_OR_RETURN_ERROR( + option.get() != nullptr, + InvalidArgument, + "Core AI compute unit option must not be null"); + static constexpr struct { + const char* name; + ETCoreAIComputeUnit value; + } 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) { @@ -97,7 +133,12 @@ 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(); + } + id loader = + ETCoreAICreateModelLoader(compute_unit.get()); ET_CHECK_OR_RETURN_ERROR( loader != nil, Internal, "Cannot create Core AI loader"); auto cache_directory = create_cache_directory(configured.get()); diff --git a/backends/apple/coreai/runtime/coreai_bridge.mm b/backends/apple/coreai/runtime/coreai_bridge.mm index 539bdb408cb..3defec3b174 100644 --- a/backends/apple/coreai/runtime/coreai_bridge.mm +++ b/backends/apple/coreai/runtime/coreai_bridge.mm @@ -24,6 +24,7 @@ BOOL ETCoreAIIsAvailable(void) { return nil; } -id ETCoreAICreateModelLoader(void) { - return [[ETCoreAISwiftModelLoader alloc] init]; +id ETCoreAICreateModelLoader( + ETCoreAIComputeUnit computeUnit) { + return [[ETCoreAISwiftModelLoader alloc] initWithComputeUnit:computeUnit]; } diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm index bd4f45a7de1..3cfe7b42a19 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.mm +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -255,7 +255,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); ET_CHECK_OR_RETURN_ERROR( loader != nil, Internal, "Cannot create Core AI model loader"); } diff --git a/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift index fbbe6274ef7..14a5b3e3ce2 100644 --- a/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift +++ b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift @@ -153,6 +153,28 @@ 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 identifiers = Set() + for (unit, options) in expected { + XCTAssertEqual(ETCoreAISwiftModelLoader.specializationOptions(for: unit), options) + let loader = try XCTUnwrap(ETCoreAISwiftModelLoader(computeUnit: unit)) + XCTAssertEqual( + loader.specializationIdentifier, try CoreAISpecialization(options: options).identifier) + identifiers.insert(loader.specializationIdentifier) + } + XCTAssertEqual(identifiers.count, expected.count) + XCTAssertEqual( + ETCoreAISwiftModelLoader(computeUnit: .default)?.specializationIdentifier, "default") + XCTAssertNil(ETCoreAISwiftModelLoader(computeUnit: ETCoreAIComputeUnit(rawValue: 99)!)) + } + func testDescriptorIsVersionedAndUsesStableComputeUnitNames() throws { var options = SpecializationOptions(preferredComputeUnitKind: .gpu) options.expectFrequentReshapes = !SpecializationOptions.default.expectFrequentReshapes diff --git a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm index c637ebd3374..84530ac9c92 100644 --- a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm @@ -927,6 +927,58 @@ TensorImpl output( } } +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; + EXPECT_EQ( + options.set_option( + "coreai_cache_dir", fixture.storage.url.path.UTF8String), + Error::Ok); + if (name != nullptr) { + EXPECT_EQ(options.set_option("coreai_compute_unit", name), Error::Ok); + } + return fixture.init_with_options( + {options.view().data(), options.view().size()}); + }; + { + SCOPED_TRACE("absent"); + ASSERT_EQ(init_with_compute_unit(nullptr), Error::Ok); + EXPECT_EQ( + bridge.state().last_compute_unit.load(), ETCoreAIComputeUnitDefault); + } + 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); + } + { + SCOPED_TRACE("integer compute unit"); + DelegateFixture fixture(backend, bridge); + BackendOption wrong_type{"coreai_compute_unit", 1}; + EXPECT_EQ( + fixture.init_with_options({&wrong_type, 1}), Error::InvalidArgument); + EXPECT_EQ(fixture.freed, 0); + } + } +} + TEST_F(CoreAIDelegateTest, ConfiguredSymlinkCacheDirectoryLoadsAndClears) { @autoreleasepool { DelegateFixture fixture(backend, bridge); diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.h b/backends/apple/coreai/runtime/test/coreai_fake_loader.h index 5fb287ed79a..72d101b1855 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.h +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.h @@ -41,6 +41,7 @@ struct FakeBridgeState { std::atomic sessions{0}; std::atomic prepared_models{0}; std::atomic loaders{0}; + std::atomic last_compute_unit{-1}; std::atomic last_input_bytes{nullptr}; std::atomic last_input_byte_count{0}; std::atomic last_input_first_byte{0}; diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm index 5eefda12f9d..4e3e69b21cb 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm @@ -537,9 +537,11 @@ BOOL ETCoreAIIsAvailable(void) { return current_context->state.device_architecture; } -id ETCoreAICreateModelLoader(void) { +id ETCoreAICreateModelLoader( + ETCoreAIComputeUnit computeUnit) { if (current_context == nil || current_context->state.fail_loader_factory) { return nil; } + current_context->state.last_compute_unit = static_cast(computeUnit); return current_context->state.bookmark_loader ?: [[FakeLoader alloc] init]; } From df52e1b8a24a0a8f21036c8b97cc8a8d89ffc00a Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 3 Nov 2026 15:09:15 -0800 Subject: [PATCH 2/3] [CoreAI] Report specialization identifiers when compute-unit identifiers collide testComputeUnitsMapToSDKSpecializationOptions fails on CI runners with fewer distinct identifiers than compute units, but passes locally with the same macOS and Xcode. Include each unit's options and identifier in the failure message to show which units collide. --- backends/apple/coreai/runtime/test/SwiftBridgeTests.swift | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift index 14a5b3e3ce2..2bbecdd6d04 100644 --- a/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift +++ b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift @@ -162,14 +162,18 @@ final class SwiftSpecializationTests: XCTestCase { (.preferNeuralEngine, SpecializationOptions(preferredComputeUnitKind: .neuralEngine)), ] var identifiers = Set() + var report: [String] = [] for (unit, options) in expected { XCTAssertEqual(ETCoreAISwiftModelLoader.specializationOptions(for: unit), options) let loader = try XCTUnwrap(ETCoreAISwiftModelLoader(computeUnit: unit)) XCTAssertEqual( loader.specializationIdentifier, try CoreAISpecialization(options: options).identifier) identifiers.insert(loader.specializationIdentifier) + report.append( + "\(unit.rawValue): \(options) -> \(loader.specializationIdentifier)") } - XCTAssertEqual(identifiers.count, expected.count) + XCTAssertEqual( + identifiers.count, expected.count, "\n" + report.joined(separator: "\n")) XCTAssertEqual( ETCoreAISwiftModelLoader(computeUnit: .default)?.specializationIdentifier, "default") XCTAssertNil(ETCoreAISwiftModelLoader(computeUnit: ETCoreAIComputeUnit(rawValue: 99)!)) From 9c2527c6f8ec2667665d8d1fbf2388ce6bf13bab Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 3 Nov 2026 16:08:04 -0800 Subject: [PATCH 3/3] [CoreAI] Check specialization identifiers pairwise against their options The SDK limits allowed compute units to the hardware, so on CPU-only CI runners .default and .cpuOnly resolve to the same options and correctly share an identifier. Instead of expecting one distinct identifier per compute unit, require two identifiers to match exactly when their options do. --- .../runtime/test/SwiftBridgeTests.swift | 26 ++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift index 2bbecdd6d04..2ebd4ce98e8 100644 --- a/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift +++ b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift @@ -161,24 +161,42 @@ final class SwiftSpecializationTests: XCTestCase { (.preferGPU, SpecializationOptions(preferredComputeUnitKind: .gpu)), (.preferNeuralEngine, SpecializationOptions(preferredComputeUnitKind: .neuralEngine)), ] - var identifiers = Set() + var variants: [(options: SpecializationOptions, identifier: String)] = [] var report: [String] = [] for (unit, options) in expected { XCTAssertEqual(ETCoreAISwiftModelLoader.specializationOptions(for: unit), options) let loader = try XCTUnwrap(ETCoreAISwiftModelLoader(computeUnit: unit)) XCTAssertEqual( loader.specializationIdentifier, try CoreAISpecialization(options: options).identifier) - identifiers.insert(loader.specializationIdentifier) + variants.append((options, loader.specializationIdentifier)) report.append( "\(unit.rawValue): \(options) -> \(loader.specializationIdentifier)") } - XCTAssertEqual( - identifiers.count, expected.count, "\n" + report.joined(separator: "\n")) + assertIdentifiersMatchExactlyWhenOptionsMatch(variants, report: report) XCTAssertEqual( ETCoreAISwiftModelLoader(computeUnit: .default)?.specializationIdentifier, "default") XCTAssertNil(ETCoreAISwiftModelLoader(computeUnit: ETCoreAIComputeUnit(rawValue: 99)!)) } + // 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