[CoreAI] Coordinate model acquisition

backends/apple/coreai/CMakeLists.txt

@@ -21,8 +21,9 @@ endforeach()
 
 enable_language(OBJCXX)
 
-set(_coreai_runtime_sources runtime/coreai_assets.mm runtime/coreai_storage.mm
-                            runtime/coreai_bookmarks.mm
+set(_coreai_runtime_sources
+    runtime/coreai_assets.mm runtime/coreai_storage.mm
+    runtime/coreai_bookmarks.mm runtime/coreai_load_coordinator.mm
 )
 
 # Keep ARC and C++ exception settings away from the private Swift module.
@@ -53,6 +54,10 @@ if(EXECUTORCH_BUILD_TESTS)
     runtime/test/coreai_bookmark_test.mm
     runtime/test/coreai_source_fixture.mm
     runtime/test/coreai_bookmark_fixture.mm
+    runtime/test/coreai_fake_loader.mm
+    runtime/test/coreai_acquisition_test.mm
+    runtime/test/coreai_acquisition_fixture.mm
+    runtime/ETCoreAITensor.mm
     ${_coreai_runtime_sources}
   )
   coreai_configure_objc_target(coreai_host_test)

backends/apple/coreai/README.md

@@ -42,6 +42,11 @@ architectures=["h17p"])` requests that compiler architecture; omitting the list
 lets the compiler emit its supported architectures for the target platform.
 These are Core AI architecture names, not CPU names such as `arm64`.
 
+Both formats can require device specialization. Every load first attempts SDK
+bookmark restoration. A missing bookmark or SDK-confirmed cache miss falls back
+to source materialization and persistent specialization, then binds the requested
+function and ordered I/O on that same acquired model.
+
 ### Asset storage
 
 The assets root must be an absolute path. The application chooses it, should
@@ -86,6 +91,15 @@ bookmark in the `bookmarks` subdirectory, a bundle directory in `staging`, and a
 empty coordination file in `locks`. There is no PTE-wide bookmark list or
 serialized lifecycle record.
 
+A healthy bookmark hit does not inspect staging or consult NDS; the acquired SDK
+model is retained through function binding. On a missing bookmark or confirmed
+SDK miss, the runtime materializes the source, specializes it and atomically
+writes the returned bookmark bytes. There is no secondary source-URL cache
+lookup. SDK errors, unreadable bookmarks and publication failures fail loading
+without automatic cache deletion. If publication fails after specialization, the
+SDK entry and source remain; the next load retries the ordinary flow. The
+backend does not track historical or unrecorded SDK entries.
+
 Raw bookmarks have an 8 MiB backend allocation limit, not an asserted SDK format
 maximum, and are published atomically because losing one orphans an SDK entry.
 Per-key disk locks coordinate loading and eviction across threads and processes;

backends/apple/coreai/runtime/coreai_load_coordinator.h

@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#pragma once
+
+#import "ETCoreAIBridge.h"
+#import "coreai_bookmarks.h"
+
+namespace executorch::backends::coreai {
+
+// The key lock outlives all SDK callbacks and raw bookmark publication.
+runtime::Result<id<ETCoreAIPreparedModel>> acquire_bookmark_model(
+    const Manifest& selected,
+    const runtime::NamedDataMap* named_data,
+    NSString* root,
+    NSString* platform,
+    NSString* architecture,
+    id<ETCoreAIModelLoading> loader);
+
+} // namespace executorch::backends::coreai

backends/apple/coreai/runtime/coreai_load_coordinator.mm

@@ -0,0 +1,96 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#import "coreai_load_coordinator.h"
+#import "coreai_storage.h"
+
+namespace executorch::backends::coreai {
+namespace {
+using runtime::Error;
+using runtime::Result;
+
+Error sdk_error(NSError* error) {
+  ET_LOG(Error, "Core AI SDK error (%s, %ld): %s", error.domain.UTF8String,
+         static_cast<long>(error.code), error.description.UTF8String);
+  return Error::Internal;
+}
+
+Result<id<ETCoreAIPreparedModel>> acquire(
+    void (^operation)(ETCoreAIAcquisitionCompletion)) {
+  ET_CHECK_OR_RETURN_ERROR(storage_fault(StorageOperation::BeforeSDK) == 0,
+                           AccessFailed, "Core AI acquisition interrupted");
+  dispatch_semaphore_t ready = dispatch_semaphore_create(0);
+  __block id<ETCoreAIPreparedModel> model = nil;
+  __block NSError* error = nil;
+  __block ETCoreAIAcquisitionStatus status = ETCoreAIAcquisitionStatusError;
+  operation(^(ETCoreAIAcquisitionStatus result, id<ETCoreAIPreparedModel> value,
+              NSError* failure) {
+    status = result;
+    model = value;
+    error = failure;
+    dispatch_semaphore_signal(ready);
+  });
+  // Do not abandon asynchronous source access on cancellation or timeout.
+  dispatch_semaphore_wait(ready, DISPATCH_TIME_FOREVER);
+  id<ETCoreAIPreparedModel> acquired = model;
+  // A loader may retain the completion block after invoking it.
+  model = nil;
+  if (error != nil) return sdk_error(error);
+  ET_CHECK_OR_RETURN_ERROR(storage_fault(StorageOperation::AfterSDK) == 0,
+                           AccessFailed, "Core AI acquisition interrupted");
+  ET_CHECK_OR_RETURN_ERROR(
+      (status == ETCoreAIAcquisitionStatusHit && acquired != nil) ||
+          (status == ETCoreAIAcquisitionStatusMiss && acquired == nil),
+      Internal, "Invalid Core AI acquisition result");
+  return acquired;
+}
+
+Result<id<ETCoreAIPreparedModel>> restore(id<ETCoreAIModelLoading> loader,
+                                          NSData* bookmark) {
+  return acquire(^(ETCoreAIAcquisitionCompletion completion) {
+    [loader restoreModelFromBookmark:bookmark completion:completion];
+  });
+}
+
+}  // namespace
+
+Result<id<ETCoreAIPreparedModel>> acquire_bookmark_model(
+    const Manifest& selected, const runtime::NamedDataMap* named_data,
+    NSString* root, NSString* platform, NSString* architecture,
+    id<ETCoreAIModelLoading> loader) {
+  ET_CHECK_OR_RETURN_ERROR(loader != nil, InvalidArgument,
+                           "Missing Core AI model loader");
+  auto key = bookmark_key(selected, platform, architecture);
+  if (!key.ok()) return key.error();
+  auto locked = lock_bookmark(root, key.get());
+  if (!locked.ok()) return locked.error();
+  auto bookmark = read_bookmark(*locked.get());
+  if (!bookmark.ok()) return bookmark.error();
+  if (bookmark.get() != nil) {
+    auto model = restore(loader, bookmark.get());
+    if (!model.ok()) return model.error();
+    if (model.get() != nil) return model;
+  }
+  auto staging = prepare_bookmark_staging(*locked.get());
+  if (!staging.ok()) return staging.error();
+  auto source =
+      prepare_source_bundle(selected, named_data, staging.get(), key.get());
+  if (!source.ok()) return source.error();
+  NSURL* url = source.get();
+  auto model = acquire(^(ETCoreAIAcquisitionCompletion completion) {
+    [loader specializeModelAtURL:url completion:completion];
+  });
+  if (!model.ok()) return model.error();
+  ET_CHECK_OR_RETURN_ERROR(model.get() != nil, Internal,
+                           "Core AI specialization returned no model");
+  ET_CHECK_OK_OR_RETURN_ERROR(
+      write_bookmark(*locked.get(), [model.get() copyBookmarkData]));
+  // SDK pins do not prove source independence; leave staged sources in place.
+  return model;
+}
+
+}  // namespace executorch::backends::coreai

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

@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include "coreai_bookmark_fixture.h"
+#include "coreai_fake_loader.h"
+#include "coreai_source_fixture.h"
+#include "coreai_storage.h"
+
+namespace executorch::backends::coreai::testing {
+
+runtime::Result<NSURL*> acquisition_staged_url(
+    NSURL* root,
+    const Manifest& manifest);
+
+struct ParallelAcquisitionObservation {
+  int creation[2] = {-1, -1};
+  int joined[2] = {-1, -1};
+  bool completed[2] = {false, false};
+  runtime::Error errors[2] = {
+      runtime::Error::Internal,
+      runtime::Error::Internal};
+  bool models[2] = {false, false};
+  bool first_entered = false;
+  bool second_observed = false;
+  int gate_timeouts = 0;
+  bool callbacks_drained = false;
+};
+ParallelAcquisitionObservation parallel_acquisitions(
+    ScopedFakeBridgeState& bridge,
+    BookmarkFakeLoader* loader,
+    NSURL* root,
+    const Manifest& first,
+    const Manifest& second,
+    TestData& data,
+    bool same_key);
+
+// Uses the existing bounded, scoped child process owner.
+void set_acquisition_test_executable(const char* path);
+::testing::AssertionResult spawn_acquisition_crash(
+    BookmarkChild& child,
+    NSURL* root,
+    StorageOperation boundary,
+    NSDictionary* manifest);
+::testing::AssertionResult spawn_parallel_acquisition(
+    BookmarkChild& child,
+    NSURL* root,
+    NSDictionary* first,
+    NSDictionary* second,
+    bool same_key);
+// Dispatch before bookmark_child_mode and InitGoogleTest. Unrelated modes
+// return -1.
+int acquisition_child_mode(int argc, char** argv);
+
+} // namespace executorch::backends::coreai::testing