[CoreAI] Add a coreai_compute_unit backend option

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,

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<ETCoreAIModelLoading> ETCoreAICreateModelLoader(void);
+id<ETCoreAIModelLoading> ETCoreAICreateModelLoader(
+    ETCoreAIComputeUnit computeUnit);
 #ifdef __cplusplus
 }
 #endif

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 }

backends/apple/coreai/runtime/coreai_backend.mm

@@ -14,6 +14,7 @@
 #import "coreai_tensor_io.h"
 
 #include <executorch/runtime/backend/interface.h>
+#include <cstring>
 #include <memory>
 #include <new>
 
@@ -42,6 +43,41 @@
   return cache_directory_path(option.get());
 }
 
+Result<ETCoreAIComputeUnit> compute_unit_option(BackendInitContext& context) {
+  auto option = context.get_runtime_spec<const char*>("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<id<ETCoreAISession>> bind_session(
     id<ETCoreAIPreparedModel> model,
     const SelectedBundle& bundle) {
@@ -97,7 +133,12 @@ bool is_available() const override {
       if (!configured.ok()) {
         return configured.error();
       }
-      id<ETCoreAIModelLoading> loader = ETCoreAICreateModelLoader();
+      auto compute_unit = compute_unit_option(context);
+      if (!compute_unit.ok()) {
+        return compute_unit.error();
+      }
+      id<ETCoreAIModelLoading> 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());

backends/apple/coreai/runtime/coreai_bridge.mm

@@ -24,6 +24,7 @@ BOOL ETCoreAIIsAvailable(void) {
   return nil;
 }
 
-id<ETCoreAIModelLoading> ETCoreAICreateModelLoader(void) {
-  return [[ETCoreAISwiftModelLoader alloc] init];
+id<ETCoreAIModelLoading> ETCoreAICreateModelLoader(
+    ETCoreAIComputeUnit computeUnit) {
+  return [[ETCoreAISwiftModelLoader alloc] initWithComputeUnit:computeUnit];
 }

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");
     }

backends/apple/coreai/runtime/test/SwiftBridgeTests.swift

@@ -153,6 +153,50 @@ final class SwiftSpecializationTests: XCTestCase {
       loader.specializationIdentifier, try CoreAISpecialization(options: options).identifier)
   }
 
+  func testComputeUnitsMapToSDKSpecializationOptions() throws {
+    let expected: [(ETCoreAIComputeUnit, SpecializationOptions)] = [
+      (.default, .default),
+      (.cpuOnly, .cpuOnly),
+      (.preferCPU, SpecializationOptions(preferredComputeUnitKind: .cpu)),
+      (.preferGPU, SpecializationOptions(preferredComputeUnitKind: .gpu)),
+      (.preferNeuralEngine, SpecializationOptions(preferredComputeUnitKind: .neuralEngine)),
+    ]
+    var variants: [(options: SpecializationOptions, identifier: String)] = []
+    var report: [String] = []
+    for (unit, options) in expected {
+      XCTAssertEqual(ETCoreAISwiftModelLoader.specializationOptions(for: unit), options)
+      let loader = try XCTUnwrap(ETCoreAISwiftModelLoader(computeUnit: unit))
+      XCTAssertEqual(
+        loader.specializationIdentifier, try CoreAISpecialization(options: options).identifier)
+      variants.append((options, loader.specializationIdentifier))
+      report.append(
+        "\(unit.rawValue): \(options) -> \(loader.specializationIdentifier)")
+    }
+    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

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<const char*, ETCoreAIComputeUnit> 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);

backends/apple/coreai/runtime/test/coreai_fake_loader.h

@@ -41,6 +41,7 @@ struct FakeBridgeState {
   std::atomic<int> sessions{0};
   std::atomic<int> prepared_models{0};
   std::atomic<int> loaders{0};
+  std::atomic<int> last_compute_unit{-1};
   std::atomic<const void*> last_input_bytes{nullptr};
   std::atomic<size_t> last_input_byte_count{0};
   std::atomic<unsigned char> last_input_first_byte{0};

backends/apple/coreai/runtime/test/coreai_fake_loader.mm

@@ -537,9 +537,11 @@ BOOL ETCoreAIIsAvailable(void) {
   return current_context->state.device_architecture;
 }
 
-id<ETCoreAIModelLoading> ETCoreAICreateModelLoader(void) {
+id<ETCoreAIModelLoading> ETCoreAICreateModelLoader(
+    ETCoreAIComputeUnit computeUnit) {
   if (current_context == nil || current_context->state.fail_loader_factory) {
     return nil;
   }
+  current_context->state.last_compute_unit = static_cast<int>(computeUnit);
   return current_context->state.bookmark_loader ?: [[FakeLoader alloc] init];
 }