From d703edf3f60713c30f18f71baf7ce11fa0a28aee Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Fri, 30 Oct 2026 17:27:46 -0700 Subject: [PATCH] [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]; }