From 04812b614baebeabbe8e556669b78ada609c7fce Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Sun, 4 Oct 2026 19:49:00 -0700 Subject: [PATCH 1/7] [CoreAI] Add atomic cache storage primitives Summary: Add the filesystem layer the Core AI runtime uses for its asset and bookmark cache (`runtime/coreai_storage.{h,mm}`, `runtime/coreai_file.h`). - Prepare an absolute storage root chosen by the application. Creating a root also excludes it from backup; inspecting an existing root changes nothing. - Enumerate directories through descriptors and return sorted names, failing on operational errors rather than returning partial listings. - Write new files relative to a directory descriptor and `fsync` each file and its parent directory. - Publish data atomically: write a temporary file, `fsync` and `F_FULLFSYNC` it, `renameat` it into place and sync the directory. Readers see either the old or the new complete bytes. A failure before the rename leaves the old bytes and removes the temporary file. - Tree removal is not included here; it lands with cache clearing. - README: the "Asset storage" root, durability and backup-exclusion text. Trust model: the application owns the cache directory. The helpers do not defend against other processes of the same user renaming or replacing paths inside it. Fault injection points use `storage_fault()`, which compiles to a constant 0 unless `COREAI_ASSETS_TESTING` is defined. Only the `coreai_host_test` target defines it. No changes outside `backends/apple/coreai`. Test Plan: ``` cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_OSX_DEPLOYMENT_TARGET=27.0 -DEXECUTORCH_BUILD_COREAI=ON \ -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON cmake --build build --target backends/apple/coreai/all ctest --test-dir build/backends/apple/coreai --output-on-failure --no-tests=error ``` CTest runs `coreai_host_test`, which adds the storage suite. AtomicPublicationPreservesOldBytesOnFailure checks that a failed rename leaves the old bytes and no `.tmp-*` entry. --- backends/apple/coreai/CMakeLists.txt | 5 +- backends/apple/coreai/README.md | 13 + backends/apple/coreai/runtime/coreai_file.h | 53 ++++ .../apple/coreai/runtime/coreai_storage.h | 60 ++++ .../apple/coreai/runtime/coreai_storage.mm | 298 ++++++++++++++++++ .../coreai/runtime/test/coreai_fault_scope.h | 57 ++++ .../runtime/test/coreai_filesystem_fixture.h | 27 ++ .../runtime/test/coreai_filesystem_fixture.mm | 61 ++++ .../runtime/test/coreai_storage_test.mm | 113 +++++++ 9 files changed, 685 insertions(+), 2 deletions(-) create mode 100644 backends/apple/coreai/runtime/coreai_file.h create mode 100644 backends/apple/coreai/runtime/coreai_storage.h create mode 100644 backends/apple/coreai/runtime/coreai_storage.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_fault_scope.h create mode 100644 backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h create mode 100644 backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_storage_test.mm diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index 800a9f098a4..6a01341c31b 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -11,7 +11,7 @@ endif() enable_language(OBJCXX) set(_coreai_runtime_sources runtime/coreai_manifest.mm - runtime/coreai_validation.mm + runtime/coreai_validation.mm runtime/coreai_storage.mm ) function(coreai_configure_objc_target target) @@ -31,7 +31,8 @@ if(EXECUTORCH_BUILD_TESTS AND CMAKE_SYSTEM_NAME STREQUAL "Darwin") add_executable( coreai_host_test runtime/test/coreai_host_test.mm runtime/test/coreai_manifest_test.mm - ${_coreai_runtime_sources} + runtime/test/coreai_storage_test.mm + runtime/test/coreai_filesystem_fixture.mm ${_coreai_runtime_sources} ) coreai_configure_objc_target(coreai_host_test) find_library(COREAI_FOUNDATION_FRAMEWORK Foundation REQUIRED) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 4b5115e1a21..623a96a14a6 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -50,6 +50,19 @@ 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`. +### Asset storage + +The assets root must be an absolute path. The application chooses it, should +reserve it for this backend, and is trusted not to rename, replace or modify its +contents while the backend uses it; the backend does not defend against other +processes of the same user rebinding paths. Written files and changed directories +are synced with `fsync`; atomic publication also uses `F_FULLFSYNC` before its +rename. + +Preparing the assets root sets `NSURLIsExcludedFromBackupKey` on it, which covers +everything beneath it. Ancestors are not modified. This is backup exclusion, not +a control for iCloud Drive synchronization. + ## Host Tests On macOS, `EXECUTORCH_BUILD_COREAI=ON` with `EXECUTORCH_BUILD_TESTS=ON` diff --git a/backends/apple/coreai/runtime/coreai_file.h b/backends/apple/coreai/runtime/coreai_file.h new file mode 100644 index 00000000000..478f03a7b91 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_file.h @@ -0,0 +1,53 @@ +/* + * 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 +#include + +namespace executorch::backends::coreai { + +class FileDescriptor { + public: + explicit FileDescriptor(int fd) : fd_(fd) {} + ~FileDescriptor() { + reset(-1); + } + FileDescriptor(const FileDescriptor&) = delete; + FileDescriptor& operator=(const FileDescriptor&) = delete; + int get() const { + return fd_; + } + int release() { + const int fd = fd_; + fd_ = -1; + return fd; + } + void reset(int fd) { + if (fd_ >= 0) { + // Retrying close can close a descriptor reused by another thread. + ::close(fd_); + } + fd_ = fd; + } + + private: + int fd_; +}; + +template +auto retry_eintr(F call) { + decltype(call()) result; + do { + result = call(); + } while (result == -1 && errno == EINTR); + return result; +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h new file mode 100644 index 00000000000..f5bd377dd70 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -0,0 +1,60 @@ +/* + * 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 +#include +#include + +namespace executorch::backends::coreai { + +// Resolves the default location without creating or modifying storage. +runtime::Result default_coreai_assets_root(); +// Requires an absolute directory path. With create, makes the directory and +// excludes it from backup; otherwise a missing directory is InvalidArgument. +runtime::Result prepare_storage_root(NSString* path, bool create); +runtime::Error ensure_excluded_from_backup(NSURL* directory); + +runtime::Error sync_storage_directory(int fd); +// Returns an independently owned descriptor for a stable coordination file. +int open_storage_shared_file(int parent_fd, const char* name); +// Creates a new file and missing parents, then syncs the file and its parent. +runtime::Error write_storage_file( + int root_fd, + NSString* relative_path, + const void* bytes, + size_t size); +// Atomically replaces name with data. On error, readers see either the old or +// the new complete bytes. +runtime::Error publish_storage_data(int root_fd, NSString* name, NSData* data); +runtime::Result*> storage_children( + int fd, + bool skip_invalid_names = false); + +enum class StorageOperation { + Rename, + BeforeSDK, + AfterSDK, + BeforeEvict, + AfterEvict, + Remove, +}; +#if defined(COREAI_ASSETS_TESTING) && COREAI_ASSETS_TESTING +int storage_fault(StorageOperation operation); +namespace testing { +using StorageFaultCallback = int (*)(StorageOperation); +extern thread_local StorageFaultCallback storage_fault_callback; +} // namespace testing +#else +constexpr int storage_fault(StorageOperation) { + return 0; +} +#endif + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm new file mode 100644 index 00000000000..66fea0e8833 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -0,0 +1,298 @@ +/* + * 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_storage.h" +#include "coreai_file.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::coreai { +namespace { +using runtime::Error; +using runtime::Result; +} // namespace + +Error ensure_excluded_from_backup(NSURL* directory) { + NSURL* fresh = [NSURL fileURLWithPath:directory.path isDirectory:YES]; + NSNumber* excluded = nil; + NSError* error = nil; + ET_CHECK_OR_RETURN_ERROR( + [fresh getResourceValue:&excluded + forKey:NSURLIsExcludedFromBackupKey + error:&error], + AccessFailed, + "Cannot read Core AI backup exclusion: %s", + error.localizedDescription.UTF8String); + if (excluded.boolValue) { + return Error::Ok; + } + ET_CHECK_OR_RETURN_ERROR( + [fresh setResourceValue:@YES + forKey:NSURLIsExcludedFromBackupKey + error:&error], + AccessFailed, + "Cannot exclude Core AI storage from backup: %s", + error.localizedDescription.UTF8String); + return Error::Ok; +} + +Result default_coreai_assets_root() { + NSURL* caches = [NSFileManager.defaultManager + URLsForDirectory:NSCachesDirectory + inDomains:NSUserDomainMask].firstObject; + ET_CHECK_OR_RETURN_ERROR(caches.isFileURL, AccessFailed, + "Cannot resolve Core AI cache directory"); + return [caches.path stringByAppendingPathComponent:@"executorch_coreai"]; +} + +Result prepare_storage_root(NSString* path, bool create) { + ET_CHECK_OR_RETURN_ERROR( + [path isKindOfClass:NSString.class] && path.isAbsolutePath && + [path rangeOfString:@"\0"].location == NSNotFound, + InvalidArgument, + "Core AI storage requires an absolute path"); + if (create) { + NSError* error = nil; + ET_CHECK_OR_RETURN_ERROR( + [NSFileManager.defaultManager + createDirectoryAtPath:path + withIntermediateDirectories:YES + attributes:@{ + NSFilePosixPermissions : @0700 + } + error:&error], + AccessFailed, + "Cannot create Core AI storage directory: %s", + error.localizedDescription.UTF8String); + } + struct stat info; + const int result = + retry_eintr([&] { return stat(path.fileSystemRepresentation, &info); }); + ET_CHECK_OR_RETURN_ERROR( + result == 0 || errno == ENOENT || errno == ENOTDIR, + AccessFailed, + "Cannot inspect Core AI storage directory"); + ET_CHECK_OR_RETURN_ERROR( + result == 0 && S_ISDIR(info.st_mode), + InvalidArgument, + "Core AI storage must be an existing directory"); + if (create) { + ET_CHECK_OK_OR_RETURN_ERROR(ensure_excluded_from_backup( + [NSURL fileURLWithPath:path isDirectory:YES])); + } + return path; +} + +#if defined(COREAI_ASSETS_TESTING) && COREAI_ASSETS_TESTING +namespace testing { +thread_local StorageFaultCallback storage_fault_callback = nullptr; +} + +int storage_fault(StorageOperation operation) { + return testing::storage_fault_callback == nullptr + ? 0 + : testing::storage_fault_callback(operation); +} +#endif + +namespace { +template +int storage_call(StorageOperation operation, F call) { + return retry_eintr([&] { + const int fault = storage_fault(operation); + if (fault != 0) { + errno = fault; + return -1; + } + return call(); + }); +} + +bool storage_component(NSString* name) { + return [name isKindOfClass:NSString.class] && name.length > 0 && + ![name isEqualToString:@"."] && ![name isEqualToString:@".."] && + ![name containsString:@"/"] && ![name containsString:@"\0"]; +} + +Error write_file( + int root_fd, + NSString* relative_path, + const void* data, + size_t size, + bool full_sync) { + NSArray* components = + [relative_path componentsSeparatedByString:@"/"]; + ET_CHECK_OR_RETURN_ERROR( + components.count > 0 && (data != nullptr || size == 0), + InvalidArgument, + "Invalid Core AI storage write"); + for (NSString* component in components) { + ET_CHECK_OR_RETURN_ERROR( + storage_component(component), InvalidArgument, "Invalid storage path"); + } + FileDescriptor parent(retry_eintr([&] { return dup(root_fd); })); + ET_CHECK_OR_RETURN_ERROR( + parent.get() >= 0, AccessFailed, "Cannot open storage parent"); + for (NSUInteger i = 0; i + 1 < components.count; ++i) { + const char* component = components[i].fileSystemRepresentation; + ET_CHECK_OR_RETURN_ERROR( + retry_eintr( + [&] { return mkdirat(parent.get(), component, 0700); }) == 0 || + errno == EEXIST, + AccessFailed, + "Cannot create Core AI source directory"); + FileDescriptor next(retry_eintr([&] { + return openat( + parent.get(), + component, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + ET_CHECK_OR_RETURN_ERROR( + next.get() >= 0, AccessFailed, "Cannot open source directory"); + ET_CHECK_OK_OR_RETURN_ERROR(sync_storage_directory(parent.get())); + parent.reset(next.release()); + } + FileDescriptor file(retry_eintr([&] { + return openat( + parent.get(), + components.lastObject.fileSystemRepresentation, + O_WRONLY | O_CREAT | O_EXCL | O_NOFOLLOW | O_CLOEXEC, + 0600); + })); + ET_CHECK_OR_RETURN_ERROR( + file.get() >= 0, AccessFailed, "Cannot create Core AI storage file"); + const auto* bytes = static_cast(data); + while (size > 0) { + const ssize_t count = + write(file.get(), bytes, std::min(size, size_t(1024 * 1024))); + if (count < 0 && errno == EINTR) continue; + ET_CHECK_OR_RETURN_ERROR( + count > 0, AccessFailed, "Cannot write Core AI storage file"); + bytes += count; + size -= count; + } + ET_CHECK_OR_RETURN_ERROR( + retry_eintr([&] { return fsync(file.get()); }) == 0 && + (!full_sync || + retry_eintr([&] { return fcntl(file.get(), F_FULLFSYNC); }) == 0), + AccessFailed, + "Cannot synchronize Core AI file"); + ET_CHECK_OR_RETURN_ERROR(::close(file.release()) == 0, AccessFailed, + "Cannot close Core AI storage file"); + return sync_storage_directory(parent.get()); +} +} // namespace + +Result*> storage_children(int fd, bool skip_invalid_names) { + FileDescriptor copy( + retry_eintr([&] { return fcntl(fd, F_DUPFD_CLOEXEC, 0); })); + ET_CHECK_OR_RETURN_ERROR(copy.get() >= 0, AccessFailed, + "Cannot duplicate storage directory"); + DIR* directory = fdopendir(copy.get()); + ET_CHECK_OR_RETURN_ERROR(directory != nullptr, AccessFailed, + "Cannot open storage directory stream"); + copy.release(); + // dup shares the enumeration offset with the original descriptor. + rewinddir(directory); + NSMutableArray* names = [NSMutableArray array]; + Error result = Error::Ok; + while (true) { + errno = 0; + dirent* entry = readdir(directory); + if (entry == nullptr) { + if (errno == EINTR) continue; + if (errno != 0) result = Error::AccessFailed; + break; + } + if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0) { + continue; + } + NSString* name = [[NSString alloc] initWithBytes:entry->d_name + length:strlen(entry->d_name) + encoding:NSUTF8StringEncoding]; + if (name == nil || + strcmp(name.fileSystemRepresentation, entry->d_name) != 0) { + if (skip_invalid_names) { + continue; + } + result = Error::InvalidExternalData; + break; + } + [names addObject:name]; + } + const int closed = closedir(directory); + ET_CHECK_OK_OR_RETURN_ERROR(result); + ET_CHECK_OR_RETURN_ERROR(closed == 0, AccessFailed, + "Cannot close storage directory stream"); + return [names sortedArrayUsingSelector:@selector(compare:)]; +} + +int open_storage_shared_file(int parent_fd, const char* name) { + // APFS can return ENOENT to a concurrent O_CREAT loser without O_EXCL. + int fd = retry_eintr([&] { + return openat(parent_fd, name, + O_RDWR | O_CREAT | O_EXCL | O_NOFOLLOW | O_CLOEXEC, 0600); + }); + if (fd < 0 && errno == EEXIST) { + fd = retry_eintr([&] { + return openat(parent_fd, name, O_RDWR | O_NOFOLLOW | O_CLOEXEC); + }); + } + return fd; +} + +Error sync_storage_directory(int fd) { + ET_CHECK_OR_RETURN_ERROR( + retry_eintr([&] { return fsync(fd); }) == 0, + AccessFailed, + "Cannot synchronize Core AI directory"); + return Error::Ok; +} + +Error write_storage_file( + int root_fd, + NSString* relative_path, + const void* data, + size_t size) { + return write_file(root_fd, relative_path, data, size, false); +} + +Error publish_storage_data(int root_fd, NSString* name, NSData* data) { + ET_CHECK_OR_RETURN_ERROR( + storage_component(name) && data != nil, + InvalidArgument, + "Invalid Core AI storage publication"); + NSString* temporary = + [@".tmp-" stringByAppendingString:NSUUID.UUID.UUIDString]; + // Bookmarks are small and costly to lose, so they get a full device flush. + auto result = + write_file(root_fd, temporary, data.bytes, data.length, true); + if (result == Error::Ok && storage_call(StorageOperation::Rename, [&] { + return renameat( + root_fd, + temporary.fileSystemRepresentation, + root_fd, + name.fileSystemRepresentation); + }) != 0) + result = Error::AccessFailed; + if (result == Error::Ok) + return sync_storage_directory(root_fd); + if (unlinkat(root_fd, temporary.fileSystemRepresentation, 0) == 0) { + (void)sync_storage_directory(root_fd); + } + return result; +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/test/coreai_fault_scope.h b/backends/apple/coreai/runtime/test/coreai_fault_scope.h new file mode 100644 index 00000000000..e55de88579a --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_fault_scope.h @@ -0,0 +1,57 @@ +/* + * 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 +#include "coreai_storage.h" + +namespace executorch::backends::coreai::testing { + +struct StorageFaultScope { + StorageOperation operation; + int error = EIO; + int match_index = 1; + int matches = 0; + int hits = 0; + bool armed = true; + void (^observe)(StorageOperation) = nil; + + explicit StorageFaultScope(StorageOperation op) + : operation(op), + previous_(current_), + previous_callback_(storage_fault_callback) { + current_ = this; + storage_fault_callback = inject; + } + ~StorageFaultScope() { + storage_fault_callback = previous_callback_; + current_ = previous_; + } + StorageFaultScope(const StorageFaultScope&) = delete; + StorageFaultScope& operator=(const StorageFaultScope&) = delete; + + private: + inline static thread_local StorageFaultScope* current_ = nullptr; + StorageFaultScope* previous_; + StorageFaultCallback previous_callback_; + + static int inject(StorageOperation operation) { + auto& scope = *current_; + if (scope.observe != nil) + scope.observe(operation); + if (scope.armed && scope.operation == operation && + ++scope.matches == scope.match_index) { + ++scope.hits; + return scope.error; + } + return 0; + } +}; + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h new file mode 100644 index 00000000000..625f1f48a8a --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h @@ -0,0 +1,27 @@ +/* + * 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 +#include + +namespace executorch::backends::coreai::testing { + +struct TestDirectory { + NSURL* url = nil; + TestDirectory(); + ~TestDirectory(); + TestDirectory(const TestDirectory&) = delete; + TestDirectory& operator=(const TestDirectory&) = delete; +}; + +::testing::AssertionResult backup_excluded(NSURL* url); +::testing::AssertionResult set_backup_excluded(NSURL* url, bool excluded); + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm new file mode 100644 index 00000000000..f9810e02aad --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm @@ -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. + */ + +#include "coreai_filesystem_fixture.h" +#include + +namespace executorch::backends::coreai::testing { + +TestDirectory::TestDirectory() { + char pattern[] = "/private/tmp/coreai-host-test-XXXXXX"; + if (mkdtemp(pattern) == nullptr) { + ADD_FAILURE() << "mkdtemp: " << strerror(errno); + return; + } + url = [NSURL fileURLWithPath:[NSString stringWithUTF8String:pattern] isDirectory:YES] + .URLByResolvingSymlinksInPath; +} + +TestDirectory::~TestDirectory() { + if (url == nil) return; + EXPECT_TRUE([url.lastPathComponent hasPrefix:@"coreai-host-test-"]); + if (![url.lastPathComponent hasPrefix:@"coreai-host-test-"]) return; + EXPECT_TRUE([NSFileManager.defaultManager removeItemAtURL:url error:nil]); +} + +::testing::AssertionResult backup_excluded(NSURL* url) { + NSURL* fresh = [NSURL fileURLWithPath:url.path]; + [fresh removeAllCachedResourceValues]; + NSNumber* value = nil; + NSError* error = nil; + if (![fresh getResourceValue:&value forKey:NSURLIsExcludedFromBackupKey error:&error] || + error != nil || value == nil) { + ADD_FAILURE() << "Cannot read backup flag: " << url.path.UTF8String; + return ::testing::AssertionFailure() << "Backup flag unavailable"; + } + return value.boolValue ? ::testing::AssertionSuccess() + : ::testing::AssertionFailure() << "Not excluded: " << url.path.UTF8String; +} + +::testing::AssertionResult set_backup_excluded(NSURL* url, bool excluded) { + // Foundation can defer metadata writes; settle deliberate fixture changes. + for (int attempt = 0; attempt < 50; ++attempt) { + NSURL* fresh = [NSURL fileURLWithPath:url.path]; + [fresh removeAllCachedResourceValues]; + if (![fresh setResourceValue:@(excluded) forKey:NSURLIsExcludedFromBackupKey error:nil]) { + return ::testing::AssertionFailure() << "Cannot set backup flag: " << url.path.UTF8String; + } + usleep(20000); + if (static_cast(backup_excluded(url)) == excluded) { + return ::testing::AssertionSuccess(); + } + } + return ::testing::AssertionFailure() << "Backup flag did not settle: " << url.path.UTF8String; +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_storage_test.mm b/backends/apple/coreai/runtime/test/coreai_storage_test.mm new file mode 100644 index 00000000000..9b1ae57c5a1 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_storage_test.mm @@ -0,0 +1,113 @@ +/* + * 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. + */ + +#include +#include "coreai_fault_scope.h" +#include "coreai_file.h" +#include "coreai_filesystem_fixture.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +TEST(CoreAIStorageTest, PreparesRootWithRootOnlyBackupPolicy) { + NSFileManager* manager = NSFileManager.defaultManager; + TestDirectory sandbox; + ASSERT_NE(sandbox.url, nil); + NSURL* parent = [sandbox.url URLByAppendingPathComponent:@"parent"]; + ASSERT_TRUE(([manager createDirectoryAtURL:parent + withIntermediateDirectories:NO + attributes:nil + error:nil])); + ASSERT_TRUE(set_backup_excluded(sandbox.url, false)); + ASSERT_TRUE(set_backup_excluded(parent, false)); + NSURL* root = [parent URLByAppendingPathComponent:@"nested/models"]; + auto prepared = prepare_storage_root(root.path, true); + ASSERT_TRUE((prepared.ok() && [prepared.get() isEqualToString:root.path])); + ASSERT_TRUE((backup_excluded(root))); + ASSERT_TRUE((!backup_excluded(parent) && !backup_excluded(sandbox.url))); + ASSERT_TRUE((!backup_excluded(root.URLByDeletingLastPathComponent))); + + ASSERT_TRUE(set_backup_excluded(root, false)); + auto inspected = prepare_storage_root(root.path, false); + ASSERT_TRUE((inspected.ok() && [inspected.get() isEqualToString:root.path])); + ASSERT_TRUE((!backup_excluded(root))); + ASSERT_TRUE((prepare_storage_root(root.path, true).ok() && backup_excluded(root))); + + NSURL* missing = [sandbox.url URLByAppendingPathComponent:@"missing"]; + ASSERT_EQ(prepare_storage_root(missing.path, false).error(), Error::InvalidArgument); + EXPECT_FALSE([manager fileExistsAtPath:missing.path]); + EXPECT_EQ(prepare_storage_root(@"relative", true).error(), Error::InvalidArgument); + EXPECT_TRUE((!backup_excluded(parent) && !backup_excluded(sandbox.url))); +} + +TEST(CoreAIStorageTest, ListsChildrenRepeatedlyInSortedOrder) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root( + open(storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_TRUE((root.get() >= 0)); + NSMutableArray* expected = [NSMutableArray array]; + for (int i = 7; i >= 0; --i) { + NSString* name = [NSString stringWithFormat:@"entry-%d", i]; + FileDescriptor file(openat(root.get(), name.fileSystemRepresentation, + O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, 0600)); + ASSERT_TRUE((file.get() >= 0)); + [expected insertObject:name atIndex:0]; + } + for (int pass = 0; pass < 3; ++pass) { + errno = EIO; + auto names = storage_children(root.get()); + ASSERT_TRUE((names.ok() && [names.get() isEqual:expected])); + ASSERT_TRUE((fcntl(root.get(), F_GETFD) >= 0)); + } +} + +TEST(CoreAIStorageTest, RetriesEintr) { + int calls = 0; + const ssize_t count = retry_eintr([&]() -> ssize_t { + if (++calls < 3) { + errno = EINTR; + return -1; + } + return 7; + }); + ASSERT_TRUE((calls == 3 && count == 7)); +} + +static ::testing::AssertionResult no_temporary_files(int root) { + auto children = storage_children(root); + if (!children.ok()) return ::testing::AssertionFailure() << "Cannot list storage"; + for (NSString* child in children.get()) { + if ([child hasPrefix:@".tmp-"]) return ::testing::AssertionFailure() << child.UTF8String; + } + return ::testing::AssertionSuccess(); +} + +TEST(CoreAIStorageTest, AtomicPublicationPreservesOldBytesOnFailure) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root( + open(storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_TRUE((root.get() >= 0)); + NSData* old = [@"old" dataUsingEncoding:NSUTF8StringEncoding]; + NSData* next = [@"next" dataUsingEncoding:NSUTF8StringEncoding]; + NSURL* published = [storage.url URLByAppendingPathComponent:@"published"]; + ASSERT_EQ(publish_storage_data(root.get(), @"published", old), Error::Ok); + { + StorageFaultScope fault(StorageOperation::Rename); + EXPECT_EQ(publish_storage_data(root.get(), @"published", next), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_TRUE(([[NSData dataWithContentsOfURL:published] isEqual:old])); + EXPECT_TRUE(no_temporary_files(root.get())); + + ASSERT_EQ(publish_storage_data(root.get(), @"published", next), Error::Ok); + EXPECT_TRUE(([[NSData dataWithContentsOfURL:published] isEqual:next])); + EXPECT_TRUE(no_temporary_files(root.get())); +} From 50f36345f4aac872dfa8e47cc79f661aa3916c10 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Mon, 26 Oct 2026 13:22:15 -0700 Subject: [PATCH 2/7] [CoreAI] Enable filesystem test fault injection The filesystem tests now own the fault-injection build setting that was removed from the parser-only base, where it had no consumers. Test Plan: Built and ran coreai_host_test locally on macOS 27. All 28 parser and filesystem tests passed, including injected atomic-publication failure. Authored with assistance from Devmate (OpenAI GPT-6-Astra). --- backends/apple/coreai/CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index 6a01341c31b..85d742adf3d 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -35,6 +35,7 @@ if(EXECUTORCH_BUILD_TESTS AND CMAKE_SYSTEM_NAME STREQUAL "Darwin") runtime/test/coreai_filesystem_fixture.mm ${_coreai_runtime_sources} ) coreai_configure_objc_target(coreai_host_test) + target_compile_definitions(coreai_host_test PRIVATE COREAI_ASSETS_TESTING=1) find_library(COREAI_FOUNDATION_FRAMEWORK Foundation REQUIRED) target_link_libraries( coreai_host_test PRIVATE executorch GTest::gtest From d3501fe87564e338b9b27fdd593f11897780581c Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Mon, 26 Oct 2026 14:23:10 -0700 Subject: [PATCH 3/7] up --- backends/apple/coreai/README.md | 11 +- .../apple/coreai/runtime/coreai_storage.h | 7 +- .../apple/coreai/runtime/coreai_storage.mm | 38 ++- .../runtime/test/coreai_storage_test.mm | 225 +++++++++++++++++- 4 files changed, 257 insertions(+), 24 deletions(-) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 623a96a14a6..b61bc86e6f1 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -56,8 +56,10 @@ The assets root must be an absolute path. The application chooses it, should reserve it for this backend, and is trusted not to rename, replace or modify its contents while the backend uses it; the backend does not defend against other processes of the same user rebinding paths. Written files and changed directories -are synced with `fsync`; atomic publication also uses `F_FULLFSYNC` before its -rename. +are synced with `fsync`. Atomic publication uses `F_FULLFSYNC` on the temporary +file before rename and on the directory afterward. A failure before rename +preserves the old contents; a failure of the final flush reports an error even +though the new complete contents are already visible. Preparing the assets root sets `NSURLIsExcludedFromBackupKey` on it, which covers everything beneath it. Ancestors are not modified. This is backup exclusion, not @@ -67,8 +69,9 @@ a control for iCloud Drive synchronization. On macOS, `EXECUTORCH_BUILD_COREAI=ON` with `EXECUTORCH_BUILD_TESTS=ON` registers the standalone `coreai_host_test` executable with CTest. It tests -manifest parsing and AOT asset selection using Foundation, without Swift or -CoreAI SDK calls. The host test target is omitted for Apple cross-builds. +manifest parsing, AOT asset selection, and filesystem storage using Foundation, +without Swift or CoreAI SDK calls. The host test target is omitted for Apple +cross-builds. CI builds and runs these tests on macOS 27. To use the same configuration from the repository root: diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h index f5bd377dd70..e3e1eb5b3a8 100644 --- a/backends/apple/coreai/runtime/coreai_storage.h +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -30,14 +30,17 @@ runtime::Error write_storage_file( NSString* relative_path, const void* bytes, size_t size); -// Atomically replaces name with data. On error, readers see either the old or -// the new complete bytes. +// Atomically replaces name with data and fully flushes the file and renamed +// entry. On error, readers see either the old or the new complete bytes. runtime::Error publish_storage_data(int root_fd, NSString* name, NSData* data); runtime::Result*> storage_children( int fd, bool skip_invalid_names = false); enum class StorageOperation { + Write, + SyncFile, + FullSync, Rename, BeforeSDK, AfterSDK, diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm index 66fea0e8833..8eed9099f63 100644 --- a/backends/apple/coreai/runtime/coreai_storage.mm +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -109,8 +109,8 @@ int storage_fault(StorageOperation operation) { namespace { template -int storage_call(StorageOperation operation, F call) { - return retry_eintr([&] { +auto storage_call(StorageOperation operation, F call) { + return retry_eintr([&]() -> decltype(call()) { const int fault = storage_fault(operation); if (fault != 0) { errno = fault; @@ -175,22 +175,27 @@ FileDescriptor file(retry_eintr([&] { file.get() >= 0, AccessFailed, "Cannot create Core AI storage file"); const auto* bytes = static_cast(data); while (size > 0) { - const ssize_t count = - write(file.get(), bytes, std::min(size, size_t(1024 * 1024))); - if (count < 0 && errno == EINTR) continue; + const ssize_t count = storage_call(StorageOperation::Write, [&] { + return write(file.get(), bytes, std::min(size, size_t(1024 * 1024))); + }); ET_CHECK_OR_RETURN_ERROR( count > 0, AccessFailed, "Cannot write Core AI storage file"); bytes += count; size -= count; } ET_CHECK_OR_RETURN_ERROR( - retry_eintr([&] { return fsync(file.get()); }) == 0 && + storage_call( + StorageOperation::SyncFile, [&] { return fsync(file.get()); }) == 0 && (!full_sync || - retry_eintr([&] { return fcntl(file.get(), F_FULLFSYNC); }) == 0), + storage_call( + StorageOperation::FullSync, + [&] { return fcntl(file.get(), F_FULLFSYNC); }) == 0), AccessFailed, "Cannot synchronize Core AI file"); - ET_CHECK_OR_RETURN_ERROR(::close(file.release()) == 0, AccessFailed, - "Cannot close Core AI storage file"); + ET_CHECK_OR_RETURN_ERROR( + ::close(file.release()) == 0, + AccessFailed, + "Cannot close Core AI storage file"); return sync_storage_directory(parent.get()); } } // namespace @@ -222,8 +227,7 @@ FileDescriptor copy( NSString* name = [[NSString alloc] initWithBytes:entry->d_name length:strlen(entry->d_name) encoding:NSUTF8StringEncoding]; - if (name == nil || - strcmp(name.fileSystemRepresentation, entry->d_name) != 0) { + if (name == nil) { if (skip_invalid_names) { continue; } @@ -287,8 +291,16 @@ Error publish_storage_data(int root_fd, NSString* name, NSData* data) { name.fileSystemRepresentation); }) != 0) result = Error::AccessFailed; - if (result == Error::Ok) - return sync_storage_directory(root_fd); + if (result == Error::Ok) { + // Flush the renamed entry as well as the file's pre-rename contents. + ET_CHECK_OR_RETURN_ERROR( + storage_call( + StorageOperation::FullSync, + [&] { return fcntl(root_fd, F_FULLFSYNC); }) == 0, + AccessFailed, + "Cannot fully synchronize Core AI publication directory"); + return Error::Ok; + } if (unlinkat(root_fd, temporary.fileSystemRepresentation, 0) == 0) { (void)sync_storage_directory(root_fd); } diff --git a/backends/apple/coreai/runtime/test/coreai_storage_test.mm b/backends/apple/coreai/runtime/test/coreai_storage_test.mm index 9b1ae57c5a1..064aab1a9d5 100644 --- a/backends/apple/coreai/runtime/test/coreai_storage_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_storage_test.mm @@ -6,7 +6,10 @@ * LICENSE file in the root directory of this source tree. */ +#include #include +#include +#include #include "coreai_fault_scope.h" #include "coreai_file.h" #include "coreai_filesystem_fixture.h" @@ -68,6 +71,141 @@ FileDescriptor file(openat(root.get(), name.fileSystemRepresentation, } } +TEST(CoreAIStorageTest, ListsComposedAndDecomposedUnicodeNames) { + for (NSString* name in @[ @"caf\u00e9", @"cafe\u0301" ]) { + SCOPED_TRACE(name.UTF8String); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + FileDescriptor file(openat( + root.get(), + name.UTF8String, + O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, + 0600)); + ASSERT_GE(file.get(), 0); + for (bool skip_invalid : {false, true}) { + auto names = storage_children(root.get(), skip_invalid); + ASSERT_TRUE(names.ok()); + ASSERT_EQ(names.get().count, 1); + EXPECT_STREQ(names.get().firstObject.UTF8String, name.UTF8String); + } + } +} + +TEST(CoreAIStorageTest, SkipsInvalidUtf8NamesOnlyWhenRequested) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + FileDescriptor invalid(openat( + root.get(), + "invalid-\xff", + O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, + 0600)); + if (invalid.get() < 0 && errno == EILSEQ) { + GTEST_SKIP() << "Filesystem rejects invalid UTF-8 filenames"; + } + ASSERT_GE(invalid.get(), 0); + FileDescriptor valid(openat( + root.get(), "valid", O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, 0600)); + ASSERT_GE(valid.get(), 0); + EXPECT_EQ(storage_children(root.get()).error(), Error::InvalidExternalData); + auto names = storage_children(root.get(), true); + ASSERT_TRUE(names.ok()); + EXPECT_TRUE(([names.get() isEqual:@[ @"valid" ]])); + EXPECT_EQ(unlinkat(root.get(), "invalid-\xff", 0), 0); +} + +TEST(CoreAIStorageTest, CreatesNestedFilesWithoutOverwritingExistingContents) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + NSData* original = [@"original" dataUsingEncoding:NSUTF8StringEncoding]; + NSData* replacement = [@"replacement" dataUsingEncoding:NSUTF8StringEncoding]; + NSString* path = @"nested/parents/model.bin"; + NSURL* file = [storage.url URLByAppendingPathComponent:path]; + ASSERT_EQ( + write_storage_file(root.get(), path, original.bytes, original.length), + Error::Ok); + EXPECT_TRUE(([[NSData dataWithContentsOfURL:file] isEqual:original])); + EXPECT_EQ( + write_storage_file( + root.get(), path, replacement.bytes, replacement.length), + Error::AccessFailed); + EXPECT_TRUE(([[NSData dataWithContentsOfURL:file] isEqual:original])); + ASSERT_EQ( + write_storage_file(root.get(), @"nested/parents/empty", nullptr, 0), + Error::Ok); + NSData* empty = [NSData + dataWithContentsOfURL: + [storage.url URLByAppendingPathComponent:@"nested/parents/empty"]]; + ASSERT_NE(empty, nil); + EXPECT_EQ(empty.length, 0); +} + +TEST(CoreAIStorageTest, RejectsUnsafeWritePathsWithoutCreatingFiles) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + for (NSString* path in @[ + @"", + @"..", + @"../escape", + @"a/../b", + @"a/./b", + @"a//b", + @"/leading", + @"trailing/" + ]) { + SCOPED_TRACE(path.UTF8String); + EXPECT_EQ( + write_storage_file(root.get(), path, "x", 1), Error::InvalidArgument); + } + EXPECT_EQ( + write_storage_file(root.get(), @"no-data", nullptr, 1), + Error::InvalidArgument); + auto names = storage_children(root.get()); + ASSERT_TRUE(names.ok()); + EXPECT_EQ(names.get().count, 0); +} + +TEST(CoreAIStorageTest, CreatesAndReopensSharedFileWithoutTruncation) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + FileDescriptor created(open_storage_shared_file(root.get(), "shared")); + ASSERT_GE(created.get(), 0); + ASSERT_EQ(write(created.get(), "data", 4), 4); + FileDescriptor reopened(open_storage_shared_file(root.get(), "shared")); + ASSERT_GE(reopened.get(), 0); + EXPECT_NE(created.get(), reopened.get()); + EXPECT_NE(fcntl(created.get(), F_GETFD) & FD_CLOEXEC, 0); + EXPECT_NE(fcntl(reopened.get(), F_GETFD) & FD_CLOEXEC, 0); + struct stat first, second; + ASSERT_EQ(fstat(created.get(), &first), 0); + ASSERT_EQ(fstat(reopened.get(), &second), 0); + EXPECT_EQ(first.st_ino, second.st_ino); + EXPECT_EQ(second.st_size, 4); + char bytes[4]; + ASSERT_EQ(read(reopened.get(), bytes, sizeof(bytes)), sizeof(bytes)); + EXPECT_EQ(memcmp(bytes, "data", sizeof(bytes)), 0); + EXPECT_EQ(lseek(created.get(), 0, SEEK_CUR), 4); +} + TEST(CoreAIStorageTest, RetriesEintr) { int calls = 0; const ssize_t count = retry_eintr([&]() -> ssize_t { @@ -99,15 +237,92 @@ FileDescriptor root( NSData* next = [@"next" dataUsingEncoding:NSUTF8StringEncoding]; NSURL* published = [storage.url URLByAppendingPathComponent:@"published"]; ASSERT_EQ(publish_storage_data(root.get(), @"published", old), Error::Ok); - { - StorageFaultScope fault(StorageOperation::Rename); - EXPECT_EQ(publish_storage_data(root.get(), @"published", next), Error::AccessFailed); + for (StorageOperation operation : + {StorageOperation::Write, + StorageOperation::SyncFile, + StorageOperation::FullSync, + StorageOperation::Rename}) { + SCOPED_TRACE(static_cast(operation)); + StorageFaultScope fault(operation); + EXPECT_EQ( + publish_storage_data(root.get(), @"published", next), + Error::AccessFailed); EXPECT_EQ(fault.hits, 1); + EXPECT_TRUE(([[NSData dataWithContentsOfURL:published] isEqual:old])); + EXPECT_TRUE(no_temporary_files(root.get())); } - EXPECT_TRUE(([[NSData dataWithContentsOfURL:published] isEqual:old])); - EXPECT_TRUE(no_temporary_files(root.get())); ASSERT_EQ(publish_storage_data(root.get(), @"published", next), Error::Ok); EXPECT_TRUE(([[NSData dataWithContentsOfURL:published] isEqual:next])); EXPECT_TRUE(no_temporary_files(root.get())); } + +TEST(CoreAIStorageTest, FullyFlushesPublicationBeforeAndAfterRename) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + __block std::vector operations; + StorageFaultScope fault(StorageOperation::FullSync); + fault.armed = false; + fault.observe = ^(StorageOperation operation) { + operations.push_back(operation); + }; + NSData* data = [@"published" dataUsingEncoding:NSUTF8StringEncoding]; + EXPECT_EQ(publish_storage_data(root.get(), @"published", data), Error::Ok); + EXPECT_EQ( + operations, + (std::vector{ + StorageOperation::Write, + StorageOperation::SyncFile, + StorageOperation::FullSync, + StorageOperation::Rename, + StorageOperation::FullSync})); +} + +TEST(CoreAIStorageTest, ReportsPostRenameFlushFailureWithNewContentsVisible) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + NSData* old = [@"old" dataUsingEncoding:NSUTF8StringEncoding]; + NSData* next = [@"next" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_EQ(publish_storage_data(root.get(), @"published", old), Error::Ok); + StorageFaultScope fault(StorageOperation::FullSync); + fault.match_index = 2; + EXPECT_EQ( + publish_storage_data(root.get(), @"published", next), + Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + NSURL* published = [storage.url URLByAppendingPathComponent:@"published"]; + EXPECT_TRUE(([[NSData dataWithContentsOfURL:published] isEqual:next])); + EXPECT_TRUE(no_temporary_files(root.get())); +} + +TEST(CoreAIStorageTest, RetriesInterruptedPublicationOperations) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + NSData* data = [@"published" dataUsingEncoding:NSUTF8StringEncoding]; + for (StorageOperation operation : + {StorageOperation::Write, + StorageOperation::SyncFile, + StorageOperation::FullSync, + StorageOperation::Rename}) { + SCOPED_TRACE(static_cast(operation)); + StorageFaultScope fault(operation); + fault.error = EINTR; + EXPECT_EQ(publish_storage_data(root.get(), @"published", data), Error::Ok); + EXPECT_EQ(fault.hits, 1); + NSURL* published = [storage.url URLByAppendingPathComponent:@"published"]; + EXPECT_TRUE(([[NSData dataWithContentsOfURL:published] isEqual:data])); + EXPECT_TRUE(no_temporary_files(root.get())); + } +} From 20d9d37d2ddb6e6f366b812f048e4eaeff9d2828 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 27 Oct 2026 12:06:17 -0700 Subject: [PATCH 4/7] [CoreAI] Harden storage failure handling and directory enumeration Summary: Keep failed writes from blocking retries and prevent directory listings from disturbing other readers. Reject invalid roots consistently, preserve strict publication durability, and cover directory-sync failures. Remove a fixture that APFS cannot create. Test Plan: All 46 native tests passed with zero skips, including partial-write cleanup and retry, independent enumeration offsets, root validation, and directory-sync ordering/failure/interruption checks. All 113 CoreAI Python tests passed. Formatting and whitespace checks passed. --- backends/apple/coreai/README.md | 3 +- .../apple/coreai/runtime/coreai_storage.h | 3 + .../apple/coreai/runtime/coreai_storage.mm | 153 ++++++++++------- .../runtime/test/coreai_filesystem_fixture.mm | 28 ++- .../runtime/test/coreai_storage_test.mm | 159 +++++++++++++----- 5 files changed, 233 insertions(+), 113 deletions(-) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index b61bc86e6f1..09361117fd6 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -59,7 +59,8 @@ processes of the same user rebinding paths. Written files and changed directorie are synced with `fsync`. Atomic publication uses `F_FULLFSYNC` on the temporary file before rename and on the directory afterward. A failure before rename preserves the old contents; a failure of the final flush reports an error even -though the new complete contents are already visible. +though the new complete contents are already visible. Use a local APFS volume; +unsupported full-device flushes are reported as errors, not downgraded to `fsync`. Preparing the assets root sets `NSURLIsExcludedFromBackupKey` on it, which covers everything beneath it. Ancestors are not modified. This is backup exclusion, not diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h index e3e1eb5b3a8..bc583e59997 100644 --- a/backends/apple/coreai/runtime/coreai_storage.h +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -25,6 +25,7 @@ runtime::Error sync_storage_directory(int fd); // Returns an independently owned descriptor for a stable coordination file. int open_storage_shared_file(int parent_fd, const char* name); // Creates a new file and missing parents, then syncs the file and its parent. +// On failure, attempts to remove the new file; created parents may remain. runtime::Error write_storage_file( int root_fd, NSString* relative_path, @@ -33,6 +34,7 @@ runtime::Error write_storage_file( // Atomically replaces name with data and fully flushes the file and renamed // entry. On error, readers see either the old or the new complete bytes. runtime::Error publish_storage_data(int root_fd, NSString* name, NSData* data); +// Lists sorted names without changing fd's enumeration offset. runtime::Result*> storage_children( int fd, bool skip_invalid_names = false); @@ -40,6 +42,7 @@ runtime::Result*> storage_children( enum class StorageOperation { Write, SyncFile, + SyncDirectory, FullSync, Rename, BeforeSDK, diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm index 8eed9099f63..5c32d432cd4 100644 --- a/backends/apple/coreai/runtime/coreai_storage.mm +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -9,9 +9,9 @@ #import "coreai_storage.h" #include "coreai_file.h" -#include -#include #include +#include +#include #include #include #include @@ -49,20 +49,32 @@ Error ensure_excluded_from_backup(NSURL* directory) { } Result default_coreai_assets_root() { - NSURL* caches = [NSFileManager.defaultManager - URLsForDirectory:NSCachesDirectory - inDomains:NSUserDomainMask].firstObject; - ET_CHECK_OR_RETURN_ERROR(caches.isFileURL, AccessFailed, - "Cannot resolve Core AI cache directory"); + NSURL* caches = + [NSFileManager.defaultManager URLsForDirectory:NSCachesDirectory + inDomains:NSUserDomainMask] + .firstObject; + ET_CHECK_OR_RETURN_ERROR( + caches.isFileURL, AccessFailed, "Cannot resolve Core AI cache directory"); return [caches.path stringByAppendingPathComponent:@"executorch_coreai"]; } Result prepare_storage_root(NSString* path, bool create) { ET_CHECK_OR_RETURN_ERROR( - [path isKindOfClass:NSString.class] && path.isAbsolutePath && + [path isKindOfClass:NSString.class] && [path hasPrefix:@"/"] && [path rangeOfString:@"\0"].location == NSNotFound, InvalidArgument, "Core AI storage requires an absolute path"); + struct stat info; + const int result = + retry_eintr([&] { return stat(path.fileSystemRepresentation, &info); }); + ET_CHECK_OR_RETURN_ERROR( + result == 0 || errno == ENOENT || errno == ENOTDIR, + AccessFailed, + "Cannot inspect Core AI storage directory"); + ET_CHECK_OR_RETURN_ERROR( + result == 0 ? S_ISDIR(info.st_mode) : create && errno == ENOENT, + InvalidArgument, + "Core AI storage must be a directory"); if (create) { NSError* error = nil; ET_CHECK_OR_RETURN_ERROR( @@ -76,19 +88,6 @@ Error ensure_excluded_from_backup(NSURL* directory) { AccessFailed, "Cannot create Core AI storage directory: %s", error.localizedDescription.UTF8String); - } - struct stat info; - const int result = - retry_eintr([&] { return stat(path.fileSystemRepresentation, &info); }); - ET_CHECK_OR_RETURN_ERROR( - result == 0 || errno == ENOENT || errno == ENOTDIR, - AccessFailed, - "Cannot inspect Core AI storage directory"); - ET_CHECK_OR_RETURN_ERROR( - result == 0 && S_ISDIR(info.st_mode), - InvalidArgument, - "Core AI storage must be an existing directory"); - if (create) { ET_CHECK_OK_OR_RETURN_ERROR(ensure_excluded_from_backup( [NSURL fileURLWithPath:path isDirectory:YES])); } @@ -142,14 +141,15 @@ Error write_file( ET_CHECK_OR_RETURN_ERROR( storage_component(component), InvalidArgument, "Invalid storage path"); } - FileDescriptor parent(retry_eintr([&] { return dup(root_fd); })); + FileDescriptor parent( + retry_eintr([&] { return fcntl(root_fd, F_DUPFD_CLOEXEC, 0); })); ET_CHECK_OR_RETURN_ERROR( parent.get() >= 0, AccessFailed, "Cannot open storage parent"); for (NSUInteger i = 0; i + 1 < components.count; ++i) { const char* component = components[i].fileSystemRepresentation; ET_CHECK_OR_RETURN_ERROR( - retry_eintr( - [&] { return mkdirat(parent.get(), component, 0700); }) == 0 || + retry_eintr([&] { return mkdirat(parent.get(), component, 0700); }) == + 0 || errno == EEXIST, AccessFailed, "Cannot create Core AI source directory"); @@ -173,52 +173,70 @@ FileDescriptor file(retry_eintr([&] { })); ET_CHECK_OR_RETURN_ERROR( file.get() >= 0, AccessFailed, "Cannot create Core AI storage file"); - const auto* bytes = static_cast(data); - while (size > 0) { - const ssize_t count = storage_call(StorageOperation::Write, [&] { - return write(file.get(), bytes, std::min(size, size_t(1024 * 1024))); - }); + const auto result = [&]() -> Error { + const auto* bytes = static_cast(data); + while (size > 0) { + const ssize_t count = storage_call(StorageOperation::Write, [&] { + return write(file.get(), bytes, std::min(size, size_t(1024 * 1024))); + }); + ET_CHECK_OR_RETURN_ERROR( + count > 0, AccessFailed, "Cannot write Core AI storage file"); + bytes += count; + size -= count; + } + ET_CHECK_OR_RETURN_ERROR( + storage_call( + StorageOperation::SyncFile, [&] { return fsync(file.get()); }) == + 0 && + (!full_sync || + storage_call( + StorageOperation::FullSync, + [&] { return fcntl(file.get(), F_FULLFSYNC); }) == 0), + AccessFailed, + "Cannot synchronize Core AI file"); ET_CHECK_OR_RETURN_ERROR( - count > 0, AccessFailed, "Cannot write Core AI storage file"); - bytes += count; - size -= count; + ::close(file.release()) == 0, + AccessFailed, + "Cannot close Core AI storage file"); + return sync_storage_directory(parent.get()); + }(); + if (result != Error::Ok) { + if (retry_eintr([&] { + return unlinkat( + parent.get(), components.lastObject.fileSystemRepresentation, 0); + }) == 0) { + (void)sync_storage_directory(parent.get()); + } else { + ET_LOG( + Error, "Cannot remove failed Core AI storage write: errno=%d", errno); + } } - ET_CHECK_OR_RETURN_ERROR( - storage_call( - StorageOperation::SyncFile, [&] { return fsync(file.get()); }) == 0 && - (!full_sync || - storage_call( - StorageOperation::FullSync, - [&] { return fcntl(file.get(), F_FULLFSYNC); }) == 0), - AccessFailed, - "Cannot synchronize Core AI file"); - ET_CHECK_OR_RETURN_ERROR( - ::close(file.release()) == 0, - AccessFailed, - "Cannot close Core AI storage file"); - return sync_storage_directory(parent.get()); + return result; } } // namespace Result*> storage_children(int fd, bool skip_invalid_names) { - FileDescriptor copy( - retry_eintr([&] { return fcntl(fd, F_DUPFD_CLOEXEC, 0); })); - ET_CHECK_OR_RETURN_ERROR(copy.get() >= 0, AccessFailed, - "Cannot duplicate storage directory"); + // Reopening avoids sharing the caller's directory offset. + FileDescriptor copy(retry_eintr( + [&] { return openat(fd, ".", O_RDONLY | O_DIRECTORY | O_CLOEXEC); })); + ET_CHECK_OR_RETURN_ERROR( + copy.get() >= 0, AccessFailed, "Cannot open storage directory"); DIR* directory = fdopendir(copy.get()); - ET_CHECK_OR_RETURN_ERROR(directory != nullptr, AccessFailed, - "Cannot open storage directory stream"); + ET_CHECK_OR_RETURN_ERROR( + directory != nullptr, + AccessFailed, + "Cannot open storage directory stream"); copy.release(); - // dup shares the enumeration offset with the original descriptor. - rewinddir(directory); NSMutableArray* names = [NSMutableArray array]; Error result = Error::Ok; while (true) { errno = 0; dirent* entry = readdir(directory); if (entry == nullptr) { - if (errno == EINTR) continue; - if (errno != 0) result = Error::AccessFailed; + if (errno == EINTR) + continue; + if (errno != 0) + result = Error::AccessFailed; break; } if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0) { @@ -238,16 +256,19 @@ FileDescriptor copy( } const int closed = closedir(directory); ET_CHECK_OK_OR_RETURN_ERROR(result); - ET_CHECK_OR_RETURN_ERROR(closed == 0, AccessFailed, - "Cannot close storage directory stream"); + ET_CHECK_OR_RETURN_ERROR( + closed == 0, AccessFailed, "Cannot close storage directory stream"); return [names sortedArrayUsingSelector:@selector(compare:)]; } int open_storage_shared_file(int parent_fd, const char* name) { // APFS can return ENOENT to a concurrent O_CREAT loser without O_EXCL. int fd = retry_eintr([&] { - return openat(parent_fd, name, - O_RDWR | O_CREAT | O_EXCL | O_NOFOLLOW | O_CLOEXEC, 0600); + return openat( + parent_fd, + name, + O_RDWR | O_CREAT | O_EXCL | O_NOFOLLOW | O_CLOEXEC, + 0600); }); if (fd < 0 && errno == EEXIST) { fd = retry_eintr([&] { @@ -259,7 +280,8 @@ int open_storage_shared_file(int parent_fd, const char* name) { Error sync_storage_directory(int fd) { ET_CHECK_OR_RETURN_ERROR( - retry_eintr([&] { return fsync(fd); }) == 0, + storage_call( + StorageOperation::SyncDirectory, [&] { return fsync(fd); }) == 0, AccessFailed, "Cannot synchronize Core AI directory"); return Error::Ok; @@ -281,16 +303,17 @@ Error publish_storage_data(int root_fd, NSString* name, NSData* data) { NSString* temporary = [@".tmp-" stringByAppendingString:NSUUID.UUID.UUIDString]; // Bookmarks are small and costly to lose, so they get a full device flush. - auto result = - write_file(root_fd, temporary, data.bytes, data.length, true); + auto result = write_file(root_fd, temporary, data.bytes, data.length, true); if (result == Error::Ok && storage_call(StorageOperation::Rename, [&] { return renameat( root_fd, temporary.fileSystemRepresentation, root_fd, name.fileSystemRepresentation); - }) != 0) + }) != 0) { + ET_LOG(Error, "Cannot publish Core AI storage data: errno=%d", errno); result = Error::AccessFailed; + } if (result == Error::Ok) { // Flush the renamed entry as well as the file's pre-rename contents. ET_CHECK_OR_RETURN_ERROR( @@ -307,4 +330,4 @@ Error publish_storage_data(int root_fd, NSString* name, NSData* data) { return result; } -} // namespace executorch::backends::coreai +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm index f9810e02aad..6c60dddd6c5 100644 --- a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm +++ b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm @@ -17,14 +17,17 @@ ADD_FAILURE() << "mkdtemp: " << strerror(errno); return; } - url = [NSURL fileURLWithPath:[NSString stringWithUTF8String:pattern] isDirectory:YES] + url = [NSURL fileURLWithPath:[NSString stringWithUTF8String:pattern] + isDirectory:YES] .URLByResolvingSymlinksInPath; } TestDirectory::~TestDirectory() { - if (url == nil) return; + if (url == nil) + return; EXPECT_TRUE([url.lastPathComponent hasPrefix:@"coreai-host-test-"]); - if (![url.lastPathComponent hasPrefix:@"coreai-host-test-"]) return; + if (![url.lastPathComponent hasPrefix:@"coreai-host-test-"]) + return; EXPECT_TRUE([NSFileManager.defaultManager removeItemAtURL:url error:nil]); } @@ -33,13 +36,16 @@ [fresh removeAllCachedResourceValues]; NSNumber* value = nil; NSError* error = nil; - if (![fresh getResourceValue:&value forKey:NSURLIsExcludedFromBackupKey error:&error] || + if (![fresh getResourceValue:&value + forKey:NSURLIsExcludedFromBackupKey + error:&error] || error != nil || value == nil) { ADD_FAILURE() << "Cannot read backup flag: " << url.path.UTF8String; return ::testing::AssertionFailure() << "Backup flag unavailable"; } return value.boolValue ? ::testing::AssertionSuccess() - : ::testing::AssertionFailure() << "Not excluded: " << url.path.UTF8String; + : ::testing::AssertionFailure() + << "Not excluded: " << url.path.UTF8String; } ::testing::AssertionResult set_backup_excluded(NSURL* url, bool excluded) { @@ -47,15 +53,19 @@ for (int attempt = 0; attempt < 50; ++attempt) { NSURL* fresh = [NSURL fileURLWithPath:url.path]; [fresh removeAllCachedResourceValues]; - if (![fresh setResourceValue:@(excluded) forKey:NSURLIsExcludedFromBackupKey error:nil]) { - return ::testing::AssertionFailure() << "Cannot set backup flag: " << url.path.UTF8String; + if (![fresh setResourceValue:@(excluded) + forKey:NSURLIsExcludedFromBackupKey + error:nil]) { + return ::testing::AssertionFailure() + << "Cannot set backup flag: " << url.path.UTF8String; } usleep(20000); if (static_cast(backup_excluded(url)) == excluded) { return ::testing::AssertionSuccess(); } } - return ::testing::AssertionFailure() << "Backup flag did not settle: " << url.path.UTF8String; + return ::testing::AssertionFailure() + << "Backup flag did not settle: " << url.path.UTF8String; } -} // namespace executorch::backends::coreai::testing +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_storage_test.mm b/backends/apple/coreai/runtime/test/coreai_storage_test.mm index 064aab1a9d5..77b23eef53f 100644 --- a/backends/apple/coreai/runtime/test/coreai_storage_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_storage_test.mm @@ -6,9 +6,10 @@ * LICENSE file in the root directory of this source tree. */ -#include #include #include +#include +#include #include #include "coreai_fault_scope.h" #include "coreai_file.h" @@ -40,34 +41,72 @@ auto inspected = prepare_storage_root(root.path, false); ASSERT_TRUE((inspected.ok() && [inspected.get() isEqualToString:root.path])); ASSERT_TRUE((!backup_excluded(root))); - ASSERT_TRUE((prepare_storage_root(root.path, true).ok() && backup_excluded(root))); + ASSERT_TRUE( + (prepare_storage_root(root.path, true).ok() && backup_excluded(root))); NSURL* missing = [sandbox.url URLByAppendingPathComponent:@"missing"]; - ASSERT_EQ(prepare_storage_root(missing.path, false).error(), Error::InvalidArgument); + ASSERT_EQ( + prepare_storage_root(missing.path, false).error(), + Error::InvalidArgument); EXPECT_FALSE([manager fileExistsAtPath:missing.path]); - EXPECT_EQ(prepare_storage_root(@"relative", true).error(), Error::InvalidArgument); + EXPECT_EQ( + prepare_storage_root(@"relative", true).error(), Error::InvalidArgument); EXPECT_TRUE((!backup_excluded(parent) && !backup_excluded(sandbox.url))); } +TEST(CoreAIStorageTest, RejectsInvalidRootsConsistently) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* regular = [storage.url URLByAppendingPathComponent:@"regular"]; + NSData* bytes = [@"unchanged" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([bytes writeToURL:regular atomically:NO]); + for (bool create : {false, true}) { + SCOPED_TRACE(create); + EXPECT_EQ( + prepare_storage_root(nil, create).error(), Error::InvalidArgument); + for (NSString* path in @[ + @"", + @"relative", + @"~", + @"~/cache", + @"~nobody/cache", + regular.path, + [regular.path stringByAppendingPathComponent:@"child"] + ]) { + SCOPED_TRACE(path.UTF8String); + EXPECT_EQ( + prepare_storage_root(path, create).error(), Error::InvalidArgument); + } + } + EXPECT_TRUE([[NSData dataWithContentsOfURL:regular] isEqual:bytes]); +} + TEST(CoreAIStorageTest, ListsChildrenRepeatedlyInSortedOrder) { TestDirectory storage; ASSERT_NE(storage.url, nil); - FileDescriptor root( - open(storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_TRUE((root.get() >= 0)); NSMutableArray* expected = [NSMutableArray array]; for (int i = 7; i >= 0; --i) { NSString* name = [NSString stringWithFormat:@"entry-%d", i]; - FileDescriptor file(openat(root.get(), name.fileSystemRepresentation, - O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, 0600)); + FileDescriptor file(openat( + root.get(), + name.fileSystemRepresentation, + O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, + 0600)); ASSERT_TRUE((file.get() >= 0)); [expected insertObject:name atIndex:0]; } + const off_t position = lseek(root.get(), 0, SEEK_CUR); + ASSERT_GE(position, 0); for (int pass = 0; pass < 3; ++pass) { errno = EIO; auto names = storage_children(root.get()); ASSERT_TRUE((names.ok() && [names.get() isEqual:expected])); ASSERT_TRUE((fcntl(root.get(), F_GETFD) >= 0)); + EXPECT_EQ(lseek(root.get(), 0, SEEK_CUR), position); } } @@ -95,32 +134,6 @@ FileDescriptor file(openat( } } -TEST(CoreAIStorageTest, SkipsInvalidUtf8NamesOnlyWhenRequested) { - TestDirectory storage; - ASSERT_NE(storage.url, nil); - FileDescriptor root(open( - storage.url.fileSystemRepresentation, - O_RDONLY | O_DIRECTORY | O_CLOEXEC)); - ASSERT_GE(root.get(), 0); - FileDescriptor invalid(openat( - root.get(), - "invalid-\xff", - O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, - 0600)); - if (invalid.get() < 0 && errno == EILSEQ) { - GTEST_SKIP() << "Filesystem rejects invalid UTF-8 filenames"; - } - ASSERT_GE(invalid.get(), 0); - FileDescriptor valid(openat( - root.get(), "valid", O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, 0600)); - ASSERT_GE(valid.get(), 0); - EXPECT_EQ(storage_children(root.get()).error(), Error::InvalidExternalData); - auto names = storage_children(root.get(), true); - ASSERT_TRUE(names.ok()); - EXPECT_TRUE(([names.get() isEqual:@[ @"valid" ]])); - EXPECT_EQ(unlinkat(root.get(), "invalid-\xff", 0), 0); -} - TEST(CoreAIStorageTest, CreatesNestedFilesWithoutOverwritingExistingContents) { TestDirectory storage; ASSERT_NE(storage.url, nil); @@ -151,6 +164,70 @@ FileDescriptor root(open( EXPECT_EQ(empty.length, 0); } +TEST(CoreAIStorageTest, RemovesFailedWritesAndAllowsRetry) { + const struct { + StorageOperation operation; + int match_index; + } rows[] = { + {StorageOperation::Write, 2}, + {StorageOperation::SyncFile, 1}, + {StorageOperation::SyncDirectory, 1}, + {StorageOperation::SyncDirectory, 2}}; + const std::vector bytes(1024 * 1024 + 1, 0x5a); + for (const auto& row : rows) { + SCOPED_TRACE(static_cast(row.operation)); + SCOPED_TRACE(row.match_index); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + NSString* path = @"nested/model.bin"; + NSURL* file = [storage.url URLByAppendingPathComponent:path]; + StorageFaultScope fault(row.operation); + fault.match_index = row.match_index; + EXPECT_EQ( + write_storage_file(root.get(), path, bytes.data(), bytes.size()), + Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:file.path]); + ASSERT_EQ( + write_storage_file(root.get(), path, bytes.data(), bytes.size()), + Error::Ok); + NSData* actual = [NSData dataWithContentsOfURL:file]; + ASSERT_NE(actual, nil); + ASSERT_EQ(actual.length, bytes.size()); + EXPECT_EQ(memcmp(actual.bytes, bytes.data(), bytes.size()), 0); + } +} + +TEST(CoreAIStorageTest, SyncsEachCreatedParentAndTheWrittenFile) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + __block std::vector operations; + StorageFaultScope fault(StorageOperation::SyncDirectory); + fault.armed = false; + fault.observe = ^(StorageOperation operation) { + operations.push_back(operation); + }; + EXPECT_EQ( + write_storage_file(root.get(), @"nested/parents/model.bin", "x", 1), + Error::Ok); + EXPECT_EQ( + operations, + (std::vector{ + StorageOperation::SyncDirectory, + StorageOperation::SyncDirectory, + StorageOperation::Write, + StorageOperation::SyncFile, + StorageOperation::SyncDirectory})); +} + TEST(CoreAIStorageTest, RejectsUnsafeWritePathsWithoutCreatingFiles) { TestDirectory storage; ASSERT_NE(storage.url, nil); @@ -220,9 +297,11 @@ FileDescriptor root(open( static ::testing::AssertionResult no_temporary_files(int root) { auto children = storage_children(root); - if (!children.ok()) return ::testing::AssertionFailure() << "Cannot list storage"; + if (!children.ok()) + return ::testing::AssertionFailure() << "Cannot list storage"; for (NSString* child in children.get()) { - if ([child hasPrefix:@".tmp-"]) return ::testing::AssertionFailure() << child.UTF8String; + if ([child hasPrefix:@".tmp-"]) + return ::testing::AssertionFailure() << child.UTF8String; } return ::testing::AssertionSuccess(); } @@ -230,8 +309,9 @@ FileDescriptor root(open( TEST(CoreAIStorageTest, AtomicPublicationPreservesOldBytesOnFailure) { TestDirectory storage; ASSERT_NE(storage.url, nil); - FileDescriptor root( - open(storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_TRUE((root.get() >= 0)); NSData* old = [@"old" dataUsingEncoding:NSUTF8StringEncoding]; NSData* next = [@"next" dataUsingEncoding:NSUTF8StringEncoding]; @@ -240,6 +320,7 @@ FileDescriptor root( for (StorageOperation operation : {StorageOperation::Write, StorageOperation::SyncFile, + StorageOperation::SyncDirectory, StorageOperation::FullSync, StorageOperation::Rename}) { SCOPED_TRACE(static_cast(operation)); @@ -278,6 +359,7 @@ FileDescriptor root(open( StorageOperation::Write, StorageOperation::SyncFile, StorageOperation::FullSync, + StorageOperation::SyncDirectory, StorageOperation::Rename, StorageOperation::FullSync})); } @@ -314,6 +396,7 @@ FileDescriptor root(open( for (StorageOperation operation : {StorageOperation::Write, StorageOperation::SyncFile, + StorageOperation::SyncDirectory, StorageOperation::FullSync, StorageOperation::Rename}) { SCOPED_TRACE(static_cast(operation)); From 94e3b32f52993f1374c2455ef399624e7d81da2c Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 28 Oct 2026 15:30:18 -0700 Subject: [PATCH 5/7] [CoreAI] Protect new caches and reuse existing directories Reduce redundant filesystem work during cache preparation, particularly on iOS, while giving newly created cache contents an explicit first-unlock protection policy. Preserve backup exclusion and durability behavior, use default directory permissions, and leave existing cache contents unmigrated. Test Plan: Native filesystem tests passed, including protection inheritance, default permissions, creation-free warm paths, concurrent creation, invalid paths, and failure/retry durability. The combined bookmark-layer suite previously passed all 95 native tests; the Python suite passed all 113 tests. --- backends/apple/coreai/README.md | 10 + .../apple/coreai/runtime/coreai_storage.h | 9 + .../apple/coreai/runtime/coreai_storage.mm | 77 ++++-- .../runtime/test/coreai_storage_test.mm | 237 ++++++++++++++++++ 4 files changed, 310 insertions(+), 23 deletions(-) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 09361117fd6..2ca9d0b8665 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -66,6 +66,16 @@ Preparing the assets root sets `NSURLIsExcludedFromBackupKey` on it, which cover everything beneath it. Ancestors are not modified. This is backup exclusion, not a control for iCloud Drive synchronization. +Newly created roots request +`NSFileProtectionCompleteUntilFirstUserAuthentication`. On iOS, this allows cache +access while the device is locked after its first unlock following a reboot. New +subdirectories and files inherit their parent's protection. Existing cache +contents are not migrated. Directory preparation skips creation when the root +already exists; child directories are opened first and created only when missing. +Directory creation uses default permissions filtered by the process umask, without +an explicit owner-only restriction. Existing permissions are not changed. +Backup exclusion and required durability flushes still apply on reused paths. + ## Host Tests On macOS, `EXECUTORCH_BUILD_COREAI=ON` with `EXECUTORCH_BUILD_TESTS=ON` diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h index bc583e59997..d807eb523de 100644 --- a/backends/apple/coreai/runtime/coreai_storage.h +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -22,6 +22,14 @@ runtime::Result prepare_storage_root(NSString* path, bool create); runtime::Error ensure_excluded_from_backup(NSURL* directory); runtime::Error sync_storage_directory(int fd); +struct StorageDirectory { + int fd; + bool created; +}; +// Opens one child component without following symlinks. The caller owns fd; +// missing with create=false returns {-1, false}. Does not flush directories. +runtime::Result +open_storage_directory(int parent_fd, NSString* name, bool create); // Returns an independently owned descriptor for a stable coordination file. int open_storage_shared_file(int parent_fd, const char* name); // Creates a new file and missing parents, then syncs the file and its parent. @@ -45,6 +53,7 @@ enum class StorageOperation { SyncDirectory, FullSync, Rename, + BeforeCreateDirectory, BeforeSDK, AfterSDK, BeforeEvict, diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm index 5c32d432cd4..550ffaaf1ea 100644 --- a/backends/apple/coreai/runtime/coreai_storage.mm +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -75,19 +75,26 @@ Error ensure_excluded_from_backup(NSURL* directory) { result == 0 ? S_ISDIR(info.st_mode) : create && errno == ENOENT, InvalidArgument, "Core AI storage must be a directory"); - if (create) { + if (create && result != 0) { + NSDictionary* attributes = @{ + NSFileProtectionKey : + NSFileProtectionCompleteUntilFirstUserAuthentication, + }; + ET_CHECK_OR_RETURN_ERROR( + storage_fault(StorageOperation::BeforeCreateDirectory) == 0, + AccessFailed, + "Cannot begin Core AI directory creation"); NSError* error = nil; ET_CHECK_OR_RETURN_ERROR( - [NSFileManager.defaultManager - createDirectoryAtPath:path - withIntermediateDirectories:YES - attributes:@{ - NSFilePosixPermissions : @0700 - } - error:&error], + [NSFileManager.defaultManager createDirectoryAtPath:path + withIntermediateDirectories:YES + attributes:attributes + error:&error], AccessFailed, "Cannot create Core AI storage directory: %s", error.localizedDescription.UTF8String); + } + if (create) { ET_CHECK_OK_OR_RETURN_ERROR(ensure_excluded_from_backup( [NSURL fileURLWithPath:path isDirectory:YES])); } @@ -146,21 +153,10 @@ FileDescriptor parent( ET_CHECK_OR_RETURN_ERROR( parent.get() >= 0, AccessFailed, "Cannot open storage parent"); for (NSUInteger i = 0; i + 1 < components.count; ++i) { - const char* component = components[i].fileSystemRepresentation; - ET_CHECK_OR_RETURN_ERROR( - retry_eintr([&] { return mkdirat(parent.get(), component, 0700); }) == - 0 || - errno == EEXIST, - AccessFailed, - "Cannot create Core AI source directory"); - FileDescriptor next(retry_eintr([&] { - return openat( - parent.get(), - component, - O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); - })); - ET_CHECK_OR_RETURN_ERROR( - next.get() >= 0, AccessFailed, "Cannot open source directory"); + auto opened = open_storage_directory(parent.get(), components[i], true); + if (!opened.ok()) + return opened.error(); + FileDescriptor next(opened.get().fd); ET_CHECK_OK_OR_RETURN_ERROR(sync_storage_directory(parent.get())); parent.reset(next.release()); } @@ -215,6 +211,41 @@ FileDescriptor file(retry_eintr([&] { } } // namespace +Result +open_storage_directory(int parent_fd, NSString* name, bool create) { + ET_CHECK_OR_RETURN_ERROR( + storage_component(name), InvalidArgument, "Invalid storage directory name"); + const auto open_directory = [&] { + return retry_eintr([&] { + return openat( + parent_fd, + name.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + }); + }; + int fd = open_directory(); + bool created = false; + if (fd < 0 && errno == ENOENT && create) { + ET_CHECK_OR_RETURN_ERROR( + storage_fault(StorageOperation::BeforeCreateDirectory) == 0, + AccessFailed, + "Cannot begin Core AI directory creation"); + const int made = retry_eintr( + [&] { return mkdirat(parent_fd, name.fileSystemRepresentation, 0777); }); + ET_CHECK_OR_RETURN_ERROR( + made == 0 || errno == EEXIST, + AccessFailed, + "Cannot create Core AI storage directory"); + created = made == 0; + fd = open_directory(); + } + if (fd < 0 && errno == ENOENT && !create) + return StorageDirectory{-1, false}; + ET_CHECK_OR_RETURN_ERROR( + fd >= 0, AccessFailed, "Cannot open Core AI storage directory"); + return StorageDirectory{fd, created}; +} + Result*> storage_children(int fd, bool skip_invalid_names) { // Reopening avoids sharing the caller's directory offset. FileDescriptor copy(retry_eintr( diff --git a/backends/apple/coreai/runtime/test/coreai_storage_test.mm b/backends/apple/coreai/runtime/test/coreai_storage_test.mm index 77b23eef53f..dd38d9386ae 100644 --- a/backends/apple/coreai/runtime/test/coreai_storage_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_storage_test.mm @@ -54,6 +54,172 @@ EXPECT_TRUE((!backup_excluded(parent) && !backup_excluded(sandbox.url))); } +TEST(CoreAIStorageTest, OpensDirectoriesAndReportsCreation) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + auto missing = open_storage_directory(root.get(), @"child", false); + ASSERT_TRUE(missing.ok()); + EXPECT_EQ(missing.get().fd, -1); + EXPECT_FALSE(missing.get().created); + + StorageFaultScope sync_fault(StorageOperation::SyncDirectory); + auto created = open_storage_directory(root.get(), @"child", true); + ASSERT_TRUE(created.ok()); + FileDescriptor first(created.get().fd); + ASSERT_GE(first.get(), 0); + EXPECT_TRUE(created.get().created); + EXPECT_NE(fcntl(first.get(), F_GETFD) & FD_CLOEXEC, 0); + EXPECT_EQ(sync_fault.matches, 0); + struct stat original; + ASSERT_EQ(fstat(first.get(), &original), 0); + EXPECT_TRUE(S_ISDIR(original.st_mode)); + + StorageFaultScope create_fault(StorageOperation::BeforeCreateDirectory); + for (bool create : {false, true}) { + auto reopened = open_storage_directory(root.get(), @"child", create); + ASSERT_TRUE(reopened.ok()); + FileDescriptor next(reopened.get().fd); + ASSERT_GE(next.get(), 0); + EXPECT_FALSE(reopened.get().created); + struct stat current; + ASSERT_EQ(fstat(next.get(), ¤t), 0); + EXPECT_EQ(current.st_ino, original.st_ino); + } + EXPECT_EQ(create_fault.matches, 0); +} + +TEST(CoreAIStorageTest, ConcurrentDirectoryCreationIsNotReportedAsOwnCreation) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + NSURL* child = [storage.url URLByAppendingPathComponent:@"child"]; + StorageFaultScope fault(StorageOperation::BeforeCreateDirectory); + fault.armed = false; + fault.observe = ^(StorageOperation operation) { + if (operation == StorageOperation::BeforeCreateDirectory) + EXPECT_EQ(mkdir(child.fileSystemRepresentation, 0777), 0); + }; + auto opened = open_storage_directory(root.get(), @"child", true); + ASSERT_TRUE(opened.ok()); + FileDescriptor directory(opened.get().fd); + EXPECT_GE(directory.get(), 0); + EXPECT_FALSE(opened.get().created); +} + +TEST(CoreAIStorageTest, RejectsInvalidDirectoryComponents) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + EXPECT_EQ( + open_storage_directory(root.get(), nil, false).error(), + Error::InvalidArgument); + for (NSString* name in @[ @"", @".", @"..", @"nested/child" ]) { + EXPECT_EQ( + open_storage_directory(root.get(), name, false).error(), + Error::InvalidArgument); + } +} + +TEST(CoreAIStorageTest, SkipsCreationForExistingRoots) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; + ASSERT_TRUE(prepare_storage_root(root.path, true).ok()); + StorageFaultScope fault(StorageOperation::BeforeCreateDirectory); + EXPECT_TRUE(prepare_storage_root(root.path, true).ok()); + EXPECT_TRUE(prepare_storage_root(root.path, false).ok()); + EXPECT_EQ(fault.matches, 0); + EXPECT_TRUE(backup_excluded(root)); + + NSURL* missing = [storage.url URLByAppendingPathComponent:@"missing"]; + EXPECT_EQ( + prepare_storage_root(missing.path, false).error(), + Error::InvalidArgument); + EXPECT_EQ(fault.matches, 0); + EXPECT_EQ( + prepare_storage_root(missing.path, true).error(), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:missing.path]); + EXPECT_TRUE(prepare_storage_root(missing.path, true).ok()); +} + +TEST(CoreAIStorageTest, CreatesDirectoriesWithDefaultPermissions) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* reference = [storage.url URLByAppendingPathComponent:@"reference"]; + ASSERT_EQ(mkdir(reference.fileSystemRepresentation, 0777), 0); + struct stat expected; + ASSERT_EQ(stat(reference.fileSystemRepresentation, &expected), 0); + NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; + ASSERT_TRUE(prepare_storage_root(root.path, true).ok()); + FileDescriptor directory(open( + root.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(directory.get(), 0); + ASSERT_EQ( + write_storage_file(directory.get(), @"nested/file", "x", 1), Error::Ok); + for (NSURL* entry in @[ root, [root URLByAppendingPathComponent:@"nested"] ]) { + struct stat actual; + ASSERT_EQ(stat(entry.fileSystemRepresentation, &actual), 0); + EXPECT_EQ(actual.st_mode & 0777, expected.st_mode & 0777); + } +} + +TEST(CoreAIStorageTest, RecreatesRemovedRoot) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; + ASSERT_TRUE(prepare_storage_root(root.path, true).ok()); + ASSERT_EQ(rmdir(root.fileSystemRepresentation), 0); + ASSERT_TRUE(prepare_storage_root(root.path, true).ok()); + EXPECT_TRUE(backup_excluded(root)); +} + +TEST(CoreAIStorageTest, NewRootAndDescendantsUseFirstUnlockProtection) { + NSFileManager* manager = NSFileManager.defaultManager; + NSURL* root = [[NSURL fileURLWithPath:NSTemporaryDirectory() isDirectory:YES] + URLByAppendingPathComponent:NSUUID.UUID.UUIDString + isDirectory:YES]; + auto prepared = prepare_storage_root(root.path, true); + EXPECT_TRUE(prepared.ok()); + if (prepared.ok()) { + FileDescriptor directory(open( + root.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + EXPECT_GE(directory.get(), 0); + if (directory.get() >= 0) { + EXPECT_EQ( + write_storage_file(directory.get(), @"nested/model.bin", "x", 1), + Error::Ok); + for (NSURL* entry in @[ + root, + [root URLByAppendingPathComponent:@"nested"], + [root URLByAppendingPathComponent:@"nested/model.bin"] + ]) { + NSDictionary* attributes = [manager attributesOfItemAtPath:entry.path + error:nil]; + EXPECT_TRUE([attributes[NSFileProtectionKey] + isEqual:NSFileProtectionCompleteUntilFirstUserAuthentication]); + } + } + EXPECT_TRUE([manager + setAttributes:@{NSFileProtectionKey : NSFileProtectionComplete} + ofItemAtPath:root.path + error:nil]); + EXPECT_TRUE(prepare_storage_root(root.path, true).ok()); + NSDictionary* existing = [manager attributesOfItemAtPath:root.path + error:nil]; + EXPECT_TRUE( + [existing[NSFileProtectionKey] isEqual:NSFileProtectionComplete]); + EXPECT_TRUE([manager removeItemAtURL:root error:nil]); + } +} + TEST(CoreAIStorageTest, RejectsInvalidRootsConsistently) { TestDirectory storage; ASSERT_NE(storage.url, nil); @@ -164,6 +330,75 @@ FileDescriptor root(open( EXPECT_EQ(empty.length, 0); } +TEST(CoreAIStorageTest, ReusesSourceParentsWithoutCreation) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + ASSERT_EQ( + write_storage_file(root.get(), @"nested/parents/first", "x", 1), + Error::Ok); + StorageFaultScope fault(StorageOperation::BeforeCreateDirectory); + EXPECT_EQ( + write_storage_file(root.get(), @"nested/parents/second", "x", 1), + Error::Ok); + EXPECT_EQ(fault.matches, 0); + EXPECT_EQ( + write_storage_file(root.get(), @"missing/file", "x", 1), + Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + EXPECT_EQ(write_storage_file(root.get(), @"missing/file", "x", 1), Error::Ok); +} + +TEST(CoreAIStorageTest, AcceptsConcurrentSourceParentCreation) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + NSURL* nested = [storage.url URLByAppendingPathComponent:@"nested"]; + StorageFaultScope fault(StorageOperation::BeforeCreateDirectory); + fault.armed = false; + __block int creations = 0; + fault.observe = ^(StorageOperation operation) { + if (operation == StorageOperation::BeforeCreateDirectory) { + ++creations; + EXPECT_EQ(mkdir(nested.fileSystemRepresentation, 0700), 0); + } + }; + EXPECT_EQ(write_storage_file(root.get(), @"nested/file", "x", 1), Error::Ok); + EXPECT_EQ(creations, 1); +} + +TEST(CoreAIStorageTest, RejectsUnsafeSourceParentsWithoutCreation) { + TestDirectory storage; + TestDirectory target; + ASSERT_NE(storage.url, nil); + ASSERT_NE(target.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + NSURL* link = [storage.url URLByAppendingPathComponent:@"link"]; + ASSERT_EQ( + symlink( + target.url.fileSystemRepresentation, link.fileSystemRepresentation), + 0); + NSURL* regular = [storage.url URLByAppendingPathComponent:@"regular"]; + ASSERT_TRUE([[NSData data] writeToURL:regular atomically:NO]); + StorageFaultScope fault(StorageOperation::BeforeCreateDirectory); + for (NSString* path in @[ @"link/file", @"regular/file" ]) { + EXPECT_EQ( + write_storage_file(root.get(), path, "x", 1), Error::AccessFailed); + } + EXPECT_EQ(fault.matches, 0); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:[target.url URLByAppendingPathComponent:@"file"].path]); +} + TEST(CoreAIStorageTest, RemovesFailedWritesAndAllowsRetry) { const struct { StorageOperation operation; @@ -221,7 +456,9 @@ FileDescriptor root(open( EXPECT_EQ( operations, (std::vector{ + StorageOperation::BeforeCreateDirectory, StorageOperation::SyncDirectory, + StorageOperation::BeforeCreateDirectory, StorageOperation::SyncDirectory, StorageOperation::Write, StorageOperation::SyncFile, From d1eb228cd0c6e0ca7f0da469ce6374c7dfaa9507 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 29 Oct 2026 13:31:16 -0700 Subject: [PATCH 6/7] [CoreAI] Separate cache-root policy and filesystem primitives Extract root policy and fault injection, name creation/enumeration/sync policies, and share no-follow descriptor and component-validation helpers. Preserve storage publication, synchronization, and failure semantics. Split protection preservation from creation/inheritance testing. Create the preexisting protected fixture directly instead of changing the protection class of a populated cache, the setup operation rejected by hosted CI. Retain protection assertions and surface NSError diagnostics. Hosted rerun remains pending; no production protection behavior changed. Test Plan: Native baseline 60/60 passed; updated suite 62/62 passed including new no-follow helper coverage. CTest, clang-format 18.1.3 across all CoreAI Objective-C++ sources, scoped lint and diff whitespace checks passed. Generic Buck owner validation remains unavailable for this checkout. --- backends/apple/coreai/CMakeLists.txt | 8 +- backends/apple/coreai/README.md | 6 + .../apple/coreai/runtime/coreai_cache_root.h | 25 ++ .../apple/coreai/runtime/coreai_cache_root.mm | 103 ++++++ .../coreai/runtime/coreai_fault_injection.h | 34 ++ .../coreai/runtime/coreai_fault_injection.mm | 24 ++ backends/apple/coreai/runtime/coreai_file.h | 22 ++ .../apple/coreai/runtime/coreai_storage.h | 45 +-- .../apple/coreai/runtime/coreai_storage.mm | 162 +++------- .../apple/coreai/runtime/coreai_validation.h | 3 + .../apple/coreai/runtime/coreai_validation.mm | 5 + .../runtime/test/coreai_cache_root_test.mm | 207 ++++++++++++ .../coreai/runtime/test/coreai_fault_scope.h | 21 +- .../runtime/test/coreai_storage_test.mm | 306 ++++++------------ 14 files changed, 583 insertions(+), 388 deletions(-) create mode 100644 backends/apple/coreai/runtime/coreai_cache_root.h create mode 100644 backends/apple/coreai/runtime/coreai_cache_root.mm create mode 100644 backends/apple/coreai/runtime/coreai_fault_injection.h create mode 100644 backends/apple/coreai/runtime/coreai_fault_injection.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_cache_root_test.mm diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index 85d742adf3d..bb63606a1a3 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -10,8 +10,10 @@ endif() enable_language(OBJCXX) -set(_coreai_runtime_sources runtime/coreai_manifest.mm - runtime/coreai_validation.mm runtime/coreai_storage.mm +set(_coreai_runtime_sources + runtime/coreai_cache_root.mm runtime/coreai_fault_injection.mm + runtime/coreai_manifest.mm runtime/coreai_validation.mm + runtime/coreai_storage.mm ) function(coreai_configure_objc_target target) @@ -31,7 +33,7 @@ if(EXECUTORCH_BUILD_TESTS AND CMAKE_SYSTEM_NAME STREQUAL "Darwin") add_executable( coreai_host_test runtime/test/coreai_host_test.mm runtime/test/coreai_manifest_test.mm - runtime/test/coreai_storage_test.mm + runtime/test/coreai_cache_root_test.mm runtime/test/coreai_storage_test.mm runtime/test/coreai_filesystem_fixture.mm ${_coreai_runtime_sources} ) coreai_configure_objc_target(coreai_host_test) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 2ca9d0b8665..1a3d96adc7a 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -76,6 +76,12 @@ Directory creation uses default permissions filtered by the process umask, witho an explicit owner-only restriction. Existing permissions are not changed. Backup exclusion and required durability flushes still apply on reused paths. +Cache-root preparation owns location and backup/protection policy; descriptor +storage owns no-follow operations, publication, and synchronization. Creation +and enumeration use explicit policies, with a missing child distinguishable +from an error. Fault injection is a separate, thread-local test interface and +compiles away in production. + ## Host Tests On macOS, `EXECUTORCH_BUILD_COREAI=ON` with `EXECUTORCH_BUILD_TESTS=ON` diff --git a/backends/apple/coreai/runtime/coreai_cache_root.h b/backends/apple/coreai/runtime/coreai_cache_root.h new file mode 100644 index 00000000000..ee536e241fc --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache_root.h @@ -0,0 +1,25 @@ +/* + * 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 +#include +#include "coreai_file.h" + +namespace executorch::backends::coreai { + +// Resolves the default location without creating or modifying storage. +runtime::Result default_coreai_assets_root(); +// Requires an absolute directory path. With Create, makes the directory and +// excludes it from backup; otherwise a missing directory is InvalidArgument. +runtime::Result prepare_storage_root( + NSString* path, + StorageCreationMode mode); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_cache_root.mm b/backends/apple/coreai/runtime/coreai_cache_root.mm new file mode 100644 index 00000000000..18ca1b6f198 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache_root.mm @@ -0,0 +1,103 @@ +/* + * 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_cache_root.h" +#include "coreai_fault_injection.h" +#include "coreai_file.h" + +#include +#include + +namespace executorch::backends::coreai { +namespace { +using runtime::Error; +using runtime::Result; + +Error ensure_excluded_from_backup(NSURL* directory) { + NSURL* fresh = [NSURL fileURLWithPath:directory.path isDirectory:YES]; + NSNumber* excluded = nil; + NSError* error = nil; + ET_CHECK_OR_RETURN_ERROR( + [fresh getResourceValue:&excluded + forKey:NSURLIsExcludedFromBackupKey + error:&error], + AccessFailed, + "Cannot read Core AI backup exclusion: %s", + error.localizedDescription.UTF8String); + if (excluded.boolValue) { + return Error::Ok; + } + ET_CHECK_OR_RETURN_ERROR( + [fresh setResourceValue:@YES + forKey:NSURLIsExcludedFromBackupKey + error:&error], + AccessFailed, + "Cannot exclude Core AI storage from backup: %s", + error.localizedDescription.UTF8String); + return Error::Ok; +} + +} // namespace + +Result default_coreai_assets_root() { + NSURL* caches = + [NSFileManager.defaultManager URLsForDirectory:NSCachesDirectory + inDomains:NSUserDomainMask] + .firstObject; + ET_CHECK_OR_RETURN_ERROR( + caches.isFileURL, AccessFailed, "Cannot resolve Core AI cache directory"); + return [caches.path stringByAppendingPathComponent:@"executorch_coreai"]; +} + +Result prepare_storage_root( + NSString* path, + StorageCreationMode mode) { + const bool create = mode == StorageCreationMode::Create; + ET_CHECK_OR_RETURN_ERROR( + [path isKindOfClass:NSString.class] && [path hasPrefix:@"/"] && + [path rangeOfString:@"\0"].location == NSNotFound, + InvalidArgument, + "Core AI storage requires an absolute path"); + struct stat info; + const int result = + retry_eintr([&] { return stat(path.fileSystemRepresentation, &info); }); + ET_CHECK_OR_RETURN_ERROR( + result == 0 || errno == ENOENT || errno == ENOTDIR, + AccessFailed, + "Cannot inspect Core AI storage directory"); + ET_CHECK_OR_RETURN_ERROR( + result == 0 ? S_ISDIR(info.st_mode) : create && errno == ENOENT, + InvalidArgument, + "Core AI storage must be a directory"); + if (create && result != 0) { + NSDictionary* attributes = @{ + NSFileProtectionKey : + NSFileProtectionCompleteUntilFirstUserAuthentication, + }; + ET_CHECK_OR_RETURN_ERROR( + fault_point(FaultPoint::BeforeCreateDirectory) == 0, + AccessFailed, + "Cannot begin Core AI directory creation"); + NSError* error = nil; + ET_CHECK_OR_RETURN_ERROR( + [NSFileManager.defaultManager createDirectoryAtPath:path + withIntermediateDirectories:YES + attributes:attributes + error:&error], + AccessFailed, + "Cannot create Core AI storage directory: %s", + error.localizedDescription.UTF8String); + } + if (create) { + ET_CHECK_OK_OR_RETURN_ERROR(ensure_excluded_from_backup( + [NSURL fileURLWithPath:path isDirectory:YES])); + } + return path; +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_fault_injection.h b/backends/apple/coreai/runtime/coreai_fault_injection.h new file mode 100644 index 00000000000..cba60c5dfb8 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_fault_injection.h @@ -0,0 +1,34 @@ +/* + * 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 + +namespace executorch::backends::coreai { + +enum class FaultPoint { + Write, + SyncFile, + SyncDirectory, + FullSync, + Rename, + BeforeCreateDirectory, +}; + +#if defined(COREAI_ASSETS_TESTING) && COREAI_ASSETS_TESTING +int fault_point(FaultPoint point); +namespace testing { +using FaultCallback = int (*)(FaultPoint); +extern thread_local FaultCallback fault_callback; +} // namespace testing +#else +constexpr int fault_point(FaultPoint) { + return 0; +} +#endif + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_fault_injection.mm b/backends/apple/coreai/runtime/coreai_fault_injection.mm new file mode 100644 index 00000000000..7bbcc0a316a --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_fault_injection.mm @@ -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. + */ + +#include "coreai_fault_injection.h" + +namespace executorch::backends::coreai { + +#if defined(COREAI_ASSETS_TESTING) && COREAI_ASSETS_TESTING +namespace testing { +thread_local FaultCallback fault_callback = nullptr; +} + +int fault_point(FaultPoint point) { + return testing::fault_callback == nullptr ? 0 + : testing::fault_callback(point); +} +#endif + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_file.h b/backends/apple/coreai/runtime/coreai_file.h index 478f03a7b91..cbc991a5fae 100644 --- a/backends/apple/coreai/runtime/coreai_file.h +++ b/backends/apple/coreai/runtime/coreai_file.h @@ -8,11 +8,18 @@ #pragma once +#include +#include #include #include namespace executorch::backends::coreai { +enum class StorageCreationMode { + OpenExisting, + Create, +}; + class FileDescriptor { public: explicit FileDescriptor(int fd) : fd_(fd) {} @@ -50,4 +57,19 @@ auto retry_eintr(F call) { return result; } +// Opens a child directory without following symlinks. Returns -1 with errno +// set, including ENOENT when the child is missing. +inline int open_directory_at(int parent_fd, const char* name) { + return retry_eintr([&] { + return openat( + parent_fd, name, O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + }); +} + +// Reads type and permission metadata for a child without following symlinks. +inline int lstat_at(int parent_fd, const char* name, struct stat* info) { + return retry_eintr( + [&] { return fstatat(parent_fd, name, info, AT_SYMLINK_NOFOLLOW); }); +} + } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h index d807eb523de..5bcfecfa4f5 100644 --- a/backends/apple/coreai/runtime/coreai_storage.h +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -11,25 +11,20 @@ #import #include #include +#include "coreai_file.h" namespace executorch::backends::coreai { -// Resolves the default location without creating or modifying storage. -runtime::Result default_coreai_assets_root(); -// Requires an absolute directory path. With create, makes the directory and -// excludes it from backup; otherwise a missing directory is InvalidArgument. -runtime::Result prepare_storage_root(NSString* path, bool create); -runtime::Error ensure_excluded_from_backup(NSURL* directory); - runtime::Error sync_storage_directory(int fd); struct StorageDirectory { int fd; bool created; }; // Opens one child component without following symlinks. The caller owns fd; -// missing with create=false returns {-1, false}. Does not flush directories. +// a missing child with OpenExisting returns {-1, false}. Does not flush +// directories. runtime::Result -open_storage_directory(int parent_fd, NSString* name, bool create); +open_storage_directory(int parent_fd, NSString* name, StorageCreationMode mode); // Returns an independently owned descriptor for a stable coordination file. int open_storage_shared_file(int parent_fd, const char* name); // Creates a new file and missing parents, then syncs the file and its parent. @@ -42,34 +37,14 @@ runtime::Error write_storage_file( // Atomically replaces name with data and fully flushes the file and renamed // entry. On error, readers see either the old or the new complete bytes. runtime::Error publish_storage_data(int root_fd, NSString* name, NSData* data); + +enum class StorageNamePolicy { + Strict, + SkipInvalid, +}; // Lists sorted names without changing fd's enumeration offset. runtime::Result*> storage_children( int fd, - bool skip_invalid_names = false); - -enum class StorageOperation { - Write, - SyncFile, - SyncDirectory, - FullSync, - Rename, - BeforeCreateDirectory, - BeforeSDK, - AfterSDK, - BeforeEvict, - AfterEvict, - Remove, -}; -#if defined(COREAI_ASSETS_TESTING) && COREAI_ASSETS_TESTING -int storage_fault(StorageOperation operation); -namespace testing { -using StorageFaultCallback = int (*)(StorageOperation); -extern thread_local StorageFaultCallback storage_fault_callback; -} // namespace testing -#else -constexpr int storage_fault(StorageOperation) { - return 0; -} -#endif + StorageNamePolicy policy = StorageNamePolicy::Strict); } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm index 550ffaaf1ea..066c60e17ba 100644 --- a/backends/apple/coreai/runtime/coreai_storage.mm +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -7,7 +7,9 @@ */ #import "coreai_storage.h" +#include "coreai_fault_injection.h" #include "coreai_file.h" +#include "coreai_validation.h" #include #include @@ -22,102 +24,11 @@ namespace { using runtime::Error; using runtime::Result; -} // namespace - -Error ensure_excluded_from_backup(NSURL* directory) { - NSURL* fresh = [NSURL fileURLWithPath:directory.path isDirectory:YES]; - NSNumber* excluded = nil; - NSError* error = nil; - ET_CHECK_OR_RETURN_ERROR( - [fresh getResourceValue:&excluded - forKey:NSURLIsExcludedFromBackupKey - error:&error], - AccessFailed, - "Cannot read Core AI backup exclusion: %s", - error.localizedDescription.UTF8String); - if (excluded.boolValue) { - return Error::Ok; - } - ET_CHECK_OR_RETURN_ERROR( - [fresh setResourceValue:@YES - forKey:NSURLIsExcludedFromBackupKey - error:&error], - AccessFailed, - "Cannot exclude Core AI storage from backup: %s", - error.localizedDescription.UTF8String); - return Error::Ok; -} - -Result default_coreai_assets_root() { - NSURL* caches = - [NSFileManager.defaultManager URLsForDirectory:NSCachesDirectory - inDomains:NSUserDomainMask] - .firstObject; - ET_CHECK_OR_RETURN_ERROR( - caches.isFileURL, AccessFailed, "Cannot resolve Core AI cache directory"); - return [caches.path stringByAppendingPathComponent:@"executorch_coreai"]; -} - -Result prepare_storage_root(NSString* path, bool create) { - ET_CHECK_OR_RETURN_ERROR( - [path isKindOfClass:NSString.class] && [path hasPrefix:@"/"] && - [path rangeOfString:@"\0"].location == NSNotFound, - InvalidArgument, - "Core AI storage requires an absolute path"); - struct stat info; - const int result = - retry_eintr([&] { return stat(path.fileSystemRepresentation, &info); }); - ET_CHECK_OR_RETURN_ERROR( - result == 0 || errno == ENOENT || errno == ENOTDIR, - AccessFailed, - "Cannot inspect Core AI storage directory"); - ET_CHECK_OR_RETURN_ERROR( - result == 0 ? S_ISDIR(info.st_mode) : create && errno == ENOENT, - InvalidArgument, - "Core AI storage must be a directory"); - if (create && result != 0) { - NSDictionary* attributes = @{ - NSFileProtectionKey : - NSFileProtectionCompleteUntilFirstUserAuthentication, - }; - ET_CHECK_OR_RETURN_ERROR( - storage_fault(StorageOperation::BeforeCreateDirectory) == 0, - AccessFailed, - "Cannot begin Core AI directory creation"); - NSError* error = nil; - ET_CHECK_OR_RETURN_ERROR( - [NSFileManager.defaultManager createDirectoryAtPath:path - withIntermediateDirectories:YES - attributes:attributes - error:&error], - AccessFailed, - "Cannot create Core AI storage directory: %s", - error.localizedDescription.UTF8String); - } - if (create) { - ET_CHECK_OK_OR_RETURN_ERROR(ensure_excluded_from_backup( - [NSURL fileURLWithPath:path isDirectory:YES])); - } - return path; -} -#if defined(COREAI_ASSETS_TESTING) && COREAI_ASSETS_TESTING -namespace testing { -thread_local StorageFaultCallback storage_fault_callback = nullptr; -} - -int storage_fault(StorageOperation operation) { - return testing::storage_fault_callback == nullptr - ? 0 - : testing::storage_fault_callback(operation); -} -#endif - -namespace { template -auto storage_call(StorageOperation operation, F call) { +auto storage_call(FaultPoint point, F call) { return retry_eintr([&]() -> decltype(call()) { - const int fault = storage_fault(operation); + const int fault = fault_point(point); if (fault != 0) { errno = fault; return -1; @@ -126,18 +37,17 @@ auto storage_call(StorageOperation operation, F call) { }); } -bool storage_component(NSString* name) { - return [name isKindOfClass:NSString.class] && name.length > 0 && - ![name isEqualToString:@"."] && ![name isEqualToString:@".."] && - ![name containsString:@"/"] && ![name containsString:@"\0"]; -} +enum class StorageSyncMode { + Normal, + Full, +}; Error write_file( int root_fd, NSString* relative_path, const void* data, size_t size, - bool full_sync) { + StorageSyncMode sync_mode) { NSArray* components = [relative_path componentsSeparatedByString:@"/"]; ET_CHECK_OR_RETURN_ERROR( @@ -153,7 +63,8 @@ FileDescriptor parent( ET_CHECK_OR_RETURN_ERROR( parent.get() >= 0, AccessFailed, "Cannot open storage parent"); for (NSUInteger i = 0; i + 1 < components.count; ++i) { - auto opened = open_storage_directory(parent.get(), components[i], true); + auto opened = open_storage_directory( + parent.get(), components[i], StorageCreationMode::Create); if (!opened.ok()) return opened.error(); FileDescriptor next(opened.get().fd); @@ -172,7 +83,7 @@ FileDescriptor file(retry_eintr([&] { const auto result = [&]() -> Error { const auto* bytes = static_cast(data); while (size > 0) { - const ssize_t count = storage_call(StorageOperation::Write, [&] { + const ssize_t count = storage_call(FaultPoint::Write, [&] { return write(file.get(), bytes, std::min(size, size_t(1024 * 1024))); }); ET_CHECK_OR_RETURN_ERROR( @@ -181,12 +92,11 @@ FileDescriptor file(retry_eintr([&] { size -= count; } ET_CHECK_OR_RETURN_ERROR( - storage_call( - StorageOperation::SyncFile, [&] { return fsync(file.get()); }) == + storage_call(FaultPoint::SyncFile, [&] { return fsync(file.get()); }) == 0 && - (!full_sync || + (sync_mode != StorageSyncMode::Full || storage_call( - StorageOperation::FullSync, + FaultPoint::FullSync, [&] { return fcntl(file.get(), F_FULLFSYNC); }) == 0), AccessFailed, "Cannot synchronize Core AI file"); @@ -211,27 +121,28 @@ FileDescriptor file(retry_eintr([&] { } } // namespace -Result -open_storage_directory(int parent_fd, NSString* name, bool create) { +Result open_storage_directory( + int parent_fd, + NSString* name, + StorageCreationMode mode) { ET_CHECK_OR_RETURN_ERROR( - storage_component(name), InvalidArgument, "Invalid storage directory name"); + storage_component(name), + InvalidArgument, + "Invalid storage directory name"); const auto open_directory = [&] { - return retry_eintr([&] { - return openat( - parent_fd, - name.fileSystemRepresentation, - O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); - }); + return open_directory_at(parent_fd, name.fileSystemRepresentation); }; + const bool create = mode == StorageCreationMode::Create; int fd = open_directory(); bool created = false; if (fd < 0 && errno == ENOENT && create) { ET_CHECK_OR_RETURN_ERROR( - storage_fault(StorageOperation::BeforeCreateDirectory) == 0, + fault_point(FaultPoint::BeforeCreateDirectory) == 0, AccessFailed, "Cannot begin Core AI directory creation"); - const int made = retry_eintr( - [&] { return mkdirat(parent_fd, name.fileSystemRepresentation, 0777); }); + const int made = retry_eintr([&] { + return mkdirat(parent_fd, name.fileSystemRepresentation, 0777); + }); ET_CHECK_OR_RETURN_ERROR( made == 0 || errno == EEXIST, AccessFailed, @@ -246,7 +157,7 @@ FileDescriptor file(retry_eintr([&] { return StorageDirectory{fd, created}; } -Result*> storage_children(int fd, bool skip_invalid_names) { +Result*> storage_children(int fd, StorageNamePolicy policy) { // Reopening avoids sharing the caller's directory offset. FileDescriptor copy(retry_eintr( [&] { return openat(fd, ".", O_RDONLY | O_DIRECTORY | O_CLOEXEC); })); @@ -277,7 +188,7 @@ FileDescriptor copy(retry_eintr( length:strlen(entry->d_name) encoding:NSUTF8StringEncoding]; if (name == nil) { - if (skip_invalid_names) { + if (policy == StorageNamePolicy::SkipInvalid) { continue; } result = Error::InvalidExternalData; @@ -311,8 +222,7 @@ int open_storage_shared_file(int parent_fd, const char* name) { Error sync_storage_directory(int fd) { ET_CHECK_OR_RETURN_ERROR( - storage_call( - StorageOperation::SyncDirectory, [&] { return fsync(fd); }) == 0, + storage_call(FaultPoint::SyncDirectory, [&] { return fsync(fd); }) == 0, AccessFailed, "Cannot synchronize Core AI directory"); return Error::Ok; @@ -323,7 +233,8 @@ Error write_storage_file( NSString* relative_path, const void* data, size_t size) { - return write_file(root_fd, relative_path, data, size, false); + return write_file( + root_fd, relative_path, data, size, StorageSyncMode::Normal); } Error publish_storage_data(int root_fd, NSString* name, NSData* data) { @@ -334,8 +245,9 @@ Error publish_storage_data(int root_fd, NSString* name, NSData* data) { NSString* temporary = [@".tmp-" stringByAppendingString:NSUUID.UUID.UUIDString]; // Bookmarks are small and costly to lose, so they get a full device flush. - auto result = write_file(root_fd, temporary, data.bytes, data.length, true); - if (result == Error::Ok && storage_call(StorageOperation::Rename, [&] { + auto result = write_file( + root_fd, temporary, data.bytes, data.length, StorageSyncMode::Full); + if (result == Error::Ok && storage_call(FaultPoint::Rename, [&] { return renameat( root_fd, temporary.fileSystemRepresentation, @@ -349,7 +261,7 @@ Error publish_storage_data(int root_fd, NSString* name, NSData* data) { // Flush the renamed entry as well as the file's pre-rename contents. ET_CHECK_OR_RETURN_ERROR( storage_call( - StorageOperation::FullSync, + FaultPoint::FullSync, [&] { return fcntl(root_fd, F_FULLFSYNC); }) == 0, AccessFailed, "Cannot fully synchronize Core AI publication directory"); diff --git a/backends/apple/coreai/runtime/coreai_validation.h b/backends/apple/coreai/runtime/coreai_validation.h index a3d060857dc..ea21885b19d 100644 --- a/backends/apple/coreai/runtime/coreai_validation.h +++ b/backends/apple/coreai/runtime/coreai_validation.h @@ -16,6 +16,9 @@ namespace executorch::backends::coreai { bool nonempty_string(id value); // Relative bundle path without escapes, backslashes, or empty components. bool relative_path(NSString* path); +// Single path component for an fd-relative open: no slashes, ".", or "..". +// Unlike relative_path, backslashes are ordinary bytes on this filesystem. +bool storage_component(NSString* name); // Lowercase hex of any positive length, as in the exporter model hash. bool lowercase_hex(NSString* value); // Lowercase hex of exactly length characters, as in bundle digests. diff --git a/backends/apple/coreai/runtime/coreai_validation.mm b/backends/apple/coreai/runtime/coreai_validation.mm index ea8b018f38f..54c611b81bc 100644 --- a/backends/apple/coreai/runtime/coreai_validation.mm +++ b/backends/apple/coreai/runtime/coreai_validation.mm @@ -38,6 +38,11 @@ bool relative_path(NSString* path) { return true; } +bool storage_component(NSString* name) { + return nonempty_string(name) && ![name isEqualToString:@"."] && + ![name isEqualToString:@".."] && ![name containsString:@"/"]; +} + bool lowercase_hex(NSString* value) { return nonempty_string(value) && [value rangeOfCharacterFromSet:invalid_hex_characters()].location == diff --git a/backends/apple/coreai/runtime/test/coreai_cache_root_test.mm b/backends/apple/coreai/runtime/test/coreai_cache_root_test.mm new file mode 100644 index 00000000000..b6cf8141065 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_cache_root_test.mm @@ -0,0 +1,207 @@ +/* + * 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. + */ + +#include +#include +#include "coreai_cache_root.h" +#include "coreai_fault_scope.h" +#include "coreai_file.h" +#include "coreai_filesystem_fixture.h" +#include "coreai_storage.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +TEST(CoreAICacheRootTest, PreparesRootWithRootOnlyBackupPolicy) { + NSFileManager* manager = NSFileManager.defaultManager; + TestDirectory sandbox; + ASSERT_NE(sandbox.url, nil); + NSURL* parent = [sandbox.url URLByAppendingPathComponent:@"parent"]; + ASSERT_TRUE(([manager createDirectoryAtURL:parent + withIntermediateDirectories:NO + attributes:nil + error:nil])); + ASSERT_TRUE(set_backup_excluded(sandbox.url, false)); + ASSERT_TRUE(set_backup_excluded(parent, false)); + NSURL* root = [parent URLByAppendingPathComponent:@"nested/models"]; + auto prepared = prepare_storage_root(root.path, StorageCreationMode::Create); + ASSERT_TRUE((prepared.ok() && [prepared.get() isEqualToString:root.path])); + ASSERT_TRUE((backup_excluded(root))); + ASSERT_TRUE((!backup_excluded(parent) && !backup_excluded(sandbox.url))); + ASSERT_TRUE((!backup_excluded(root.URLByDeletingLastPathComponent))); + + ASSERT_TRUE(set_backup_excluded(root, false)); + auto inspected = + prepare_storage_root(root.path, StorageCreationMode::OpenExisting); + ASSERT_TRUE((inspected.ok() && [inspected.get() isEqualToString:root.path])); + ASSERT_TRUE((!backup_excluded(root))); + ASSERT_TRUE( + (prepare_storage_root(root.path, StorageCreationMode::Create).ok() && + backup_excluded(root))); + + NSURL* missing = [sandbox.url URLByAppendingPathComponent:@"missing"]; + ASSERT_EQ( + prepare_storage_root(missing.path, StorageCreationMode::OpenExisting) + .error(), + Error::InvalidArgument); + EXPECT_FALSE([manager fileExistsAtPath:missing.path]); + EXPECT_EQ( + prepare_storage_root(@"relative", StorageCreationMode::Create).error(), + Error::InvalidArgument); + EXPECT_TRUE((!backup_excluded(parent) && !backup_excluded(sandbox.url))); +} + +TEST(CoreAICacheRootTest, SkipsCreationForExistingRoots) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; + ASSERT_TRUE( + prepare_storage_root(root.path, StorageCreationMode::Create).ok()); + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); + EXPECT_TRUE( + prepare_storage_root(root.path, StorageCreationMode::Create).ok()); + EXPECT_TRUE( + prepare_storage_root(root.path, StorageCreationMode::OpenExisting).ok()); + EXPECT_EQ(fault.matches, 0); + EXPECT_TRUE(backup_excluded(root)); + + NSURL* missing = [storage.url URLByAppendingPathComponent:@"missing"]; + EXPECT_EQ( + prepare_storage_root(missing.path, StorageCreationMode::OpenExisting) + .error(), + Error::InvalidArgument); + EXPECT_EQ(fault.matches, 0); + EXPECT_EQ( + prepare_storage_root(missing.path, StorageCreationMode::Create).error(), + Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:missing.path]); + EXPECT_TRUE( + prepare_storage_root(missing.path, StorageCreationMode::Create).ok()); +} + +TEST(CoreAICacheRootTest, CreatesDirectoriesWithDefaultPermissions) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* reference = [storage.url URLByAppendingPathComponent:@"reference"]; + ASSERT_EQ(mkdir(reference.fileSystemRepresentation, 0777), 0); + struct stat expected; + ASSERT_EQ(stat(reference.fileSystemRepresentation, &expected), 0); + NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; + ASSERT_TRUE( + prepare_storage_root(root.path, StorageCreationMode::Create).ok()); + FileDescriptor directory( + open(root.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(directory.get(), 0); + ASSERT_EQ( + write_storage_file(directory.get(), @"nested/file", "x", 1), Error::Ok); + for (NSURL* entry in + @[ root, [root URLByAppendingPathComponent:@"nested"] ]) { + struct stat actual; + ASSERT_EQ(stat(entry.fileSystemRepresentation, &actual), 0); + EXPECT_EQ(actual.st_mode & 0777, expected.st_mode & 0777); + } +} + +TEST(CoreAICacheRootTest, RecreatesRemovedRoot) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; + ASSERT_TRUE( + prepare_storage_root(root.path, StorageCreationMode::Create).ok()); + ASSERT_EQ(rmdir(root.fileSystemRepresentation), 0); + ASSERT_TRUE( + prepare_storage_root(root.path, StorageCreationMode::Create).ok()); + EXPECT_TRUE(backup_excluded(root)); +} + +TEST(CoreAICacheRootTest, NewRootAndDescendantsUseFirstUnlockProtection) { + TestDirectory sandbox; + ASSERT_NE(sandbox.url, nil); + NSURL* root = [sandbox.url URLByAppendingPathComponent:@"cache"]; + ASSERT_TRUE( + prepare_storage_root(root.path, StorageCreationMode::Create).ok()); + FileDescriptor directory( + open_directory_at(AT_FDCWD, root.fileSystemRepresentation)); + ASSERT_GE(directory.get(), 0); + ASSERT_EQ( + write_storage_file(directory.get(), @"nested/model.bin", "x", 1), + Error::Ok); + for (NSURL* entry in @[ + root, + [root URLByAppendingPathComponent:@"nested"], + [root URLByAppendingPathComponent:@"nested/model.bin"] + ]) { + SCOPED_TRACE(entry.path.UTF8String); + NSError* error = nil; + NSDictionary* attributes = + [NSFileManager.defaultManager attributesOfItemAtPath:entry.path + error:&error]; + ASSERT_NE(attributes, nil) << error.description.UTF8String; + EXPECT_TRUE([attributes[NSFileProtectionKey] + isEqual:NSFileProtectionCompleteUntilFirstUserAuthentication]); + } +} + +TEST(CoreAICacheRootTest, PreservesExistingRootProtection) { + TestDirectory sandbox; + ASSERT_NE(sandbox.url, nil); + NSURL* root = [sandbox.url URLByAppendingPathComponent:@"existing"]; + NSFileManager* manager = NSFileManager.defaultManager; + NSError* error = nil; + // Set the fixture policy at creation, without reclassifying a populated + // cache. + ASSERT_TRUE([manager + createDirectoryAtURL:root + withIntermediateDirectories:NO + attributes:@{ + NSFileProtectionKey : NSFileProtectionComplete + } + error:&error]) + << error.description.UTF8String; + NSDictionary* before = [manager attributesOfItemAtPath:root.path + error:&error]; + ASSERT_NE(before, nil) << error.description.UTF8String; + ASSERT_TRUE([before[NSFileProtectionKey] isEqual:NSFileProtectionComplete]); + for (StorageCreationMode mode : + {StorageCreationMode::OpenExisting, StorageCreationMode::Create}) { + ASSERT_TRUE(prepare_storage_root(root.path, mode).ok()); + NSDictionary* after = [manager attributesOfItemAtPath:root.path + error:&error]; + ASSERT_NE(after, nil) << error.description.UTF8String; + EXPECT_TRUE([after[NSFileProtectionKey] isEqual:NSFileProtectionComplete]); + } +} + +TEST(CoreAICacheRootTest, RejectsInvalidRootsConsistently) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* regular = [storage.url URLByAppendingPathComponent:@"regular"]; + NSData* bytes = [@"unchanged" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([bytes writeToURL:regular atomically:NO]); + for (StorageCreationMode mode : + {StorageCreationMode::OpenExisting, StorageCreationMode::Create}) { + SCOPED_TRACE(static_cast(mode)); + EXPECT_EQ(prepare_storage_root(nil, mode).error(), Error::InvalidArgument); + for (NSString* path in @[ + @"", + @"relative", + @"~", + @"~/cache", + @"~nobody/cache", + regular.path, + [regular.path stringByAppendingPathComponent:@"child"] + ]) { + SCOPED_TRACE(path.UTF8String); + EXPECT_EQ( + prepare_storage_root(path, mode).error(), Error::InvalidArgument); + } + } + EXPECT_TRUE([[NSData dataWithContentsOfURL:regular] isEqual:bytes]); +} diff --git a/backends/apple/coreai/runtime/test/coreai_fault_scope.h b/backends/apple/coreai/runtime/test/coreai_fault_scope.h index e55de88579a..f60092cd91a 100644 --- a/backends/apple/coreai/runtime/test/coreai_fault_scope.h +++ b/backends/apple/coreai/runtime/test/coreai_fault_scope.h @@ -8,29 +8,28 @@ #pragma once +#import #include -#include "coreai_storage.h" +#include "coreai_fault_injection.h" namespace executorch::backends::coreai::testing { struct StorageFaultScope { - StorageOperation operation; + FaultPoint operation; int error = EIO; int match_index = 1; int matches = 0; int hits = 0; bool armed = true; - void (^observe)(StorageOperation) = nil; + void (^observe)(FaultPoint) = nil; - explicit StorageFaultScope(StorageOperation op) - : operation(op), - previous_(current_), - previous_callback_(storage_fault_callback) { + explicit StorageFaultScope(FaultPoint op) + : operation(op), previous_(current_), previous_callback_(fault_callback) { current_ = this; - storage_fault_callback = inject; + fault_callback = inject; } ~StorageFaultScope() { - storage_fault_callback = previous_callback_; + fault_callback = previous_callback_; current_ = previous_; } StorageFaultScope(const StorageFaultScope&) = delete; @@ -39,9 +38,9 @@ struct StorageFaultScope { private: inline static thread_local StorageFaultScope* current_ = nullptr; StorageFaultScope* previous_; - StorageFaultCallback previous_callback_; + FaultCallback previous_callback_; - static int inject(StorageOperation operation) { + static int inject(FaultPoint operation) { auto& scope = *current_; if (scope.observe != nil) scope.observe(operation); diff --git a/backends/apple/coreai/runtime/test/coreai_storage_test.mm b/backends/apple/coreai/runtime/test/coreai_storage_test.mm index dd38d9386ae..266dfbd34bf 100644 --- a/backends/apple/coreai/runtime/test/coreai_storage_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_storage_test.mm @@ -14,59 +14,49 @@ #include "coreai_fault_scope.h" #include "coreai_file.h" #include "coreai_filesystem_fixture.h" +#include "coreai_storage.h" using namespace executorch::runtime; using namespace executorch::backends::coreai; using namespace executorch::backends::coreai::testing; -TEST(CoreAIStorageTest, PreparesRootWithRootOnlyBackupPolicy) { - NSFileManager* manager = NSFileManager.defaultManager; - TestDirectory sandbox; - ASSERT_NE(sandbox.url, nil); - NSURL* parent = [sandbox.url URLByAppendingPathComponent:@"parent"]; - ASSERT_TRUE(([manager createDirectoryAtURL:parent - withIntermediateDirectories:NO - attributes:nil - error:nil])); - ASSERT_TRUE(set_backup_excluded(sandbox.url, false)); - ASSERT_TRUE(set_backup_excluded(parent, false)); - NSURL* root = [parent URLByAppendingPathComponent:@"nested/models"]; - auto prepared = prepare_storage_root(root.path, true); - ASSERT_TRUE((prepared.ok() && [prepared.get() isEqualToString:root.path])); - ASSERT_TRUE((backup_excluded(root))); - ASSERT_TRUE((!backup_excluded(parent) && !backup_excluded(sandbox.url))); - ASSERT_TRUE((!backup_excluded(root.URLByDeletingLastPathComponent))); - - ASSERT_TRUE(set_backup_excluded(root, false)); - auto inspected = prepare_storage_root(root.path, false); - ASSERT_TRUE((inspected.ok() && [inspected.get() isEqualToString:root.path])); - ASSERT_TRUE((!backup_excluded(root))); - ASSERT_TRUE( - (prepare_storage_root(root.path, true).ok() && backup_excluded(root))); - - NSURL* missing = [sandbox.url URLByAppendingPathComponent:@"missing"]; - ASSERT_EQ( - prepare_storage_root(missing.path, false).error(), - Error::InvalidArgument); - EXPECT_FALSE([manager fileExistsAtPath:missing.path]); - EXPECT_EQ( - prepare_storage_root(@"relative", true).error(), Error::InvalidArgument); - EXPECT_TRUE((!backup_excluded(parent) && !backup_excluded(sandbox.url))); +TEST(CoreAIStorageTest, DirectoryHelpersDoNotFollowSymlinks) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root( + open_directory_at(AT_FDCWD, storage.url.fileSystemRepresentation)); + ASSERT_GE(root.get(), 0); + EXPECT_NE(fcntl(root.get(), F_GETFD) & FD_CLOEXEC, 0); + ASSERT_EQ(mkdirat(root.get(), "directory", 0700), 0); + ASSERT_EQ(symlinkat("directory", root.get(), "link"), 0); + struct stat info; + ASSERT_EQ(lstat_at(root.get(), "link", &info), 0); + EXPECT_TRUE(S_ISLNK(info.st_mode)); + ASSERT_EQ(lstat_at(root.get(), "directory", &info), 0); + EXPECT_TRUE(S_ISDIR(info.st_mode)); + EXPECT_EQ(open_directory_at(root.get(), "link"), -1); + EXPECT_EQ(open_directory_at(root.get(), "missing"), -1); + EXPECT_EQ(errno, ENOENT); + EXPECT_EQ(lstat_at(root.get(), "missing", &info), -1); + EXPECT_EQ(errno, ENOENT); } TEST(CoreAIStorageTest, OpensDirectoriesAndReportsCreation) { TestDirectory storage; ASSERT_NE(storage.url, nil); FileDescriptor root(open( - storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_GE(root.get(), 0); - auto missing = open_storage_directory(root.get(), @"child", false); + auto missing = open_storage_directory( + root.get(), @"child", StorageCreationMode::OpenExisting); ASSERT_TRUE(missing.ok()); EXPECT_EQ(missing.get().fd, -1); EXPECT_FALSE(missing.get().created); - StorageFaultScope sync_fault(StorageOperation::SyncDirectory); - auto created = open_storage_directory(root.get(), @"child", true); + StorageFaultScope sync_fault(FaultPoint::SyncDirectory); + auto created = + open_storage_directory(root.get(), @"child", StorageCreationMode::Create); ASSERT_TRUE(created.ok()); FileDescriptor first(created.get().fd); ASSERT_GE(first.get(), 0); @@ -77,9 +67,10 @@ FileDescriptor root(open( ASSERT_EQ(fstat(first.get(), &original), 0); EXPECT_TRUE(S_ISDIR(original.st_mode)); - StorageFaultScope create_fault(StorageOperation::BeforeCreateDirectory); - for (bool create : {false, true}) { - auto reopened = open_storage_directory(root.get(), @"child", create); + StorageFaultScope create_fault(FaultPoint::BeforeCreateDirectory); + for (StorageCreationMode mode : + {StorageCreationMode::OpenExisting, StorageCreationMode::Create}) { + auto reopened = open_storage_directory(root.get(), @"child", mode); ASSERT_TRUE(reopened.ok()); FileDescriptor next(reopened.get().fd); ASSERT_GE(next.get(), 0); @@ -95,16 +86,18 @@ FileDescriptor root(open( TestDirectory storage; ASSERT_NE(storage.url, nil); FileDescriptor root(open( - storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_GE(root.get(), 0); NSURL* child = [storage.url URLByAppendingPathComponent:@"child"]; - StorageFaultScope fault(StorageOperation::BeforeCreateDirectory); + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); fault.armed = false; - fault.observe = ^(StorageOperation operation) { - if (operation == StorageOperation::BeforeCreateDirectory) + fault.observe = ^(FaultPoint operation) { + if (operation == FaultPoint::BeforeCreateDirectory) EXPECT_EQ(mkdir(child.fileSystemRepresentation, 0777), 0); }; - auto opened = open_storage_directory(root.get(), @"child", true); + auto opened = + open_storage_directory(root.get(), @"child", StorageCreationMode::Create); ASSERT_TRUE(opened.ok()); FileDescriptor directory(opened.get().fd); EXPECT_GE(directory.get(), 0); @@ -115,138 +108,22 @@ FileDescriptor root(open( TestDirectory storage; ASSERT_NE(storage.url, nil); FileDescriptor root(open( - storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_GE(root.get(), 0); EXPECT_EQ( - open_storage_directory(root.get(), nil, false).error(), + open_storage_directory(root.get(), nil, StorageCreationMode::OpenExisting) + .error(), Error::InvalidArgument); for (NSString* name in @[ @"", @".", @"..", @"nested/child" ]) { EXPECT_EQ( - open_storage_directory(root.get(), name, false).error(), + open_storage_directory( + root.get(), name, StorageCreationMode::OpenExisting) + .error(), Error::InvalidArgument); } } -TEST(CoreAIStorageTest, SkipsCreationForExistingRoots) { - TestDirectory storage; - ASSERT_NE(storage.url, nil); - NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; - ASSERT_TRUE(prepare_storage_root(root.path, true).ok()); - StorageFaultScope fault(StorageOperation::BeforeCreateDirectory); - EXPECT_TRUE(prepare_storage_root(root.path, true).ok()); - EXPECT_TRUE(prepare_storage_root(root.path, false).ok()); - EXPECT_EQ(fault.matches, 0); - EXPECT_TRUE(backup_excluded(root)); - - NSURL* missing = [storage.url URLByAppendingPathComponent:@"missing"]; - EXPECT_EQ( - prepare_storage_root(missing.path, false).error(), - Error::InvalidArgument); - EXPECT_EQ(fault.matches, 0); - EXPECT_EQ( - prepare_storage_root(missing.path, true).error(), Error::AccessFailed); - EXPECT_EQ(fault.hits, 1); - EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:missing.path]); - EXPECT_TRUE(prepare_storage_root(missing.path, true).ok()); -} - -TEST(CoreAIStorageTest, CreatesDirectoriesWithDefaultPermissions) { - TestDirectory storage; - ASSERT_NE(storage.url, nil); - NSURL* reference = [storage.url URLByAppendingPathComponent:@"reference"]; - ASSERT_EQ(mkdir(reference.fileSystemRepresentation, 0777), 0); - struct stat expected; - ASSERT_EQ(stat(reference.fileSystemRepresentation, &expected), 0); - NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; - ASSERT_TRUE(prepare_storage_root(root.path, true).ok()); - FileDescriptor directory(open( - root.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); - ASSERT_GE(directory.get(), 0); - ASSERT_EQ( - write_storage_file(directory.get(), @"nested/file", "x", 1), Error::Ok); - for (NSURL* entry in @[ root, [root URLByAppendingPathComponent:@"nested"] ]) { - struct stat actual; - ASSERT_EQ(stat(entry.fileSystemRepresentation, &actual), 0); - EXPECT_EQ(actual.st_mode & 0777, expected.st_mode & 0777); - } -} - -TEST(CoreAIStorageTest, RecreatesRemovedRoot) { - TestDirectory storage; - ASSERT_NE(storage.url, nil); - NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; - ASSERT_TRUE(prepare_storage_root(root.path, true).ok()); - ASSERT_EQ(rmdir(root.fileSystemRepresentation), 0); - ASSERT_TRUE(prepare_storage_root(root.path, true).ok()); - EXPECT_TRUE(backup_excluded(root)); -} - -TEST(CoreAIStorageTest, NewRootAndDescendantsUseFirstUnlockProtection) { - NSFileManager* manager = NSFileManager.defaultManager; - NSURL* root = [[NSURL fileURLWithPath:NSTemporaryDirectory() isDirectory:YES] - URLByAppendingPathComponent:NSUUID.UUID.UUIDString - isDirectory:YES]; - auto prepared = prepare_storage_root(root.path, true); - EXPECT_TRUE(prepared.ok()); - if (prepared.ok()) { - FileDescriptor directory(open( - root.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); - EXPECT_GE(directory.get(), 0); - if (directory.get() >= 0) { - EXPECT_EQ( - write_storage_file(directory.get(), @"nested/model.bin", "x", 1), - Error::Ok); - for (NSURL* entry in @[ - root, - [root URLByAppendingPathComponent:@"nested"], - [root URLByAppendingPathComponent:@"nested/model.bin"] - ]) { - NSDictionary* attributes = [manager attributesOfItemAtPath:entry.path - error:nil]; - EXPECT_TRUE([attributes[NSFileProtectionKey] - isEqual:NSFileProtectionCompleteUntilFirstUserAuthentication]); - } - } - EXPECT_TRUE([manager - setAttributes:@{NSFileProtectionKey : NSFileProtectionComplete} - ofItemAtPath:root.path - error:nil]); - EXPECT_TRUE(prepare_storage_root(root.path, true).ok()); - NSDictionary* existing = [manager attributesOfItemAtPath:root.path - error:nil]; - EXPECT_TRUE( - [existing[NSFileProtectionKey] isEqual:NSFileProtectionComplete]); - EXPECT_TRUE([manager removeItemAtURL:root error:nil]); - } -} - -TEST(CoreAIStorageTest, RejectsInvalidRootsConsistently) { - TestDirectory storage; - ASSERT_NE(storage.url, nil); - NSURL* regular = [storage.url URLByAppendingPathComponent:@"regular"]; - NSData* bytes = [@"unchanged" dataUsingEncoding:NSUTF8StringEncoding]; - ASSERT_TRUE([bytes writeToURL:regular atomically:NO]); - for (bool create : {false, true}) { - SCOPED_TRACE(create); - EXPECT_EQ( - prepare_storage_root(nil, create).error(), Error::InvalidArgument); - for (NSString* path in @[ - @"", - @"relative", - @"~", - @"~/cache", - @"~nobody/cache", - regular.path, - [regular.path stringByAppendingPathComponent:@"child"] - ]) { - SCOPED_TRACE(path.UTF8String); - EXPECT_EQ( - prepare_storage_root(path, create).error(), Error::InvalidArgument); - } - } - EXPECT_TRUE([[NSData dataWithContentsOfURL:regular] isEqual:bytes]); -} - TEST(CoreAIStorageTest, ListsChildrenRepeatedlyInSortedOrder) { TestDirectory storage; ASSERT_NE(storage.url, nil); @@ -291,8 +168,9 @@ FileDescriptor file(openat( O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, 0600)); ASSERT_GE(file.get(), 0); - for (bool skip_invalid : {false, true}) { - auto names = storage_children(root.get(), skip_invalid); + for (StorageNamePolicy policy : + {StorageNamePolicy::Strict, StorageNamePolicy::SkipInvalid}) { + auto names = storage_children(root.get(), policy); ASSERT_TRUE(names.ok()); ASSERT_EQ(names.get().count, 1); EXPECT_STREQ(names.get().firstObject.UTF8String, name.UTF8String); @@ -340,7 +218,7 @@ FileDescriptor root(open( ASSERT_EQ( write_storage_file(root.get(), @"nested/parents/first", "x", 1), Error::Ok); - StorageFaultScope fault(StorageOperation::BeforeCreateDirectory); + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); EXPECT_EQ( write_storage_file(root.get(), @"nested/parents/second", "x", 1), Error::Ok); @@ -360,11 +238,11 @@ FileDescriptor root(open( O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_GE(root.get(), 0); NSURL* nested = [storage.url URLByAppendingPathComponent:@"nested"]; - StorageFaultScope fault(StorageOperation::BeforeCreateDirectory); + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); fault.armed = false; __block int creations = 0; - fault.observe = ^(StorageOperation operation) { - if (operation == StorageOperation::BeforeCreateDirectory) { + fault.observe = ^(FaultPoint operation) { + if (operation == FaultPoint::BeforeCreateDirectory) { ++creations; EXPECT_EQ(mkdir(nested.fileSystemRepresentation, 0700), 0); } @@ -389,7 +267,7 @@ FileDescriptor root(open( 0); NSURL* regular = [storage.url URLByAppendingPathComponent:@"regular"]; ASSERT_TRUE([[NSData data] writeToURL:regular atomically:NO]); - StorageFaultScope fault(StorageOperation::BeforeCreateDirectory); + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); for (NSString* path in @[ @"link/file", @"regular/file" ]) { EXPECT_EQ( write_storage_file(root.get(), path, "x", 1), Error::AccessFailed); @@ -401,13 +279,13 @@ FileDescriptor root(open( TEST(CoreAIStorageTest, RemovesFailedWritesAndAllowsRetry) { const struct { - StorageOperation operation; + FaultPoint operation; int match_index; } rows[] = { - {StorageOperation::Write, 2}, - {StorageOperation::SyncFile, 1}, - {StorageOperation::SyncDirectory, 1}, - {StorageOperation::SyncDirectory, 2}}; + {FaultPoint::Write, 2}, + {FaultPoint::SyncFile, 1}, + {FaultPoint::SyncDirectory, 1}, + {FaultPoint::SyncDirectory, 2}}; const std::vector bytes(1024 * 1024 + 1, 0x5a); for (const auto& row : rows) { SCOPED_TRACE(static_cast(row.operation)); @@ -444,10 +322,10 @@ FileDescriptor root(open( storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_GE(root.get(), 0); - __block std::vector operations; - StorageFaultScope fault(StorageOperation::SyncDirectory); + __block std::vector operations; + StorageFaultScope fault(FaultPoint::SyncDirectory); fault.armed = false; - fault.observe = ^(StorageOperation operation) { + fault.observe = ^(FaultPoint operation) { operations.push_back(operation); }; EXPECT_EQ( @@ -455,14 +333,14 @@ FileDescriptor root(open( Error::Ok); EXPECT_EQ( operations, - (std::vector{ - StorageOperation::BeforeCreateDirectory, - StorageOperation::SyncDirectory, - StorageOperation::BeforeCreateDirectory, - StorageOperation::SyncDirectory, - StorageOperation::Write, - StorageOperation::SyncFile, - StorageOperation::SyncDirectory})); + (std::vector{ + FaultPoint::BeforeCreateDirectory, + FaultPoint::SyncDirectory, + FaultPoint::BeforeCreateDirectory, + FaultPoint::SyncDirectory, + FaultPoint::Write, + FaultPoint::SyncFile, + FaultPoint::SyncDirectory})); } TEST(CoreAIStorageTest, RejectsUnsafeWritePathsWithoutCreatingFiles) { @@ -554,12 +432,12 @@ FileDescriptor root(open( NSData* next = [@"next" dataUsingEncoding:NSUTF8StringEncoding]; NSURL* published = [storage.url URLByAppendingPathComponent:@"published"]; ASSERT_EQ(publish_storage_data(root.get(), @"published", old), Error::Ok); - for (StorageOperation operation : - {StorageOperation::Write, - StorageOperation::SyncFile, - StorageOperation::SyncDirectory, - StorageOperation::FullSync, - StorageOperation::Rename}) { + for (FaultPoint operation : + {FaultPoint::Write, + FaultPoint::SyncFile, + FaultPoint::SyncDirectory, + FaultPoint::FullSync, + FaultPoint::Rename}) { SCOPED_TRACE(static_cast(operation)); StorageFaultScope fault(operation); EXPECT_EQ( @@ -582,23 +460,23 @@ FileDescriptor root(open( storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_GE(root.get(), 0); - __block std::vector operations; - StorageFaultScope fault(StorageOperation::FullSync); + __block std::vector operations; + StorageFaultScope fault(FaultPoint::FullSync); fault.armed = false; - fault.observe = ^(StorageOperation operation) { + fault.observe = ^(FaultPoint operation) { operations.push_back(operation); }; NSData* data = [@"published" dataUsingEncoding:NSUTF8StringEncoding]; EXPECT_EQ(publish_storage_data(root.get(), @"published", data), Error::Ok); EXPECT_EQ( operations, - (std::vector{ - StorageOperation::Write, - StorageOperation::SyncFile, - StorageOperation::FullSync, - StorageOperation::SyncDirectory, - StorageOperation::Rename, - StorageOperation::FullSync})); + (std::vector{ + FaultPoint::Write, + FaultPoint::SyncFile, + FaultPoint::FullSync, + FaultPoint::SyncDirectory, + FaultPoint::Rename, + FaultPoint::FullSync})); } TEST(CoreAIStorageTest, ReportsPostRenameFlushFailureWithNewContentsVisible) { @@ -611,7 +489,7 @@ FileDescriptor root(open( NSData* old = [@"old" dataUsingEncoding:NSUTF8StringEncoding]; NSData* next = [@"next" dataUsingEncoding:NSUTF8StringEncoding]; ASSERT_EQ(publish_storage_data(root.get(), @"published", old), Error::Ok); - StorageFaultScope fault(StorageOperation::FullSync); + StorageFaultScope fault(FaultPoint::FullSync); fault.match_index = 2; EXPECT_EQ( publish_storage_data(root.get(), @"published", next), @@ -630,12 +508,12 @@ FileDescriptor root(open( O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_GE(root.get(), 0); NSData* data = [@"published" dataUsingEncoding:NSUTF8StringEncoding]; - for (StorageOperation operation : - {StorageOperation::Write, - StorageOperation::SyncFile, - StorageOperation::SyncDirectory, - StorageOperation::FullSync, - StorageOperation::Rename}) { + for (FaultPoint operation : + {FaultPoint::Write, + FaultPoint::SyncFile, + FaultPoint::SyncDirectory, + FaultPoint::FullSync, + FaultPoint::Rename}) { SCOPED_TRACE(static_cast(operation)); StorageFaultScope fault(operation); fault.error = EINTR; From 6853ed1cd2d520ea7124038ebbf2a396ec74e938 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Fri, 30 Oct 2026 09:24:51 -0700 Subject: [PATCH 7/7] [CoreAI] Persist raw bookmarks (#23389) Persists the raw SDK bookmark for a selected Core AI bundle, so later loads can restore a specialized model without re-reading the payload (`runtime/coreai_bookmarks.{h,mm}`). * Store raw opaque bookmarks (at most 8 MiB) with atomic publication. * Keys include the bundle digest, platform, architecture and persistence policy, but not function bindings or delivery location. * Per-key OS locks coordinate threads and processes; independent keys proceed in parallel. * SDK acquisition and cache removal are not included here. * Tests: the bookmark suite in `coreai_host_test`, including a cross-process lock test that spawns child processes. * README: the "Asset storage" key derivation, bookmark bound and locking text. No changes outside `backends/apple/coreai`. Stack: 3 of 9, based on #23388. Review only this PR's commit. Next: #23390. Test plan: built locally for macOS 27.0; `ctest -N` lists `coreai_host_test`, which the Core AI workflow runs on the macOS 27 runner. `lintrunner` is clean. --- CMakeLists.txt | 10 + backends/apple/coreai/CMakeLists.txt | 436 ++++++++- backends/apple/coreai/README.md | 390 +++++++- .../coreai/cmake/PublishSwiftHeader.cmake | 17 + backends/apple/coreai/compiler/preprocess.py | 42 +- .../apple/coreai/coreai-dependencies.cmake | 118 +++ .../apple/coreai/runtime/ETCoreAIBridge.h | 142 +++ .../apple/coreai/runtime/ETCoreAIError.swift | 30 + .../coreai/runtime/ETCoreAIModelLoader.swift | 279 ++++++ .../apple/coreai/runtime/ETCoreAITensor.mm | 63 ++ .../coreai/runtime/ETCoreAITensors.swift | 214 ++++ .../apple/coreai/runtime/coreai_backend.mm | 156 +++ .../apple/coreai/runtime/coreai_backend_id.h | 13 + .../apple/coreai/runtime/coreai_bridge.mm | 29 + backends/apple/coreai/runtime/coreai_cache.h | 58 ++ backends/apple/coreai/runtime/coreai_cache.mm | 141 +++ .../apple/coreai/runtime/coreai_cache_entry.h | 123 +++ .../coreai/runtime/coreai_cache_entry.mm | 370 +++++++ .../apple/coreai/runtime/coreai_cache_key.h | 42 + .../apple/coreai/runtime/coreai_cache_key.mm | 185 ++++ .../apple/coreai/runtime/coreai_cache_root.h | 28 +- .../apple/coreai/runtime/coreai_cache_root.mm | 191 +++- backends/apple/coreai/runtime/coreai_error.h | 20 + backends/apple/coreai/runtime/coreai_error.mm | 48 + .../coreai/runtime/coreai_fault_injection.h | 24 + backends/apple/coreai/runtime/coreai_file.h | 14 +- .../coreai/runtime/coreai_load_coordinator.h | 40 + .../coreai/runtime/coreai_load_coordinator.mm | 275 ++++++ .../apple/coreai/runtime/coreai_manifest.h | 10 +- .../apple/coreai/runtime/coreai_manifest.mm | 45 +- backends/apple/coreai/runtime/coreai_pte.cpp | 256 +++++ backends/apple/coreai/runtime/coreai_pte.h | 42 + .../coreai/runtime/coreai_source_bundle.h | 33 + .../coreai/runtime/coreai_source_bundle.mm | 328 +++++++ .../apple/coreai/runtime/coreai_storage.h | 44 +- .../apple/coreai/runtime/coreai_storage.mm | 248 ++++- .../apple/coreai/runtime/coreai_tensor_io.h | 22 + .../apple/coreai/runtime/coreai_tensor_io.mm | 403 ++++++++ .../apple/coreai/runtime/coreai_validation.h | 14 +- .../apple/coreai/runtime/coreai_validation.mm | 18 +- .../runtime/test/SDKSourceLifetimeTest.swift | 180 ++++ .../runtime/test/SwiftBridgeTests.swift | 754 +++++++++++++++ .../runtime/test/coreai_acquisition_fixture.h | 68 ++ .../test/coreai_acquisition_fixture.mm | 364 +++++++ .../runtime/test/coreai_acquisition_test.mm | 804 +++++++++++++++ .../coreai_cache_entry_contention_test.mm | 289 ++++++ .../runtime/test/coreai_cache_entry_fixture.h | 70 ++ .../test/coreai_cache_entry_fixture.mm | 442 +++++++++ .../runtime/test/coreai_cache_entry_test.mm | 521 ++++++++++ .../runtime/test/coreai_cache_key_test.mm | 209 ++++ .../runtime/test/coreai_cache_root_test.mm | 148 ++- .../coreai/runtime/test/coreai_cache_test.mm | 841 ++++++++++++++++ .../runtime/test/coreai_delegate_fixture.h | 123 +++ .../runtime/test/coreai_delegate_fixture.mm | 87 ++ .../runtime/test/coreai_delegate_test.mm | 912 ++++++++++++++++++ .../runtime/test/coreai_eviction_test.mm | 470 +++++++++ .../coreai/runtime/test/coreai_fake_loader.h | 128 +++ .../coreai/runtime/test/coreai_fake_loader.mm | 545 +++++++++++ .../runtime/test/coreai_filesystem_fixture.h | 23 + .../runtime/test/coreai_filesystem_fixture.mm | 98 ++ .../coreai/runtime/test/coreai_host_test.mm | 13 + .../runtime/test/coreai_manifest_fixture.h | 17 +- .../runtime/test/coreai_manifest_test.mm | 67 +- .../runtime/test/coreai_pte_export_test.cpp | 55 ++ .../coreai/runtime/test/coreai_pte_fixture.h | 76 ++ .../coreai/runtime/test/coreai_pte_fixture.mm | 254 +++++ .../coreai/runtime/test/coreai_pte_test.mm | 538 +++++++++++ .../runtime/test/coreai_source_fixture.h | 56 ++ .../runtime/test/coreai_source_fixture.mm | 112 +++ .../test/coreai_source_integrity_test.mm | 304 ++++++ .../test/coreai_source_publication_test.mm | 295 ++++++ .../coreai/runtime/test/coreai_source_test.mm | 428 ++++++++ .../runtime/test/coreai_storage_test.mm | 202 +++- .../coreai/runtime/test/coreai_test_utils.h | 20 + .../runtime/test/export_bridge_fixture.py | 125 +++ .../runtime/test/run_sdk_source_lifetime.py | 114 +++ .../coreai/runtime/test/runtime_smoke.cpp | 48 + .../test/test_export_bridge_fixture.py | 72 ++ .../test/test_run_sdk_source_lifetime.py | 97 ++ backends/apple/coreai/test/test_preprocess.py | 8 +- tools/cmake/executorch-config.cmake | 1 + tools/cmake/preset/default.cmake | 6 +- 82 files changed, 14670 insertions(+), 275 deletions(-) create mode 100644 backends/apple/coreai/cmake/PublishSwiftHeader.cmake create mode 100644 backends/apple/coreai/coreai-dependencies.cmake create mode 100644 backends/apple/coreai/runtime/ETCoreAIBridge.h create mode 100644 backends/apple/coreai/runtime/ETCoreAIError.swift create mode 100644 backends/apple/coreai/runtime/ETCoreAIModelLoader.swift create mode 100644 backends/apple/coreai/runtime/ETCoreAITensor.mm create mode 100644 backends/apple/coreai/runtime/ETCoreAITensors.swift create mode 100644 backends/apple/coreai/runtime/coreai_backend.mm create mode 100644 backends/apple/coreai/runtime/coreai_backend_id.h create mode 100644 backends/apple/coreai/runtime/coreai_bridge.mm create mode 100644 backends/apple/coreai/runtime/coreai_cache.h create mode 100644 backends/apple/coreai/runtime/coreai_cache.mm create mode 100644 backends/apple/coreai/runtime/coreai_cache_entry.h create mode 100644 backends/apple/coreai/runtime/coreai_cache_entry.mm create mode 100644 backends/apple/coreai/runtime/coreai_cache_key.h create mode 100644 backends/apple/coreai/runtime/coreai_cache_key.mm create mode 100644 backends/apple/coreai/runtime/coreai_error.h create mode 100644 backends/apple/coreai/runtime/coreai_error.mm create mode 100644 backends/apple/coreai/runtime/coreai_load_coordinator.h create mode 100644 backends/apple/coreai/runtime/coreai_load_coordinator.mm create mode 100644 backends/apple/coreai/runtime/coreai_pte.cpp create mode 100644 backends/apple/coreai/runtime/coreai_pte.h create mode 100644 backends/apple/coreai/runtime/coreai_source_bundle.h create mode 100644 backends/apple/coreai/runtime/coreai_source_bundle.mm create mode 100644 backends/apple/coreai/runtime/coreai_tensor_io.h create mode 100644 backends/apple/coreai/runtime/coreai_tensor_io.mm create mode 100644 backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift create mode 100644 backends/apple/coreai/runtime/test/SwiftBridgeTests.swift create mode 100644 backends/apple/coreai/runtime/test/coreai_acquisition_fixture.h create mode 100644 backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_acquisition_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_cache_entry_contention_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_cache_entry_fixture.h create mode 100644 backends/apple/coreai/runtime/test/coreai_cache_entry_fixture.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_cache_entry_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_cache_key_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_cache_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_delegate_fixture.h create mode 100644 backends/apple/coreai/runtime/test/coreai_delegate_fixture.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_delegate_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_eviction_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_fake_loader.h create mode 100644 backends/apple/coreai/runtime/test/coreai_fake_loader.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_pte_export_test.cpp create mode 100644 backends/apple/coreai/runtime/test/coreai_pte_fixture.h create mode 100644 backends/apple/coreai/runtime/test/coreai_pte_fixture.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_pte_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_source_fixture.h create mode 100644 backends/apple/coreai/runtime/test/coreai_source_fixture.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_source_integrity_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_source_publication_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_source_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_test_utils.h create mode 100644 backends/apple/coreai/runtime/test/export_bridge_fixture.py create mode 100644 backends/apple/coreai/runtime/test/run_sdk_source_lifetime.py create mode 100644 backends/apple/coreai/runtime/test/runtime_smoke.cpp create mode 100644 backends/apple/coreai/runtime/test/test_export_bridge_fixture.py create mode 100644 backends/apple/coreai/runtime/test/test_run_sdk_source_lifetime.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 0620bd5ad27..9abc10ff8bc 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -800,6 +800,15 @@ endif() if(EXECUTORCH_BUILD_VULKAN) string(APPEND _executorch_backend_dependencies "find_dependency(Threads)\n") endif() +if(EXECUTORCH_BUILD_COREAI) + string(APPEND _executorch_backend_dependencies + "include(\"\${CMAKE_CURRENT_LIST_DIR}/coreai-dependencies.cmake\")\n" + "executorch_coreai_dependencies()\n" + ) + install(FILES backends/apple/coreai/coreai-dependencies.cmake + DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/ExecuTorch + ) +endif() configure_file( tools/cmake/executorch-backend-dependencies.cmake.in ${CMAKE_CURRENT_BINARY_DIR}/executorch-backend-dependencies.cmake @ONLY @@ -845,6 +854,7 @@ endif() if(EXECUTORCH_BUILD_COREAI) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/apple/coreai) + list(APPEND _executorch_backends coreaidelegate) endif() if(EXECUTORCH_BUILD_MLX) diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index bb63606a1a3..33d0d2d1e33 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -8,14 +8,54 @@ if(NOT APPLE) message(FATAL_ERROR "Core AI requires Apple platforms") endif() +if(CMAKE_VERSION VERSION_LESS 3.31) + message(FATAL_ERROR "Core AI requires CMake 3.31 or newer for Swift support") +endif() + +# Inherited Clang flags and numeric definitions are not valid Swift options. +foreach(property IN ITEMS COMPILE_DEFINITIONS COMPILE_OPTIONS) + get_directory_property(_coreai_inherited ${property}) + set_property(DIRECTORY PROPERTY ${property} "") + foreach(value IN LISTS _coreai_inherited) + set_property( + DIRECTORY + APPEND + PROPERTY ${property} "$<$>:${value}>" + ) + endforeach() +endforeach() + enable_language(OBJCXX) +# The consolidated shared runtime target is created after this directory. +if(EXECUTORCH_BUILD_SHARED) + set(_coreai_library_type SHARED) + set(_coreai_runtime executorch_shared) + set(_coreai_data_loader "") +else() + set(_coreai_library_type STATIC) + set(_coreai_runtime executorch) + set(_coreai_data_loader extension_data_loader) +endif() + set(_coreai_runtime_sources - runtime/coreai_cache_root.mm runtime/coreai_fault_injection.mm - runtime/coreai_manifest.mm runtime/coreai_validation.mm + runtime/coreai_backend.mm + runtime/coreai_cache_entry.mm + runtime/coreai_cache.mm + runtime/coreai_cache_key.mm + runtime/coreai_cache_root.mm + runtime/coreai_error.mm + runtime/coreai_fault_injection.mm + runtime/coreai_load_coordinator.mm + runtime/coreai_manifest.mm + runtime/coreai_pte.cpp + runtime/coreai_source_bundle.mm + runtime/coreai_tensor_io.mm + runtime/coreai_validation.mm runtime/coreai_storage.mm ) +# Keep ARC and C++ exception settings away from the private Swift module. function(coreai_configure_objc_target target) target_include_directories( ${target} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/runtime @@ -32,17 +72,399 @@ endfunction() if(EXECUTORCH_BUILD_TESTS AND CMAKE_SYSTEM_NAME STREQUAL "Darwin") add_executable( coreai_host_test - runtime/test/coreai_host_test.mm runtime/test/coreai_manifest_test.mm - runtime/test/coreai_cache_root_test.mm runtime/test/coreai_storage_test.mm - runtime/test/coreai_filesystem_fixture.mm ${_coreai_runtime_sources} + runtime/test/coreai_host_test.mm + runtime/test/coreai_manifest_test.mm + runtime/test/coreai_cache_root_test.mm + runtime/test/coreai_storage_test.mm + runtime/test/coreai_filesystem_fixture.mm + runtime/test/coreai_cache_key_test.mm + runtime/test/coreai_cache_entry_test.mm + runtime/test/coreai_cache_entry_contention_test.mm + runtime/test/coreai_source_test.mm + runtime/test/coreai_source_integrity_test.mm + runtime/test/coreai_source_publication_test.mm + runtime/test/coreai_cache_entry_fixture.mm + runtime/test/coreai_source_fixture.mm + runtime/test/coreai_fake_loader.mm + runtime/test/coreai_acquisition_test.mm + runtime/test/coreai_acquisition_fixture.mm + runtime/test/coreai_delegate_fixture.mm + runtime/test/coreai_delegate_test.mm + runtime/test/coreai_pte_test.mm + runtime/test/coreai_pte_fixture.mm + runtime/test/coreai_cache_test.mm + runtime/test/coreai_eviction_test.mm + runtime/ETCoreAITensor.mm + ${_coreai_runtime_sources} ) coreai_configure_objc_target(coreai_host_test) target_compile_definitions(coreai_host_test PRIVATE COREAI_ASSETS_TESTING=1) find_library(COREAI_FOUNDATION_FRAMEWORK Foundation REQUIRED) target_link_libraries( - coreai_host_test PRIVATE executorch GTest::gtest - ${COREAI_FOUNDATION_FRAMEWORK} + coreai_host_test + PRIVATE ${_coreai_runtime} ${_coreai_data_loader} program_schema + GTest::gtest ${COREAI_FOUNDATION_FRAMEWORK} ) add_test(NAME coreai_host_test COMMAND coreai_host_test) set_tests_properties(coreai_host_test PROPERTIES TIMEOUT 120) endif() + +if(NOT CMAKE_GENERATOR MATCHES "Ninja|Xcode") + message(FATAL_ERROR "Core AI Swift compilation requires Ninja or Xcode") +endif() +if(NOT CMAKE_SYSTEM_NAME MATCHES "^(Darwin|iOS)$") + message(FATAL_ERROR "Core AI currently supports macOS and iOS only") +endif() + +include(${CMAKE_CURRENT_LIST_DIR}/coreai-dependencies.cmake) +executorch_coreai_dependencies() +set(_coreai_deployment 27.0) +if(CMAKE_OSX_DEPLOYMENT_TARGET VERSION_GREATER _coreai_deployment) + set(_coreai_deployment "${CMAKE_OSX_DEPLOYMENT_TARGET}") +endif() +# CMake 3.31 does not infer an iOS target for its Swift compiler probe. +if(CMAKE_SYSTEM_NAME STREQUAL "iOS" + AND NOT CMAKE_GENERATOR STREQUAL "Xcode" + AND NOT CMAKE_Swift_COMPILER_TARGET +) + if(CMAKE_OSX_ARCHITECTURES) + list(GET CMAKE_OSX_ARCHITECTURES 0 _coreai_probe_arch) + else() + set(_coreai_probe_arch "${CMAKE_SYSTEM_PROCESSOR}") + endif() + set(_coreai_probe_environment "") + if(CMAKE_OSX_SYSROOT MATCHES "[Ss]imulator") + set(_coreai_probe_environment "-simulator") + endif() + set(CMAKE_Swift_COMPILER_TARGET + "${_coreai_probe_arch}-apple-ios${_coreai_deployment}${_coreai_probe_environment}" + ) +endif() +enable_language(Swift) + +# CMake has no per-target OSX_DEPLOYMENT_TARGET property for Ninja. +function(coreai_require_os27 target) + set(deployment "${_coreai_deployment}") + set(target_environment "") + if(CMAKE_SYSTEM_NAME STREQUAL "iOS") + set(platform ios) + if(CMAKE_OSX_SYSROOT MATCHES "[Ss]imulator") + set(minimum_flag "-mios-simulator-version-min=${deployment}") + set(target_environment -simulator) + else() + set(minimum_flag "-miphoneos-version-min=${deployment}") + endif() + else() + set(platform macosx) + set(minimum_flag "-mmacosx-version-min=${deployment}") + endif() + set_target_properties( + ${target} + PROPERTIES XCODE_ATTRIBUTE_MACOSX_DEPLOYMENT_TARGET "${deployment}" + XCODE_ATTRIBUTE_IPHONEOS_DEPLOYMENT_TARGET "${deployment}" + ) + if(NOT CMAKE_GENERATOR STREQUAL "Xcode") + set(arch "${CMAKE_OSX_ARCHITECTURES}") + if(NOT arch) + set(arch "${CMAKE_SYSTEM_PROCESSOR}") + endif() + list(LENGTH arch arch_count) + if(NOT arch_count EQUAL 1) + message( + FATAL_ERROR + "Core AI Ninja builds require one CMAKE_OSX_ARCHITECTURES value" + ) + endif() + set(triple "${arch}-apple-${platform}${deployment}${target_environment}") + target_compile_options( + ${target} + PRIVATE "$<$:${minimum_flag}>" + "$<$:-target;${triple}>" + ) + target_link_options( + ${target} PRIVATE "$<$:${minimum_flag}>" + "$<$:-target;${triple}>" + ) + endif() +endfunction() + +function(coreai_configure_swift_target target module_name) + set_target_properties( + ${target} PROPERTIES Swift_MODULE_NAME "${module_name}" + Swift_LANGUAGE_VERSION 6 + ) + coreai_require_os27(${target}) + target_link_libraries(${target} PRIVATE executorch::coreai_dependencies) + if(CMAKE_GENERATOR STREQUAL "Xcode") + set_target_properties( + ${target} + PROPERTIES XCODE_ATTRIBUTE_PRODUCT_MODULE_NAME "${module_name}" + XCODE_ATTRIBUTE_SWIFT_MODULE_NAME "${module_name}" + XCODE_ATTRIBUTE_SWIFT_OBJC_BRIDGING_HEADER + "${CMAKE_CURRENT_SOURCE_DIR}/runtime/ETCoreAIBridge.h" + ) + else() + target_compile_options( + ${target} + PRIVATE + "$<$:-import-objc-header;${CMAKE_CURRENT_SOURCE_DIR}/runtime/ETCoreAIBridge.h>" + ) + endif() +endfunction() + +set(_coreai_generated_dir "${CMAKE_CURRENT_BINARY_DIR}/generated") +if(CMAKE_GENERATOR STREQUAL "Xcode") + string(APPEND _coreai_generated_dir + "/$(CONFIGURATION)$(EFFECTIVE_PLATFORM_NAME)" + ) +else() + file(MAKE_DIRECTORY "${_coreai_generated_dir}") +endif() +set(_coreai_swift_header "${_coreai_generated_dir}/CoreAIBridge-Swift.h") +set(_coreai_swift_sources + runtime/ETCoreAIError.swift runtime/ETCoreAITensors.swift + runtime/ETCoreAIModelLoader.swift +) +add_library(coreai_swift STATIC ${_coreai_swift_sources}) +coreai_configure_swift_target(coreai_swift CoreAIBridge) +set_target_properties( + coreai_swift PROPERTIES Swift_MODULE_DIRECTORY "${_coreai_generated_dir}" + POSITION_INDEPENDENT_CODE ON +) +if(CMAKE_GENERATOR STREQUAL "Xcode") + set_target_properties( + coreai_swift PROPERTIES XCODE_ATTRIBUTE_SWIFT_OBJC_INTERFACE_HEADER_NAME + CoreAIBridge-Swift.h + ) + # Xcode merges architecture-specific headers into DERIVED_SOURCES_DIR. + add_custom_command( + TARGET coreai_swift + POST_BUILD + COMMAND + ${CMAKE_COMMAND} "-DPUBLISH_ROOT=${CMAKE_CURRENT_BINARY_DIR}/generated" -P + "${CMAKE_CURRENT_SOURCE_DIR}/cmake/PublishSwiftHeader.cmake" + BYPRODUCTS "${_coreai_swift_header}" + COMMENT "Publish the merged Swift Objective-C header" + VERBATIM + ) +else() + target_compile_options( + coreai_swift + PRIVATE + "$<$:-emit-objc-header;-emit-objc-header-path;${_coreai_swift_header}>" + ) + add_custom_command( + TARGET coreai_swift + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E true + BYPRODUCTS "${_coreai_swift_header}" + COMMENT "Track the generated Swift Objective-C header" + ) +endif() +# The bridge header is an input to Swift even when only its declarations change. +set_property( + SOURCE ${_coreai_swift_sources} + APPEND + PROPERTY OBJECT_DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/runtime/ETCoreAIBridge.h" +) + +# Keep SDK entry points separate so delegate tests can substitute a fake. +add_library( + coreai_bridge_obj OBJECT runtime/ETCoreAITensor.mm runtime/coreai_bridge.mm +) +coreai_configure_objc_target(coreai_bridge_obj) +coreai_require_os27(coreai_bridge_obj) +set_target_properties(coreai_bridge_obj PROPERTIES POSITION_INDEPENDENT_CODE ON) +target_include_directories(coreai_bridge_obj PRIVATE "${_coreai_generated_dir}") +add_dependencies(coreai_bridge_obj coreai_swift) +set_property( + SOURCE runtime/coreai_bridge.mm + APPEND + PROPERTY OBJECT_DEPENDS "${_coreai_swift_header}" +) + +set(_coreai_run_sdk_tests FALSE) +if(EXECUTORCH_BUILD_TESTS + AND CMAKE_SYSTEM_NAME STREQUAL "Darwin" + AND NOT CMAKE_CROSSCOMPILING +) + execute_process( + COMMAND sw_vers -productVersion + OUTPUT_VARIABLE _coreai_host_version + OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY + ) + if(_coreai_host_version VERSION_GREATER_EQUAL _coreai_deployment) + set(_coreai_run_sdk_tests TRUE) + endif() +endif() + +# The Swift tests use the bridge's internal declarations, so they compile into +# one module with it and run under xctest. Ninja on macOS only. +if(EXECUTORCH_BUILD_TESTS + AND CMAKE_SYSTEM_NAME STREQUAL "Darwin" + AND NOT CMAKE_GENERATOR STREQUAL "Xcode" +) + execute_process( + COMMAND xcrun --sdk macosx --show-sdk-platform-path + OUTPUT_VARIABLE _coreai_platform + OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY + ) + set(_coreai_xctest_frameworks + "${_coreai_platform}/Developer/Library/Frameworks" + ) + set(_coreai_xctest_libs "${_coreai_platform}/Developer/usr/lib") + add_library( + coreai_swift_bridge_test MODULE + ${_coreai_swift_sources} runtime/test/SwiftBridgeTests.swift + runtime/ETCoreAITensor.mm + ) + set_target_properties( + coreai_swift_bridge_test PROPERTIES BUNDLE TRUE BUNDLE_EXTENSION xctest + ) + coreai_configure_objc_target(coreai_swift_bridge_test) + coreai_configure_swift_target(coreai_swift_bridge_test CoreAIBridgeTests) + target_compile_options( + coreai_swift_bridge_test + PRIVATE + "$<$:-F;${_coreai_xctest_frameworks};-I;${_coreai_xctest_libs}>" + ) + target_link_options( + coreai_swift_bridge_test PRIVATE "-F${_coreai_xctest_frameworks}" + "-L${_coreai_xctest_libs}" "LINKER:-rpath,${_coreai_xctest_frameworks}" + "LINKER:-rpath,${_coreai_xctest_libs}" + ) + target_link_libraries( + coreai_swift_bridge_test PRIVATE "-framework XCTest" XCTestSwiftSupport + ) + file(GLOB_RECURSE _coreai_fixture_sources CONFIGURE_DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/*.py" + ) + set(_coreai_fixture_dir "${CMAKE_CURRENT_BINARY_DIR}/bridge_fixture") + add_custom_command( + OUTPUT "${_coreai_fixture_dir}/manifest.json" + "${_coreai_fixture_dir}/dynamic/manifest.json" + "${_coreai_fixture_dir}/model.pte" + BYPRODUCTS "${_coreai_fixture_dir}/model.aimodel/main.hash" + "${_coreai_fixture_dir}/model.aimodel/main.mlirb" + "${_coreai_fixture_dir}/model.aimodel/metadata.json" + "${_coreai_fixture_dir}/dynamic/model.aimodel/main.hash" + "${_coreai_fixture_dir}/dynamic/model.aimodel/main.mlirb" + "${_coreai_fixture_dir}/dynamic/model.aimodel/metadata.json" + COMMAND + "${PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/runtime/test/export_bridge_fixture.py" + "${_coreai_fixture_dir}" + DEPENDS ${_coreai_fixture_sources} + WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}" + COMMENT "Export the Core AI session test model" + VERBATIM + ) + add_custom_target( + coreai_bridge_fixture + DEPENDS "${_coreai_fixture_dir}/manifest.json" + "${_coreai_fixture_dir}/dynamic/manifest.json" + "${_coreai_fixture_dir}/model.pte" + ) + add_dependencies(coreai_swift_bridge_test coreai_bridge_fixture) + add_executable( + coreai_pte_export_test runtime/test/coreai_pte_export_test.cpp + runtime/coreai_pte.cpp + ) + coreai_configure_objc_target(coreai_pte_export_test) + target_link_libraries( + coreai_pte_export_test PRIVATE ${_coreai_runtime} ${_coreai_data_loader} + program_schema + ) + add_dependencies(coreai_pte_export_test coreai_bridge_fixture) + if(_coreai_run_sdk_tests) + add_test(NAME coreai_pte_export_test + COMMAND coreai_pte_export_test "${_coreai_fixture_dir}/model.pte" + "${_coreai_fixture_dir}/manifest.json" + ) + set_tests_properties(coreai_pte_export_test PROPERTIES TIMEOUT 30) + add_test(NAME coreai_swift_bridge_test + COMMAND xcrun xctest + "$" + ) + set_tests_properties( + coreai_swift_bridge_test + PROPERTIES TIMEOUT 600 ENVIRONMENT + "COREAI_BRIDGE_FIXTURE_DIR=${_coreai_fixture_dir}" + ) + endif() + add_executable( + coreai_sdk_source_lifetime_test + runtime/test/SDKSourceLifetimeTest.swift ${_coreai_swift_sources} + runtime/ETCoreAITensor.mm + ) + coreai_configure_objc_target(coreai_sdk_source_lifetime_test) + coreai_configure_swift_target( + coreai_sdk_source_lifetime_test coreai_sdk_source_lifetime_test + ) + target_compile_options( + coreai_sdk_source_lifetime_test + PRIVATE "$<$:-parse-as-library>" + ) + if(_coreai_run_sdk_tests) + add_test( + NAME coreai_sdk_source_lifetime_test + COMMAND + "${Python3_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/runtime/test/run_sdk_source_lifetime.py" + "$" + ) + set_tests_properties(coreai_sdk_source_lifetime_test PROPERTIES TIMEOUT 240) + endif() +endif() + +add_library( + coreaidelegate ${_coreai_library_type} ${_coreai_runtime_sources} + $ +) +coreai_configure_objc_target(coreaidelegate) +# The delegate contains no Swift sources; its archive is assembled by C++. +set_target_properties(coreaidelegate PROPERTIES LINKER_LANGUAGE CXX) +coreai_require_os27(coreaidelegate) +target_include_directories(coreaidelegate PUBLIC ${_common_include_directories}) +add_dependencies(coreaidelegate coreai_swift) +target_link_libraries( + coreaidelegate PRIVATE coreai_swift ${_coreai_runtime} ${_coreai_data_loader} + program_schema executorch::coreai_dependencies +) +executorch_target_link_options_shared_lib(coreaidelegate) +if(EXECUTORCH_BUILD_SHARED) + set_target_properties( + coreaidelegate PROPERTIES OUTPUT_NAME executorch_backend_coreai + ) + executorch_target_soname_policy(coreaidelegate) + executorch_target_shipped_runtime_path(coreaidelegate) +endif() + +# A consumer is necessary to check Swift linkage and registration retention. +add_executable(coreai_runtime_smoke runtime/test/runtime_smoke.cpp) +if(_coreai_run_sdk_tests) + add_test(NAME coreai_runtime_smoke COMMAND coreai_runtime_smoke) +endif() +if(NOT EXECUTORCH_BUILD_TESTS) + set_target_properties(coreai_runtime_smoke PROPERTIES EXCLUDE_FROM_ALL TRUE) +endif() +coreai_require_os27(coreai_runtime_smoke) +set_target_properties(coreai_runtime_smoke PROPERTIES LINKER_LANGUAGE CXX) +target_link_libraries( + coreai_runtime_smoke PRIVATE coreaidelegate ${_coreai_runtime} +) + +install( + FILES runtime/coreai_cache.h + DESTINATION + ${CMAKE_INSTALL_INCLUDEDIR}/executorch/backends/apple/coreai/runtime +) + +install( + TARGETS coreaidelegate coreai_swift + EXPORT ExecuTorchTargets + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} +) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 1a3d96adc7a..e76c209991e 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -30,18 +30,62 @@ The delegate manifest describes `.aimodel` source bundles, which require `inline` packaging, and AOT-compiled `.aimodelc` bundles, which require `aot_compiled_inline`. Both formats require a `function` name, and ordered `input_names` and `output_names` matching the converter's bindings. Dictionary -iteration order is not a binding contract. The `files` object maps relative +iteration order is not a binding contract. `export_hash` identifies the exported +program used to namespace embedded assets. The `files` object maps relative filenames to byte sizes; `bundle_digests` maps each bundle basename to its export-time SHA-256 digest. Asset paths, including bundle and directory names, must remain distinct under case-insensitive, canonically normalized Unicode comparison. Accepted names are preserved, not rewritten. Parsing returns an immutable `Manifest` covering all exported bundles. Pure -platform and architecture selection produces an immutable `SelectedBundle` -with one bundle's path, digest, and files. Neither type can be constructed +platform and architecture selection through `select_bundle` produces an immutable +`SelectedBundle` with one bundle's path, digest, and files. Host probing returns a +`SelectedHostBundle` containing that selection and its platform/architecture. without validation; filesystem materialization and SDK host probing remain outside the manifest parser. +The experimental `CoreAIBackend` runtime supports stateless FP32/FP16 tensor +models in both formats. ExecuTorch tensors must be contiguous. Output shapes +must fit the executor's declared resize and capacity constraints. Stateful +functions, image values, other dtypes, and interleaved layouts are not supported. +Calls on one session must be serialized by the caller. + +Supported inputs are borrowed directly from ExecuTorch storage through Core AI +raw views. The bridge copies shape metadata, not tensor bytes, and does not +allocate intermediate input `NSData`, `Data`, or `NDArray` buffers. Each call +wraps `const_data_ptr()` and the current shape's `nbytes()`, not the allocation's +maximum capacity. Outputs use the SDK's actual returned shape; the backend +resizes them within ExecuTorch's declared bounds before copying only the logical +result bytes. Unused capacity is neither wrapped as input nor written as output. +The backend waits for inference and all input access to finish before returning +from `execute`, including on errors. This uses ExecuTorch's normal input-lifetime +contract; callers must not concurrently mutate, release, or rebind the input +storage. Output copiers retain the SDK arrays and copy directly into validated +ExecuTorch destinations, without an intermediate `Data` allocation or payload +copy. If any destination overlaps an input or another destination, all outputs +are staged before writing destinations in order. This preserves input/output +aliasing and overlapping-output semantics. Staging is not zero-initialized. +This does not guarantee that Core AI itself avoids device transfers or internal +copies. + +After an execution error, do not use the outputs as inference results. Output +shapes may already have changed, but no result bytes are copied unless every +output passes validation and resizing. Malformed SDK outputs return `Internal`. +Output dtype/rank mismatches and incompatible shapes or capacity return +`InvalidArgument`. + +### Calling from Swift + +When using the Core AI backend, run ExecuTorch model loading, inference and +cache operations on a background `DispatchQueue` or a dedicated worker thread. +Do not call these synchronous APIs directly inside Swift `Task` or +`Task.detached`, because doing so can hang your application. + +From async Swift code, submit the entire ExecuTorch call with +`DispatchQueue.async` and await its result using a continuation. Do not block +the Swift task with `DispatchQueue.sync` or a semaphore. Keep input tensors alive +and unchanged through all output copies, even if the waiting task is cancelled. + ### AOT architecture selection Configure AOT export with `AOTCompileConfig`, including the target `platform` and @@ -50,19 +94,72 @@ 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`. -### Asset storage +At load time, the runtime queries `AIModel.deviceArchitectureName` and selects +the exact matching entry in the manifest's `archs` map. No runtime architecture +option is required, and there is no fallback to another architecture. Platform, +deployment floor, and architecture checks happen before asset reads or storage +creation. Missing-architecture errors identify the device architecture, list +available architectures, and request a matching export. Missing selected files +are reported separately from an absent architecture entry. + +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. -The assets root must be an absolute path. The application chooses it, should -reserve it for this backend, and is trusted not to rename, replace or modify its -contents while the backend uses it; the backend does not defend against other -processes of the same user rebinding paths. Written files and changed directories -are synced with `fsync`. Atomic publication uses `F_FULLFSYNC` on the temporary -file before rename and on the directory afterward. A failure before rename -preserves the old contents; a failure of the final flush reports an error even -though the new complete contents are already visible. Use a local APFS volume; -unsupported full-device flushes are reported as errors, not downgraded to `fsync`. +### Cache directory -Preparing the assets root sets `NSURLIsExcludedFromBackupKey` on it, which covers +Backend initialization reads the string runtime spec `coreai_cache_dir`. If +absent, it sets the directory to `executorch_coreai` beneath user-domain +`NSCachesDirectory`. The coordinator and storage helpers always receive that +explicit directory. It contains raw bookmarks, source bundles and lock files. + +The cache directory must be an absolute path reserved for this backend. A +symlink in the final component of the configured path is resolved before +opening the cache directory. Entry and source operations retain directory descriptors and do not follow internal +symlinks. The application must still prevent concurrent renaming, replacement, +or modification while SDK calls use source URLs. Written files and changed directories +are synced with `fsync`. Bookmark publication uses `F_FULLFSYNC` on the temporary +file before rename and on the directory afterward. Source-bundle publication +fully flushes the pending build directory after syncing its files and subdirectories, +then fully flushes the destination directory after rename. A failure before +rename preserves the old contents; a failure of the final flush reports an error +even though the new complete contents are already visible. Use a local APFS +volume; unsupported full-device flushes are reported as errors, not downgraded +to `fsync`. + +Inline bundles are read through `NamedDataMap` and materialized as unchanged +files under the cache directory. Named data can be supplied externally; inline +packaging does not require all bytes to reside in one physical PTE. The exporter +computes each bundle digest from all delivered relative filenames and file bytes, +not just the SDK's bytecode-only `main.hash`. Each AOT architecture has an +independent digest; only the selected architecture is materialized. NDS key names +and raw asset bytes are unchanged. + +When source recovery is needed, existing bundles are checked for the expected +file set, sizes, entry types and directory structure without reading asset +contents or requesting NDS payloads. Unicode paths are compared in canonical +decomposed form; NDS keys and bytes remain unchanged. Reuse fully flushes the +`sources` directory, including after a previous publication's final flush failed. +Missing bundles are written to a `.pending-` directory and published with an +exclusive rename; each selected NDS payload is fetched once and checked for size +before writing. If another loader publishes first, its bundle is validated and +used. Incomplete, size-mismatched or otherwise malformed existing bundles fail +normal source preparation without replacement, since an SDK model may still be +using them. The acquisition coordinator can explicitly replace a damaged bundle +under the same entry lock and rebuild it from available named data. Replacement +preflights the full descriptor-bound tree, then revalidates entries while +removing them without following links. Named-data and SDK failures are not +retried. Temporary bundles are removed on ordinary failures before publication; +a caller can explicitly retry after addressing the failure. + +These are completeness checks, not full content-integrity verification. Same-size +content changes are not detected, and a complete stored bundle does not cause its +NDS source to be reread. The export-time digest is a trusted artifact identifier; +the runtime does not rehash source or stored bytes. Core AI may reject invalid +contents, but SDK load failures do not automatically evict the extracted copy. + +Preparing the cache directory sets `NSURLIsExcludedFromBackupKey` on it, which covers everything beneath it. Ancestors are not modified. This is backup exclusion, not a control for iCloud Drive synchronization. @@ -76,19 +173,270 @@ Directory creation uses default permissions filtered by the process umask, witho an explicit owner-only restriction. Existing permissions are not changed. Backup exclusion and required durability flushes still apply on reused paths. -Cache-root preparation owns location and backup/protection policy; descriptor +Cache-directory preparation owns location and backup/protection policy; descriptor storage owns no-follow operations, publication, and synchronization. Creation and enumeration use explicit policies, with a missing child distinguishable from an error. Fault injection is a separate, thread-local test interface and compiles away in production. +Delegate construction and registration do not create storage directories or load +models. Per-model initialization performs SDK acquisition and function loading, +validating source assets only when recovery requires them. There is no synthetic +zero-input inference at initialization. SDK specialization and function resource +loading can still be substantial. + +The runtime uses `AIModelCache.default` with storage-pressure-only purging and +stores an opaque bookmark for direct restoration. `coreai_cache_dir` does not +choose the SDK artifact directory. The backend does not inspect SDK-private +files, alter their backup flags, or use an App Group cache. + +Each selected bundle has a `bundle_identity` derived from its bundle digest, +name, and source/AOT mode. This identifies the delivered bundle independently of +`export_hash`, which namespaces embedded assets. A readable cache key combines +that bundle identity with +the platform and Core AI device architecture, plus a hash of a versioned +specialization identity supplied as plain data by the SDK boundary. Function +bindings and delivery location are excluded, and no model bytes are hashed at +runtime. Each variant key determines a raw bookmark in `bookmarks`, a bundle +directory in `sources`, and one of 256 shared coordination files in `locks`. +There is no PTE-wide bookmark list or serialized lifecycle record. + +A healthy bookmark hit does not consult NDS and attempts to clean the leftover +source bundle. On a missing bookmark or confirmed SDK miss, the coordinator prepares +the source, specializes it and atomically writes the returned bookmark bytes. +After durable publication, it attempts to remove the entire keyed source +directory while holding the key's stripe lock. The SDK cache supports +inference and bookmark restoration without that extracted copy. +Keep the packaged named data available: a later SDK cache miss requires source +materialization again. 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. Cleanup after a bookmark hit or durable publication is +best-effort: failures are logged without discarding the usable model or bookmark. +A later load can finish cleanup or rebuild a partially removed source after an +SDK miss. No cleanup failure triggers SDK eviction or retries within that load. +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. +A fixed pool of at most 256 disk locks coordinates loading and eviction across +threads and processes. A SHA-256 hash of the full validated cache key selects a +stripe modulo the configured pool size. Files are created on demand, so new +model versions do not grow the pool beyond 256. Different keys on the same +stripe serialize; different stripes remain independent. Bookmarks and source bundles +remain separate for every key. Runtime operations hold only one stripe at a +time. + +Lock ownership is maintained by the OS, and acquiring a lock does not flush its +file or directory. Stripe files are never removed, including during cache +clearing or after process exit, so waiters and future openers keep using the +same inode. Removing the cache directory externally requires all loads, sessions +and maintenance to stop. + +### Lifecycle logging + +Build with `EXECUTORCH_ENABLE_LOGGING=ON` and `EXECUTORCH_LOG_LEVEL=Info` to +observe cache lookup, bookmark restoration hits, missing-bookmark versus +SDK-confirmed misses, specialization start/success, durable bookmark saves, and +cache-clearing milestones. Each message carries the cache key so events from +multiple partitions can be correlated. These messages do not include bookmark +contents, model data, or per-inference logs. + +A cache hit means that the SDK successfully restored a model from its bookmark. +Specialization start means a request to the SDK's default cache; the +SDK does not report whether that request compiled afresh or reused an internal +artifact. The durable-save message confirms local bookmark publication, +including its final directory flush, not independent verification of SDK-private +artifact storage. During clearing, SDK entry absence and successful local +cleanup are reported separately. Failed operations do not emit their success +messages, and clearing an absent entry remains quiet. + +### Clearing cached assets + +The public C++ maintenance APIs are in `executorch::backends::coreai`: + +```cpp +runtime::Error clear_cache(const char* coreai_cache_dir = nullptr); +runtime::Error clear_cache_for_pte( + runtime::DataLoader& pte, const char* coreai_cache_dir = nullptr); +runtime::Error clear_cache_for_pte( + const char* pte_path, const char* coreai_cache_dir = nullptr); +``` + +All three APIs use the same default cache directory as backend initialization +when the directory is omitted or null. A supplied path must be absolute; invalid +paths fail without falling back to the default. If loading used a custom cache +directory, pass it when clearing; the PTE does not store runtime overrides. + +Cache clearing inventories keyed bookmarks and staged directories. PTE clearing +finds every Core AI delegate in every method and derives each bundle's identity. +Matching that identity in cache keys clears every recorded platform, +device, and specialization variant without querying the host or SDK. All +manifests are validated before any entry is cleared; their minimum OS and runtime +compatibility do not restrict maintenance. Discovery does not load models or +materialize assets. Actual SDK eviction still requires a usable bookmark and an +available loader. + +These synchronous calls require callers to unload affected models and prevent +concurrent loads, inference and maintenance across processes. Stripe locks are +not a whole-root barrier. Clearing first deletes the recorded SDK entry, then +its source bundle, then its bookmark. Busy or unknown SDK errors preserve the files; +SDK-confirmed absence permits cleanup. Filesystem errors may leave partial +cleanup, but the bookmark is retained until source removal succeeds so cleanup +can be retried. Independent keys are still attempted in sorted order, and the +first key error is returned. Each removal syncs the directory it changed. A +missing root or absent entry is a noncreating no-op. + +Everything under a keyed `sources/` directory belongs to the backend and is +removed with that key. Outside keyed entries, clearing does not delete PTEs, +other files or the stable lock files. Unkeyed interrupted `.pending-*` directories +also remain untouched. Bookmarks are SDK eviction handles, not weights. For an +existing entry with a missing or structurally invalid bookmark, clearing logs an +Info message with the cache key, attempts local cleanup, and continues to other keys. +The message states that SDK eviction could not be confirmed; bookmark contents +are not logged. Absent entries remain quiet no-ops. I/O failures and SDK eviction +errors still preserve bookmarks for retry. Lost bookmarks can leave SDK entries +that these APIs cannot target. No history of removed entries is retained. +Copied or reused exported artifacts can share keys; clearing them also removes +the other copy's warm-cache benefit. Subsequent loads can specialize again. + + +### Persistent SDK Rejection + +SDK errors are not classified as transient or permanent. +Busy errors and all unclassified eviction/restore errors retain the local +bookmark for retry; only successful eviction or a confirmed SDK miss permits +normal removal. If retries keep failing, first stop every process using the +cache directory. Back up or quarantine only the affected key's bookmark and keyed +source bundle outside that directory, leaving the stable lock files in place. A later load +can attempt cold acquisition and publish a new bookmark. +This manual recovery does not confirm SDK eviction or guarantee disk-space +reclamation: keep the backup because it may contain the only remaining eviction +handle. Do not perform manual recovery while loading, inference, or maintenance +is running. + +### Runtime options + +Pass options through the public `Module` API before loading. In an error-returning +loader using the `executorch::runtime` namespace: + +```cpp +BackendOptions<1> options; +ET_CHECK_OK_OR_RETURN_ERROR(options.set_option("coreai_cache_dir", cache_dir)); +LoadBackendOptionsMap map; +ET_CHECK_OK_OR_RETURN_ERROR(map.set_options("CoreAIBackend", options.view())); +ET_CHECK_OK_OR_RETURN_ERROR(module.load(map)); +``` + +No options are required when using the default cache directory. `coreai_cache_dir` +selects an application-chosen cache directory reserved for this backend. The +supplied path is validated during preflight, but warm bookmark hits need no +source. + +Use separate option maps to choose different cache directories for different models. +The path names the cache directory, not a bundle. ExecuTorch option strings +have a 255-byte limit. + +## SDK Bridge + +The private Swift module `CoreAIBridge` separates error/completion support, +tensor layout/copying, and model/session lifecycle into three source files. It +wraps the `CoreAI` framework behind the Objective-C protocols in +`runtime/ETCoreAIBridge.h`. It restores models from bookmarks, specializes source +bundles in the default SDK cache with storage-pressure-only purging, and binds +named functions with ordered inputs and outputs. Each loader owns an immutable +snapshot of its actual specialization options and exposes a stable, versioned +identity derived from that snapshot. Asynchronous SDK calls run in detached +tasks and report through completion blocks; function binding completes +synchronously. + +Inputs are borrowed from caller storage through Core AI raw views; only shape +metadata is copied. Outputs retain their SDK storage and expose a copier after +all fallible layout validation finishes. Dense row-major outputs copy with one +`memcpy`, while padded or transposed outputs use strided copies. Only FP16 and +FP32 NDArrays are supported. Callers keep inputs live through output copying and +stage copies before writing destinations that could alias inputs or other output +destinations. + +Bridge-owned Objective-C types live in `runtime/ETCoreAITensor.mm`, separately +from the production SDK entry points in `runtime/coreai_bridge.mm`. Both build in +the private `coreai_bridge_obj` object library, while SDK-free tests can link the +types and substitute fake entry points. + +`coreai_swift_bridge_test` runs the bridge's XCTests under `xctest` with +`EXECUTORCH_BUILD_TESTS=ON`. It links CoreAI and is registered with CTest only +when the macOS host meets the configured Core AI deployment target. +The test build exports static- and dynamic-shape models for real session tests. +The SDK source-lifetime regression checks live inference and fresh-process +bookmark restoration through the bridge after removing disposable portable and +AOT source copies. The AOT case is skipped when the SDK cannot identify the +device architecture. A killed probe can leave an SDK entry if it dies before +saving its bookmark; cleanup never purges unrelated default-cache entries. +CMake's `PYTHON_EXECUTABLE` must use the Python environment with the pinned Core +AI requirements above. Backend Python changes regenerate the dedicated fixture +directory, replacing previous contents. CTest supplies its location to XCTest; +direct XCTest runs without it skip the session tests. Building only +`coreai_host_test` does not export fixtures or build the Swift test bundle. + +## Building + +Core AI builds require an Apple SDK 27 or newer, Swift from Xcode 27, and CMake +3.31 or newer with the Ninja or Xcode generator. Execution requires macOS 27 or +iOS 27. `EXECUTORCH_BUILD_COREAI` defaults to `OFF`. Enabling it builds the +private Swift implementation, module `CoreAIBridge`. Only Core AI targets get an +OS 27 minimum; unrelated runtime targets keep the configured deployment target. +Ninja builds use one architecture per build directory. Current runtime support +covers arm64 macOS and iOS device builds. x86_64 builds are blocked by Swift +`Float16` availability, and the tested iOS simulator SDKs do not contain Core AI. + +Enable `EXECUTORCH_BUILD_EXTENSION_DATA_LOADER` for the public PTE-file +maintenance API and consumers such as `Module`. +Link the CMake target `coreaidelegate`, not just its archive filename. Its +transitive dependencies supply Swift and Foundation/CoreAI linkage, and its +link interface retains static backend registration. With +`EXECUTORCH_BUILD_SHARED=ON`, the delegate is a shared library linked to the +single consolidated `executorch_shared` runtime. Otherwise it is static. +Consumers of the shared delegate must use that same shared runtime. +Installed CMake targets resolve framework and Swift runtime dependencies using +the consumer's selected SDK and toolchain, not the producer's Xcode paths. +Select the Apple SDK/toolchain when configuring the consumer project. +SwiftPM and XCFramework distribution are not integrated yet. + +The `coreai_runtime_smoke` target checks final linkage of a C++ consumer, +including Swift dependencies, backend registration retention, and the public +cache-clearing overloads. With +`EXECUTORCH_BUILD_TESTS=ON` it is built; otherwise it is excluded from the +default build. SDK-linked tests are registered with CTest only for native macOS +builds when the host meets the target's deployment minimum (at least 27.0). +Cross-builds and older hosts can build these tests without registering them. +Running the smoke test checks registration and availability; it does not perform +model inference. Do not execute SDK27 binaries on an older host. + +## Backend-local PTE Inspection + +The private `inspect_coreai_pte(DataLoader&)` reader extracts Core AI processed +buffers without loading methods, external assets, or SDK models. Inspection does +not derive cache keys or mutate cache state. It verifies FlatBuffer structure +and program semantics, skips delegates without a matching backend ID, and +returns `NotFound` for absent segments or out-of-range references. Legacy +headers may omit the segment extent; selected segment bounds are still checked. + ## Host Tests On macOS, `EXECUTORCH_BUILD_COREAI=ON` with `EXECUTORCH_BUILD_TESTS=ON` registers the standalone `coreai_host_test` executable with CTest. It tests -manifest parsing, AOT asset selection, and filesystem storage using Foundation, -without Swift or CoreAI SDK calls. The host test target is omitted for Apple -cross-builds. +manifest parsing, PTE inspection, AOT asset selection, cache identity, bookmark +I/O, acquisition coordination, lock contention, filesystem storage, and delegate +initialization and execution against a fake bridge. It uses Foundation without +Swift or CoreAI SDK calls. Acquisition contention tests run two threads inside a +bounded child process; bookmark lock tests separately cover contention between +processes. The host test target is omitted for Apple cross-builds. On supported +macOS hosts, CTest also runs the Swift bridge tests, source-lifetime probes, and +`coreai_runtime_smoke`, which link Swift and the Core AI framework. +`coreai_pte_export_test` compares inspected metadata with the manifest from an +actual exported PTE. It uses the generated bridge fixture without loading SDK +models. CI builds and runs these tests on macOS 27. To use the same configuration from the repository root: @@ -101,3 +449,9 @@ cmake -S . -B -G Ninja \ cmake --build --target backends/apple/coreai/all ctest --test-dir /backends/apple/coreai --output-on-failure --no-tests=error ``` + +For a compile-only production smoke build, use a fresh build directory without +`EXECUTORCH_BUILD_TESTS`, set `CMAKE_OSX_ARCHITECTURES=arm64`, and build +`coreai_runtime_smoke`. For iOS also set `CMAKE_SYSTEM_NAME=iOS` and +`CMAKE_OSX_SYSROOT=iphoneos`. Do not run the iOS binary on the host or install +a simulator as part of this build. diff --git a/backends/apple/coreai/cmake/PublishSwiftHeader.cmake b/backends/apple/coreai/cmake/PublishSwiftHeader.cmake new file mode 100644 index 00000000000..19c4006acb4 --- /dev/null +++ b/backends/apple/coreai/cmake/PublishSwiftHeader.cmake @@ -0,0 +1,17 @@ +# 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. + +if("$ENV{DERIVED_SOURCES_DIR}" STREQUAL "" OR "$ENV{CONFIGURATION}" STREQUAL "") + message(FATAL_ERROR "Swift header publication requires Xcode build settings") +endif() +set(_header "$ENV{DERIVED_SOURCES_DIR}/$ENV{SWIFT_OBJC_INTERFACE_HEADER_NAME}") +set(_destination + "${PUBLISH_ROOT}/$ENV{CONFIGURATION}$ENV{EFFECTIVE_PLATFORM_NAME}" +) +file(MAKE_DIRECTORY "${_destination}") +file(COPY_FILE "${_header}" "${_destination}/CoreAIBridge-Swift.h" + ONLY_IF_DIFFERENT +) diff --git a/backends/apple/coreai/compiler/preprocess.py b/backends/apple/coreai/compiler/preprocess.py index cf2f86224a8..7ff320accb1 100644 --- a/backends/apple/coreai/compiler/preprocess.py +++ b/backends/apple/coreai/compiler/preprocess.py @@ -157,10 +157,9 @@ def _get_compile_spec( return None -def _nds_key(model_hash: str, relative_path: str) -> str: - # Namespaced by content hash so identical assets dedup to one entry and the - # hash doubles as an on-device cache key. - return f"coreai/{model_hash}/{relative_path}" +def _nds_key(export_hash: str, relative_path: str) -> str: + # Identical exports share embedded assets under their export hash. + return f"coreai/{export_hash}/{relative_path}" def _asset_metadata(staging: Path) -> Dict[str, Any]: @@ -194,23 +193,23 @@ def _asset_metadata(staging: Path) -> Dict[str, Any]: def _deliver( staging: Path, - model_hash: str, + export_hash: str, packaging: AssetPackaging, manifest_extra: Dict[str, Any], ) -> PreprocessResult: """Embed a staging directory of bundles in the .pte with its manifest. ``staging`` holds the ``.aimodel`` / ``.aimodelc`` bundles, never their - contents, so assets land under ``coreai///`` in the + contents, so assets land under ``coreai///`` in the NamedDataStore. """ manifest = { "packaging": packaging.value, - "hash": model_hash, + "export_hash": export_hash, **manifest_extra, **_asset_metadata(staging), } - return _embed_dir_inline(staging, model_hash, manifest) + return _embed_dir_inline(staging, export_hash, manifest) def _prepare_program_for_conversion(edge_program: ExportedProgram) -> ExportedProgram: @@ -255,21 +254,21 @@ def _convert_to_aiprogram( # Asset embedding helpers. def _embed_dir_inline( - root_dir: Path, model_hash: str, manifest_extra: Dict[str, Any] + root_dir: Path, export_hash: str, manifest_extra: Dict[str, Any] ) -> PreprocessResult: """Flatten every file under ``root_dir`` into the NamedDataStore. ``root_dir`` is the directory *containing* the ``.aimodel`` / ``.aimodelc`` bundles, never a bundle itself, so the bundle name survives into the keys. - Keys are ``coreai/{hash}/{relpath}``; ``manifest_extra`` becomes the + Keys are ``coreai/{export_hash}/{relpath}``; ``manifest_extra`` becomes the manifest (packaging, files, archs, ...). """ store = NamedDataStore() for path in sorted(root_dir.rglob("*")): if path.is_file(): store.add_named_data( - _nds_key(model_hash, path.relative_to(root_dir).as_posix()), + _nds_key(export_hash, path.relative_to(root_dir).as_posix()), path.read_bytes(), alignment=_ASSET_ALIGNMENT, ) @@ -426,15 +425,15 @@ def _compiled_arch_bundles(out_dir: Path) -> List[Tuple[str, Path]]: def _compile_aot(program, opts: Dict[str, Any], tmp_dir: Path) -> Tuple[str, Path]: """Save the intermediate ``.aimodel`` and AOT-compile it to ``.aimodelc``. - Returns ``(model_hash, compiled_out_dir)``. The intermediate ``.aimodel`` + Returns ``(export_hash, compiled_out_dir)``. The intermediate ``.aimodel`` keeps the default OS floor; ``min_deployment_version`` targets the compiled ``.aimodelc`` via ``coreai-build``, so it is not double-specified. """ - model_hash = _save_and_hash(program, tmp_dir / "model.aimodel") + export_hash = _save_and_hash(program, tmp_dir / "model.aimodel") out = tmp_dir / "compiled" out.mkdir() _run_coreai_build(tmp_dir / "model.aimodel", out, opts) - return model_hash, out + return export_hash, out @final @@ -446,8 +445,7 @@ class CoreAIBackend(BackendDetails): ``xcrun coreai-build``; architecture selection is one / a list / all). The bundles are embedded in the ``.pte`` via NamedDataStore. - ``processed_bytes`` is always a small JSON manifest naming what/where. - Runtime execution is not wired up yet. + ``processed_bytes`` contains the manifest used by the runtime delegate. """ @staticmethod @@ -514,15 +512,15 @@ def _preprocess_portable( program, min_os, bindings: Dict[str, Any] ) -> PreprocessResult: with TemporaryDirectory() as tmp: - model_hash = _save_and_hash(program, Path(tmp) / "model.aimodel", min_os) + export_hash = _save_and_hash(program, Path(tmp) / "model.aimodel", min_os) return _deliver( Path(tmp), - model_hash, + export_hash, AssetPackaging.INLINE, { **bindings, # relative path the runtime resolves against its base - "path": f"{model_hash}/model.aimodel", + "path": f"{export_hash}/model.aimodel", "min_deployment_version": _os_version_text(min_os), }, ) @@ -533,17 +531,17 @@ def _preprocess_aot_compiled( program, opts: Dict[str, Any], bindings: Dict[str, Any] ) -> PreprocessResult: with TemporaryDirectory() as tmp: - model_hash, out = _compile_aot(program, opts, Path(tmp)) + export_hash, out = _compile_aot(program, opts, Path(tmp)) return _deliver( out, - model_hash, + export_hash, AssetPackaging.AOT_COMPILED_INLINE, { **bindings, "platform": opts["platform"], "min_deployment_version": opts["min_deployment_version"], "archs": { - arch: f"{model_hash}/{bundle.name}" + arch: f"{export_hash}/{bundle.name}" for arch, bundle in _compiled_arch_bundles(out) }, }, diff --git a/backends/apple/coreai/coreai-dependencies.cmake b/backends/apple/coreai/coreai-dependencies.cmake new file mode 100644 index 00000000000..98f1c202a79 --- /dev/null +++ b/backends/apple/coreai/coreai-dependencies.cmake @@ -0,0 +1,118 @@ +# 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. + +# Resolve dependencies on the consuming machine, including C++-only projects. +function(executorch_coreai_dependencies) + if(TARGET executorch::coreai_dependencies) + return() + endif() + if(NOT APPLE OR NOT CMAKE_SYSTEM_NAME MATCHES "^(Darwin|iOS)$") + message(FATAL_ERROR "Core AI requires macOS or iOS") + endif() + + set(sdk "${CMAKE_OSX_SYSROOT}") + if(NOT sdk) + if(CMAKE_SYSTEM_NAME STREQUAL "iOS") + set(sdk iphoneos) + else() + set(sdk macosx) + endif() + endif() + execute_process( + COMMAND xcrun --sdk "${sdk}" --show-sdk-version + OUTPUT_VARIABLE sdk_version + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE sdk_result + ERROR_VARIABLE sdk_error + ) + if(NOT sdk_result EQUAL 0 OR sdk_version VERSION_LESS 27.0) + message( + FATAL_ERROR + "Core AI requires an Apple SDK 27 or newer; selected '${sdk}' " + "reports '${sdk_version}'. Select Xcode 27 with DEVELOPER_DIR. " + "${sdk_error}" + ) + endif() + execute_process( + COMMAND xcrun --sdk "${sdk}" --show-sdk-path + OUTPUT_VARIABLE sdk_path + OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY + ) + find_library( + foundation Foundation + PATHS "${sdk_path}/System/Library/Frameworks" + NO_DEFAULT_PATH NO_CMAKE_FIND_ROOT_PATH NO_CACHE REQUIRED + ) + find_library( + coreai CoreAI + PATHS "${sdk_path}/System/Library/Frameworks" + NO_DEFAULT_PATH NO_CMAKE_FIND_ROOT_PATH NO_CACHE REQUIRED + ) + + # C++ package consumers need Swift's autolink paths, not the Swift language. + set(swift_compiler "${CMAKE_Swift_COMPILER}") + if(NOT swift_compiler) + execute_process( + COMMAND xcrun --sdk "${sdk}" --find swiftc + OUTPUT_VARIABLE swift_compiler + OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY + ) + endif() + set(arch "${CMAKE_OSX_ARCHITECTURES}") + if(NOT arch) + set(arch "${CMAKE_SYSTEM_PROCESSOR}") + endif() + if(NOT arch) + set(arch "${CMAKE_HOST_SYSTEM_PROCESSOR}") + endif() + list(GET arch 0 arch) + set(deployment 27.0) + if(CMAKE_OSX_DEPLOYMENT_TARGET VERSION_GREATER deployment) + set(deployment "${CMAKE_OSX_DEPLOYMENT_TARGET}") + endif() + set(target_environment "") + if(CMAKE_SYSTEM_NAME STREQUAL "iOS") + set(platform ios) + if(sdk_path MATCHES "[Ss]imulator") + set(target_environment -simulator) + endif() + else() + set(platform macosx) + endif() + execute_process( + COMMAND "${swift_compiler}" -print-target-info -sdk "${sdk_path}" -target + "${arch}-apple-${platform}${deployment}${target_environment}" + OUTPUT_VARIABLE target_info COMMAND_ERROR_IS_FATAL ANY + ) + string(JSON path_count LENGTH "${target_info}" paths runtimeLibraryPaths) + set(runtime_paths "") + if(path_count GREATER 0) + math(EXPR path_last "${path_count} - 1") + foreach(index RANGE ${path_last}) + string( + JSON + runtime_path + GET + "${target_info}" + paths + runtimeLibraryPaths + ${index} + ) + list(APPEND runtime_paths "${runtime_path}") + endforeach() + endif() + + # GLOBAL keeps the target visible to runners outside the backend directory. + add_library(executorch::coreai_dependencies INTERFACE IMPORTED GLOBAL) + set_target_properties( + executorch::coreai_dependencies + PROPERTIES + INTERFACE_LINK_LIBRARIES "${foundation};${coreai}" + INTERFACE_LINK_DIRECTORIES "${runtime_paths}" + INTERFACE_LINK_OPTIONS + "SHELL:-F \"${sdk_path}/System/Library/SubFrameworks\";LINKER:-rpath,/usr/lib/swift" + ) +endfunction() diff --git a/backends/apple/coreai/runtime/ETCoreAIBridge.h b/backends/apple/coreai/runtime/ETCoreAIBridge.h new file mode 100644 index 00000000000..900aa000aa2 --- /dev/null +++ b/backends/apple/coreai/runtime/ETCoreAIBridge.h @@ -0,0 +1,142 @@ +/* + * 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 + +NS_ASSUME_NONNULL_BEGIN + +// Bridge completions can require Swift's cooperative thread pool. Callers that +// synchronously wait for them must use ordinary worker threads, not Swift Task +// or Task.detached workers, which can starve the work needed for completion. + +typedef NS_ENUM(NSInteger, ETCoreAIScalarType) { + ETCoreAIScalarTypeFloat16 = 0, + ETCoreAIScalarTypeFloat32 = 1, +}; + +typedef NS_ENUM(NSInteger, ETCoreAIErrorCode) { + // Invalid model descriptors, functions or bindings. + ETCoreAIErrorInvalidModel = 1, + // Unsupported tensor types, layouts or model features. + ETCoreAIErrorUnsupported = 2, + // Invalid caller arguments or input tensor payloads. + ETCoreAIErrorInvalidArgument = 3, + // SDK or system failures, or inconsistent SDK outputs. Does not mean cached + // files are corrupt. + ETCoreAIErrorRuntime = 4, +}; + +FOUNDATION_EXPORT NSErrorDomain const ETCoreAIErrorDomain; + +// Input storage is borrowed until output copying finishes. The bridge never +// retains or frees it. +@interface ETCoreAIInputTensor : NSObject +@property(nonatomic, readonly) const void* _Nullable bytes; +@property(nonatomic, readonly) NSUInteger byteCount; +@property(nonatomic, readonly, copy) NSArray* shape; +@property(nonatomic, readonly) ETCoreAIScalarType scalarType; +- (instancetype)init NS_UNAVAILABLE; ++ (instancetype)new NS_UNAVAILABLE; +- (instancetype)initWithBytes:(const void* _Nullable)bytes + byteCount:(NSUInteger)byteCount + shape:(NSArray*)shape + scalarType:(ETCoreAIScalarType)scalarType + NS_DESIGNATED_INITIALIZER; +@end + +// Holds the SDK output until copyBytesTo: runs; raw SDK pointers are used only +// inside that call. Construct only after validating the source layout, because +// copying cannot report errors. +@interface ETCoreAITensor : NSObject +@property(nonatomic, readonly) NSUInteger byteCount; +@property(nonatomic, readonly, copy) NSArray* shape; +@property(nonatomic, readonly) ETCoreAIScalarType scalarType; +- (instancetype)init NS_UNAVAILABLE; ++ (instancetype)new NS_UNAVAILABLE; +- (instancetype)initWithByteCount:(NSUInteger)byteCount + shape:(NSArray*)shape + scalarType:(ETCoreAIScalarType)scalarType + copyBytes:(void (^)(void* bytes))copyBytes + NS_DESIGNATED_INITIALIZER; +- (instancetype)initWithData:(NSData*)data + shape:(NSArray*)shape + scalarType:(ETCoreAIScalarType)scalarType; +// Caller supplies byteCount writable bytes, disjoint from live source storage. +// A null destination is allowed only for an empty output. +- (void)copyBytesTo:(void* _Nullable)bytes NS_SWIFT_NAME(copyBytes(to:)); +@end + +// Keep inputs valid until output copying finishes or completion reports +// failure. Outputs may share input storage. Stage copies before writing any +// destination that aliases an input or another output destination. +// +// Completion runs exactly once, possibly before executeInputs returns, on an +// unspecified thread or queue. +@protocol ETCoreAISession +- (void)executeInputs:(NSArray*)inputs + completion:(void (^)( + NSArray* _Nullable outputs, + NSError* _Nullable error))completion; +@end + +// Holds the SDK model, keeping its cache entry in use. Bound sessions hold the +// model too, so they stay valid after the prepared model is released. +@protocol ETCoreAIPreparedModel +- (NSData*)copyBookmarkData; +// Completes once with either a session or an error, never both. Completion may +// run synchronously on the calling thread. +- (void)loadFunctionNamed:(NSString*)functionName + inputNames:(NSArray*)inputNames + outputNames:(NSArray*)outputNames + completion:(void (^)( + id _Nullable session, + NSError* _Nullable error))completion; +@end + +typedef NS_ENUM(NSInteger, ETCoreAIAcquisitionStatus) { + ETCoreAIAcquisitionStatusHit = 0, + ETCoreAIAcquisitionStatusMiss = 1, + ETCoreAIAcquisitionStatusError = 2, +}; + +// Hit: nonnull model, nil error. +// Miss: nil model and error, only for an SDK nil result. +// Error: nil model, nonnull error. Function errors never report Miss. +// +// Acquisition completions run exactly once off the main thread. +typedef void (^ETCoreAIAcquisitionCompletion)( + ETCoreAIAcquisitionStatus status, + id _Nullable preparedModel, + NSError* _Nullable error); + +@protocol ETCoreAIModelLoading +// Stable for this loader's lifetime; derived from its specialization options. +@property(nonatomic, readonly, copy) NSString* specializationIdentifier; +- (void)restoreModelFromBookmark:(NSData*)bookmark + completion:(ETCoreAIAcquisitionCompletion)completion; +// Uses the SDK default cache with storage-pressure-only purging. +// Specialization never reports Miss. +- (void)specializeModelAtURL:(NSURL*)url + completion:(ETCoreAIAcquisitionCompletion)completion; +// Nil error means SDK-confirmed deletion; the bridge adds no +// busy/absent classification. +- (void)evictModelWithBookmark:(NSData*)bookmark + completion:(void (^)(NSError* _Nullable error))completion; +@end + +#ifdef __cplusplus +extern "C" { +#endif +BOOL ETCoreAIIsAvailable(void); +NSString* _Nullable ETCoreAIDeviceArchitectureName(void); +id ETCoreAICreateModelLoader(void); +#ifdef __cplusplus +} +#endif + +NS_ASSUME_NONNULL_END diff --git a/backends/apple/coreai/runtime/ETCoreAIError.swift b/backends/apple/coreai/runtime/ETCoreAIError.swift new file mode 100644 index 00000000000..65c20c5f472 --- /dev/null +++ b/backends/apple/coreai/runtime/ETCoreAIError.swift @@ -0,0 +1,30 @@ +// 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 Foundation + +func bridgeError( + _ code: ETCoreAIErrorCode, _ message: String, underlying: Error? = nil +) -> NSError { + var info: [String: Any] = [NSLocalizedDescriptionKey: message] + if let underlying { info[NSUnderlyingErrorKey] = underlying } + return NSError(domain: ETCoreAIErrorDomain, code: code.rawValue, userInfo: info) +} + +func mappedError(_ error: Error, code: ETCoreAIErrorCode, context: String) -> NSError { + let error = error as NSError + return error.domain == ETCoreAIErrorDomain + ? error : bridgeError(code, "\(context): \(error.localizedDescription)", underlying: error) +} + +// Objective-C blocks have no Sendable annotation. Each box is invoked by one detached task only. +final class Completion: @unchecked Sendable { + let call: Callback + + init(_ call: Callback) { + self.call = call + } +} diff --git a/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift b/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift new file mode 100644 index 00000000000..adf80fbc91f --- /dev/null +++ b/backends/apple/coreai/runtime/ETCoreAIModelLoader.swift @@ -0,0 +1,279 @@ +// 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 +import Foundation + +// Only a successful SDK nil is a miss; thrown errors retain their underlying cause. +func acquireCoreAIModel( + context: String, completion: @escaping ETCoreAIAcquisitionCompletion, + operation: @escaping @Sendable () async throws -> (any ETCoreAIPreparedModel)? +) { + let callback = Completion(completion) + Task.detached { + do { + if let model = try await operation() { + callback.call(.hit, model, nil) + } else { + callback.call(.miss, nil, nil) + } + } catch { + callback.call(.error, nil, mappedError(error, code: .runtime, context: context)) + } + } +} + +@available(macOS 27.0, iOS 27.0, *) +private final class CoreAISwiftSession: NSObject, ETCoreAISession, Sendable { + let model: AIModel + let function: InferenceFunction + let inputNames: [String] + let outputNames: [String] + let inputDescriptors: [NDArrayDescriptor] + let outputDescriptors: [NDArrayDescriptor] + + init(model: AIModel, function: InferenceFunction, inputNames: [String], outputNames: [String]) + throws + { + let descriptor = function.descriptor + guard descriptor.stateNames.isEmpty else { + throw bridgeError(.unsupported, "Stateful functions are not supported") + } + guard inputNames.count == descriptor.inputCount, + outputNames.count == descriptor.outputCount, + Set(inputNames).count == inputNames.count, Set(outputNames).count == outputNames.count, + descriptor.inputNames.count == inputNames.count, + descriptor.outputNames.count == outputNames.count, + Set(inputNames) == Set(descriptor.inputNames), Set(outputNames) == Set(descriptor.outputNames) + else { + throw bridgeError(.invalidModel, "Ordered bindings do not match the function descriptor") + } + self.model = model + self.function = function + self.inputNames = inputNames + self.outputNames = outputNames + inputDescriptors = try inputNames.map { + try tensorDescriptor(descriptor.inputDescriptor(of: $0)) + } + outputDescriptors = try outputNames.map { + try tensorDescriptor(descriptor.outputDescriptor(of: $0)) + } + super.init() + } + + // All nonescapable input views end here on both return and throw, before completion. + private func runBorrowing(_ payloads: [CoreAITensorPayload]) async throws + -> InferenceFunction.Outputs + { + var arguments = InferenceFunction.Inputs() + for index in payloads.indices { + let payload = payloads[index] + let descriptor = inputDescriptors[index] + try validateShape(payload.shape, descriptor: descriptor) + let expectedType = try CoreAIScalarInfo.from(sdkType: descriptor.scalarType) + guard payload.scalarType == expectedType.bridgeType else { + throw bridgeError(.invalidArgument, "Input dtype does not match the function descriptor") + } + // Unsafe boundary: the ET caller, not this metadata snapshot, owns the storage. + let bytes = RawSpan( + _unsafeBytes: UnsafeRawBufferPointer( + start: payload.bytes, count: payload.byteCount)) + let view = NDArray.RawView( + bytes: bytes, scalarType: descriptor.scalarType, shape: payload.shape) + arguments.insert(view, for: inputNames[index]) + } + return try await function.run(inputs: arguments) + } + + func executeInputs( + _ inputs: [ETCoreAIInputTensor], completion: @escaping ([ETCoreAITensor]?, Error?) -> Void + ) { + let payloads: [CoreAITensorPayload] + do { + guard inputs.count == inputNames.count else { + throw bridgeError(.invalidArgument, "Input count does not match the function") + } + payloads = try inputs.map(CoreAITensorPayload.init) + } catch { + completion(nil, mappedError(error, code: .invalidArgument, context: "Invalid inputs")) + return + } + let callback = Completion(completion) + Task.detached { [self, payloads, callback] in + do { + var results = try await runBorrowing(payloads) + guard results.count == outputNames.count, Set(results.names) == Set(outputNames) else { + throw bridgeError(.runtime, "Inference returned unexpected output names") + } + var tensors: [ETCoreAITensor] = [] + for index in outputNames.indices { + guard let value = results.remove(outputNames[index]) else { + throw bridgeError(.runtime, "Inference output is missing") + } + guard let array = value.ndArray else { + throw bridgeError(.unsupported, "Inference returned an image output") + } + let descriptor = outputDescriptors[index] + try validateShape(array.shape, descriptor: descriptor, errorCode: .runtime) + guard array.scalarType == descriptor.scalarType else { + throw bridgeError(.runtime, "Output dtype does not match the function descriptor") + } + tensors.append(try prepareOutput(array)) + } + callback.call(tensors, nil) + } catch { + callback.call(nil, mappedError(error, code: .runtime, context: "Core AI inference failed")) + } + } + } +} + +@available(macOS 27.0, iOS 27.0, *) +private func bindFunction( + model: AIModel, functionName: String, inputNames: [String], outputNames: [String] +) throws -> CoreAISwiftSession { + guard !functionName.isEmpty else { + throw bridgeError(.invalidArgument, "A nonempty function name is required") + } + guard let function = try model.loadFunction(named: functionName) else { + throw bridgeError(.invalidModel, "Model has no function named '\(functionName)'") + } + return try CoreAISwiftSession( + model: model, function: function, inputNames: inputNames, outputNames: outputNames) +} + +@available(macOS 27.0, iOS 27.0, *) +private final class CoreAISwiftPreparedModel: NSObject, ETCoreAIPreparedModel, Sendable { + let model: AIModel + + init(model: AIModel) { + self.model = model + super.init() + } + + func copyBookmarkData() -> Data { + model.bookmarkData + } + + // Binding never awaits, so it completes on the calling thread. + func loadFunctionNamed( + _ functionName: String, inputNames: [String], outputNames: [String], + completion: @escaping ((any ETCoreAISession)?, Error?) -> Void + ) { + do { + completion( + try bindFunction( + model: model, functionName: functionName, inputNames: inputNames, + outputNames: outputNames), nil) + } catch { + completion( + nil, mappedError(error, code: .runtime, context: "Core AI function binding failed")) + } + } +} + +@available(macOS 27.0, iOS 27.0, *) +struct CoreAISpecialization: Sendable { + let options: SpecializationOptions + let identifier: String + + static let `default` = Self() + + private init() { + options = .default + identifier = "default" + } + + // New SDK option fields require a schema review and versioned identity changes. + private struct Descriptor: Encodable { + let version: Int + let allowedComputeUnits: [String] + let preferredComputeUnit: String? + let expectFrequentReshapes: Bool + } + + init(options: SpecializationOptions) throws { + if options == .default { + self = Self.default + return + } + let descriptor = Descriptor( + version: 1, + allowedComputeUnits: try options.allowedComputeUnitKinds.map(Self.name).sorted(), + preferredComputeUnit: try options.preferredComputeUnitKind.map(Self.name), + expectFrequentReshapes: options.expectFrequentReshapes) + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + identifier = String(decoding: try encoder.encode(descriptor), as: UTF8.self) + self.options = options + } + + private static func name(_ kind: ComputeUnitKind) throws -> String { + switch kind { + case .cpu: return "cpu" + case .gpu: return "gpu" + case .neuralEngine: return "neuralEngine" + @unknown default: + throw bridgeError(.unsupported, "Unknown specialization compute unit") + } + } +} + +@available(macOS 27.0, iOS 27.0, *) +@objc(ETCoreAISwiftModelLoader) +public final class ETCoreAISwiftModelLoader: NSObject, ETCoreAIModelLoading { + static let cachePolicy = AIModelCache.Policy(purgeConditions: [.storagePressure]) + private let specialization: CoreAISpecialization + + public override init() { + specialization = .default + super.init() + } + + init(specialization: CoreAISpecialization) { + self.specialization = specialization + super.init() + } + + public var specializationIdentifier: String { specialization.identifier } + + @objc public static var deviceArchitectureName: String { AIModel.deviceArchitectureName } + + public func restoreModel( + fromBookmark bookmark: Data, completion: @escaping ETCoreAIAcquisitionCompletion + ) { + acquireCoreAIModel(context: "Core AI bookmark restoration failed", completion: completion) { + try AIModel(resolvingBookmark: bookmark).map { CoreAISwiftPreparedModel(model: $0) } + } + } + + public func specializeModel(at url: URL, completion: @escaping ETCoreAIAcquisitionCompletion) { + let options = specialization.options + acquireCoreAIModel(context: "Core AI specialization failed", completion: completion) { + guard url.isFileURL else { + throw bridgeError(.invalidArgument, "A file URL is required") + } + let model = try await AIModel.specialize( + contentsOf: url, options: options, cache: .default, cachePolicy: Self.cachePolicy) + return CoreAISwiftPreparedModel(model: model) + } + } + + public func evictModel( + withBookmark bookmark: Data, completion: @escaping (Error?) -> Void + ) { + let callback = Completion(completion) + Task.detached { + do { + try AIModelCache.deleteEntry(referencedBy: bookmark) + callback.call(nil) + } catch { + callback.call( + mappedError(error, code: .runtime, context: "Core AI cache eviction failed")) + } + } + } +} diff --git a/backends/apple/coreai/runtime/ETCoreAITensor.mm b/backends/apple/coreai/runtime/ETCoreAITensor.mm new file mode 100644 index 00000000000..611812e3d39 --- /dev/null +++ b/backends/apple/coreai/runtime/ETCoreAITensor.mm @@ -0,0 +1,63 @@ +/* + * 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. + */ + +#include +#import "ETCoreAIBridge.h" + +NSErrorDomain const ETCoreAIErrorDomain = @"org.pytorch.executorch.coreai"; + +@implementation ETCoreAIInputTensor +- (instancetype)initWithBytes:(const void*)bytes + byteCount:(NSUInteger)byteCount + shape:(NSArray*)shape + scalarType:(ETCoreAIScalarType)scalarType { + self = [super init]; + if (self) { + _bytes = bytes; + _byteCount = byteCount; + _shape = [shape copy]; + _scalarType = scalarType; + } + return self; +} +@end + +@implementation ETCoreAITensor { + void (^_copyBytes)(void*); +} +- (instancetype)initWithByteCount:(NSUInteger)byteCount + shape:(NSArray*)shape + scalarType:(ETCoreAIScalarType)scalarType + copyBytes:(void (^)(void*))copyBytes { + NSParameterAssert(copyBytes != nil); + self = [super init]; + if (self) { + _byteCount = byteCount; + _shape = [shape copy]; + _scalarType = scalarType; + _copyBytes = [copyBytes copy]; + } + return self; +} +- (instancetype)initWithData:(NSData*)data + shape:(NSArray*)shape + scalarType:(ETCoreAIScalarType)scalarType { + NSData* storage = [data copy]; + return [self initWithByteCount:storage.length + shape:shape + scalarType:scalarType + copyBytes:^(void* bytes) { + std::memcpy(bytes, storage.bytes, storage.length); + }]; +} +- (void)copyBytesTo:(void*)bytes { + if (_byteCount > 0) { + _copyBytes(bytes); + } +} +@end diff --git a/backends/apple/coreai/runtime/ETCoreAITensors.swift b/backends/apple/coreai/runtime/ETCoreAITensors.swift new file mode 100644 index 00000000000..950cc705ffc --- /dev/null +++ b/backends/apple/coreai/runtime/ETCoreAITensors.swift @@ -0,0 +1,214 @@ +// 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 +import Foundation + +@available(macOS 27.0, iOS 27.0, *) +struct CoreAIScalarInfo: Sendable { + let sdkType: NDArray.ScalarType + let bridgeType: ETCoreAIScalarType + let byteWidth: Int + + // Only byte-addressable scalar types belong here; packed types need a separate layout policy. + static let supported: [Self] = [ + Self(sdkType: .float16, bridgeType: .float16, byteWidth: 2), + Self(sdkType: .float32, bridgeType: .float32, byteWidth: 4), + ] + + static func from(sdkType: NDArray.ScalarType) throws -> Self { + guard let info = supported.first(where: { $0.sdkType == sdkType }) else { + throw bridgeError(.unsupported, "Unsupported Core AI scalar type: \(sdkType)") + } + return info + } + + static func from(bridgeType: ETCoreAIScalarType) throws -> Self { + guard let info = supported.first(where: { $0.bridgeType == bridgeType }) else { + throw bridgeError(.unsupported, "Unsupported bridge scalar type: \(bridgeType.rawValue)") + } + return info + } + + func byteCount(shape: [Int], errorCode: ETCoreAIErrorCode = .invalidArgument) throws -> Int { + guard shape.allSatisfy({ $0 >= 0 }) else { + throw bridgeError(errorCode, "Negative tensor dimension") + } + // Check suffix products too: SDK dense-stride construction must not overflow on empty tensors. + var count = byteWidth + for dimension in shape.reversed() { + let (next, overflow) = count.multipliedReportingOverflow(by: max(dimension, 1)) + guard !overflow else { + throw bridgeError(errorCode, "Tensor byte count overflows Int") + } + count = next + } + return shape.contains(0) ? 0 : count + } +} + +// ExecuTorch waits synchronously and keeps input storage live through output copying. +// This snapshot transfers only metadata and a pointer to the task; it does not own storage. +@available(macOS 27.0, iOS 27.0, *) +struct CoreAITensorPayload: @unchecked Sendable { + let bytes: UnsafeRawPointer? + let byteCount: Int + let shape: [Int] + let scalarType: ETCoreAIScalarType + + init(_ tensor: ETCoreAIInputTensor) throws { + shape = try tensor.shape.map { number in + guard let dimension = number as? Int, dimension >= 0 else { + throw bridgeError(.invalidArgument, "Tensor dimensions must be nonnegative integers") + } + return dimension + } + scalarType = tensor.scalarType + byteCount = try Self.byteCount(shape: shape, scalarType: scalarType) + guard tensor.byteCount == byteCount else { + throw bridgeError( + .invalidArgument, "Tensor payload byte count does not match shape and dtype") + } + bytes = tensor.bytes + guard byteCount == 0 || bytes != nil else { + throw bridgeError(.invalidArgument, "Nonempty tensor has null data") + } + } + + static func byteCount( + shape: [Int], scalarType: ETCoreAIScalarType, + errorCode: ETCoreAIErrorCode = .invalidArgument + ) throws -> Int { + try CoreAIScalarInfo.from(bridgeType: scalarType).byteCount(shape: shape, errorCode: errorCode) + } +} + +@available(macOS 27.0, iOS 27.0, *) +func tensorDescriptor(_ value: InferenceValue.Descriptor?) throws -> NDArrayDescriptor { + guard let value else { throw bridgeError(.invalidModel, "Missing tensor descriptor") } + guard case .ndArray(let descriptor) = value else { + throw bridgeError(.unsupported, "Image values are not supported") + } + _ = try CoreAIScalarInfo.from(sdkType: descriptor.scalarType) + guard descriptor.interleaveLayout == nil else { + throw bridgeError(.unsupported, "Interleaved tensor descriptors are not supported") + } + guard descriptor.shape.allSatisfy({ $0 >= -1 }) else { + throw bridgeError(.invalidModel, "Invalid descriptor dimension") + } + return descriptor +} + +@available(macOS 27.0, iOS 27.0, *) +func validateShape( + _ shape: [Int], descriptor: NDArrayDescriptor, + errorCode: ETCoreAIErrorCode = .invalidArgument +) throws { + guard shape.count == descriptor.shape.count, + zip(shape, descriptor.shape).allSatisfy({ actual, expected in + actual >= 0 && (expected == -1 || actual == expected) + }) + else { + throw bridgeError(errorCode, "Tensor shape does not match the function descriptor") + } +} + +// Offsets are in elements, including padded or transposed SDK layouts. +struct CoreAITensorLayout { + let shape: [Int] + let strides: [Int] + let count: Int + // Row-major contiguous, ignoring strides of size-1 dimensions. + let isDense: Bool + + init(shape: [Int], strides: [Int], width: Int, storageByteCount: Int) throws { + guard shape.count == strides.count, shape.allSatisfy({ $0 >= 0 }), + strides.allSatisfy({ $0 >= 0 }), width > 0, storageByteCount >= 0 + else { throw bridgeError(.unsupported, "Invalid or negative NDArray strides") } + self.shape = shape + self.strides = strides + var count = 1 + var maximumOffset = 0 + for (dimension, stride) in zip(shape, strides) { + let (next, countOverflow) = count.multipliedReportingOverflow(by: dimension) + let (extent, extentOverflow) = max(dimension - 1, 0).multipliedReportingOverflow(by: stride) + let (offset, offsetOverflow) = maximumOffset.addingReportingOverflow(extent) + guard !countOverflow, !extentOverflow, !offsetOverflow else { + throw bridgeError(.runtime, "NDArray layout overflows Int") + } + count = next + maximumOffset = offset + } + self.count = count + var expectedStride = 1 + var isDense = true + for axis in shape.indices.reversed() where shape[axis] != 1 { + isDense = isDense && strides[axis] == expectedStride + expectedStride = expectedStride &* shape[axis] + } + self.isDense = isDense + if count > 0 { + let (span, spanOverflow) = maximumOffset.addingReportingOverflow(1) + let (bytes, byteOverflow) = span.multipliedReportingOverflow(by: width) + guard !spanOverflow, !byteOverflow, bytes <= storageByteCount else { + throw bridgeError(.runtime, "NDArray layout exceeds its storage") + } + } + } + + func offset(for linearIndex: Int) -> Int { + var remainder = linearIndex + var offset = 0 + for axis in shape.indices.reversed() { + offset += (remainder % shape[axis]) * strides[axis] + remainder /= shape[axis] + } + return offset + } +} + +@available(macOS 27.0, iOS 27.0, *) +func prepareOutput(_ array: NDArray) throws -> ETCoreAITensor { + let type = try CoreAIScalarInfo.from(sdkType: array.scalarType) + let byteCount = try type.byteCount(shape: array.shape, errorCode: .runtime) + guard array.interleaveLayout == nil else { + throw bridgeError(.unsupported, "Interleaved output storage is not supported") + } + let view = array.rawView() + let layout = try CoreAITensorLayout( + shape: array.shape, strides: array.strides, width: type.byteWidth, + storageByteCount: view.bytes.byteCount) + let shape = view.shape + let strides = view.strides + guard view.scalarType == type.sdkType, view.interleaveLayout == nil, + shape.count == layout.shape.count, strides.count == layout.strides.count + else { + throw bridgeError(.runtime, "Output view layout changed") + } + for axis in layout.shape.indices { + guard shape[axis] == layout.shape[axis], strides[axis] == layout.strides[axis] else { + throw bridgeError(.runtime, "Output view layout changed") + } + } + // Retain the immutable array, not its scoped pointer. All fallible work is complete. + return ETCoreAITensor( + byteCount: UInt(byteCount), shape: layout.shape.map(NSNumber.init(value:)), + scalarType: type.bridgeType + ) { destination in + let view = array.rawView() + view.bytes.withUnsafeBytes { source in + if layout.isDense { + destination.copyMemory(from: source.baseAddress!, byteCount: byteCount) + return + } + for index in 0.. +#include +#include + +namespace executorch::backends::coreai { +namespace { +using namespace executorch::runtime; + +struct Handle final { + size_t input_count = 0; + size_t output_count = 0; + id session = nil; +}; + +Result cache_directory_option(BackendInitContext& context) { + auto option = context.get_runtime_spec("coreai_cache_dir"); + if (!option.ok()) { + if (option.error() == Error::NotFound) { + return cache_directory_path(nullptr); + } + return option.error(); + } + ET_CHECK_OR_RETURN_ERROR( + option.get() != nullptr, + InvalidArgument, + "Core AI cache directory option must not be null"); + return cache_directory_path(option.get()); +} + +Result> bind_session( + id model, + const SelectedBundle& bundle) { + dispatch_semaphore_t ready = dispatch_semaphore_create(0); + __block id session = nil; + __block NSError* error = nil; + [model loadFunctionNamed:bundle.function + inputNames:bundle.inputs + outputNames:bundle.outputs + completion:^(id result, NSError* failure) { + session = result; + error = failure; + dispatch_semaphore_signal(ready); + }]; + dispatch_semaphore_wait(ready, DISPATCH_TIME_FOREVER); + id bound = session; + NSError* failure = error; + // The SDK may retain the completion block after invoking it. + session = nil; + error = nil; + if (failure != nil) { + return execution_error(failure); + } + ET_CHECK_OR_RETURN_ERROR( + bound != nil, Internal, "Core AI returned no session"); + return bound; +} + +class CoreAIBackend final : public BackendInterface { + public: + bool is_available() const override { + return ETCoreAIIsAvailable(); + } + + Result init( + BackendInitContext& context, + FreeableBuffer* processed, + ArrayRef) const override { + @autoreleasepool { + ET_CHECK_OR_RETURN_ERROR( + processed && processed->data() && processed->size() > 0, + InvalidProgram, + "Missing Core AI manifest"); + auto manifest = parse_manifest(processed->data(), processed->size()); + if (!manifest.ok()) { + return manifest.error(); + } + auto host_bundle = select_host_bundle(manifest.get()); + if (!host_bundle.ok()) { + return host_bundle.error(); + } + auto configured = cache_directory_option(context); + if (!configured.ok()) { + return configured.error(); + } + id loader = ETCoreAICreateModelLoader(); + ET_CHECK_OR_RETURN_ERROR( + loader != nil, Internal, "Cannot create Core AI loader"); + auto cache_directory = create_cache_directory(configured.get()); + if (!cache_directory.ok()) { + return cache_directory.error(); + } + auto prepared = acquire_model( + host_bundle.get(), + context.get_named_data_map(), + cache_directory.get(), + loader); + if (!prepared.ok()) { + return prepared.error(); + } + auto session = bind_session(prepared.get(), 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"); + handle->input_count = host_bundle->bundle.inputs.count; + handle->output_count = host_bundle->bundle.outputs.count; + handle->session = session.get(); + processed->Free(); + return static_cast(handle.release()); + } + } + + Error execute( + BackendExecutionContext&, + DelegateHandle* opaque, + Span args) const override { + @autoreleasepool { + ET_CHECK_OR_RETURN_ERROR( + opaque != nullptr, DelegateInvalidHandle, "Null Core AI handle"); + auto& handle = *static_cast(opaque); + return execute_session( + handle.session, args, handle.input_count, handle.output_count); + } + } + + void destroy(DelegateHandle* handle) const override { + @autoreleasepool { + delete static_cast(handle); + } + } +}; + +CoreAIBackend backend; +[[maybe_unused]] const auto registration = + register_backend({kCoreAIBackendId, &backend}); +} // namespace +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_backend_id.h b/backends/apple/coreai/runtime/coreai_backend_id.h new file mode 100644 index 00000000000..8fbd72951e1 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_backend_id.h @@ -0,0 +1,13 @@ +/* + * 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 + +namespace executorch::backends::coreai { +inline constexpr char kCoreAIBackendId[] = "CoreAIBackend"; +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_bridge.mm b/backends/apple/coreai/runtime/coreai_bridge.mm new file mode 100644 index 00000000000..539bdb408cb --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_bridge.mm @@ -0,0 +1,29 @@ +/* + * 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 "ETCoreAIBridge.h" + +#import "CoreAIBridge-Swift.h" + +BOOL ETCoreAIIsAvailable(void) { + if (@available(macOS 27.0, iOS 27.0, *)) { + return YES; + } + return NO; +} + +NSString* ETCoreAIDeviceArchitectureName(void) { + if (@available(macOS 27.0, iOS 27.0, *)) { + return ETCoreAISwiftModelLoader.deviceArchitectureName; + } + return nil; +} + +id ETCoreAICreateModelLoader(void) { + return [[ETCoreAISwiftModelLoader alloc] init]; +} diff --git a/backends/apple/coreai/runtime/coreai_cache.h b/backends/apple/coreai/runtime/coreai_cache.h new file mode 100644 index 00000000000..0d2ea862772 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache.h @@ -0,0 +1,58 @@ +/* + * 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 + +namespace executorch::runtime { +class DataLoader; +} + +namespace executorch::backends::coreai { + +/** + * Clears tracked SDK entries, per-key source bundles, and bookmarks in a + * dedicated Core AI cache directory. Null selects the same default as loading; + * supplied paths must be absolute. Missing directories are successful no-ops. + * PTEs, stable lock files, and unkeyed temporaries are not removed. + * SDK entries without saved bookmarks cannot be targeted. + * + * Callers must unload affected models and prevent concurrent loading, + * inference, and maintenance across processes until this call returns. + * Independent keys are attempted in sorted order, returning the first error; + * clearing is not a transaction. SDK failures preserve an entry's files unless + * the SDK confirms that entry is absent. + */ +[[nodiscard]] runtime::Error clear_cache( + const char* coreai_cache_dir = nullptr); + +/** + * Clears all recorded cache variants referenced by every CoreAI delegate in + * the PTE after validating its manifests. Each exported bundle's identity + * selects every device and specialization variant. Discovery does not require + * host compatibility or SDK availability; evicting an entry with a saved + * bookmark does. Does not initialize delegates or materialize Core AI assets. + * + * A null directory selects the same Caches default as backend initialization. + * Supplied paths must resolve to valid explicit cache directories; they never + * fall back to the default. The directory is never derived from the PTE's + * location. + * Copied PTEs can share cache keys. The quiescence and failure contract of + * clear_cache also applies. Keep the loader alive and its data unchanged until + * the call returns. + */ +[[nodiscard]] runtime::Error clear_cache_for_pte( + runtime::DataLoader& loader, + const char* coreai_cache_dir = nullptr); + +/** FileDataLoader convenience overload with the same contract. */ +[[nodiscard]] runtime::Error clear_cache_for_pte( + const char* pte_path, + const char* coreai_cache_dir = nullptr); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_cache.mm b/backends/apple/coreai/runtime/coreai_cache.mm new file mode 100644 index 00000000000..2f6c092eb81 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache.mm @@ -0,0 +1,141 @@ +/* + * 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. + */ +#include "coreai_cache.h" + +#import "coreai_cache_entry.h" +#import "coreai_cache_key.h" +#import "coreai_cache_root.h" +#import "coreai_load_coordinator.h" +#import "coreai_manifest.h" +#include "coreai_pte.h" + +#include +#include + +namespace executorch::backends::coreai { +namespace { +using runtime::Error; +using runtime::Result; + +Result> cache_directory(const char* path) { + auto configured = cache_directory_path(path); + if (!configured.ok()) { + return configured.error(); + } + return open_cache_directory(configured.get()); +} + +Error clear_keys(const CacheDirectory& directory, NSArray* keys) { + NSArray* ordered = [[[NSSet setWithArray:keys] allObjects] + sortedArrayUsingSelector:@selector(compare:)]; + Error first_error = Error::Ok; + for (NSString* key in ordered) { + const Error error = clear_cache_entry(directory, key, nil); + if (error != Error::Ok) { + ET_LOG( + Error, + "Core AI cache clear failed for key %s: 0x%x", + key.UTF8String, + static_cast(error)); + if (first_error == Error::Ok) { + first_error = error; + } + } + } + return first_error; +} + +Error clear_pte( + runtime::DataLoader& loader, + const std::optional& directory) { + // Keep the verified image and selected segment owners alive through eviction. + auto pte = inspect_coreai_pte(loader); + if (!pte.ok()) { + return pte.error(); + } + NSMutableSet* bundle_identities = [NSMutableSet set]; + for (const auto& buffer : pte->processed) { + auto data = buffer.data_safe(); + if (!data.ok()) { + return data.error(); + } + ET_CHECK_OR_RETURN_ERROR( + data.get() != nullptr && buffer.size() > 0, + InvalidProgram, + "Missing Core AI manifest"); + auto manifest = parse_manifest(data.get(), buffer.size()); + if (!manifest.ok()) { + return manifest.error(); + } + for (NSString* bundle_name in manifest->bundle_digests) { + auto identity = bundle_identity(manifest.get(), bundle_name); + if (!identity.ok()) { + return identity.error(); + } + [bundle_identities addObject:identity.get()]; + } + } + if (bundle_identities.count == 0 || !directory.has_value()) { + return Error::Ok; + } + auto inventory = inventory_cache_keys(*directory); + if (!inventory.ok()) { + return inventory.error(); + } + NSMutableArray* keys = [NSMutableArray array]; + for (NSString* key in inventory.get()) { + if ([bundle_identities containsObject:cache_key_bundle_identity(key)]) { + [keys addObject:key]; + } + } + return clear_keys(*directory, keys); +} +} // namespace + +Error clear_cache(const char* coreai_cache_dir) { + @autoreleasepool { + auto directory = cache_directory(coreai_cache_dir); + if (!directory.ok()) { + return directory.error(); + } + if (!directory->has_value()) { + return Error::Ok; + } + auto keys = inventory_cache_keys(directory->value()); + if (!keys.ok()) { + return keys.error(); + } + return clear_keys(directory->value(), keys.get()); + } +} + +Error clear_cache_for_pte( + runtime::DataLoader& loader, + const char* coreai_cache_dir) { + @autoreleasepool { + auto directory = cache_directory(coreai_cache_dir); + if (!directory.ok()) { + return directory.error(); + } + return clear_pte(loader, directory.get()); + } +} + +Error clear_cache_for_pte(const char* pte_path, const char* coreai_cache_dir) { + ET_CHECK_OR_RETURN_ERROR( + pte_path != nullptr && pte_path[0] != '\0', + InvalidArgument, + "Missing PTE path"); + auto loader = extension::FileDataLoader::from(pte_path); + if (!loader.ok()) { + return loader.error(); + } + return clear_cache_for_pte(loader.get(), coreai_cache_dir); +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_cache_entry.h b/backends/apple/coreai/runtime/coreai_cache_entry.h new file mode 100644 index 00000000000..f3b8230f5b1 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache_entry.h @@ -0,0 +1,123 @@ +/* + * 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 +#include +#include +#include +#include "coreai_cache_root.h" + +namespace executorch::backends::coreai { + +// Invalid and unowned names are omitted from the sorted inventory. +runtime::Result*> inventory_cache_keys( + const CacheDirectory& directory); + +enum class CacheBookmarkState { + Missing, + Corrupt, + Present, +}; + +struct CacheBookmarkInspection final { + CacheBookmarkState state; + NSData* data; +}; + +enum class CacheEntryPersistence { + BookmarkCapable, + // Only from acquire_existing when bookmarks/ is missing; never creates it. + SourceOnly, +}; + +struct CacheSourceDirectory final { + FileDescriptor descriptor; + NSURL* url; +}; + +class CacheEntryMaintenance; + +class LockedCacheEntry final { + public: + LockedCacheEntry(const LockedCacheEntry&) = delete; + LockedCacheEntry& operator=(const LockedCacheEntry&) = delete; + + // Each acquisition owns an independent descriptor on one stable stripe. + static runtime::Result> acquire( + const CacheDirectory& directory, + NSString* key); + // Missing entries return a null handle without changing entry storage. + static runtime::Result> acquire_existing( + const CacheDirectory& directory, + NSString* key); + + CacheEntryPersistence persistence() const { + return persistence_; + } + NSString* key() const { + return key_; + } + // Missing is nil. Empty, oversized, inaccessible, or nonregular files fail + // closed. + runtime::Result read_bookmark() const; + runtime::Error write_bookmark(NSData* data) const; + runtime::Result prepare_source_directory() const; + // A missing source bundle is a no-op; unexpected entry types fail before + // deletion. + runtime::Error remove_source() const; + + // The view must not outlive this locked entry. + CacheEntryMaintenance maintenance() const; + + private: + friend class CacheEntryMaintenance; + + LockedCacheEntry( + FileDescriptor root, + FileDescriptor bookmarks, + FileDescriptor lock_file, + NSString* root_path, + NSString* key, + CacheEntryPersistence persistence); + static runtime::Result> acquire_impl( + const CacheDirectory& directory, + NSString* key, + StorageCreationMode mode); + runtime::Result> open_source_directory( + StorageCreationMode mode) const; + // Applies a storage operation to this key's source bundle. A missing + // sources/ directory is a no-op. + runtime::Error apply_to_source( + runtime::Error (*operation)(int, NSString*, StorageRemovalPolicy), + StorageRemovalPolicy policy) const; + + FileDescriptor root_; + FileDescriptor bookmarks_; + FileDescriptor lock_file_; + NSString* root_path_; + NSString* key_; + CacheEntryPersistence persistence_; +}; + +class CacheEntryMaintenance final { + public: + runtime::Result inspect_bookmark() const; + runtime::Error preflight() const; + runtime::Error remove_source() const; + runtime::Error remove_bookmark() const; + + private: + friend class LockedCacheEntry; + explicit CacheEntryMaintenance(const LockedCacheEntry& entry) + : entry_(entry) {} + + const LockedCacheEntry& entry_; +}; + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_cache_entry.mm b/backends/apple/coreai/runtime/coreai_cache_entry.mm new file mode 100644 index 00000000000..ce363e04760 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache_entry.mm @@ -0,0 +1,370 @@ +/* + * 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_cache_entry.h" + +#import "coreai_cache_key.h" +#import "coreai_fault_injection.h" +#import "coreai_storage.h" +#import "coreai_validation.h" + +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::coreai { +namespace { +using runtime::Error; +using runtime::Result; + +constexpr size_t kBookmarkByteLimit = 8 * 1024 * 1024; +NSString* const kBookmarksDirectory = @"bookmarks"; +NSString* const kSourcesDirectory = @"sources"; +NSString* const kLocksDirectory = @"locks"; +NSString* const kBookmarkSuffix = @(kBookmarkExtension); + +NSString* bookmark_name(NSString* key) { + return [key stringByAppendingString:kBookmarkSuffix]; +} + +Result> +open_child(int root_fd, NSString* name, StorageCreationMode mode) { + auto opened = open_storage_directory(root_fd, name, mode); + if (!opened.ok()) { + return opened.error(); + } + if (opened->created) { + ET_CHECK_OK_OR_RETURN_ERROR(sync_storage_directory(root_fd)); + } + return std::move(opened->descriptor); +} + +Result entry_exists(int root_fd, NSString* key) { + for (NSString* directory_name in + @[ kBookmarksDirectory, kSourcesDirectory ]) { + auto opened = + open_child(root_fd, directory_name, StorageCreationMode::OpenExisting); + if (!opened.ok()) { + return opened.error(); + } + if (!opened->has_value()) { + continue; + } + NSString* entry = + [directory_name isEqual:kBookmarksDirectory] ? bookmark_name(key) : key; + struct stat info; + if (lstat_at( + opened->value().get(), entry.fileSystemRepresentation, &info) == + 0) { + return true; + } + ET_CHECK_OR_RETURN_ERROR( + errno == ENOENT, AccessFailed, "Cannot inspect Core AI cache entry"); + } + return false; +} + +Result duplicate_directory(int fd) { + FileDescriptor duplicate( + retry_eintr([&] { return fcntl(fd, F_DUPFD_CLOEXEC, 0); })); + ET_CHECK_OR_RETURN_ERROR( + duplicate.get() >= 0, + AccessFailed, + "Cannot retain Core AI cache directory"); + return duplicate; +} + +} // namespace + +Result*> inventory_cache_keys( + const CacheDirectory& directory) { + NSMutableSet* keys = [NSMutableSet set]; + for (NSString* child_name in @[ kBookmarksDirectory, kSourcesDirectory ]) { + auto child = open_child( + directory.descriptor.get(), + child_name, + StorageCreationMode::OpenExisting); + if (!child.ok()) { + return child.error(); + } + if (!child->has_value()) { + continue; + } + auto names = + storage_children(child->value().get(), StorageNamePolicy::SkipInvalid); + if (!names.ok()) { + return names.error(); + } + for (NSString* name in names.get()) { + NSString* key = name; + if ([child_name isEqual:kBookmarksDirectory]) { + if (![name hasSuffix:kBookmarkSuffix] || + name.length <= kBookmarkSuffix.length) { + continue; + } + key = [name substringToIndex:name.length - kBookmarkSuffix.length]; + } + if (is_cache_key(key)) { + [keys addObject:key]; + } + } + } + return [[keys allObjects] sortedArrayUsingSelector:@selector(compare:)]; +} + +LockedCacheEntry::LockedCacheEntry( + FileDescriptor root, + FileDescriptor bookmarks, + FileDescriptor lock_file, + NSString* root_path, + NSString* key, + CacheEntryPersistence persistence) + : root_(std::move(root)), + bookmarks_(std::move(bookmarks)), + lock_file_(std::move(lock_file)), + root_path_(root_path), + key_(key), + persistence_(persistence) {} + +Result> LockedCacheEntry::acquire( + const CacheDirectory& directory, + NSString* key) { + return acquire_impl(directory, key, StorageCreationMode::Create); +} + +Result> LockedCacheEntry::acquire_existing( + const CacheDirectory& directory, + NSString* key) { + return acquire_impl(directory, key, StorageCreationMode::OpenExisting); +} + +Result> LockedCacheEntry::acquire_impl( + const CacheDirectory& cache_directory, + NSString* key, + StorageCreationMode mode) { + auto lock_name = cache_entry_lock_name(key); + if (!lock_name.ok()) { + return lock_name.error(); + } + auto duplicated = duplicate_directory(cache_directory.descriptor.get()); + if (!duplicated.ok()) { + return duplicated.error(); + } + FileDescriptor directory = std::move(duplicated.get()); + const bool create = mode == StorageCreationMode::Create; + // Checking before locking avoids creating locks/ for a missing entry. + if (!create) { + auto exists = entry_exists(directory.get(), key); + if (!exists.ok()) { + return exists.error(); + } + if (!exists.get()) { + return std::unique_ptr(); + } + } + auto locks = + open_child(directory.get(), kLocksDirectory, StorageCreationMode::Create); + if (!locks.ok()) { + return locks.error(); + } + FileDescriptor lock_file(open_storage_shared_file( + locks->value().get(), lock_name.get().fileSystemRepresentation)); + struct stat info; + ET_CHECK_OR_RETURN_ERROR( + lock_file.get() >= 0 && fstat(lock_file.get(), &info) == 0 && + S_ISREG(info.st_mode), + AccessFailed, + "Invalid Core AI cache entry lock"); + ET_CHECK_OR_RETURN_ERROR( + fault_point(FaultPoint::BeforeCacheEntryLock) == 0, + AccessFailed, + "Cannot begin Core AI cache entry lock acquisition"); + ET_CHECK_OR_RETURN_ERROR( + retry_eintr([&] { return flock(lock_file.get(), LOCK_EX); }) == 0, + AccessFailed, + "Cannot lock Core AI cache entry"); + // Another process may have cleared the entry while this one waited. + if (!create) { + auto exists = entry_exists(directory.get(), key); + if (!exists.ok()) { + return exists.error(); + } + if (!exists.get()) { + return std::unique_ptr(); + } + } + auto bookmarks = open_child(directory.get(), kBookmarksDirectory, mode); + if (!bookmarks.ok()) { + return bookmarks.error(); + } + const CacheEntryPersistence persistence = bookmarks->has_value() + ? CacheEntryPersistence::BookmarkCapable + : CacheEntryPersistence::SourceOnly; + FileDescriptor bookmark_directory = + bookmarks->has_value() ? std::move(bookmarks->value()) : FileDescriptor(); + return std::unique_ptr(new LockedCacheEntry( + std::move(directory), + std::move(bookmark_directory), + std::move(lock_file), + [cache_directory.path copy], + [key copy], + persistence)); +} + +Result LockedCacheEntry::read_bookmark() const { + if (persistence_ == CacheEntryPersistence::SourceOnly) { + return static_cast(nil); + } + FileDescriptor file(retry_eintr([&] { + return openat( + bookmarks_.get(), + bookmark_name(key_).fileSystemRepresentation, + O_RDONLY | O_NOFOLLOW | O_NONBLOCK | O_CLOEXEC); + })); + if (file.get() < 0 && errno == ENOENT) { + return static_cast(nil); + } + ET_CHECK_OR_RETURN_ERROR( + file.get() >= 0, AccessFailed, "Cannot open Core AI bookmark"); + struct stat info; + ET_CHECK_OR_RETURN_ERROR( + fstat(file.get(), &info) == 0 && S_ISREG(info.st_mode) && + info.st_size > 0 && + static_cast(info.st_size) <= kBookmarkByteLimit, + InvalidExternalData, + "Invalid Core AI bookmark file"); + NSMutableData* bytes = [NSMutableData dataWithLength:info.st_size]; + ET_CHECK_OR_RETURN_ERROR( + bytes != nil, MemoryAllocationFailed, "Cannot allocate Core AI bookmark"); + size_t offset = 0; + while (offset < bytes.length) { + const ssize_t count = retry_eintr([&] { + return ::read( + file.get(), + static_cast(bytes.mutableBytes) + offset, + bytes.length - offset); + }); + ET_CHECK_OR_RETURN_ERROR( + count > 0, AccessFailed, "Cannot read Core AI bookmark"); + offset += count; + } + return bytes; +} + +Error LockedCacheEntry::write_bookmark(NSData* data) const { + ET_CHECK_OR_RETURN_ERROR( + [data isKindOfClass:NSData.class] && data.length > 0 && + data.length <= kBookmarkByteLimit, + InvalidExternalData, + "Invalid Core AI bookmark data"); + ET_CHECK_OR_RETURN_ERROR( + persistence_ == CacheEntryPersistence::BookmarkCapable, + InvalidArgument, + "Cannot write through a source-only Core AI cache entry lock"); + return publish_storage_data(bookmarks_.get(), bookmark_name(key_), data); +} + +Result> LockedCacheEntry::open_source_directory( + StorageCreationMode mode) const { + return open_child(root_.get(), kSourcesDirectory, mode); +} + +Error LockedCacheEntry::apply_to_source( + Error (*operation)(int, NSString*, StorageRemovalPolicy), + StorageRemovalPolicy policy) const { + auto sources = open_source_directory(StorageCreationMode::OpenExisting); + if (!sources.ok()) { + return sources.error(); + } + if (!sources->has_value()) { + return Error::Ok; + } + return operation(sources->value().get(), key_, policy); +} + +Result LockedCacheEntry::prepare_source_directory() + const { + auto sources = open_source_directory(StorageCreationMode::Create); + if (!sources.ok()) { + return sources.error(); + } + NSURL* url = [[NSURL fileURLWithPath:root_path_ isDirectory:YES] + URLByAppendingPathComponent:kSourcesDirectory + isDirectory:YES]; + return CacheSourceDirectory{std::move(sources->value()), url}; +} + +Error LockedCacheEntry::remove_source() const { + return apply_to_source( + remove_storage_directory, StorageRemovalPolicy::Strict); +} + +CacheEntryMaintenance LockedCacheEntry::maintenance() const { + return CacheEntryMaintenance(*this); +} + +Result CacheEntryMaintenance::inspect_bookmark() + const { + if (entry_.persistence_ == CacheEntryPersistence::SourceOnly) { + return CacheBookmarkInspection{CacheBookmarkState::Missing, nil}; + } + struct stat info; + const int inspected = lstat_at( + entry_.bookmarks_.get(), + bookmark_name(entry_.key_).fileSystemRepresentation, + &info); + if (inspected < 0 && errno == ENOENT) { + return CacheBookmarkInspection{CacheBookmarkState::Missing, nil}; + } + ET_CHECK_OR_RETURN_ERROR( + inspected == 0, AccessFailed, "Cannot inspect Core AI bookmark"); + if (!S_ISREG(info.st_mode)) { + return CacheBookmarkInspection{CacheBookmarkState::Corrupt, nil}; + } + auto bookmark = entry_.read_bookmark(); + if (!bookmark.ok()) { + if (bookmark.error() == Error::InvalidExternalData) { + return CacheBookmarkInspection{CacheBookmarkState::Corrupt, nil}; + } + return bookmark.error(); + } + return CacheBookmarkInspection{ + bookmark.get() == nil ? CacheBookmarkState::Missing + : CacheBookmarkState::Present, + bookmark.get()}; +} + +Error CacheEntryMaintenance::preflight() const { + if (entry_.persistence_ == CacheEntryPersistence::BookmarkCapable) { + ET_CHECK_OK_OR_RETURN_ERROR(preflight_storage_directory( + entry_.bookmarks_.get(), + bookmark_name(entry_.key_), + StorageRemovalPolicy::Maintenance)); + } + return entry_.apply_to_source( + preflight_storage_directory, StorageRemovalPolicy::Maintenance); +} + +Error CacheEntryMaintenance::remove_source() const { + return entry_.apply_to_source( + remove_storage_directory, StorageRemovalPolicy::Maintenance); +} + +Error CacheEntryMaintenance::remove_bookmark() const { + if (entry_.persistence_ == CacheEntryPersistence::SourceOnly) { + return Error::Ok; + } + return remove_storage_directory( + entry_.bookmarks_.get(), + bookmark_name(entry_.key_), + StorageRemovalPolicy::Maintenance); +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_cache_key.h b/backends/apple/coreai/runtime/coreai_cache_key.h new file mode 100644 index 00000000000..b252821cbb3 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache_key.h @@ -0,0 +1,42 @@ +/* + * 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 +#include +#include + +namespace executorch::backends::coreai { + +class Manifest; +class SelectedBundle; + +bool is_cache_key(NSString* key); +NSString* cache_key_bundle_identity(NSString* key); + +// Bookmark files are named key + extension, which bounds the key length. +inline constexpr char kBookmarkExtension[] = ".bookmark"; + +// Bundle identity excludes device and specialization so maintenance can find +// every runtime variant of one exported bundle. +runtime::Result bundle_identity( + const Manifest& manifest, + NSString* bundle_name); +runtime::Result bundle_identity(const SelectedBundle& bundle); +// specialization_identity is stable, versioned plain data supplied by the SDK +// boundary. This module does not depend on SDK types. +runtime::Result cache_variant_key( + const SelectedBundle& bundle, + NSString* platform, + NSString* architecture, + NSString* specialization_identity); + +constexpr size_t kCacheEntryLockStripeCount = 256; +runtime::Result cache_entry_lock_name(NSString* key); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_cache_key.mm b/backends/apple/coreai/runtime/coreai_cache_key.mm new file mode 100644 index 00000000000..8f63ce7fd25 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache_key.mm @@ -0,0 +1,185 @@ +/* + * 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_cache_key.h" + +#include +#include +#import "coreai_manifest.h" +#import "coreai_validation.h" + +namespace executorch::backends::coreai { +namespace { +using runtime::Result; + +constexpr NSUInteger kHashLength = 2 * CC_SHA256_DIGEST_LENGTH; +constexpr NSUInteger kFixedLength = 2 * (kHashLength + 1); +constexpr NSUInteger kMaxKeyLength = + NAME_MAX - (sizeof(kBookmarkExtension) - 1); + +static_assert(kCacheEntryLockStripeCount > 0); +static_assert(kCacheEntryLockStripeCount <= 256); + +bool valid_device(NSString* device) { + if (![device isKindOfClass:NSString.class] || + device.length > kMaxKeyLength - kFixedLength) { + return false; + } + const NSUInteger separator = [device rangeOfString:@"_"].location; + if (separator == NSNotFound || separator == 0 || + separator + 1 == device.length) { + return false; + } + for (NSUInteger i = 0; i < device.length; ++i) { + const unichar c = [device characterAtIndex:i]; + if (i == separator || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || + (c >= '0' && c <= '9')) { + continue; + } + if (i < separator || (c != '_' && c != '-')) { + return false; + } + } + return true; +} + +Result encode_json(id value) { + NSError* error = nil; + NSData* data = [NSJSONSerialization dataWithJSONObject:value + options:0 + error:&error]; + ET_CHECK_OR_RETURN_ERROR( + data != nil, + InvalidExternalData, + "Cannot encode Core AI cache identity: %s", + error.localizedDescription.UTF8String); + return data; +} + +Result hash_data(NSData* data) { + ET_CHECK_OR_RETURN_ERROR( + data != nil && data.length <= UINT_MAX, + InvalidExternalData, + "Cannot hash Core AI cache identity"); + unsigned char bytes[CC_SHA256_DIGEST_LENGTH]; + CC_SHA256(data.bytes, static_cast(data.length), bytes); + NSMutableString* result = [NSMutableString stringWithCapacity:kHashLength]; + for (unsigned char byte : bytes) { + [result appendFormat:@"%02x", byte]; + } + return [result copy]; +} + +Result compute_bundle_identity( + NSString* digest, + NSString* bundle_name, + bool aot_compiled) { + // Bumping the leading version makes every existing cache entry unreachable. + auto encoded = encode_json( + @[ @"1", digest, aot_compiled ? @"aot" : @"source", bundle_name ]); + if (!encoded.ok()) { + return encoded.error(); + } + return hash_data(encoded.get()); +} + +Result utf8_data(NSString* value) { + ET_CHECK_OR_RETURN_ERROR( + is_nonempty_string(value), + DelegateInvalidCompatibility, + "Missing Core AI cache specialization"); + NSData* data = [value dataUsingEncoding:NSUTF8StringEncoding]; + ET_CHECK_OR_RETURN_ERROR( + data != nil, InvalidExternalData, "Cannot encode Core AI specialization"); + return data; +} +} // namespace + +bool is_cache_key(NSString* key) { + if (![key isKindOfClass:NSString.class] || key.length <= kFixedLength || + key.length > kMaxKeyLength) { + return false; + } + const NSUInteger options_start = key.length - kHashLength; + if ([key characterAtIndex:kHashLength] != '-' || + [key characterAtIndex:options_start - 1] != '-') { + return false; + } + NSString* model = [key substringToIndex:kHashLength]; + NSString* options = [key substringFromIndex:options_start]; + if (!is_lowercase_hex(model, kHashLength) || + !is_lowercase_hex(options, kHashLength)) { + return false; + } + return valid_device( + [key substringWithRange:NSMakeRange( + kHashLength + 1, key.length - kFixedLength)]); +} + +NSString* cache_key_bundle_identity(NSString* key) { + return is_cache_key(key) ? [key substringToIndex:kHashLength] : nil; +} + +Result bundle_identity( + const Manifest& manifest, + NSString* bundle_name) { + NSString* digest = manifest.bundle_digests[bundle_name]; + ET_CHECK_OR_RETURN_ERROR( + digest != nil, InvalidArgument, "Unknown Core AI bundle name"); + return compute_bundle_identity(digest, bundle_name, manifest.aot_compiled); +} + +Result bundle_identity(const SelectedBundle& bundle) { + return compute_bundle_identity( + bundle.digest, bundle.bundle_name, bundle.aot_compiled); +} + +Result cache_variant_key( + const SelectedBundle& bundle, + NSString* platform, + NSString* architecture, + NSString* specialization_identity) { + auto model = bundle_identity(bundle); + if (!model.ok()) { + return model.error(); + } + ET_CHECK_OR_RETURN_ERROR( + is_nonempty_string(platform) && is_nonempty_string(architecture) && + ![platform containsString:@"_"], + DelegateInvalidCompatibility, + "Invalid Core AI cache context"); + NSString* device = + [NSString stringWithFormat:@"%@_%@", platform, architecture]; + ET_CHECK_OR_RETURN_ERROR( + valid_device(device), + DelegateInvalidCompatibility, + "Invalid Core AI cache device"); + auto specialization = utf8_data(specialization_identity); + if (!specialization.ok()) { + return specialization.error(); + } + auto options = hash_data(specialization.get()); + if (!options.ok()) { + return options.error(); + } + return [NSString + stringWithFormat:@"%@-%@-%@", model.get(), device, options.get()]; +} + +Result cache_entry_lock_name(NSString* key) { + ET_CHECK_OR_RETURN_ERROR( + is_cache_key(key), InvalidArgument, "Invalid Core AI cache key"); + NSData* encoded = [key dataUsingEncoding:NSUTF8StringEncoding]; + unsigned char digest[CC_SHA256_DIGEST_LENGTH]; + CC_SHA256(encoded.bytes, static_cast(encoded.length), digest); + const size_t stripe = digest[0] % kCacheEntryLockStripeCount; + // A different stripe scheme needs a new prefix so its files never alias + // these. + return [NSString stringWithFormat:@"stripe-v1-%02zx.lock", stripe]; +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_cache_root.h b/backends/apple/coreai/runtime/coreai_cache_root.h index ee536e241fc..f2c3f696202 100644 --- a/backends/apple/coreai/runtime/coreai_cache_root.h +++ b/backends/apple/coreai/runtime/coreai_cache_root.h @@ -10,16 +10,30 @@ #import #include -#include "coreai_file.h" +#include +#include "coreai_storage.h" namespace executorch::backends::coreai { +struct CacheDirectory final { + NSString* path; + FileDescriptor descriptor; +}; + // Resolves the default location without creating or modifying storage. -runtime::Result default_coreai_assets_root(); -// Requires an absolute directory path. With Create, makes the directory and -// excludes it from backup; otherwise a missing directory is InvalidArgument. -runtime::Result prepare_storage_root( - NSString* path, - StorageCreationMode mode); +runtime::Result default_cache_directory(); +// Null selects the default; invalid explicit paths never fall back to it. +runtime::Result cache_directory_path(const char* configured_path); + +// Both functions resolve a symlink in the final path component once and hold +// the directory open, so later path changes cannot redirect the cache. +// +// Creates the directory if missing and excludes it from backup. A path that is +// not a directory is the caller's configuration error: InvalidArgument. +runtime::Result create_cache_directory(NSString* path); +// Returns no value when the directory is missing. A path that is not a +// directory means existing storage is unusable: AccessFailed. +runtime::Result> open_cache_directory( + NSString* path); } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_cache_root.mm b/backends/apple/coreai/runtime/coreai_cache_root.mm index 18ca1b6f198..abc039bed0f 100644 --- a/backends/apple/coreai/runtime/coreai_cache_root.mm +++ b/backends/apple/coreai/runtime/coreai_cache_root.mm @@ -8,24 +8,29 @@ #import "coreai_cache_root.h" #include "coreai_fault_injection.h" -#include "coreai_file.h" +#include "coreai_validation.h" +#include #include +#include #include +#include +#include +#include namespace executorch::backends::coreai { namespace { using runtime::Error; using runtime::Result; -Error ensure_excluded_from_backup(NSURL* directory) { - NSURL* fresh = [NSURL fileURLWithPath:directory.path isDirectory:YES]; +Error ensure_excluded_from_backup(NSString* path) { + NSURL* directory = [NSURL fileURLWithPath:path isDirectory:YES]; NSNumber* excluded = nil; NSError* error = nil; ET_CHECK_OR_RETURN_ERROR( - [fresh getResourceValue:&excluded - forKey:NSURLIsExcludedFromBackupKey - error:&error], + [directory getResourceValue:&excluded + forKey:NSURLIsExcludedFromBackupKey + error:&error], AccessFailed, "Cannot read Core AI backup exclusion: %s", error.localizedDescription.UTF8String); @@ -33,71 +38,157 @@ Error ensure_excluded_from_backup(NSURL* directory) { return Error::Ok; } ET_CHECK_OR_RETURN_ERROR( - [fresh setResourceValue:@YES - forKey:NSURLIsExcludedFromBackupKey - error:&error], + [directory setResourceValue:@YES + forKey:NSURLIsExcludedFromBackupKey + error:&error], AccessFailed, - "Cannot exclude Core AI storage from backup: %s", + "Cannot exclude Core AI cache directory from backup: %s", error.localizedDescription.UTF8String); return Error::Ok; } -} // namespace +Result canonical_path(NSString* path) { + char* resolved = realpath(path.fileSystemRepresentation, nullptr); + ET_CHECK_OR_RETURN_ERROR( + resolved != nullptr, + AccessFailed, + "Cannot resolve Core AI cache directory"); + NSString* canonical = [NSFileManager.defaultManager + stringWithFileSystemRepresentation:resolved + length:std::strlen(resolved)]; + std::free(resolved); + ET_CHECK_OR_RETURN_ERROR( + canonical != nil, + AccessFailed, + "Cannot decode Core AI cache directory path"); + return canonical; +} -Result default_coreai_assets_root() { - NSURL* caches = - [NSFileManager.defaultManager URLsForDirectory:NSCachesDirectory - inDomains:NSUserDomainMask] - .firstObject; +// isAbsolutePath accepts "~" paths, which the filesystem would treat literally. +bool is_cache_directory_path(NSString* path) { + return is_nonempty_string(path) && path.isAbsolutePath && + ![path hasPrefix:@"~"]; +} + +Error create_directory(NSString* path) { + NSDictionary* attributes = @{ + NSFileProtectionKey : NSFileProtectionCompleteUntilFirstUserAuthentication, + }; ET_CHECK_OR_RETURN_ERROR( - caches.isFileURL, AccessFailed, "Cannot resolve Core AI cache directory"); - return [caches.path stringByAppendingPathComponent:@"executorch_coreai"]; + fault_point(FaultPoint::BeforeCreateDirectory) == 0, + AccessFailed, + "Cannot begin Core AI cache directory creation"); + NSError* error = nil; + ET_CHECK_OR_RETURN_ERROR( + [NSFileManager.defaultManager createDirectoryAtPath:path + withIntermediateDirectories:YES + attributes:attributes + error:&error], + AccessFailed, + "Cannot create Core AI cache directory: %s", + error.localizedDescription.UTF8String); + return Error::Ok; } -Result prepare_storage_root( +Error not_a_directory(StorageCreationMode mode) { + ET_LOG(Error, "Core AI cache path is not a directory"); + return mode == StorageCreationMode::Create ? Error::InvalidArgument + : Error::AccessFailed; +} + +Result> resolve_cache_directory( NSString* path, StorageCreationMode mode) { const bool create = mode == StorageCreationMode::Create; ET_CHECK_OR_RETURN_ERROR( - [path isKindOfClass:NSString.class] && [path hasPrefix:@"/"] && - [path rangeOfString:@"\0"].location == NSNotFound, + is_cache_directory_path(path), InvalidArgument, - "Core AI storage requires an absolute path"); + "Core AI cache directory requires an absolute path"); struct stat info; - const int result = - retry_eintr([&] { return stat(path.fileSystemRepresentation, &info); }); - ET_CHECK_OR_RETURN_ERROR( - result == 0 || errno == ENOENT || errno == ENOTDIR, - AccessFailed, - "Cannot inspect Core AI storage directory"); - ET_CHECK_OR_RETURN_ERROR( - result == 0 ? S_ISDIR(info.st_mode) : create && errno == ENOENT, - InvalidArgument, - "Core AI storage must be a directory"); - if (create && result != 0) { - NSDictionary* attributes = @{ - NSFileProtectionKey : - NSFileProtectionCompleteUntilFirstUserAuthentication, - }; + if (retry_eintr([&] { return stat(path.fileSystemRepresentation, &info); }) != + 0) { + if (errno == ENOTDIR) { + return not_a_directory(mode); + } ET_CHECK_OR_RETURN_ERROR( - fault_point(FaultPoint::BeforeCreateDirectory) == 0, + errno == ENOENT, AccessFailed, - "Cannot begin Core AI directory creation"); - NSError* error = nil; - ET_CHECK_OR_RETURN_ERROR( - [NSFileManager.defaultManager createDirectoryAtPath:path - withIntermediateDirectories:YES - attributes:attributes - error:&error], - AccessFailed, - "Cannot create Core AI storage directory: %s", - error.localizedDescription.UTF8String); + "Cannot inspect Core AI cache directory"); + if (!create) { + return std::optional(); + } + ET_CHECK_OK_OR_RETURN_ERROR(create_directory(path)); + } else if (!S_ISDIR(info.st_mode)) { + return not_a_directory(mode); + } + NSString* resolved_path = path; + struct stat link_info; + if (retry_eintr([&] { + return lstat(path.fileSystemRepresentation, &link_info); + }) == 0 && + S_ISLNK(link_info.st_mode)) { + auto canonical = canonical_path(path); + if (!canonical.ok()) { + return canonical.error(); + } + resolved_path = canonical.get(); } + FileDescriptor directory(retry_eintr([&] { + return open( + resolved_path.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + if (directory.get() < 0 && errno == ENOENT && !create) { + return std::optional(); + } + ET_CHECK_OR_RETURN_ERROR( + directory.get() >= 0 && fstat(directory.get(), &info) == 0 && + S_ISDIR(info.st_mode), + AccessFailed, + "Cannot open Core AI cache directory"); if (create) { - ET_CHECK_OK_OR_RETURN_ERROR(ensure_excluded_from_backup( - [NSURL fileURLWithPath:path isDirectory:YES])); + ET_CHECK_OK_OR_RETURN_ERROR(ensure_excluded_from_backup(resolved_path)); + } + return std::optional( + std::in_place, + CacheDirectory{[resolved_path copy], std::move(directory)}); +} + +} // namespace + +Result default_cache_directory() { + NSURL* caches = + [NSFileManager.defaultManager URLsForDirectory:NSCachesDirectory + inDomains:NSUserDomainMask] + .firstObject; + ET_CHECK_OR_RETURN_ERROR( + caches.isFileURL, AccessFailed, "Cannot resolve Core AI cache directory"); + return [caches.path stringByAppendingPathComponent:@"executorch_coreai"]; +} + +Result cache_directory_path(const char* configured_path) { + if (configured_path == nullptr) { + return default_cache_directory(); } + NSString* path = [NSString stringWithUTF8String:configured_path]; + ET_CHECK_OR_RETURN_ERROR( + is_cache_directory_path(path), + InvalidArgument, + "Core AI cache directory requires an absolute UTF-8 path"); return path; } +Result create_cache_directory(NSString* path) { + auto directory = resolve_cache_directory(path, StorageCreationMode::Create); + if (!directory.ok()) { + return directory.error(); + } + // Create mode reports a missing directory as an error, never as no value. + return std::move(**directory); +} + +Result> open_cache_directory(NSString* path) { + return resolve_cache_directory(path, StorageCreationMode::OpenExisting); +} + } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_error.h b/backends/apple/coreai/runtime/coreai_error.h new file mode 100644 index 00000000000..09c276bb43b --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_error.h @@ -0,0 +1,20 @@ +/* + * 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 +#include + +namespace executorch::backends::coreai { + +// Binding and execution preserve bridge errors that callers can act on. +runtime::Error execution_error(NSError* error); +// Acquisition failures describe SDK/cache internals, not caller arguments. +runtime::Error acquisition_error(NSError* error); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_error.mm b/backends/apple/coreai/runtime/coreai_error.mm new file mode 100644 index 00000000000..e63e66fd477 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_error.mm @@ -0,0 +1,48 @@ +/* + * 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_error.h" + +#import "ETCoreAIBridge.h" + +#include + +namespace executorch::backends::coreai { + +runtime::Error execution_error(NSError* error) { + ET_LOG( + Error, + "Core AI execution error (%s, %ld): %s", + error.domain.UTF8String, + static_cast(error.code), + error.description.UTF8String); + if ([error.domain isEqualToString:ETCoreAIErrorDomain]) { + switch (error.code) { + case ETCoreAIErrorInvalidModel: + return runtime::Error::InvalidProgram; + case ETCoreAIErrorUnsupported: + return runtime::Error::NotSupported; + case ETCoreAIErrorInvalidArgument: + return runtime::Error::InvalidArgument; + default: + break; + } + } + return runtime::Error::Internal; +} + +runtime::Error acquisition_error(NSError* error) { + ET_LOG( + Error, + "Core AI acquisition error (%s, %ld): %s", + error.domain.UTF8String, + static_cast(error.code), + error.description.UTF8String); + return runtime::Error::Internal; +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_fault_injection.h b/backends/apple/coreai/runtime/coreai_fault_injection.h index cba60c5dfb8..79b9f38da6b 100644 --- a/backends/apple/coreai/runtime/coreai_fault_injection.h +++ b/backends/apple/coreai/runtime/coreai_fault_injection.h @@ -8,6 +8,10 @@ #pragma once +#include +#include +#include "coreai_file.h" + namespace executorch::backends::coreai { enum class FaultPoint { @@ -16,7 +20,15 @@ enum class FaultPoint { SyncDirectory, FullSync, Rename, + Remove, + CreatePendingDirectory, + OpenPendingDirectory, BeforeCreateDirectory, + BeforeCacheEntryLock, + BeforeSDK, + AfterSDK, + BeforeEvict, + AfterEvict, }; #if defined(COREAI_ASSETS_TESTING) && COREAI_ASSETS_TESTING @@ -31,4 +43,16 @@ constexpr int fault_point(FaultPoint) { } #endif +template +auto fault_aware_retry(FaultPoint point, F&& call) { + return retry_eintr([&]() -> decltype(call()) { + const int fault = fault_point(point); + if (fault != 0) { + errno = fault; + return -1; + } + return call(); + }); +} + } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_file.h b/backends/apple/coreai/runtime/coreai_file.h index cbc991a5fae..473c8d5bb7d 100644 --- a/backends/apple/coreai/runtime/coreai_file.h +++ b/backends/apple/coreai/runtime/coreai_file.h @@ -15,19 +15,21 @@ namespace executorch::backends::coreai { -enum class StorageCreationMode { - OpenExisting, - Create, -}; - class FileDescriptor { public: - explicit FileDescriptor(int fd) : fd_(fd) {} + explicit FileDescriptor(int fd = -1) : fd_(fd) {} ~FileDescriptor() { reset(-1); } FileDescriptor(const FileDescriptor&) = delete; FileDescriptor& operator=(const FileDescriptor&) = delete; + FileDescriptor(FileDescriptor&& other) noexcept : fd_(other.release()) {} + FileDescriptor& operator=(FileDescriptor&& other) noexcept { + if (this != &other) { + reset(other.release()); + } + return *this; + } int get() const { return fd_; } diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.h b/backends/apple/coreai/runtime/coreai_load_coordinator.h new file mode 100644 index 00000000000..0cff9087e7d --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.h @@ -0,0 +1,40 @@ +/* + * 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 +#import "ETCoreAIBridge.h" +#import "coreai_cache_root.h" +#import "coreai_manifest.h" + +namespace executorch::backends::coreai { + +struct SelectedHostBundle final { + SelectedBundle bundle; + NSString* const platform; + NSString* const architecture; +}; + +runtime::Result select_host_bundle( + const Manifest& manifest); + +// The cache entry lock outlives all SDK callbacks and bookmark publication. +// Cleanup after a restore hit or durable publication is best-effort. +runtime::Result> acquire_model( + const SelectedHostBundle& host_bundle, + const runtime::NamedDataMap* named_data, + const CacheDirectory& cache_directory, + id loader); + +// Callers must keep affected loads and sessions quiescent across processes. +runtime::Error clear_cache_entry( + const CacheDirectory& cache_directory, + NSString* key, + id loader); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm new file mode 100644 index 00000000000..bd4f45a7de1 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -0,0 +1,275 @@ +/* + * 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_cache_entry.h" +#import "coreai_cache_key.h" +#import "coreai_error.h" +#import "coreai_fault_injection.h" +#import "coreai_source_bundle.h" + +#include +#include + +namespace executorch::backends::coreai { +namespace { +using runtime::Error; +using runtime::Result; + +Error require_coreai_available() { + ET_CHECK_OR_RETURN_ERROR( + ETCoreAIIsAvailable(), + NotSupported, + "Core AI requires macOS 27 or iOS 27"); + return Error::Ok; +} + +Result> wait_for_acquisition( + void (^operation)(ETCoreAIAcquisitionCompletion)) { + ET_CHECK_OR_RETURN_ERROR( + fault_point(FaultPoint::BeforeSDK) == 0, + AccessFailed, + "Core AI acquisition interrupted"); + dispatch_semaphore_t ready = dispatch_semaphore_create(0); + __block id model = nil; + __block NSError* error = nil; + __block ETCoreAIAcquisitionStatus status = ETCoreAIAcquisitionStatusError; + operation( + ^(ETCoreAIAcquisitionStatus result, + id value, + NSError* failure) { + status = result; + model = value; + error = failure; + dispatch_semaphore_signal(ready); + }); + // Wait without a timeout; the SDK may still be reading the staged source. + dispatch_semaphore_wait(ready, DISPATCH_TIME_FOREVER); + id acquired = model; + NSError* failure = error; + // A loader may retain the completion block after invoking it. + model = nil; + error = nil; + if (failure != nil) { + return acquisition_error(failure); + } + ET_CHECK_OR_RETURN_ERROR( + fault_point(FaultPoint::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> restore( + id loader, + NSData* bookmark) { + return wait_for_acquisition(^(ETCoreAIAcquisitionCompletion completion) { + [loader restoreModelFromBookmark:bookmark completion:completion]; + }); +} + +void cleanup_source(const LockedCacheEntry& cache_entry) { + const Error error = cache_entry.remove_source(); + if (error != Error::Ok) { + ET_LOG( + Error, + "Cannot remove Core AI source bundle: error=%u", + static_cast(error)); + } +} + +Error evict_one(id loader, NSData* bookmark) { + ET_CHECK_OR_RETURN_ERROR( + fault_point(FaultPoint::BeforeEvict) == 0, + AccessFailed, + "Core AI eviction interrupted"); + dispatch_semaphore_t ready = dispatch_semaphore_create(0); + __block NSError* failure = nil; + [loader evictModelWithBookmark:bookmark + completion:^(NSError* error) { + failure = error; + dispatch_semaphore_signal(ready); + }]; + dispatch_semaphore_wait(ready, DISPATCH_TIME_FOREVER); + NSError* eviction_failure = failure; + // A loader may retain the completion block after invoking it. + failure = nil; + if (eviction_failure != nil) { + // A confirmed restore miss proves that a failed eviction is already done. + @autoreleasepool { + auto probe = restore(loader, bookmark); + if (probe.ok() && probe.get() == nil) { + return Error::Ok; + } + } + return acquisition_error(eviction_failure); + } + ET_CHECK_OR_RETURN_ERROR( + fault_point(FaultPoint::AfterEvict) == 0, + AccessFailed, + "Core AI eviction interrupted"); + return Error::Ok; +} +} // namespace + +Result select_host_bundle(const Manifest& manifest) { + ET_CHECK_OR_RETURN_ERROR( + [NSProcessInfo.processInfo + isOperatingSystemAtLeastVersion:manifest.minimum_version], + DelegateInvalidCompatibility, + "Core AI model requires a newer operating system"); + ET_CHECK_OK_OR_RETURN_ERROR(require_coreai_available()); +#if TARGET_OS_OSX + NSString* platform = @"macOS"; +#elif TARGET_OS_IOS + NSString* platform = @"iOS"; +#else + NSString* platform = @"unsupported"; +#endif + NSString* architecture = ETCoreAIDeviceArchitectureName(); + auto selected = select_bundle(manifest, architecture, platform); + if (!selected.ok()) { + return selected.error(); + } + return SelectedHostBundle{ + std::move(selected.get()), [platform copy], [architecture copy]}; +} + +Result> acquire_model( + const SelectedHostBundle& host_bundle, + const runtime::NamedDataMap* named_data, + const CacheDirectory& cache_directory, + id loader) { + ET_CHECK_OR_RETURN_ERROR( + loader != nil, InvalidArgument, "Missing Core AI model loader"); + NSString* specialization = [loader.specializationIdentifier copy]; + auto key = cache_variant_key( + host_bundle.bundle, + host_bundle.platform, + host_bundle.architecture, + specialization); + if (!key.ok()) { + return key.error(); + } + NSString* cache_key = key.get(); + ET_LOG(Info, "Core AI cache lookup: key=%s", cache_key.UTF8String); + auto locked = LockedCacheEntry::acquire(cache_directory, cache_key); + if (!locked.ok()) { + return locked.error(); + } + auto bookmark = locked.get()->read_bookmark(); + 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) { + ET_LOG( + Info, + "Core AI cache hit: restored model from bookmark, key=%s", + cache_key.UTF8String); + cleanup_source(*locked.get()); + return model; + } + ET_LOG( + Info, + "Core AI cache miss: SDK could not restore bookmark, key=%s", + cache_key.UTF8String); + } else { + ET_LOG( + Info, "Core AI cache miss: no bookmark, key=%s", cache_key.UTF8String); + } + auto source = prepare_source_bundle( + host_bundle.bundle, + named_data, + *locked.get(), + CorruptBundlePolicy::Replace); + if (!source.ok()) { + return source.error(); + } + NSURL* url = source.get(); + auto model = + wait_for_acquisition(^(ETCoreAIAcquisitionCompletion completion) { + ET_LOG( + Info, + "Core AI specialization started: kind=%s, key=%s", + host_bundle.bundle.aot_compiled ? "aot" : "source", + cache_key.UTF8String); + [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_LOG( + Info, "Core AI specialization succeeded: key=%s", cache_key.UTF8String); + ET_CHECK_OK_OR_RETURN_ERROR( + locked.get()->write_bookmark([model.get() copyBookmarkData])); + ET_LOG( + Info, + "Core AI cache bookmark saved durably: key=%s", + cache_key.UTF8String); + cleanup_source(*locked.get()); + return model; +} + +Error clear_cache_entry( + const CacheDirectory& cache_directory, + NSString* key, + id loader) { + auto locked = LockedCacheEntry::acquire_existing(cache_directory, key); + if (!locked.ok()) { + return locked.error(); + } + if (locked.get() == nullptr) { + return Error::Ok; + } + ET_LOG(Info, "Core AI cache clear started: key=%s", key.UTF8String); + auto maintenance = locked.get()->maintenance(); + auto bookmark = maintenance.inspect_bookmark(); + if (!bookmark.ok()) { + return bookmark.error(); + } + ET_CHECK_OK_OR_RETURN_ERROR(maintenance.preflight()); + if (bookmark->state != CacheBookmarkState::Present) { + ET_LOG( + Info, + "Core AI cache key %s has a %s bookmark; SDK eviction could not be " + "confirmed; attempting local cleanup", + key.UTF8String, + bookmark->state == CacheBookmarkState::Missing ? "missing" : "corrupt"); + } else { + if (loader == nil) { + ET_CHECK_OK_OR_RETURN_ERROR(require_coreai_available()); + loader = ETCoreAICreateModelLoader(); + ET_CHECK_OR_RETURN_ERROR( + loader != nil, Internal, "Cannot create Core AI model loader"); + } + ET_LOG(Info, "Core AI SDK eviction requested: key=%s", key.UTF8String); + ET_CHECK_OK_OR_RETURN_ERROR(evict_one(loader, bookmark->data)); + ET_LOG( + Info, + "Core AI SDK cache entry absence confirmed: key=%s", + key.UTF8String); + } + ET_CHECK_OK_OR_RETURN_ERROR(maintenance.remove_source()); + ET_CHECK_OK_OR_RETURN_ERROR(maintenance.remove_bookmark()); + ET_LOG(Info, "Core AI local cache entry cleared: key=%s", key.UTF8String); + return Error::Ok; +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_manifest.h b/backends/apple/coreai/runtime/coreai_manifest.h index c929e9cf632..9627761c0ce 100644 --- a/backends/apple/coreai/runtime/coreai_manifest.h +++ b/backends/apple/coreai/runtime/coreai_manifest.h @@ -15,7 +15,7 @@ namespace executorch::backends::coreai { class Manifest final { public: - NSString* const model_hash; + NSString* const export_hash; NSString* const function; NSArray* const inputs; NSArray* const outputs; @@ -34,7 +34,7 @@ class Manifest final { private: Manifest( - NSString* model_hash, + NSString* export_hash, NSString* function, NSArray* inputs, NSArray* outputs, @@ -50,7 +50,7 @@ class Manifest final { class SelectedBundle final { public: - NSString* const model_hash; + NSString* const export_hash; NSString* const path; NSString* const bundle_name; NSString* const digest; @@ -71,7 +71,7 @@ class SelectedBundle final { NSString* path, NSDictionary* files); - friend runtime::Result select_assets( + friend runtime::Result select_bundle( const Manifest& manifest, NSString* device_architecture, NSString* platform); @@ -79,7 +79,7 @@ class SelectedBundle final { runtime::Result parse_manifest(NSData* data); runtime::Result parse_manifest(const void* data, size_t size); -runtime::Result select_assets( +runtime::Result select_bundle( const Manifest& manifest, NSString* device_architecture, NSString* platform); diff --git a/backends/apple/coreai/runtime/coreai_manifest.mm b/backends/apple/coreai/runtime/coreai_manifest.mm index 06ffff5e832..af139c8f09d 100644 --- a/backends/apple/coreai/runtime/coreai_manifest.mm +++ b/backends/apple/coreai/runtime/coreai_manifest.mm @@ -14,7 +14,7 @@ namespace executorch::backends::coreai { Manifest::Manifest( - NSString* model_hash, + NSString* export_hash, NSString* function, NSArray* inputs, NSArray* outputs, @@ -24,7 +24,7 @@ NSString* platform, NSDictionary* architectures, NSOperatingSystemVersion minimum_version) - : model_hash(model_hash), + : export_hash(export_hash), function(function), inputs(inputs), outputs(outputs), @@ -39,10 +39,10 @@ const Manifest& manifest, NSString* path, NSDictionary* files) - : model_hash(manifest.model_hash), + : export_hash(manifest.export_hash), path(path), bundle_name(path.lastPathComponent), - digest(manifest.bundle_digests[bundle_name]), + digest(manifest.bundle_digests[path.lastPathComponent]), function(manifest.function), inputs(manifest.inputs), outputs(manifest.outputs), @@ -59,7 +59,7 @@ bool names_array(id value) { } NSMutableSet* seen = [NSMutableSet set]; for (id name in value) { - if (!nonempty_string(name) || [seen containsObject:name]) { + if (!is_nonempty_string(name) || [seen containsObject:name]) { return false; } [seen addObject:name]; @@ -91,7 +91,7 @@ Error validate_files( NSLocale* locale = [NSLocale localeWithLocaleIdentifier:@"en_US_POSIX"]; for (NSString* file in filenames) { ET_CHECK_OR_RETURN_ERROR( - relative_path(file) && file_size(filenames[file]) && + is_relative_path(file) && file_size(filenames[file]) && [bundles containsObject:[file componentsSeparatedByString:@"/"] .firstObject] && [file containsString:@"/"] && @@ -132,7 +132,7 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { version = {}; return true; } - if (!nonempty_string(value)) { + if (!is_nonempty_string(value)) { return false; } NSArray* parts = [value componentsSeparatedByString:@"."]; @@ -179,21 +179,23 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { NotSupported, "Unsupported Core AI packaging; expected .aimodel or AOT .aimodelc"); const bool aot_compiled = [packaging hasPrefix:@"aot_compiled_"]; - NSString* model_hash = dict[@"hash"]; + NSString* export_hash = dict[@"export_hash"]; NSString* function = dict[@"function"]; NSArray* inputs = dict[@"input_names"]; NSArray* outputs = dict[@"output_names"]; NSString* platform = nil; NSDictionary* architectures = nil; ET_CHECK_OR_RETURN_ERROR( - lowercase_hex(model_hash), InvalidProgram, "Invalid Core AI model hash"); + is_lowercase_hex(export_hash), + InvalidProgram, + "Invalid Core AI export hash"); NSMutableSet* bundles = [NSMutableSet set]; NSString* source_path = nil; if (aot_compiled) { platform = dict[@"platform"]; architectures = dict[@"archs"]; ET_CHECK_OR_RETURN_ERROR( - nonempty_string(platform), + is_nonempty_string(platform), InvalidProgram, "Core AI AOT manifest requires a target platform"); ET_CHECK_OR_RETURN_ERROR( @@ -203,15 +205,15 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { "Core AI AOT manifest requires a nonempty architecture map"); for (id arch in architectures) { ET_CHECK_OR_RETURN_ERROR( - relative_path(arch) && ![arch containsString:@"/"] && + is_relative_path(arch) && ![arch containsString:@"/"] && [arch canBeConvertedToEncoding:NSUTF8StringEncoding], InvalidProgram, "Invalid Core AI AOT architecture name"); NSString* bundle = [NSString stringWithFormat:@"model.%@.aimodelc", arch]; id path = architectures[arch]; ET_CHECK_OR_RETURN_ERROR( - relative_path(path) && - [path isEqualToString:[model_hash + is_relative_path(path) && + [path isEqualToString:[export_hash stringByAppendingPathComponent:bundle]], InvalidProgram, "Invalid Core AI AOT bundle path for architecture %s", @@ -221,16 +223,17 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { } else { source_path = dict[@"path"]; ET_CHECK_OR_RETURN_ERROR( - relative_path(source_path) && + is_relative_path(source_path) && [source_path - isEqualToString:[model_hash + isEqualToString:[export_hash stringByAppendingString:@"/model.aimodel"]], InvalidProgram, "Invalid Core AI source bundle path"); [bundles addObject:@"model.aimodel"]; } ET_CHECK_OR_RETURN_ERROR( - nonempty_string(function) && names_array(inputs) && names_array(outputs), + is_nonempty_string(function) && names_array(inputs) && + names_array(outputs), InvalidProgram, "Core AI manifest requires ordered input/output names"); NSOperatingSystemVersion minimum_version{}; @@ -249,12 +252,12 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { for (NSString* bundle in bundles) { id digest = bundle_digests[bundle]; ET_CHECK_OR_RETURN_ERROR( - lowercase_hex(digest, 64), + is_lowercase_hex(digest, 64), InvalidProgram, "Invalid Core AI exported bundle digest"); } return Manifest( - model_hash, + export_hash, function, inputs, outputs, @@ -272,14 +275,14 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { return parse_manifest([NSData dataWithBytes:data length:size]); } -Result select_assets( +Result select_bundle( const Manifest& manifest, NSString* device_architecture, NSString* platform) { if (!manifest.aot_compiled) { return SelectedBundle( manifest, - [manifest.model_hash stringByAppendingString:@"/model.aimodel"], + [manifest.export_hash stringByAppendingString:@"/model.aimodel"], manifest.files); } ET_CHECK_OR_RETURN_ERROR( @@ -295,7 +298,7 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { }]; NSString* available = [architectures componentsJoinedByString:@", "]; ET_CHECK_OR_RETURN_ERROR( - nonempty_string(device_architecture), + is_nonempty_string(device_architecture), DelegateInvalidCompatibility, "Core AI SDK device architecture is unavailable (%s); available " "architectures: [%s]; re-export the model as .aimodel or for this device", diff --git a/backends/apple/coreai/runtime/coreai_pte.cpp b/backends/apple/coreai/runtime/coreai_pte.cpp new file mode 100644 index 00000000000..512ae6c31e5 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_pte.cpp @@ -0,0 +1,256 @@ +/* + * 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. + */ + +#include "coreai_pte.h" +#include "coreai_backend_id.h" + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace executorch::backends::coreai { +namespace { +namespace fb = executorch_flatbuffer; +using runtime::DataLoader; +using runtime::Error; +using runtime::ExtendedHeader; +using runtime::FreeableBuffer; +using runtime::Result; +using SegmentInfo = DataLoader::SegmentInfo; + +struct ProgramExtent { + size_t file_size; + size_t program_size; + ExtendedHeader header{}; + bool has_header = false; +}; + +struct VerifiedProgram { + FreeableBuffer storage; + const fb::Program* program; +}; + +Result cpu_data(const FreeableBuffer& buffer, size_t size) { + auto data = buffer.data_safe(); + if (!data.ok()) { + return data.error(); + } + if ((data.get() == nullptr && size != 0) || buffer.size() < size) { + return Error::InvalidProgram; + } + return data.get(); +} + +// Mirror Program::load's header and buffer checks without loading any methods. +Result parse_header(DataLoader& loader) { + auto source_size = loader.size(); + if (!source_size.ok()) { + return source_size.error(); + } + ProgramExtent extent{source_size.get(), source_size.get()}; + if (extent.file_size >= ExtendedHeader::kNumHeadBytes) { + auto prefix = loader.load( + 0, + ExtendedHeader::kNumHeadBytes, + SegmentInfo(SegmentInfo::Type::Program)); + if (!prefix.ok()) { + return prefix.error(); + } + auto data = cpu_data(prefix.get(), ExtendedHeader::kNumHeadBytes); + if (!data.ok()) { + return data.error(); + } + auto parsed = + ExtendedHeader::Parse(data.get(), ExtendedHeader::kNumHeadBytes); + if (parsed.ok()) { + extent.header = parsed.get(); + extent.has_header = true; + const auto& header = extent.header; + if (header.program_size > std::numeric_limits::max() || + header.segment_base_offset > std::numeric_limits::max() || + header.segment_data_size > std::numeric_limits::max() || + header.program_size < ExtendedHeader::kNumHeadBytes || + header.program_size > extent.file_size) { + return Error::InvalidProgram; + } + extent.program_size = static_cast(header.program_size); + // Legacy headers do not specify the segment extent. + if (header.segment_data_size != 0 && + (header.segment_base_offset < extent.program_size || + header.segment_base_offset > extent.file_size || + header.segment_data_size > + extent.file_size - header.segment_base_offset)) { + return Error::InvalidProgram; + } + } else if (parsed.error() != Error::NotFound) { + return Error::InvalidProgram; + } + } + if (extent.program_size == 0 || + extent.program_size >= FLATBUFFERS_MAX_BUFFER_SIZE) { + return Error::InvalidProgram; + } + return extent; +} + +Result load_program( + DataLoader& loader, + const ProgramExtent& extent) { + auto loaded = loader.load( + 0, extent.program_size, SegmentInfo(SegmentInfo::Type::Program)); + if (!loaded.ok()) { + return loaded.error(); + } + auto data = cpu_data(loaded.get(), extent.program_size); + if (!data.ok()) { + return data.error(); + } + if (reinterpret_cast(data.get()) % alignof(std::max_align_t) != + 0) { + return Error::InvalidArgument; + } + flatbuffers::Verifier verifier( + static_cast(data.get()), extent.program_size); + if (!fb::VerifyProgramBuffer(verifier)) { + return Error::InvalidProgram; + } + const auto* program = fb::GetProgram(data.get()); + if (program->version() > runtime::Program::kMaxSupportedSchemaVersion || + runtime::validate_program(program) != Error::Ok) { + return Error::InvalidProgram; + } + return VerifiedProgram{std::move(loaded.get()), program}; +} + +Result load_inline(const fb::Program& program, uint32_t index) { + const auto* buffers = program.backend_delegate_data(); + if (buffers == nullptr) { + return Error::InvalidProgram; + } + if (index >= buffers->size()) { + return Error::NotFound; + } + const auto* bytes = buffers->Get(index)->data(); + if (bytes == nullptr) { + return Error::InvalidProgram; + } + return FreeableBuffer(bytes->data(), bytes->size(), nullptr, nullptr); +} + +Result load_segment( + DataLoader& loader, + const fb::Program& program, + const ProgramExtent& extent, + uint32_t index, + const char* backend_id) { + const auto& header = extent.header; + const auto* segments = program.segments(); + if (!extent.has_header || header.segment_base_offset == 0) { + return Error::NotFound; + } + if (segments == nullptr) { + return Error::InvalidProgram; + } + if (index >= segments->size()) { + return Error::NotFound; + } + const auto* segment = segments->Get(index); + if (segment->offset() > std::numeric_limits::max() || + segment->size() > std::numeric_limits::max()) { + return Error::InvalidProgram; + } + const size_t base = static_cast(header.segment_base_offset); + const size_t offset = static_cast(segment->offset()); + const size_t size = static_cast(segment->size()); + if (base > extent.file_size || offset > extent.file_size - base || + size > extent.file_size - base - offset || + (header.segment_data_size != 0 && + (offset > header.segment_data_size || + size > header.segment_data_size - offset))) { + return Error::InvalidProgram; + } + auto buffer = loader.load( + base + offset, + size, + SegmentInfo(SegmentInfo::Type::Backend, index, backend_id)); + if (!buffer.ok()) { + return buffer.error(); + } + auto bytes = cpu_data(buffer.get(), size); + if (!bytes.ok()) { + return bytes.error(); + } + // Keep the loader's owner while rejecting data beyond the declared extent. + if (buffer->size() != size) { + return Error::InvalidProgram; + } + return std::move(buffer.get()); +} + +Result load_processed( + DataLoader& loader, + const fb::Program& program, + const ProgramExtent& extent, + const fb::BackendDelegate& delegate) { + const auto* reference = delegate.processed(); + if (reference == nullptr) { + return Error::InvalidProgram; + } + switch (reference->location()) { + case fb::DataLocation::INLINE: + return load_inline(program, reference->index()); + case fb::DataLocation::SEGMENT: + return load_segment( + loader, program, extent, reference->index(), delegate.id()->c_str()); + default: + return Error::InvalidProgram; + } +} +} // namespace + +Result inspect_coreai_pte(DataLoader& loader) { + auto extent = parse_header(loader); + if (!extent.ok()) { + return extent.error(); + } + auto loaded = load_program(loader, extent.get()); + if (!loaded.ok()) { + return loaded.error(); + } + const auto* program = loaded->program; + CoreAIPteData result(std::move(loaded->storage)); + for (const auto* plan : *program->execution_plan()) { + if (plan->delegates() == nullptr) { + continue; + } + for (const auto* delegate : *plan->delegates()) { + if (delegate->id() == nullptr) { + continue; + } + const std::string_view id( + delegate->id()->c_str(), delegate->id()->size()); + if (id != kCoreAIBackendId) { + continue; + } + auto buffer = load_processed(loader, *program, extent.get(), *delegate); + if (!buffer.ok()) { + return buffer.error(); + } + result.processed.emplace_back(std::move(buffer.get())); + } + } + return result; +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_pte.h b/backends/apple/coreai/runtime/coreai_pte.h new file mode 100644 index 00000000000..e65bf154b84 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_pte.h @@ -0,0 +1,42 @@ +/* + * 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 +#include +#include + +#include +#include + +namespace executorch::backends::coreai { + +struct CoreAIPteData { + // Member order keeps program_storage alive until processed's inline borrows + // die. + runtime::FreeableBuffer program_storage; + std::vector processed; + + explicit CoreAIPteData(runtime::FreeableBuffer&& storage) + : program_storage(std::move(storage)) {} + CoreAIPteData(CoreAIPteData&&) = default; + CoreAIPteData(const CoreAIPteData&) = delete; + CoreAIPteData& operator=(const CoreAIPteData&) = delete; + CoreAIPteData& operator=(CoreAIPteData&&) = delete; +}; + +// Inspects verified CoreAIBackend processed data in method/delegate order, +// without loading methods, external constants or named assets, or calling the +// SDK. Inline data borrows program_storage; selected segments retain their +// loader callbacks and must match their declared size. Payload contents are not +// parsed. Keep the returned owner and loader alive, and the loader's data +// stable, while using any processed buffer. +runtime::Result inspect_coreai_pte(runtime::DataLoader& loader); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_source_bundle.h b/backends/apple/coreai/runtime/coreai_source_bundle.h new file mode 100644 index 00000000000..6dae0d68447 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_source_bundle.h @@ -0,0 +1,33 @@ +/* + * 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 +#include +#include + +namespace executorch::backends::coreai { + +class LockedCacheEntry; +class SelectedBundle; + +enum class CorruptBundlePolicy { + Reject, + Replace, +}; + +// cache_entry remains locked for this call. Reuse validates the exact tree +// structure and sizes, not file contents. +runtime::Result prepare_source_bundle( + const SelectedBundle& bundle, + const runtime::NamedDataMap* named_data, + const LockedCacheEntry& cache_entry, + CorruptBundlePolicy corrupt_policy = CorruptBundlePolicy::Reject); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_source_bundle.mm b/backends/apple/coreai/runtime/coreai_source_bundle.mm new file mode 100644 index 00000000000..f52f4e2a38c --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_source_bundle.mm @@ -0,0 +1,328 @@ +/* + * 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_source_bundle.h" + +#import "coreai_cache_entry.h" +#import "coreai_fault_injection.h" +#import "coreai_manifest.h" +#import "coreai_storage.h" + +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::coreai { +namespace { +using runtime::Error; +using runtime::Result; + +enum class EntryState { + Missing, + Valid, + Corrupt, +}; + +Error path_error(int error) { + return error == ENOENT || error == ENOTDIR || error == ELOOP + ? Error::InvalidExternalData + : Error::AccessFailed; +} + +Error validate_directory( + int fd, + NSString* relative, + NSDictionary* files, + NSSet* expected_directories, + NSMutableSet* actual_files) { + auto children = storage_children(fd); + if (!children.ok()) { + return children.error(); + } + for (NSString* name in children.get()) { + @autoreleasepool { + NSString* path = [relative stringByAppendingPathComponent:name] + .decomposedStringWithCanonicalMapping; + struct stat info; + if (lstat_at(fd, name.fileSystemRepresentation, &info) != 0) { + return path_error(errno); + } + const bool directory = S_ISDIR(info.st_mode); + ET_CHECK_OR_RETURN_ERROR( + (directory && [expected_directories containsObject:path]) || + (S_ISREG(info.st_mode) && files[path] != nil), + InvalidExternalData, + "Unexpected Core AI bundle entry: %s", + path.UTF8String); + if (directory) { + FileDescriptor child( + open_directory_at(fd, name.fileSystemRepresentation)); + if (child.get() < 0) { + return path_error(errno); + } + ET_CHECK_OK_OR_RETURN_ERROR(validate_directory( + child.get(), path, files, expected_directories, actual_files)); + } else { + ET_CHECK_OR_RETURN_ERROR( + static_cast(info.st_size) == + files[path].unsignedLongLongValue, + InvalidExternalData, + "Core AI asset size does not match manifest"); + [actual_files addObject:path]; + } + } + } + return Error::Ok; +} + +// Filesystem conversion decomposes Unicode without changing NamedDataMap keys. +NSDictionary* normalized_files( + NSDictionary* files) { + NSMutableDictionary* normalized = + [NSMutableDictionary dictionaryWithCapacity:files.count]; + for (NSString* file in files) { + normalized[file.decomposedStringWithCanonicalMapping] = files[file]; + } + return normalized; +} + +Error validate_tree( + int fd, + NSDictionary* selected_files) { + NSDictionary* files = normalized_files(selected_files); + NSMutableSet* expected_directories = [NSMutableSet set]; + for (NSString* file in files) { + NSString* directory = file.stringByDeletingLastPathComponent; + while (directory.length > 0) { + [expected_directories addObject:directory]; + directory = directory.stringByDeletingLastPathComponent; + } + } + NSMutableSet* actual_files = [NSMutableSet set]; + ET_CHECK_OK_OR_RETURN_ERROR( + validate_directory(fd, @"", files, expected_directories, actual_files)); + ET_CHECK_OR_RETURN_ERROR( + [actual_files isEqualToSet:[NSSet setWithArray:files.allKeys]], + InvalidExternalData, + "Core AI bundle file set does not match manifest"); + return Error::Ok; +} + +Result inspect_source_bundle( + int sources_fd, + NSString* key, + NSDictionary* files) { + FileDescriptor entry( + open_directory_at(sources_fd, key.fileSystemRepresentation)); + if (entry.get() < 0) { + if (errno == ENOENT) { + return EntryState::Missing; + } + ET_CHECK_OR_RETURN_ERROR( + errno == ENOTDIR || errno == ELOOP, + AccessFailed, + "Cannot inspect Core AI stored bundle"); + return EntryState::Corrupt; + } + const Error error = validate_tree(entry.get(), files); + if (error == Error::Ok) { + return EntryState::Valid; + } + if (error == Error::InvalidExternalData) { + return EntryState::Corrupt; + } + return error; +} + +// Removes the unpublished build directory on every early return. +struct PendingDirectory { + int parent_fd; + NSString* name; + bool published = false; + + PendingDirectory(int parent_fd, NSString* name) + : parent_fd(parent_fd), name(name) {} + PendingDirectory(const PendingDirectory&) = delete; + PendingDirectory& operator=(const PendingDirectory&) = delete; + PendingDirectory(PendingDirectory&& other) + : parent_fd(other.parent_fd), + name(other.name), + published(other.published) { + other.published = true; + } + + ~PendingDirectory() { + if (published) { + return; + } + const Error error = + remove_storage_directory(parent_fd, name, StorageRemovalPolicy::Strict); + if (error != Error::Ok) { + ET_LOG(Error, "Cannot remove pending Core AI source directory"); + } + } +}; + +Result create_pending_directory(int sources_fd) { + NSString* name = + [@".pending-" stringByAppendingString:NSUUID.UUID.UUIDString]; + ET_CHECK_OR_RETURN_ERROR( + fault_aware_retry( + FaultPoint::CreatePendingDirectory, + [&] { + return mkdirat(sources_fd, name.fileSystemRepresentation, 0777); + }) == 0, + AccessFailed, + "Cannot create pending Core AI source directory"); + return PendingDirectory{sources_fd, [name copy]}; +} + +Result open_pending_directory(int sources_fd, NSString* name) { + FileDescriptor fd(fault_aware_retry(FaultPoint::OpenPendingDirectory, [&] { + return open_directory_at(sources_fd, name.fileSystemRepresentation); + })); + ET_CHECK_OR_RETURN_ERROR( + fd.get() >= 0, + AccessFailed, + "Cannot open pending Core AI source directory"); + return fd; +} + +Error write_pending_bundle( + int pending_fd, + const SelectedBundle& bundle, + const runtime::NamedDataMap& named_data) { + for (NSString* file in bundle.files) { + @autoreleasepool { + NSString* data_key = + [NSString stringWithFormat:@"coreai/%@/%@", bundle.export_hash, file]; + auto buffer = named_data.get_data(data_key.UTF8String); + if (!buffer.ok()) { + ET_CHECK_OR_RETURN_ERROR( + buffer.error() != Error::NotFound, + InvalidExternalData, + "Missing selected Core AI named data: %s", + data_key.UTF8String); + return buffer.error(); + } + ET_CHECK_OR_RETURN_ERROR( + buffer->size() == bundle.files[file].unsignedLongLongValue && + (buffer->size() == 0 || buffer->data() != nullptr), + InvalidExternalData, + "Invalid Core AI named data buffer"); + ET_CHECK_OK_OR_RETURN_ERROR( + write_storage_file(pending_fd, file, buffer->data(), buffer->size())); + } + } + ET_CHECK_OK_OR_RETURN_ERROR(validate_tree(pending_fd, bundle.files)); + return full_sync_storage_directory(pending_fd); +} + +Result publish_pending_directory( + int sources_fd, + PendingDirectory& pending, + NSString* key) { + const int result = fault_aware_retry(FaultPoint::Rename, [&] { + return renameatx_np( + sources_fd, + pending.name.fileSystemRepresentation, + sources_fd, + key.fileSystemRepresentation, + RENAME_EXCL); + }); + if (result == 0) { + pending.published = true; + ET_CHECK_OK_OR_RETURN_ERROR(full_sync_storage_directory(sources_fd)); + return true; + } + ET_CHECK_OR_RETURN_ERROR( + errno == EEXIST, AccessFailed, "Cannot publish Core AI stored bundle"); + return false; +} + +Error recover_lost_publication_race( + int sources_fd, + NSString* key, + NSDictionary* files) { + auto winner = inspect_source_bundle(sources_fd, key, files); + if (!winner.ok()) { + return winner.error(); + } + ET_CHECK_OR_RETURN_ERROR( + winner.get() == EntryState::Valid, + InvalidExternalData, + "Concurrently published Core AI source is incomplete"); + return full_sync_storage_directory(sources_fd); +} + +} // namespace + +Result prepare_source_bundle( + const SelectedBundle& bundle, + const runtime::NamedDataMap* named_data, + const LockedCacheEntry& cache_entry, + CorruptBundlePolicy corrupt_policy) { + NSString* key = cache_entry.key(); + auto sources = cache_entry.prepare_source_directory(); + if (!sources.ok()) { + return sources.error(); + } + NSURL* model_url = [[sources->url URLByAppendingPathComponent:key + isDirectory:YES] + URLByAppendingPathComponent:bundle.bundle_name + isDirectory:YES]; + FileDescriptor sources_fd = std::move(sources->descriptor); + + auto state = inspect_source_bundle(sources_fd.get(), key, bundle.files); + if (!state.ok()) { + return state.error(); + } + if (state.get() == EntryState::Valid) { + ET_CHECK_OK_OR_RETURN_ERROR(full_sync_storage_directory(sources_fd.get())); + return model_url; + } + if (state.get() == EntryState::Corrupt) { + ET_CHECK_OR_RETURN_ERROR( + corrupt_policy == CorruptBundlePolicy::Replace && named_data != nullptr, + InvalidExternalData, + "Cannot replace a damaged Core AI stored bundle"); + ET_CHECK_OK_OR_RETURN_ERROR(remove_storage_directory( + sources_fd.get(), key, StorageRemovalPolicy::Maintenance)); + } + ET_CHECK_OR_RETURN_ERROR( + named_data != nullptr, + InvalidProgram, + "Missing NamedDataMap for Core AI source recovery"); + + auto created = create_pending_directory(sources_fd.get()); + if (!created.ok()) { + return created.error(); + } + PendingDirectory pending = std::move(created.get()); + auto pending_fd = open_pending_directory(sources_fd.get(), pending.name); + if (!pending_fd.ok()) { + return pending_fd.error(); + } + ET_CHECK_OK_OR_RETURN_ERROR( + write_pending_bundle(pending_fd->get(), bundle, *named_data)); + + auto published = publish_pending_directory(sources_fd.get(), pending, key); + if (!published.ok()) { + return published.error(); + } + if (!published.get()) { + ET_CHECK_OK_OR_RETURN_ERROR( + recover_lost_publication_race(sources_fd.get(), key, bundle.files)); + } + return model_url; +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h index 5bcfecfa4f5..b61699b6ff0 100644 --- a/backends/apple/coreai/runtime/coreai_storage.h +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -11,22 +11,32 @@ #import #include #include +#include #include "coreai_file.h" namespace executorch::backends::coreai { +enum class StorageCreationMode { + OpenExisting, + Create, +}; + runtime::Error sync_storage_directory(int fd); -struct StorageDirectory { - int fd; +// Also flushes the drive's write cache, which fsync alone does not on Apple +// platforms. +runtime::Error full_sync_storage_directory(int fd); + +struct StorageDirectory final { + std::optional descriptor; bool created; }; -// Opens one child component without following symlinks. The caller owns fd; -// a missing child with OpenExisting returns {-1, false}. Does not flush -// directories. +// Opens one child component without following symlinks. A missing child with +// OpenExisting has no descriptor. Does not flush directories. runtime::Result open_storage_directory(int parent_fd, NSString* name, StorageCreationMode mode); // Returns an independently owned descriptor for a stable coordination file. -int open_storage_shared_file(int parent_fd, const char* name); +// On failure the descriptor is invalid and errno is set. +FileDescriptor open_storage_shared_file(int parent_fd, const char* name); // Creates a new file and missing parents, then syncs the file and its parent. // On failure, attempts to remove the new file; created parents may remain. runtime::Error write_storage_file( @@ -47,4 +57,26 @@ runtime::Result*> storage_children( int fd, StorageNamePolicy policy = StorageNamePolicy::Strict); +enum class StorageRemovalPolicy { + Strict, + Maintenance, +}; + +// Checks one child tree without following links or changing storage. A missing +// entry counts as removed. Maintenance accepts malformed backend-owned +// non-directory entries. +runtime::Error preflight_storage_directory( + int parent_fd, + NSString* name, + StorageRemovalPolicy policy = StorageRemovalPolicy::Strict); + +// Removes one child directory without following links. The complete tree is +// checked before removal and every entry is checked again during removal. +// A missing entry counts as removed. Maintenance also removes malformed +// backend-owned non-directory entries. +runtime::Error remove_storage_directory( + int parent_fd, + NSString* name, + StorageRemovalPolicy policy = StorageRemovalPolicy::Strict); + } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm index 066c60e17ba..585152356d7 100644 --- a/backends/apple/coreai/runtime/coreai_storage.mm +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -19,29 +19,23 @@ #include #include #include +#include namespace executorch::backends::coreai { namespace { using runtime::Error; using runtime::Result; -template -auto storage_call(FaultPoint point, F call) { - return retry_eintr([&]() -> decltype(call()) { - const int fault = fault_point(point); - if (fault != 0) { - errno = fault; - return -1; - } - return call(); - }); -} - enum class StorageSyncMode { Normal, Full, }; +bool full_fsync(int fd) { + return fault_aware_retry( + FaultPoint::FullSync, [&] { return fcntl(fd, F_FULLFSYNC); }) == 0; +} + Error write_file( int root_fd, NSString* relative_path, @@ -56,7 +50,9 @@ Error write_file( "Invalid Core AI storage write"); for (NSString* component in components) { ET_CHECK_OR_RETURN_ERROR( - storage_component(component), InvalidArgument, "Invalid storage path"); + is_storage_component(component), + InvalidArgument, + "Invalid storage path"); } FileDescriptor parent( retry_eintr([&] { return fcntl(root_fd, F_DUPFD_CLOEXEC, 0); })); @@ -65,11 +61,12 @@ FileDescriptor parent( for (NSUInteger i = 0; i + 1 < components.count; ++i) { auto opened = open_storage_directory( parent.get(), components[i], StorageCreationMode::Create); - if (!opened.ok()) + if (!opened.ok()) { return opened.error(); - FileDescriptor next(opened.get().fd); + } + FileDescriptor next = std::move(*opened->descriptor); ET_CHECK_OK_OR_RETURN_ERROR(sync_storage_directory(parent.get())); - parent.reset(next.release()); + parent = std::move(next); } FileDescriptor file(retry_eintr([&] { return openat( @@ -83,7 +80,7 @@ FileDescriptor file(retry_eintr([&] { const auto result = [&]() -> Error { const auto* bytes = static_cast(data); while (size > 0) { - const ssize_t count = storage_call(FaultPoint::Write, [&] { + const ssize_t count = fault_aware_retry(FaultPoint::Write, [&] { return write(file.get(), bytes, std::min(size, size_t(1024 * 1024))); }); ET_CHECK_OR_RETURN_ERROR( @@ -92,12 +89,9 @@ FileDescriptor file(retry_eintr([&] { size -= count; } ET_CHECK_OR_RETURN_ERROR( - storage_call(FaultPoint::SyncFile, [&] { return fsync(file.get()); }) == - 0 && - (sync_mode != StorageSyncMode::Full || - storage_call( - FaultPoint::FullSync, - [&] { return fcntl(file.get(), F_FULLFSYNC); }) == 0), + fault_aware_retry( + FaultPoint::SyncFile, [&] { return fsync(file.get()); }) == 0 && + (sync_mode != StorageSyncMode::Full || full_fsync(file.get())), AccessFailed, "Cannot synchronize Core AI file"); ET_CHECK_OR_RETURN_ERROR( @@ -119,6 +113,139 @@ FileDescriptor file(retry_eintr([&] { } return result; } + +enum class StorageTreePass { + Preflight, + Remove, +}; + +enum class StorageEntryState { + Missing, + NonDirectory, + Directory, +}; + +struct InspectedStorageEntry final { + StorageEntryState state; + std::optional directory; + struct stat info; +}; + +Result inspect_storage_entry( + int parent_fd, + NSString* name, + StorageRemovalPolicy policy) { + ET_CHECK_OR_RETURN_ERROR( + is_storage_component(name), + InvalidArgument, + "Invalid Core AI entry name"); + struct stat info {}; + const int inspected = + lstat_at(parent_fd, name.fileSystemRepresentation, &info); + if (inspected < 0 && errno == ENOENT) { + return InspectedStorageEntry{ + StorageEntryState::Missing, std::nullopt, info}; + } + ET_CHECK_OR_RETURN_ERROR( + inspected == 0, AccessFailed, "Cannot inspect Core AI storage entry"); + if (!S_ISDIR(info.st_mode)) { + ET_CHECK_OR_RETURN_ERROR( + policy == StorageRemovalPolicy::Maintenance, + InvalidExternalData, + "Core AI storage entry is not a directory"); + return InspectedStorageEntry{ + StorageEntryState::NonDirectory, std::nullopt, info}; + } + FileDescriptor entry( + open_directory_at(parent_fd, name.fileSystemRepresentation)); + if (entry.get() < 0 && errno == ENOENT) { + return InspectedStorageEntry{ + StorageEntryState::Missing, std::nullopt, info}; + } + ET_CHECK_OR_RETURN_ERROR( + entry.get() >= 0, AccessFailed, "Cannot open Core AI storage directory"); + return InspectedStorageEntry{ + StorageEntryState::Directory, + std::optional(std::move(entry)), + info}; +} + +bool same_storage_entry(const struct stat& left, const struct stat& right) { + return left.st_dev == right.st_dev && left.st_ino == right.st_ino && + (left.st_mode & S_IFMT) == (right.st_mode & S_IFMT); +} + +// Removes name only while it is still the entry described by expected. A +// missing entry counts as removed. Does not sync parent_fd. +Error remove_if_unchanged( + int parent_fd, + NSString* name, + const struct stat& expected) { + struct stat current; + const int inspected = + lstat_at(parent_fd, name.fileSystemRepresentation, ¤t); + if (inspected != 0 && errno == ENOENT) { + return Error::Ok; + } + ET_CHECK_OR_RETURN_ERROR( + inspected == 0 && same_storage_entry(expected, current), + AccessFailed, + "Core AI storage entry changed during removal"); + const int removed = fault_aware_retry(FaultPoint::Remove, [&] { + return unlinkat( + parent_fd, + name.fileSystemRepresentation, + S_ISDIR(expected.st_mode) ? AT_REMOVEDIR : 0); + }); + ET_CHECK_OR_RETURN_ERROR( + removed == 0 || errno == ENOENT, + AccessFailed, + "Cannot remove Core AI storage entry"); + return Error::Ok; +} + +Error walk_storage_tree( + int fd, + StorageTreePass pass, + StorageRemovalPolicy policy) { + auto children = storage_children(fd); + if (!children.ok()) { + return children.error(); + } + for (NSString* name in children.get()) { + struct stat info; + const int inspected = lstat_at(fd, name.fileSystemRepresentation, &info); + if (inspected < 0 && errno == ENOENT) { + continue; + } + ET_CHECK_OR_RETURN_ERROR( + inspected == 0, AccessFailed, "Cannot inspect Core AI storage entry"); + const bool directory = S_ISDIR(info.st_mode); + if (policy == StorageRemovalPolicy::Strict) { + ET_CHECK_OR_RETURN_ERROR( + directory || S_ISREG(info.st_mode), + InvalidExternalData, + "Unexpected Core AI storage entry type"); + } + if (directory) { + FileDescriptor child( + open_directory_at(fd, name.fileSystemRepresentation)); + if (child.get() < 0 && errno == ENOENT) { + continue; + } + ET_CHECK_OR_RETURN_ERROR( + child.get() >= 0, + AccessFailed, + "Cannot open Core AI storage directory"); + ET_CHECK_OK_OR_RETURN_ERROR(walk_storage_tree(child.get(), pass, policy)); + } + if (pass == StorageTreePass::Remove) { + ET_CHECK_OK_OR_RETURN_ERROR(remove_if_unchanged(fd, name, info)); + } + } + return pass == StorageTreePass::Remove ? sync_storage_directory(fd) + : Error::Ok; +} } // namespace Result open_storage_directory( @@ -126,7 +253,7 @@ FileDescriptor file(retry_eintr([&] { NSString* name, StorageCreationMode mode) { ET_CHECK_OR_RETURN_ERROR( - storage_component(name), + is_storage_component(name), InvalidArgument, "Invalid storage directory name"); const auto open_directory = [&] { @@ -150,11 +277,12 @@ FileDescriptor file(retry_eintr([&] { created = made == 0; fd = open_directory(); } - if (fd < 0 && errno == ENOENT && !create) - return StorageDirectory{-1, false}; + if (fd < 0 && errno == ENOENT && !create) { + return StorageDirectory{std::nullopt, false}; + } ET_CHECK_OR_RETURN_ERROR( fd >= 0, AccessFailed, "Cannot open Core AI storage directory"); - return StorageDirectory{fd, created}; + return StorageDirectory{FileDescriptor(fd), created}; } Result*> storage_children(int fd, StorageNamePolicy policy) { @@ -175,10 +303,12 @@ FileDescriptor copy(retry_eintr( errno = 0; dirent* entry = readdir(directory); if (entry == nullptr) { - if (errno == EINTR) + if (errno == EINTR) { continue; - if (errno != 0) + } + if (errno != 0) { result = Error::AccessFailed; + } break; } if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0) { @@ -203,7 +333,43 @@ FileDescriptor copy(retry_eintr( return [names sortedArrayUsingSelector:@selector(compare:)]; } -int open_storage_shared_file(int parent_fd, const char* name) { +Error preflight_storage_directory( + int parent_fd, + NSString* name, + StorageRemovalPolicy policy) { + auto entry = inspect_storage_entry(parent_fd, name, policy); + if (!entry.ok()) { + return entry.error(); + } + return entry->state == StorageEntryState::Directory + ? walk_storage_tree( + entry->directory->get(), StorageTreePass::Preflight, policy) + : Error::Ok; +} + +Error remove_storage_directory( + int parent_fd, + NSString* name, + StorageRemovalPolicy policy) { + auto entry = inspect_storage_entry(parent_fd, name, policy); + if (!entry.ok()) { + return entry.error(); + } + if (entry->state == StorageEntryState::Missing) { + return Error::Ok; + } + if (entry->state == StorageEntryState::Directory) { + ET_CHECK_OK_OR_RETURN_ERROR(walk_storage_tree( + entry->directory->get(), StorageTreePass::Preflight, policy)); + ET_CHECK_OK_OR_RETURN_ERROR(walk_storage_tree( + entry->directory->get(), StorageTreePass::Remove, policy)); + } + ET_CHECK_OK_OR_RETURN_ERROR( + remove_if_unchanged(parent_fd, name, entry->info)); + return sync_storage_directory(parent_fd); +} + +FileDescriptor open_storage_shared_file(int parent_fd, const char* name) { // APFS can return ENOENT to a concurrent O_CREAT loser without O_EXCL. int fd = retry_eintr([&] { return openat( @@ -217,17 +383,27 @@ int open_storage_shared_file(int parent_fd, const char* name) { return openat(parent_fd, name, O_RDWR | O_NOFOLLOW | O_CLOEXEC); }); } - return fd; + return FileDescriptor(fd); } Error sync_storage_directory(int fd) { ET_CHECK_OR_RETURN_ERROR( - storage_call(FaultPoint::SyncDirectory, [&] { return fsync(fd); }) == 0, + fault_aware_retry(FaultPoint::SyncDirectory, [&] { return fsync(fd); }) == + 0, AccessFailed, "Cannot synchronize Core AI directory"); return Error::Ok; } +Error full_sync_storage_directory(int fd) { + ET_CHECK_OK_OR_RETURN_ERROR(sync_storage_directory(fd)); + ET_CHECK_OR_RETURN_ERROR( + full_fsync(fd), + AccessFailed, + "Cannot fully synchronize Core AI directory"); + return Error::Ok; +} + Error write_storage_file( int root_fd, NSString* relative_path, @@ -239,7 +415,7 @@ Error write_storage_file( Error publish_storage_data(int root_fd, NSString* name, NSData* data) { ET_CHECK_OR_RETURN_ERROR( - storage_component(name) && data != nil, + is_storage_component(name) && data != nil, InvalidArgument, "Invalid Core AI storage publication"); NSString* temporary = @@ -247,7 +423,7 @@ Error publish_storage_data(int root_fd, NSString* name, NSData* data) { // Bookmarks are small and costly to lose, so they get a full device flush. auto result = write_file( root_fd, temporary, data.bytes, data.length, StorageSyncMode::Full); - if (result == Error::Ok && storage_call(FaultPoint::Rename, [&] { + if (result == Error::Ok && fault_aware_retry(FaultPoint::Rename, [&] { return renameat( root_fd, temporary.fileSystemRepresentation, @@ -260,9 +436,7 @@ Error publish_storage_data(int root_fd, NSString* name, NSData* data) { if (result == Error::Ok) { // Flush the renamed entry as well as the file's pre-rename contents. ET_CHECK_OR_RETURN_ERROR( - storage_call( - FaultPoint::FullSync, - [&] { return fcntl(root_fd, F_FULLFSYNC); }) == 0, + full_fsync(root_fd), AccessFailed, "Cannot fully synchronize Core AI publication directory"); return Error::Ok; diff --git a/backends/apple/coreai/runtime/coreai_tensor_io.h b/backends/apple/coreai/runtime/coreai_tensor_io.h new file mode 100644 index 00000000000..1a49564face --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_tensor_io.h @@ -0,0 +1,22 @@ +/* + * 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 +#include +#import "ETCoreAIBridge.h" + +namespace executorch::backends::coreai { + +runtime::Error execute_session( + id session, + runtime::Span args, + size_t input_count, + size_t output_count); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_tensor_io.mm b/backends/apple/coreai/runtime/coreai_tensor_io.mm new file mode 100644 index 00000000000..d7971c8fc2d --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_tensor_io.mm @@ -0,0 +1,403 @@ +/* + * 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_tensor_io.h" + +#import "coreai_error.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::coreai { +namespace { +using namespace executorch::runtime; +using executorch::aten::ScalarType; +using executorch::aten::SizesType; +using executorch::aten::Tensor; + +struct TensorMetadata final { + ETCoreAIScalarType scalar_type; + size_t bytes; + size_t rank; +}; + +struct ValidatedArguments final { + std::vector tensors; +}; + +struct OutputPlan final { + std::vector shape; + size_t bytes = 0; + void* destination = nullptr; +}; + +Result scalar_type(ScalarType type) { + switch (type) { + case ScalarType::Half: + return ETCoreAIScalarTypeFloat16; + case ScalarType::Float: + return ETCoreAIScalarTypeFloat32; + default: + ET_LOG(Error, "Core AI runtime supports FP16 and FP32 tensors only"); + return Error::NotSupported; + } +} + +Result bridge_scalar_width(ETCoreAIScalarType type) { + switch (type) { + case ETCoreAIScalarTypeFloat16: + return size_t(2); + case ETCoreAIScalarTypeFloat32: + return size_t(4); + default: + return Error::Internal; + } +} + +Result validate_tensor(const Tensor& tensor) { + auto type = scalar_type(tensor.scalar_type()); + if (!type.ok()) { + return type.error(); + } + ET_CHECK_OR_RETURN_ERROR( + tensor.device_type() == executorch::aten::DeviceType::CPU, + NotSupported, + "Core AI requires CPU-accessible tensor storage"); + ET_CHECK_OR_RETURN_ERROR( + tensor.dim() >= 0 && + static_cast(tensor.dim()) <= kTensorDimensionLimit, + InvalidArgument, + "Invalid Core AI tensor rank"); + ET_CHECK_OR_RETURN_ERROR( + tensor_is_default_dim_order(tensor), + NotSupported, + "Core AI requires default tensor dimension order"); + size_t elements = 1; + size_t expected_stride = 1; + for (size_t i = tensor.sizes().size(); i > 0; --i) { + const auto size = tensor.sizes()[i - 1]; + ET_CHECK_OR_RETURN_ERROR( + size >= 0, InvalidArgument, "Negative Core AI tensor dimension"); + ET_CHECK_OR_RETURN_ERROR( + tensor.strides()[i - 1] >= 0 && + static_cast(tensor.strides()[i - 1]) == expected_stride, + NotSupported, + "Core AI requires contiguous tensor storage"); + ET_CHECK_OR_RETURN_ERROR( + !c10::mul_overflows(elements, static_cast(size), &elements) && + !c10::mul_overflows( + expected_stride, + std::max(static_cast(size), size_t(1)), + &expected_stride), + InvalidArgument, + "Core AI tensor size overflow"); + } + size_t bytes = 0; + ET_CHECK_OR_RETURN_ERROR( + !c10::mul_overflows( + elements, static_cast(tensor.element_size()), &bytes) && + bytes <= static_cast(std::numeric_limits::max()) && + bytes == tensor.nbytes(), + InvalidArgument, + "Invalid Core AI tensor byte count"); + return TensorMetadata{type.get(), bytes, tensor.sizes().size()}; +} + +Result validate_arguments( + Span args, + size_t input_count, + size_t output_count) { + ET_CHECK_OR_RETURN_ERROR( + args.size() >= input_count && args.size() - input_count == output_count, + InvalidArgument, + "Core AI argument count mismatch"); + ValidatedArguments validated; + validated.tensors.reserve(args.size()); + for (EValue* arg : args) { + ET_CHECK_OR_RETURN_ERROR( + arg && arg->isTensor(), + InvalidArgument, + "Core AI arguments must be tensors"); + auto metadata = validate_tensor(arg->toTensor()); + if (!metadata.ok()) { + return metadata.error(); + } + validated.tensors.push_back(metadata.get()); + } + return validated; +} + +Result*> borrow_inputs( + Span args, + const ValidatedArguments& validated, + size_t input_count) { + NSMutableArray* inputs = + [NSMutableArray arrayWithCapacity:input_count]; + for (size_t i = 0; i < input_count; ++i) { + const auto& tensor = args[i]->toTensor(); + const auto& metadata = validated.tensors[i]; + ET_CHECK_OR_RETURN_ERROR( + metadata.bytes == 0 || tensor.const_data_ptr() != nullptr, + InvalidArgument, + "Core AI input has no storage"); + NSMutableArray* shape = + [NSMutableArray arrayWithCapacity:tensor.dim()]; + for (auto size : tensor.sizes()) { + [shape addObject:@(size)]; + } + [inputs addObject:[[ETCoreAIInputTensor alloc] + initWithBytes:tensor.const_data_ptr() + byteCount:metadata.bytes + shape:shape + scalarType:metadata.scalar_type]]; + } + return inputs; +} + +Result*> wait_for_outputs( + id session, + NSArray* inputs) { + dispatch_semaphore_t ready = dispatch_semaphore_create(0); + __block NSArray* outputs = nil; + __block NSError* error = nil; + [session executeInputs:inputs + completion:^(NSArray* result, NSError* failure) { + outputs = result; + error = failure; + dispatch_semaphore_signal(ready); + }]; + // Returning earlier would let the SDK outlive borrowed input storage. + dispatch_semaphore_wait(ready, DISPATCH_TIME_FOREVER); + NSArray* completed = outputs; + NSError* failure = error; + // The SDK may retain the completion block after invoking it. + outputs = nil; + error = nil; + if (failure != nil) { + return execution_error(failure); + } + return completed; +} + +Result> validate_bridge_outputs( + NSArray* outputs, + const ValidatedArguments& validated, + size_t input_count, + size_t output_count) { + ET_CHECK_OR_RETURN_ERROR( + outputs != nil && outputs.count == output_count, + Internal, + "Core AI output count mismatch"); + std::vector plans(output_count); + for (size_t i = 0; i < output_count; ++i) { + ETCoreAITensor* output = outputs[i]; + ET_CHECK_OR_RETURN_ERROR( + [output isKindOfClass:ETCoreAITensor.class], + Internal, + "Invalid Core AI output tensor"); + auto width = bridge_scalar_width(output.scalarType); + if (!width.ok()) { + return width.error(); + } + const auto& expected = validated.tensors[input_count + i]; + ET_CHECK_OR_RETURN_ERROR( + output.scalarType == expected.scalar_type && + output.shape.count == expected.rank, + InvalidArgument, + "Core AI output type or rank mismatch"); + OutputPlan& plan = plans[i]; + plan.bytes = width.get(); + plan.shape.reserve(output.shape.count); + for (id size in output.shape) { + ET_CHECK_OR_RETURN_ERROR( + [size isKindOfClass:NSNumber.class] && [size longLongValue] >= 0 && + [size longLongValue] <= std::numeric_limits::max() && + [size isEqualToNumber:@([size longLongValue])], + Internal, + "Invalid Core AI output dimension"); + const auto dimension = static_cast([size longLongValue]); + plan.shape.push_back(dimension); + ET_CHECK_OR_RETURN_ERROR( + !c10::mul_overflows( + plan.bytes, static_cast(dimension), &plan.bytes), + Internal, + "Core AI output size overflow"); + } + ET_CHECK_OR_RETURN_ERROR( + plan.bytes == output.byteCount, + Internal, + "Core AI output byte count mismatch"); + } + return plans; +} + +Error resize_outputs( + std::vector& plans, + Span args, + size_t input_count) { + for (size_t i = 0; i < plans.size(); ++i) { + auto& tensor = args[input_count + i]->toTensor(); + auto& shape = plans[i].shape; + const Error resized = + resize_tensor(tensor, ArrayRef(shape.data(), shape.size())); + if (resized != Error::Ok) { + return resized == Error::NotSupported ? Error::InvalidArgument : resized; + } + } + return Error::Ok; +} + +Error validate_resized_outputs( + std::vector& plans, + Span args, + const ValidatedArguments& validated, + size_t input_count) { + for (size_t i = 0; i < plans.size(); ++i) { + auto& tensor = args[input_count + i]->toTensor(); + auto& plan = plans[i]; + ET_CHECK_OR_RETURN_ERROR( + tensor.dim() >= 0 && + static_cast(tensor.dim()) == plan.shape.size() && + std::equal( + plan.shape.begin(), plan.shape.end(), tensor.sizes().begin()), + InvalidArgument, + "Core AI output shape changed during resize"); + auto metadata = validate_tensor(tensor); + if (!metadata.ok()) { + return metadata.error(); + } + ET_CHECK_OR_RETURN_ERROR( + metadata->scalar_type == + validated.tensors[input_count + i].scalar_type && + metadata->bytes == plan.bytes && + (plan.bytes == 0 || tensor.mutable_data_ptr() != nullptr), + InvalidArgument, + "Core AI output has insufficient storage"); + plan.destination = tensor.mutable_data_ptr(); + } + return Error::Ok; +} + +bool storage_overlaps( + const void* left, + size_t left_size, + const void* right, + size_t right_size) { + if (left_size == 0 || right_size == 0) { + return false; + } + const auto a = reinterpret_cast(left); + const auto b = reinterpret_cast(right); + return a <= b ? b - a < left_size : a - b < right_size; +} + +bool outputs_need_staging( + const std::vector& plans, + NSArray* inputs) { + for (size_t i = 0; i < plans.size(); ++i) { + const auto& plan = plans[i]; + for (ETCoreAIInputTensor* input in inputs) { + if (storage_overlaps( + plan.destination, plan.bytes, input.bytes, input.byteCount)) { + return true; + } + } + for (size_t j = 0; j < i; ++j) { + if (storage_overlaps( + plan.destination, + plan.bytes, + plans[j].destination, + plans[j].bytes)) { + return true; + } + } + } + return false; +} + +Error copy_direct( + NSArray* outputs, + const std::vector& plans) { + for (size_t i = 0; i < plans.size(); ++i) { + [outputs[i] copyBytesTo:plans[i].destination]; + } + return Error::Ok; +} + +Error copy_staged( + NSArray* outputs, + const std::vector& plans) { + size_t total_bytes = 0; + for (const auto& plan : plans) { + ET_CHECK_OR_RETURN_ERROR( + plan.bytes <= std::numeric_limits::max() - total_bytes, + MemoryAllocationFailed, + "Core AI output staging size overflow"); + total_bytes += plan.bytes; + } + std::unique_ptr staging(new (std::nothrow) uint8_t[total_bytes]); + ET_CHECK_OR_RETURN_ERROR( + staging != nullptr, + MemoryAllocationFailed, + "Cannot stage Core AI outputs"); + size_t offset = 0; + for (size_t i = 0; i < plans.size(); ++i) { + [outputs[i] copyBytesTo:staging.get() + offset]; + offset += plans[i].bytes; + } + offset = 0; + for (const auto& plan : plans) { + if (plan.bytes > 0) { + memcpy(plan.destination, staging.get() + offset, plan.bytes); + } + offset += plan.bytes; + } + return Error::Ok; +} + +} // 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"); + auto validated = validate_arguments(args, input_count, output_count); + if (!validated.ok()) { + return validated.error(); + } + auto inputs = borrow_inputs(args, validated.get(), input_count); + if (!inputs.ok()) { + return inputs.error(); + } + auto outputs = wait_for_outputs(session, inputs.get()); + if (!outputs.ok()) { + return outputs.error(); + } + auto plans = validate_bridge_outputs( + outputs.get(), validated.get(), input_count, output_count); + if (!plans.ok()) { + return plans.error(); + } + ET_CHECK_OK_OR_RETURN_ERROR(resize_outputs(plans.get(), args, input_count)); + // Finish every resize before copying so a late failure cannot write results. + ET_CHECK_OK_OR_RETURN_ERROR(validate_resized_outputs( + plans.get(), args, validated.get(), input_count)); + return outputs_need_staging(plans.get(), inputs.get()) + ? copy_staged(outputs.get(), plans.get()) + : copy_direct(outputs.get(), plans.get()); +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_validation.h b/backends/apple/coreai/runtime/coreai_validation.h index ea21885b19d..14f5e9fd63e 100644 --- a/backends/apple/coreai/runtime/coreai_validation.h +++ b/backends/apple/coreai/runtime/coreai_validation.h @@ -13,15 +13,15 @@ namespace executorch::backends::coreai { // Nonempty text without embedded NULs. -bool nonempty_string(id value); +bool is_nonempty_string(id value); // Relative bundle path without escapes, backslashes, or empty components. -bool relative_path(NSString* path); +bool is_relative_path(NSString* path); // Single path component for an fd-relative open: no slashes, ".", or "..". -// Unlike relative_path, backslashes are ordinary bytes on this filesystem. -bool storage_component(NSString* name); -// Lowercase hex of any positive length, as in the exporter model hash. -bool lowercase_hex(NSString* value); +// Unlike is_relative_path, backslashes are ordinary bytes on this filesystem. +bool is_storage_component(NSString* name); +// Lowercase hex of any positive length, as in the export hash. +bool is_lowercase_hex(NSString* value); // Lowercase hex of exactly length characters, as in bundle digests. -bool lowercase_hex(NSString* value, NSUInteger length); +bool is_lowercase_hex(NSString* value, NSUInteger length); } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_validation.mm b/backends/apple/coreai/runtime/coreai_validation.mm index 54c611b81bc..4b635772561 100644 --- a/backends/apple/coreai/runtime/coreai_validation.mm +++ b/backends/apple/coreai/runtime/coreai_validation.mm @@ -19,13 +19,13 @@ } // namespace -bool nonempty_string(id value) { +bool is_nonempty_string(id value) { return [value isKindOfClass:NSString.class] && [value length] > 0 && [value rangeOfString:@"\0"].location == NSNotFound; } -bool relative_path(NSString* path) { - if (!nonempty_string(path) || path.isAbsolutePath || +bool is_relative_path(NSString* path) { + if (!is_nonempty_string(path) || path.isAbsolutePath || [path containsString:@"\\"]) { return false; } @@ -38,19 +38,19 @@ bool relative_path(NSString* path) { return true; } -bool storage_component(NSString* name) { - return nonempty_string(name) && ![name isEqualToString:@"."] && +bool is_storage_component(NSString* name) { + return is_nonempty_string(name) && ![name isEqualToString:@"."] && ![name isEqualToString:@".."] && ![name containsString:@"/"]; } -bool lowercase_hex(NSString* value) { - return nonempty_string(value) && +bool is_lowercase_hex(NSString* value) { + return is_nonempty_string(value) && [value rangeOfCharacterFromSet:invalid_hex_characters()].location == NSNotFound; } -bool lowercase_hex(NSString* value, NSUInteger length) { - return lowercase_hex(value) && [value length] == length; +bool is_lowercase_hex(NSString* value, NSUInteger length) { + return is_lowercase_hex(value) && [value length] == length; } } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift b/backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift new file mode 100644 index 00000000000..b2c24401297 --- /dev/null +++ b/backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift @@ -0,0 +1,180 @@ +// 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 Foundation +import Synchronization + +private struct Fixture: Decodable { + let function: String + let inputs: [String] + let outputs: [String] + let expected: [[Float]] +} + +private enum Failure: Error { + case invalidArguments + case invalidResult + case incorrectOutput(String) + case unexpectedSource +} + +// Bridge results are immutable and consumed only after completion. +private struct BridgeResult: @unchecked Sendable { + let value: Result +} + +// The executable waits on an ordinary thread, not Swift's cooperative pool. +private func waitForResult( + _ operation: (@escaping @Sendable (Result) -> Void) -> Void +) throws -> Value { + let result = Mutex?>(nil) + let ready = DispatchSemaphore(value: 0) + operation { value in + result.withLock { $0 = BridgeResult(value: value) } + ready.signal() + } + ready.wait() + let value = result.withLock { value in + let snapshot = value + value = nil + return snapshot + } + guard let value else { throw Failure.invalidResult } + return try value.value.get() +} + +private func acquire( + _ operation: (@escaping ETCoreAIAcquisitionCompletion) -> Void +) throws -> (any ETCoreAIPreparedModel)? { + try waitForResult { completed in + operation { status, model, error in + if let error { + completed(.failure(error)) + } else if (status == .hit && model != nil) || (status == .miss && model == nil) { + completed(.success(model)) + } else { + completed(.failure(Failure.invalidResult)) + } + } + } +} + +private func checkInference(_ model: any ETCoreAIPreparedModel, fixture: Fixture) throws { + guard fixture.inputs.count == 2, fixture.outputs.count == fixture.expected.count else { + throw Failure.invalidArguments + } + let session: any ETCoreAISession = try waitForResult { completed in + model.loadFunctionNamed( + fixture.function, inputNames: fixture.inputs, outputNames: fixture.outputs + ) { session, error in + if let error { + completed(.failure(error)) + } else if let session { + completed(.success(session)) + } else { + completed(.failure(Failure.invalidResult)) + } + } + } + let storage = [[Float(5), 9], [Float(1), 2]].map { values in + values.withUnsafeBytes { NSData(bytes: $0.baseAddress, length: $0.count) } + } + let inputs = storage.map { + ETCoreAIInputTensor( + bytes: $0.bytes, byteCount: UInt($0.length), shape: [2], scalarType: .float32) + } + try withExtendedLifetime(storage) { + let outputs: [ETCoreAITensor] = try waitForResult { completed in + session.executeInputs(inputs) { outputs, error in + if let error { + completed(.failure(error)) + } else if let outputs { + completed(.success(outputs)) + } else { + completed(.failure(Failure.invalidResult)) + } + } + } + guard outputs.count == fixture.outputs.count else { + throw Failure.incorrectOutput("count") + } + for (index, output) in outputs.enumerated() { + let name = fixture.outputs[index] + guard output.shape == [2], output.scalarType == .float32, + output.byteCount == 8, fixture.expected[index].count == 2 + else { throw Failure.incorrectOutput(name) } + var values = [Float](repeating: 0, count: 2) + values.withUnsafeMutableBytes { output.copyBytes(to: $0.baseAddress) } + guard zip(values, fixture.expected[index]).allSatisfy({ abs($0 - $1) < 0.001 }) else { + throw Failure.incorrectOutput("\(name): \(values), expected \(fixture.expected[index])") + } + } + } +} + +@main +private struct SDKSourceLifetimeTest { + static func main() { + do { + let args = CommandLine.arguments + if args.count == 2, args[1] == "architecture" { + print(ETCoreAISwiftModelLoader.deviceArchitectureName) + return + } + guard args.count == 5 else { throw Failure.invalidArguments } + let mode = args[1] + let source = URL(fileURLWithPath: args[2]).standardizedFileURL + let fixtureURL = URL(fileURLWithPath: args[3]).standardizedFileURL + let bookmarkURL = URL(fileURLWithPath: args[4]) + let directory = source.deletingLastPathComponent() + guard ["model.aimodel", "model.aimodelc"].contains(source.lastPathComponent), + directory == fixtureURL.deletingLastPathComponent(), + directory.deletingLastPathComponent().lastPathComponent.hasPrefix("coreai-lifetime-") + else { throw Failure.invalidArguments } + let loader = ETCoreAISwiftModelLoader() + if mode == "evict" { + let bookmark = try Data(contentsOf: bookmarkURL) + let _: Void = try waitForResult { completed in + loader.evictModel(withBookmark: bookmark) { error in + if let error { completed(.failure(error)) } else { completed(.success(())) } + } + } + return + } + let fixture = try JSONDecoder().decode(Fixture.self, from: Data(contentsOf: fixtureURL)) + if mode == "specialize" { + guard let model = try acquire({ loader.specializeModel(at: source, completion: $0) }) else { + throw Failure.invalidResult + } + try model.copyBookmarkData().write(to: bookmarkURL, options: .atomic) + try checkInference(model, fixture: fixture) + try FileManager.default.removeItem(at: source) + guard !FileManager.default.fileExists(atPath: source.path) else { + throw Failure.unexpectedSource + } + try checkInference(model, fixture: fixture) + print("Bridge inference passed after source removal") + } else if mode == "restore" { + guard !FileManager.default.fileExists(atPath: source.path) else { + throw Failure.unexpectedSource + } + let bookmark = try Data(contentsOf: bookmarkURL) + guard + let model = try acquire({ loader.restoreModel(fromBookmark: bookmark, completion: $0) }) + else { + exit(3) + } + try checkInference(model, fixture: fixture) + print("Bridge restoration and inference passed without source") + } else { + throw Failure.invalidArguments + } + } catch { + print("SDK source-lifetime test failed: \(error)") + exit(1) + } + } +} diff --git a/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift new file mode 100644 index 00000000000..da3dce548a5 --- /dev/null +++ b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift @@ -0,0 +1,754 @@ +// 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 +import Foundation +import Synchronization +import XCTest + +private func copiedData(_ output: ETCoreAITensor) -> Data { + var data = Data(count: Int(output.byteCount)) + data.withUnsafeMutableBytes { output.copyBytes(to: $0.baseAddress) } + return data +} + +@available(macOS 27.0, iOS 27.0, *) +private func copyOutput(_ array: NDArray) throws -> Data { + copiedData(try prepareOutput(array)) +} + +private final class ProtocolPreparedModel: NSObject, ETCoreAIPreparedModel, Sendable { + let released: XCTestExpectation + + init(released: XCTestExpectation) { + self.released = released + super.init() + } + + deinit { released.fulfill() } + + func copyBookmarkData() -> Data { Data([1, 2, 3]) } + + func loadFunctionNamed( + _ functionName: String, inputNames: [String], outputNames: [String], + completion: @escaping ((any ETCoreAISession)?, Error?) -> Void + ) { + completion( + nil, NSError(domain: ETCoreAIErrorDomain, code: ETCoreAIErrorCode.invalidModel.rawValue)) + } +} + +// These contract tests never invoke the SDK; the linked test bundle still requires OS 27. +final class SwiftAcquisitionTests: XCTestCase { + func testHitPinsPreparedModelThroughCompletion() { + let completed = expectation(description: "acquired") + let released = expectation(description: "prepared model released") + completed.assertForOverFulfill = true + released.assertForOverFulfill = true + acquireCoreAIModel(context: "test hit") { status, model, error in + XCTAssertFalse(Thread.isMainThread) + XCTAssertEqual(status, .hit) + XCTAssertNil(error) + XCTAssertEqual(model?.copyBookmarkData(), Data([1, 2, 3])) + completed.fulfill() + } operation: { + ProtocolPreparedModel(released: released) + } + wait(for: [completed, released], timeout: 5, enforceOrder: true) + } + + func testOnlyNilIsMiss() { + let completed = expectation(description: "miss") + completed.assertForOverFulfill = true + acquireCoreAIModel(context: "test miss") { status, model, error in + XCTAssertFalse(Thread.isMainThread) + XCTAssertEqual(status, .miss) + XCTAssertNil(model) + XCTAssertNil(error) + completed.fulfill() + } operation: { + nil + } + wait(for: [completed], timeout: 5) + } + + func testErrorsAreWrappedOnceAndPreserveTheUnderlyingError() { + let sdk = NSError(domain: "SDK.test", code: 91, userInfo: ["detail": "opaque"]) + let bridge = NSError( + domain: ETCoreAIErrorDomain, code: ETCoreAIErrorCode.invalidArgument.rawValue) + for underlying in [sdk, bridge] { + let completed = expectation(description: underlying.domain) + completed.assertForOverFulfill = true + acquireCoreAIModel(context: "test restoration") { status, model, error in + XCTAssertFalse(Thread.isMainThread) + XCTAssertEqual(status, .error) + XCTAssertNil(model) + let error = error as NSError? + if underlying === bridge { + XCTAssertTrue(error === bridge) + } else { + XCTAssertEqual(error?.domain, ETCoreAIErrorDomain) + XCTAssertEqual(error?.code, ETCoreAIErrorCode.runtime.rawValue) + XCTAssertTrue(error?.localizedDescription.hasPrefix("test restoration:") == true) + XCTAssertTrue((error?.userInfo[NSUnderlyingErrorKey] as? NSError) === sdk) + } + completed.fulfill() + } operation: { + throw underlying + } + wait(for: [completed], timeout: 5) + } + } +} + +@available(macOS 27.0, iOS 27.0, *) +final class SwiftSpecializationTests: XCTestCase { + func testCachePolicyAllowsOnlyStoragePressure() { + XCTAssertEqual(ETCoreAISwiftModelLoader.cachePolicy.purgeConditions, [.storagePressure]) + } + func testDefaultPreservesIdentifier() throws { + let snapshot = try CoreAISpecialization(options: .default) + XCTAssertEqual(snapshot.options, .default) + XCTAssertEqual(snapshot.identifier, "default") + XCTAssertEqual(ETCoreAISwiftModelLoader().specializationIdentifier, snapshot.identifier) + } + + func testAllOptionFieldsParticipateInIdentity() throws { + let choices: [SpecializationOptions] = [ + .default, .cpuOnly, + SpecializationOptions(preferredComputeUnitKind: .cpu), + SpecializationOptions(preferredComputeUnitKind: .gpu), + SpecializationOptions(preferredComputeUnitKind: .neuralEngine), + ] + let options = + choices + + choices.map { original in + var changed = original + changed.expectFrequentReshapes.toggle() + return changed + } + let snapshots = try options.map { try CoreAISpecialization(options: $0) } + for first in snapshots { + XCTAssertEqual(first.identifier, try CoreAISpecialization(options: first.options).identifier) + XCTAssertEqual( + ETCoreAISwiftModelLoader(specialization: first).specializationIdentifier, first.identifier) + for second in snapshots { + XCTAssertEqual(first.options == second.options, first.identifier == second.identifier) + } + } + } + + func testSnapshotDoesNotFollowSubsequentOptionMutation() throws { + var options = SpecializationOptions.cpuOnly + let original = options + let snapshot = try CoreAISpecialization(options: options) + let loader = ETCoreAISwiftModelLoader(specialization: snapshot) + options.expectFrequentReshapes.toggle() + XCTAssertEqual(snapshot.options, original) + XCTAssertEqual(loader.specializationIdentifier, snapshot.identifier) + XCTAssertNotEqual( + loader.specializationIdentifier, try CoreAISpecialization(options: options).identifier) + } + + func testDescriptorIsVersionedAndUsesStableComputeUnitNames() throws { + var options = SpecializationOptions(preferredComputeUnitKind: .gpu) + options.expectFrequentReshapes = !SpecializationOptions.default.expectFrequentReshapes + let snapshot = try CoreAISpecialization(options: options) + let json = try XCTUnwrap( + JSONSerialization.jsonObject(with: Data(snapshot.identifier.utf8)) as? [String: Any]) + XCTAssertEqual(json["version"] as? Int, 1) + XCTAssertEqual(json["preferredComputeUnit"] as? String, "gpu") + XCTAssertEqual(json["expectFrequentReshapes"] as? Bool, options.expectFrequentReshapes) + let allowed = try XCTUnwrap(json["allowedComputeUnits"] as? [String]) + XCTAssertEqual(allowed, allowed.sorted()) + XCTAssertTrue(Set(allowed).isSubset(of: ["cpu", "gpu", "neuralEngine"])) + XCTAssertEqual(allowed.count, options.allowedComputeUnitKinds.count) + } +} + +@available(macOS 27.0, iOS 27.0, *) +final class SwiftSessionTests: XCTestCase { + private struct Manifest: Decodable { + let function: String + let inputNames: [String] + let outputNames: [String] + + enum CodingKeys: String, CodingKey { + case function + case inputNames = "input_names" + case outputNames = "output_names" + } + } + + private struct Input: @unchecked Sendable { + let storage: NSData + let tensor: ETCoreAIInputTensor + + init(_ data: Data, shape: [Int], type: ETCoreAIScalarType = .float32) { + storage = data as NSData + tensor = ETCoreAIInputTensor( + bytes: storage.bytes, byteCount: UInt(storage.length), + shape: shape.map(NSNumber.init(value:)), scalarType: type) + } + + init(_ values: [Float]) { + self.init(values.withUnsafeBytes { Data($0) }, shape: [values.count]) + } + } + + private var loader = ETCoreAISwiftModelLoader() + private var manifest: Manifest! + private var fixtureDirectory: URL! + private var model: (any ETCoreAIPreparedModel)? + private var session: (any ETCoreAISession)? + private var bookmark: Data? + + override func setUpWithError() throws { + guard let root = ProcessInfo.processInfo.environment["COREAI_BRIDGE_FIXTURE_DIR"] else { + throw XCTSkip("COREAI_BRIDGE_FIXTURE_DIR is required for SDK session tests") + } + try useFixture(at: URL(fileURLWithPath: root, isDirectory: true)) + } + + private func useFixture(at directory: URL) throws { + fixtureDirectory = directory + manifest = try JSONDecoder().decode( + Manifest.self, from: Data(contentsOf: directory.appendingPathComponent("manifest.json"))) + XCTAssertEqual(manifest.inputNames.count, 2) + XCTAssertEqual(manifest.outputNames.count, 2) + } + + func testSpecializationUsesTheSnapshotOptions() throws { + let snapshot = try CoreAISpecialization(options: .cpuOnly) + loader = ETCoreAISwiftModelLoader(specialization: snapshot) + try prepareModel() + XCTAssertNotNil( + try AIModelCache.default.model( + for: fixtureDirectory.appendingPathComponent("model.aimodel"), options: snapshot.options)) + } + + private func prepareModel() throws { + guard model == nil else { return } + let completed = expectation(description: "specialized session fixture") + completed.assertForOverFulfill = true + let result = Mutex<(model: (any ETCoreAIPreparedModel)?, error: Error?)>((nil, nil)) + loader.specializeModel(at: fixtureDirectory.appendingPathComponent("model.aimodel")) { + status, model, error in + XCTAssertEqual(status, .hit) + result.withLock { $0 = (model, error) } + completed.fulfill() + } + wait(for: [completed], timeout: 30) + let value = result.withLock { value in + let snapshot = value + value = (nil, nil) + return snapshot + } + if let failure = value.error { throw failure } + model = try XCTUnwrap(value.model) + bookmark = model?.copyBookmarkData() + } + + override func tearDownWithError() throws { + session = nil + model = nil + if let bookmark { + let completed = expectation(description: "evicted session fixture") + completed.assertForOverFulfill = true + loader.evictModel(withBookmark: bookmark) { error in + XCTAssertNil(error) + completed.fulfill() + } + wait(for: [completed], timeout: 10) + } + bookmark = nil + } + + private func bind(inputNames: [String]? = nil, outputNames: [String]? = nil) throws { + try prepareModel() + let result = Mutex<(session: (any ETCoreAISession)?, error: Error?, calls: Int)>((nil, nil, 0)) + try XCTUnwrap(model).loadFunctionNamed( + manifest.function, inputNames: inputNames ?? manifest.inputNames, + outputNames: outputNames ?? manifest.outputNames + ) { session, error in + result.withLock { $0 = (session, error, $0.calls + 1) } + } + let value = result.withLock { value in + let snapshot = value + value = (nil, nil, 0) + return snapshot + } + XCTAssertEqual(value.calls, 1) + if let failure = value.error { throw failure } + session = try XCTUnwrap(value.session) + } + + private func execute(_ inputs: [Input]) throws -> [ETCoreAITensor] { + let completed = expectation(description: "executed session fixture") + completed.assertForOverFulfill = true + let result = Mutex<(outputs: [ETCoreAITensor]?, error: Error?)>((nil, nil)) + try XCTUnwrap(session).executeInputs(inputs.map(\.tensor)) { tensors, error in + let owned = tensors?.map { + ETCoreAITensor(data: copiedData($0), shape: $0.shape, scalarType: $0.scalarType) + } + result.withLock { $0 = (owned, error) } + // Retain the borrowed allocations even if the XCTest wait times out. + withExtendedLifetime(inputs) { completed.fulfill() } + } + wait(for: [completed], timeout: 10) + let value = result.withLock { value in + let snapshot = value + value = (nil, nil) + return snapshot + } + if let failure = value.error { throw failure } + return try XCTUnwrap(value.outputs) + } + + func testExecutesWithOrderedAndReversedBindings() throws { + for reversed in [false, true] { + try bind( + inputNames: reversed ? Array(manifest.inputNames.reversed()) : manifest.inputNames, + outputNames: reversed ? Array(manifest.outputNames.reversed()) : manifest.outputNames) + let inputs = [Input([5, 9]), Input([1, 2])] + let outputs = try execute(reversed ? Array(inputs.reversed()) : inputs) + let expected: [[Float]] = reversed ? [[4, 7], [6, 11]] : [[6, 11], [4, 7]] + XCTAssertEqual(outputs.count, expected.count) + for (output, values) in zip(outputs, expected) { + XCTAssertEqual(output.scalarType, .float32) + XCTAssertEqual(output.shape, [2]) + XCTAssertEqual(copiedData(output), values.withUnsafeBytes { Data($0) }) + } + } + } + + func testDynamicDescriptorsTrackRuntimeShapes() throws { + try useFixture(at: fixtureDirectory.appendingPathComponent("dynamic")) + try bind() + let sdkModel = try XCTUnwrap(AIModel(resolvingBookmark: XCTUnwrap(bookmark))) + let function = try XCTUnwrap(sdkModel.loadFunction(named: manifest.function)) + for name in manifest.inputNames { + guard case .ndArray(let descriptor) = function.descriptor.inputDescriptor(of: name) else { + return XCTFail("Expected an NDArray input descriptor") + } + XCTAssertEqual(descriptor.shape, [-1]) + } + for count in [2, 4, 1] { + let left = (0..(_ values: [T], scalarType: ETCoreAIScalarType) throws { + try values.withUnsafeBytes { storage in + try assertBorrowedInput( + UnsafeRawBufferPointer(start: storage.baseAddress, count: 6 * MemoryLayout.stride), + shape: [2, 3], scalarType: scalarType) + try assertBorrowedInput( + UnsafeRawBufferPointer(start: storage.baseAddress, count: MemoryLayout.stride), + shape: [], scalarType: scalarType) + try assertBorrowedInput( + UnsafeRawBufferPointer(start: nil, count: 0), shape: [2, 0, 3], scalarType: scalarType) + // A dynamic-shape subrange starts mid-allocation and tracks the current rows. + for rows in [4, 0] { + try assertBorrowedInput( + UnsafeRawBufferPointer( + start: storage.baseAddress!.advanced(by: 2 * MemoryLayout.stride), + count: rows * 2 * MemoryLayout.stride), + shape: [rows, 2], scalarType: scalarType) + } + } + } + + private func assertBorrowedInput( + _ source: UnsafeRawBufferPointer, shape: [Int], scalarType: ETCoreAIScalarType, + file: StaticString = #filePath, line: UInt = #line + ) throws { + let payload = try CoreAITensorPayload( + ETCoreAIInputTensor( + bytes: source.baseAddress, byteCount: UInt(source.count), + shape: shape.map(NSNumber.init(value:)), scalarType: scalarType)) + XCTAssertEqual(payload.bytes, source.baseAddress, file: file, line: line) + XCTAssertEqual(payload.byteCount, source.count, file: file, line: line) + XCTAssertEqual(payload.shape, shape, file: file, line: line) + XCTAssertEqual(payload.scalarType, scalarType, file: file, line: line) + // The source's withUnsafeBytes scope is still active for every view access. + let bytes = RawSpan( + _unsafeBytes: UnsafeRawBufferPointer( + start: payload.bytes, count: payload.byteCount)) + let expectedType = try CoreAIScalarInfo.from(bridgeType: scalarType).sdkType + let view = NDArray.RawView( + bytes: bytes, scalarType: expectedType, shape: payload.shape) + view.bytes.withUnsafeBytes { borrowed in + XCTAssertEqual(borrowed.baseAddress, source.baseAddress, file: file, line: line) + XCTAssertEqual(borrowed.count, source.count, file: file, line: line) + } + let viewShape = view.shape + XCTAssertEqual(viewShape.count, shape.count, file: file, line: line) + for axis in shape.indices { + XCTAssertEqual(viewShape[axis], shape[axis], file: file, line: line) + } + XCTAssertEqual(view.scalarType, expectedType, file: file, line: line) + } + + func testPreparedOutputsRetainSDKStorageAndCopyIntoCallerSubranges() throws { + try assertOwnedOutput([Float(3), 7], scalarType: .float32) + try assertOwnedOutput([Float16(3), 7], scalarType: .float16) + } + + private func assertOwnedOutput( + _ values: [T], scalarType: NDArray.ScalarType + ) throws { + let output: ETCoreAITensor = try { + var array = NDArray(shape: [values.count], scalarType: scalarType) + do { + var view = array.mutableView(as: T.self) + view.withUnsafeMutablePointer { pointer, _, _ in + for index in values.indices { pointer[index] = values[index] } + } + } + return try prepareOutput(array) + }() + XCTAssertEqual(output.byteCount, UInt(values.count * MemoryLayout.stride)) + for _ in 0..<2 { + var destination = [T](repeating: 99, count: values.count + 2) + destination.withUnsafeMutableBytes { bytes in + output.copyBytes(to: bytes.baseAddress!.advanced(by: MemoryLayout.stride)) + } + XCTAssertEqual(destination, [T(99)] + values + [T(99)]) + } + } + + func testEmptyOutputDoesNotInvokeCopier() { + var calls = 0 + let output = ETCoreAITensor(byteCount: 0, shape: [0], scalarType: .float32) { _ in + calls += 1 + } + output.copyBytes(to: nil) + XCTAssertEqual(calls, 0) + } + + func testOutputCopierIsDeferredUntilCopyAndRetainsItsOwner() { + final class Owner { var calls = 0 } + weak var observed: Owner? + var output: ETCoreAITensor? = { + let owner = Owner() + observed = owner + return ETCoreAITensor(byteCount: 4, shape: [1], scalarType: .float32) { destination in + owner.calls += 1 + var value: Float = 7 + withUnsafeBytes(of: &value) { source in + destination.copyMemory(from: source.baseAddress!, byteCount: source.count) + } + } + }() + XCTAssertEqual(observed?.calls, 0) + XCTAssertEqual(copiedData(output!), [Float(7)].withUnsafeBytes { Data($0) }) + XCTAssertEqual(observed?.calls, 1) + output = nil + XCTAssertNil(observed) + } + + func testOutputCopiesTrackActualShape() throws { + for rows in [2, 0] { + let values = (0..<(rows * 2)).map(Float.init) + var array = NDArray(shape: [rows, 2], scalarType: .float32, strides: [3, 1]) + do { + var view = array.mutableView(as: Float.self) + view.withUnsafeMutablePointer { pointer, _, strides in + for row in 0.. acquisition_staged_url( + NSURL* root, + const SelectedBundle& bundle); +runtime::Result> acquire_test_model( + const SelectedHostBundle& bundle, + const runtime::NamedDataMap* data, + NSString* cache_path, + id loader); +runtime::Error clear_test_cache_entry( + NSString* cache_path, + NSString* key, + id loader); + +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 SelectedBundle& first, + const SelectedBundle& second, + TestData& data, + bool same_key); + +::testing::AssertionResult spawn_acquisition_crash( + BookmarkChild& child, + NSURL* root, + FaultPoint boundary, + NSDictionary* manifest); +::testing::AssertionResult spawn_parallel_acquisition( + BookmarkChild& child, + NSURL* root, + NSDictionary* first, + NSDictionary* second, + bool same_key); +// Dispatch before InitGoogleTest. Unrelated modes return -1. +int acquisition_child_mode(int argc, char** argv); + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm b/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm new file mode 100644 index 00000000000..bb2c5708d35 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm @@ -0,0 +1,364 @@ +/* + * 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. + */ + +#include "coreai_acquisition_fixture.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "coreai_fault_scope.h" +#include "coreai_load_coordinator.h" +#include "coreai_test_utils.h" + +namespace executorch::backends::coreai::testing { +namespace { +using runtime::Error; +using runtime::Result; + +struct AcquisitionWorker { + const SelectedBundle* bundle = nullptr; + TestData* data = nullptr; + NSString* root = nil; + BookmarkFakeLoader* loader = nil; + Error error = Error::Internal; + bool model = false; + dispatch_semaphore_t started = dispatch_semaphore_create(0); + dispatch_semaphore_t done = dispatch_semaphore_create(0); + + static void* run(void* opaque) { + auto* worker = static_cast(opaque); + dispatch_semaphore_signal(worker->started); + @autoreleasepool { + SelectedHostBundle assets{*worker->bundle, @"macOS", @"arch_b"}; + auto acquired = acquire_test_model( + assets, worker->data, worker->root, worker->loader); + worker->error = acquired.ok() ? Error::Ok : acquired.error(); + worker->model = acquired.ok() && acquired.get() != nil; + } + dispatch_semaphore_signal(worker->done); + return nullptr; + } +}; + +struct AcquisitionWorkers { + ScopedFakeBridgeState& bridge; + BookmarkFakeLoader* loader; + dispatch_semaphore_t proceed = dispatch_semaphore_create(0); + pthread_t threads[2]{}; + bool started[2] = {false, false}; + AcquisitionWorker workers[2]; + ParallelAcquisitionObservation& observation; + bool finished = false; + + AcquisitionWorkers( + ScopedFakeBridgeState& state, + BookmarkFakeLoader* fake, + ParallelAcquisitionObservation& result) + : bridge(state), loader(fake), observation(result) {} + + void finish() { + if (finished) + return; + for (int i = 0; i < 2; ++i) + dispatch_semaphore_signal(proceed); + for (int i = 0; i < 2; ++i) { + if (!started[i]) + continue; + observation.completed[i] = + dispatch_semaphore_wait(workers[i].done, test_deadline()) == 0; + observation.joined[i] = pthread_join(threads[i], nullptr); + observation.errors[i] = workers[i].error; + observation.models[i] = workers[i].model; + } + observation.callbacks_drained = bridge.wait_for_callbacks(test_deadline()); + if (observation.callbacks_drained) + loader.onSpecialize = nil; + finished = true; + } + ~AcquisitionWorkers() { + finish(); + } +}; + +} // namespace + +Result> acquire_test_model( + const SelectedHostBundle& bundle, + const runtime::NamedDataMap* data, + NSString* cache_path, + id loader) { + auto directory = test_cache_directory(cache_path); + if (!directory.ok()) { + return directory.error(); + } + return acquire_model(bundle, data, directory.get(), loader); +} + +Error clear_test_cache_entry( + NSString* cache_path, + NSString* key, + id loader) { + auto directory = + test_cache_directory(cache_path, StorageCreationMode::OpenExisting); + if (!directory.ok()) { + return directory.error() == Error::NotFound ? Error::Ok : directory.error(); + } + return clear_cache_entry(directory.get(), key, loader); +} + +Result acquisition_staged_url( + NSURL* root, + const SelectedBundle& bundle) { + auto key = test_cache_key(bundle); + if (!key.ok()) + return key.error(); + return [[[root URLByAppendingPathComponent:@"sources"] + URLByAppendingPathComponent:key.get()] + URLByAppendingPathComponent:bundle.bundle_name]; +} + +ParallelAcquisitionObservation parallel_acquisitions( + ScopedFakeBridgeState& bridge, + BookmarkFakeLoader* loader, + NSURL* root, + const SelectedBundle& first, + const SelectedBundle& second, + TestData& data, + bool same_key) { + ParallelAcquisitionObservation observation; + auto timeouts = std::make_shared>(0); + dispatch_semaphore_t entered = dispatch_semaphore_create(0); + AcquisitionWorkers group(bridge, loader, observation); + dispatch_semaphore_t proceed = group.proceed; + loader.onSpecialize = ^(NSURL*) { + dispatch_semaphore_signal(entered); + if (dispatch_semaphore_wait(proceed, test_deadline()) != 0) + ++*timeouts; + }; + for (int i = 0; i < 2; ++i) { + auto& worker = group.workers[i]; + worker.bundle = i == 0 ? &first : &second; + worker.data = &data; + worker.root = root.path; + worker.loader = loader; + } + observation.creation[0] = pthread_create( + &group.threads[0], nullptr, AcquisitionWorker::run, &group.workers[0]); + group.started[0] = observation.creation[0] == 0; + if (group.started[0]) { + observation.first_entered = + dispatch_semaphore_wait(entered, test_deadline()) == 0; + } + if (observation.first_entered) { + observation.creation[1] = pthread_create( + &group.threads[1], nullptr, AcquisitionWorker::run, &group.workers[1]); + group.started[1] = observation.creation[1] == 0; + if (group.started[1] && + dispatch_semaphore_wait(group.workers[1].started, test_deadline()) == + 0) { + const auto second_deadline = same_key + ? dispatch_time(DISPATCH_TIME_NOW, 50 * NSEC_PER_MSEC) + : test_deadline(); + const bool second_entered = + dispatch_semaphore_wait(entered, second_deadline) == 0; + observation.second_observed = same_key ? !second_entered : second_entered; + } + } + group.finish(); + observation.gate_timeouts = timeouts->load(); + return observation; +} + +::testing::AssertionResult spawn_acquisition_crash( + BookmarkChild& child, + NSURL* root, + FaultPoint boundary, + NSDictionary* manifest) { + // Serialize the shared fixture in the parent; its builders contain GTest + // checks. + NSData* payload = encode( + @{@"boundary" : @(static_cast(boundary)), + @"manifest" : manifest}); + if (payload == nil) + return ::testing::AssertionFailure() << "Cannot encode crash child input"; + NSString* argument = [[NSString alloc] initWithData:payload + encoding:NSUTF8StringEncoding]; + return child.spawn(@"--acquisition-crash", root.path, argument); +} + +::testing::AssertionResult spawn_parallel_acquisition( + BookmarkChild& child, + NSURL* root, + NSDictionary* first, + NSDictionary* second, + bool same_key) { + NSData* payload = encode( + @{@"manifest" : first, + @"second" : second, + @"same_key" : @(same_key)}); + if (payload == nil) + return ::testing::AssertionFailure() + << "Cannot encode parallel child input"; + NSString* argument = [[NSString alloc] initWithData:payload + encoding:NSUTF8StringEncoding]; + return child.spawn(@"--acquisition-parallel", root.path, argument); +} + +namespace { +[[noreturn]] void parallel_child( + NSString* root, + const SelectedBundle& first, + const SelectedBundle& second, + bool same_key, + int fd) { + ScopedFakeBridgeState bridge; + @autoreleasepool { + BookmarkFakeLoader* loader = [[BookmarkFakeLoader alloc] init]; + if (loader == nil) + _exit(8); + bridge.state().bookmark_loader = loader; + TestData data; + auto result = parallel_acquisitions( + bridge, + loader, + [NSURL fileURLWithPath:root], + first, + second, + data, + same_key); + if (!result.callbacks_drained) + _exit(11); + // These bytes are observations, not child assertions. The parent checks + // each one. + const bool harness_ok = result.creation[0] == 0 && + result.creation[1] == 0 && result.first_entered && + result.joined[0] == 0 && result.joined[1] == 0 && result.completed[0] && + result.completed[1] && result.gate_timeouts == 0; + const bool both_ok = result.errors[0] == Error::Ok && result.models[0] && + result.errors[1] == Error::Ok && result.models[1]; + const char observations[] = { + static_cast(harness_ok), + static_cast(result.second_observed), + static_cast(both_ok), + static_cast(loader->specializations.load()), + static_cast(data.attempts.load()), + static_cast(loader->restores.load())}; + for (char observation : observations) { + if (!send_child_byte(fd, observation)) + _exit(12); + } + if (loader->evictions.load() != 0 || + data.requests.load() != data.releases.load()) + _exit(14); + [loader clearRetainedState]; + bridge.state().bookmark_loader = nil; + loader = nil; + } + auto& state = bridge.state(); + const bool none_live = state.sessions.load() == 0 && + state.prepared_models.load() == 0 && state.loaders.load() == 0 && + state.binding_mismatches.load() == 0 && + state.missing_bundles.load() == 0 && + state.input_wait_timeouts.load() == 0; + if (!send_child_byte(fd, static_cast(none_live))) + _exit(12); + // Avoid GTest-bearing destructors before InitGoogleTest. + _exit(close(fd) == 0 ? 0 : 13); +} +} // namespace + +int acquisition_child_mode(int argc, char** argv) { + if (argc < 2) + return -1; + const bool crash = std::strcmp(argv[1], "--acquisition-crash") == 0; + const bool parallel = std::strcmp(argv[1], "--acquisition-parallel") == 0; + if (!crash && !parallel) + return -1; + if (argc != 5) + return 1; + NSString* root = [NSString stringWithUTF8String:argv[2]]; + NSString* argument = [NSString stringWithUTF8String:argv[3]]; + if (root == nil || argument == nil) + return 2; + id payload = [NSJSONSerialization + JSONObjectWithData:[argument dataUsingEncoding:NSUTF8StringEncoding] + options:0 + error:nil]; + if (![payload isKindOfClass:NSDictionary.class]) + return 3; + id number = payload[@"boundary"]; + id dictionary = payload[@"manifest"]; + if (![dictionary isKindOfClass:NSDictionary.class]) + return 3; + FaultPoint boundary = FaultPoint::AfterSDK; + if (crash) { + if (![number isKindOfClass:NSNumber.class]) + return 3; + boundary = static_cast([number intValue]); + if (boundary != FaultPoint::BeforeSDK && boundary != FaultPoint::AfterSDK) + return 4; + } + auto parsed = parse_manifest(encode(dictionary)); + if (!parsed.ok()) + return 5; + auto selected = select_bundle(parsed.get(), @"arch_b", @"macOS"); + if (!selected.ok()) + return 5; + const int fd = parse_child_fd(argv[4]); + if (fd < 0 || fcntl(fd, F_GETFD) < 0) + return 6; + signal(SIGPIPE, SIG_IGN); + if (!send_child_byte(fd, 'C')) + return 7; + if (parallel) { + id second_dictionary = payload[@"second"]; + id same_key = payload[@"same_key"]; + if (![second_dictionary isKindOfClass:NSDictionary.class] || + ![same_key isKindOfClass:NSNumber.class]) + return 3; + auto second_parsed = parse_manifest(encode(second_dictionary)); + if (!second_parsed.ok()) + return 5; + auto second = select_bundle(second_parsed.get(), @"arch_b", @"macOS"); + if (!second.ok()) + return 5; + parallel_child( + root, selected.get(), second.get(), [same_key boolValue], fd); + } + if (close(fd) != 0) + return 7; + + ScopedFakeBridgeState bridge; + @autoreleasepool { + // Terminal child paths deliberately bypass assertion-bearing scope + // destructors. + BookmarkFakeLoader* loader = [[BookmarkFakeLoader alloc] init]; + if (loader == nil) + _exit(8); + bridge.state().bookmark_loader = loader; + TestData data; + StorageFaultScope fault(boundary); + fault.armed = false; + fault.observe = ^(FaultPoint event) { + if (event == boundary) + _exit(73); + }; + SelectedHostBundle assets{selected.get(), @"macOS", @"arch_b"}; + auto result = acquire_test_model(assets, &data, root, loader); + _exit(result.ok() ? 9 : 10); + } +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm new file mode 100644 index 00000000000..5b1df356756 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm @@ -0,0 +1,804 @@ +/* + * 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. + */ + +#include "coreai_acquisition_fixture.h" + +#include +#include +#include +#include +#include +#include +#include +#include "coreai_fault_scope.h" +#include "coreai_filesystem_fixture.h" +#include "coreai_load_coordinator.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +@interface InvalidAcquisitionLoader : BookmarkFakeLoader +@property(nonatomic) ETCoreAIAcquisitionStatus restoreStatus; +@property(nonatomic, strong) id restoreModel; +@end + +@implementation InvalidAcquisitionLoader +- (void)restoreModelFromBookmark:(NSData*)bookmark + completion:(ETCoreAIAcquisitionCompletion)completion { + ++restores; + completion(self.restoreStatus, self.restoreModel, nil); +} +- (void)specializeModelAtURL:(NSURL*)url + completion:(ETCoreAIAcquisitionCompletion)completion { + ++specializations; + completion(ETCoreAIAcquisitionStatusMiss, nil, nil); +} +@end + +namespace { + +TEST(CoreAISelectedHostBundleTest, ProbesHostAndSelectsImmutableBundle) { + ScopedFakeBridgeState bridge; + auto parsed = parse_manifest(encode(aot_manifest_dict())); + ASSERT_TRUE(parsed.ok()); + auto assets = select_host_bundle(parsed.get()); + ASSERT_TRUE(assets.ok()); + EXPECT_TRUE([assets->platform isEqual:@"macOS"]); + EXPECT_TRUE([assets->architecture isEqual:@"arch_b"]); + EXPECT_TRUE([assets->bundle.bundle_name isEqual:@"model.arch_b.aimodelc"]); + EXPECT_TRUE([assets->bundle.export_hash isEqual:parsed->export_hash]); +} + +class CoreAIAcquisitionTest : public ::testing::Test { + protected: + // Declared first so each case's loader, models and data die before it. + ScopedFakeBridgeState bridge; + TestDirectory root; + TestData data; + std::optional assets; + NSString* key = nil; + + void SetUp() override { + ASSERT_NE(root.url, nil); + auto selected = selected_bundle(manifest_dict()); + ASSERT_TRUE(selected.ok()); + assets.emplace(SelectedHostBundle{selected.get(), @"macOS", @"arch_b"}); + auto identity = test_cache_key(assets->bundle); + ASSERT_TRUE(identity.ok()); + key = identity.get(); + } + + Error acquire(BookmarkFakeLoader* loader, const NamedDataMap* named_data) { + @autoreleasepool { + auto model = + acquire_test_model(*assets, named_data, root.url.path, loader); + if (!model.ok()) + return model.error(); + return model.get() != nil ? Error::Ok : Error::Internal; + } + } + + NSData* saved() { + auto bookmark = saved_bookmark(root.url, key); + EXPECT_TRUE(bookmark.ok()); + return bookmark.ok() ? bookmark.get() : nil; + } + + NSURL* source_key_url() { + return [[[root.url URLByAppendingPathComponent:@"sources"] + URLByAppendingPathComponent:key] copy]; + } + + BOOL source_key_exists() { + return + [NSFileManager.defaultManager fileExistsAtPath:source_key_url().path]; + } + + NSURL* lock_url() { + return test_cache_lock_url(root.url, key); + } + + ino_t lock_inode() { + struct stat info; + if (lstat(lock_url().fileSystemRepresentation, &info) != 0) + return 0; + return info.st_ino; + } +}; + +TEST_F(CoreAIAcquisitionTest, ColdSpecializesStagedSourceAndPersistsBookmark) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(data.attempts.load(), 1); + auto staged = acquisition_staged_url(root.url, assets->bundle); + ASSERT_TRUE(staged.ok()); + EXPECT_TRUE([fake.loader.lastSpecializeURL isEqual:staged.get().path]); + EXPECT_TRUE([saved() isEqual:fake.loader.lastAcquired.bookmark]); + EXPECT_FALSE(source_key_exists()); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:staged.get().path]); +} + +TEST_F(CoreAIAcquisitionTest, WarmHitRestoresWithoutNamedData) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + NSData* original = saved(); + EXPECT_FALSE(source_key_exists()); + const int attempts = data.attempts.load(); + ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->restores.load(), 1); + EXPECT_EQ(data.attempts.load(), attempts); + EXPECT_TRUE([saved() isEqual:original]); + EXPECT_FALSE(source_key_exists()); +} + +TEST_F(CoreAIAcquisitionTest, RestoreErrorFailsClosedWithoutSourceFallback) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + NSData* original = saved(); + const int attempts = data.attempts; + fake.loader.restoreError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EIO + userInfo:nil]; + EXPECT_EQ(acquire(fake.loader, &data), Error::Internal); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(data.attempts.load(), attempts); + EXPECT_TRUE([saved() isEqual:original]); +} + +TEST_F( + CoreAIAcquisitionTest, + RestoreMissRespecializesStagedSourceAndReplacesBookmark) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + NSData* original = saved(); + const int attempts = data.attempts; + fake.loader.restoreMiss = YES; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 2); + EXPECT_EQ(data.attempts.load(), attempts + 1); + NSData* replaced = saved(); + EXPECT_NE(replaced, nil); + EXPECT_FALSE([replaced isEqual:original]); + EXPECT_FALSE(source_key_exists()); +} + +TEST_F( + CoreAIAcquisitionTest, + SpecializationErrorFailsClosedPreservingBookmark) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + NSData* original = saved(); + ASSERT_NE(original, nil); + const int attempts = data.attempts.load(); + const int specializations = fake.loader->specializations.load(); + fake.loader.restoreMiss = YES; + fake.loader.specializeError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EIO + userInfo:nil]; + EXPECT_EQ(acquire(fake.loader, &data), Error::Internal); + EXPECT_EQ(fake.loader->specializations.load(), specializations + 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(data.attempts.load(), attempts + 1); + EXPECT_TRUE([saved() isEqual:original]); + EXPECT_TRUE([NSFileManager.defaultManager + fileExistsAtPath:fake.loader.lastSpecializeURL]); +} + +TEST_F(CoreAIAcquisitionTest, EmptyBookmarkFailsClosedPreservingPrior) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + NSData* original = saved(); + ASSERT_NE(original, nil); + const int specializations = fake.loader->specializations.load(); + fake.loader.restoreMiss = YES; + fake.loader.emptyBookmark = YES; + EXPECT_EQ(acquire(fake.loader, &data), Error::InvalidExternalData); + EXPECT_EQ(fake.loader->specializations.load(), specializations + 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_TRUE([saved() isEqual:original]); + EXPECT_TRUE([NSFileManager.defaultManager + fileExistsAtPath:fake.loader.lastSpecializeURL]); +} + +TEST_F(CoreAIAcquisitionTest, NilLoaderRejectsWithoutFilesystemEffects) { + NSDictionary* before = nil; + ASSERT_TRUE(asset_tree_snapshot(root.url, before)); + Error error = Error::Ok; + @autoreleasepool { + auto result = acquire_test_model(*assets, &data, root.url.path, nil); + ASSERT_FALSE(result.ok()); + error = result.error(); + } + EXPECT_EQ(error, Error::InvalidArgument); + EXPECT_TRUE(snapshot_matches(root.url, before)); + EXPECT_EQ(saved(), nil); + EXPECT_FALSE(source_key_exists()); +} + +TEST_F(CoreAIAcquisitionTest, RestoreHitWithoutModelFailsClosed) { + NSData* original = nil; + { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + original = saved(); + ASSERT_NE(original, nil); + } + InvalidAcquisitionLoader* loader = [[InvalidAcquisitionLoader alloc] init]; + ASSERT_NE(loader, nil); + loader.restoreStatus = ETCoreAIAcquisitionStatusHit; + Error error = Error::Ok; + @autoreleasepool { + auto result = acquire_test_model(*assets, &data, root.url.path, loader); + ASSERT_FALSE(result.ok()); + error = result.error(); + } + EXPECT_EQ(error, Error::Internal); + loader = nil; + EXPECT_TRUE([saved() isEqual:original]); +} + +TEST_F(CoreAIAcquisitionTest, RestoreMissWithModelFailsClosed) { + NSData* original = nil; + { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + original = saved(); + ASSERT_NE(original, nil); + } + InvalidAcquisitionLoader* loader = [[InvalidAcquisitionLoader alloc] init]; + ASSERT_NE(loader, nil); + BookmarkFakeModel* model = [[BookmarkFakeModel alloc] init]; + ASSERT_NE(model, nil); + loader.restoreStatus = ETCoreAIAcquisitionStatusMiss; + loader.restoreModel = model; + Error error = Error::Ok; + @autoreleasepool { + auto result = acquire_test_model(*assets, &data, root.url.path, loader); + ASSERT_FALSE(result.ok()); + error = result.error(); + } + EXPECT_EQ(error, Error::Internal); + loader.restoreModel = nil; + loader = nil; + EXPECT_TRUE([saved() isEqual:original]); +} + +TEST_F(CoreAIAcquisitionTest, RestoreErrorWithoutErrorFailsClosed) { + NSData* original = nil; + { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + original = saved(); + ASSERT_NE(original, nil); + } + InvalidAcquisitionLoader* loader = [[InvalidAcquisitionLoader alloc] init]; + ASSERT_NE(loader, nil); + loader.restoreStatus = ETCoreAIAcquisitionStatusError; + Error error = Error::Ok; + @autoreleasepool { + auto result = acquire_test_model(*assets, &data, root.url.path, loader); + ASSERT_FALSE(result.ok()); + error = result.error(); + } + EXPECT_EQ(error, Error::Internal); + loader = nil; + EXPECT_TRUE([saved() isEqual:original]); +} + +TEST_F(CoreAIAcquisitionTest, SpecializeMissFailsClosed) { + InvalidAcquisitionLoader* loader = [[InvalidAcquisitionLoader alloc] init]; + ASSERT_NE(loader, nil); + Error error = Error::Ok; + @autoreleasepool { + auto result = acquire_test_model(*assets, &data, root.url.path, loader); + ASSERT_FALSE(result.ok()); + error = result.error(); + } + EXPECT_EQ(error, Error::Internal); + EXPECT_EQ(loader->specializations.load(), 1); + EXPECT_EQ(loader->restores.load(), 0); + loader = nil; + EXPECT_EQ(saved(), nil); +} + +TEST_F(CoreAIAcquisitionTest, ColdSuccessRemovesWholeSourceAfterBookmarkFlush) { + BookmarkLoaderScope fake; + __block std::vector ops; + Error error = Error::Internal; + { + StorageFaultScope fault(FaultPoint::Remove); + fault.armed = false; + fault.observe = ^(FaultPoint operation) { + ops.push_back(operation); + }; + error = acquire(fake.loader, &data); + } + ASSERT_EQ(error, Error::Ok); + EXPECT_TRUE([saved() isEqual:fake.loader.lastAcquired.bookmark]); + EXPECT_FALSE(source_key_exists()); + auto staged = acquisition_staged_url(root.url, assets->bundle); + ASSERT_TRUE(staged.ok()); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:staged.get().path]); + size_t last_flush = 0; + size_t removes = 0; + size_t first_remove = ops.size(); + for (size_t i = 0; i < ops.size(); ++i) { + if (ops[i] == FaultPoint::FullSync) + last_flush = i; + if (ops[i] == FaultPoint::Remove) { + ++removes; + if (first_remove == ops.size()) + first_remove = i; + } + } + EXPECT_EQ(removes, 3u); + EXPECT_GT(first_remove, last_flush); +} + +TEST_F(CoreAIAcquisitionTest, SpecializeErrorPerformsNoCleanup) { + BookmarkLoaderScope fake; + __block std::vector ops; + StorageFaultScope fault(FaultPoint::Remove); + fault.armed = false; + fault.observe = ^(FaultPoint operation) { + ops.push_back(operation); + }; + fake.loader.specializeError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EIO + userInfo:nil]; + EXPECT_EQ(acquire(fake.loader, &data), Error::Internal); + size_t removes = 0; + for (auto operation : ops) { + if (operation == FaultPoint::Remove) + ++removes; + } + EXPECT_EQ(removes, 0u); + EXPECT_EQ(saved(), nil); + EXPECT_EQ(fake.loader->evictions.load(), 0); +} + +TEST_F(CoreAIAcquisitionTest, BookmarkFinalFlushFailurePerformsNoCleanup) { + BookmarkLoaderScope fake; + __block std::vector ops; + Error error = Error::Internal; + { + // Cold source contributes two FullSyncs and the bookmark two more, so + // index four is the post-rename bookmark flush. + StorageFaultScope fault(FaultPoint::FullSync); + fault.match_index = 4; + fault.observe = ^(FaultPoint operation) { + ops.push_back(operation); + }; + error = acquire(fake.loader, &data); + } + EXPECT_EQ(error, Error::AccessFailed); + EXPECT_TRUE([saved() isEqual:fake.loader.lastAcquired.bookmark]); + size_t removes = 0; + for (auto operation : ops) { + if (operation == FaultPoint::Remove) + ++removes; + } + EXPECT_EQ(removes, 0u); + auto staged = acquisition_staged_url(root.url, assets->bundle); + ASSERT_TRUE(staged.ok()); + EXPECT_TRUE( + [NSFileManager.defaultManager fileExistsAtPath:staged.get().path]); +} + +TEST_F( + CoreAIAcquisitionTest, + CleanupFirstRemoveFailurePreservesDurableBookmark) { + BookmarkLoaderScope fake; + NSData* bookmark = nil; + int specializations = 0; + { + StorageFaultScope fault(FaultPoint::Remove); + fault.match_index = 1; + EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fault.hits, 1); + bookmark = saved(); + ASSERT_NE(bookmark, nil); + EXPECT_TRUE([bookmark isEqual:fake.loader.lastAcquired.bookmark]); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(fake.loader->evictions.load(), 0); + specializations = fake.loader->specializations.load(); + } + EXPECT_EQ(fake.loader->specializations.load(), specializations); + const int attempts = data.attempts.load(); + const int restores = fake.loader->restores.load(); + EXPECT_TRUE(source_key_exists()); + ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), specializations); + EXPECT_EQ(fake.loader->restores.load(), restores + 1); + EXPECT_EQ(data.attempts.load(), attempts); + EXPECT_TRUE([saved() isEqual:bookmark]); + EXPECT_FALSE(source_key_exists()); +} + +TEST_F( + CoreAIAcquisitionTest, + CleanupPartialRemoveFailurePreservesDurableBookmark) { + BookmarkLoaderScope fake; + NSData* bookmark = nil; + int specializations = 0; + { + StorageFaultScope fault(FaultPoint::Remove); + fault.match_index = 2; + EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fault.hits, 1); + bookmark = saved(); + ASSERT_NE(bookmark, nil); + EXPECT_TRUE([bookmark isEqual:fake.loader.lastAcquired.bookmark]); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + specializations = fake.loader->specializations.load(); + } + EXPECT_EQ(fake.loader->specializations.load(), specializations); + const int attempts = data.attempts.load(); + const int restores = fake.loader->restores.load(); + EXPECT_TRUE(source_key_exists()); + ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), specializations); + EXPECT_EQ(fake.loader->restores.load(), restores + 1); + EXPECT_EQ(data.attempts.load(), attempts); + EXPECT_TRUE([saved() isEqual:bookmark]); + EXPECT_FALSE(source_key_exists()); +} + +TEST_F( + CoreAIAcquisitionTest, + CleanupDirectorySyncFailurePreservesDurableBookmark) { + BookmarkLoaderScope fake; + NSData* bookmark = nil; + int specializations = 0; + { + StorageFaultScope scope(FaultPoint::SyncDirectory); + scope.armed = false; + StorageFaultScope* arming = &scope; + scope.observe = ^(FaultPoint operation) { + if (operation == FaultPoint::Remove) + arming->armed = true; + }; + EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(scope.hits, 1); + bookmark = saved(); + ASSERT_NE(bookmark, nil); + EXPECT_TRUE([bookmark isEqual:fake.loader.lastAcquired.bookmark]); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + specializations = fake.loader->specializations.load(); + } + const int attempts = data.attempts.load(); + const int restores = fake.loader->restores.load(); + EXPECT_TRUE(source_key_exists()); + ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), specializations); + EXPECT_EQ(fake.loader->restores.load(), restores + 1); + EXPECT_EQ(data.attempts.load(), attempts); + EXPECT_TRUE([saved() isEqual:bookmark]); + EXPECT_FALSE(source_key_exists()); +} + +TEST_F(CoreAIAcquisitionTest, WarmHitCleanupFailureStillReturnsModel) { + BookmarkLoaderScope fake; + { + StorageFaultScope fault(FaultPoint::Remove); + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + } + NSData* original = saved(); + StorageFaultScope fault(FaultPoint::Remove); + EXPECT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fault.hits, 1); + EXPECT_EQ(fake.loader->restores.load(), 1); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(data.attempts.load(), 1); + EXPECT_TRUE([saved() isEqual:original]); +} + +TEST_F(CoreAIAcquisitionTest, PartialCleanupThenMissRebuildsSource) { + BookmarkLoaderScope fake; + { + StorageFaultScope fault(FaultPoint::Remove); + fault.match_index = 2; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fault.hits, 1); + } + ASSERT_TRUE(source_key_exists()); + NSData* original = saved(); + fake.loader.restoreMiss = YES; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fake.loader->restores.load(), 1); + EXPECT_EQ(fake.loader->specializations.load(), 2); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(data.attempts.load(), 2); + EXPECT_FALSE([saved() isEqual:original]); + EXPECT_FALSE(source_key_exists()); +} + +TEST_F( + CoreAIAcquisitionTest, + DamagedSourceCleanupFailureDoesNotReadOrSpecialize) { + BookmarkLoaderScope fake; + { + StorageFaultScope fault(FaultPoint::Remove); + fault.match_index = 2; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + } + NSData* original = saved(); + fake.loader.restoreMiss = YES; + { + StorageFaultScope fault(FaultPoint::Remove); + EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_EQ(data.attempts.load(), 1); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_TRUE([saved() isEqual:original]); + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_FALSE(source_key_exists()); +} + +TEST_F(CoreAIAcquisitionTest, MissingNamedDataIsNotRetried) { + BookmarkLoaderScope fake; + data.files.clear(); + EXPECT_EQ(acquire(fake.loader, &data), Error::InvalidExternalData); + EXPECT_EQ(data.attempts.load(), 1); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(saved(), nil); +} + +TEST_F(CoreAIAcquisitionTest, BeforeSdkFaultNeverInvokesLoader) { + BookmarkLoaderScope fake; + StorageFaultScope fault(FaultPoint::BeforeSDK); + EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(fault.hits, 1); + EXPECT_EQ(saved(), nil); +} + +TEST_F(CoreAIAcquisitionTest, KeyLockIdentitySurvivesSuccessfulAcquisition) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + ino_t first = lock_inode(); + ASSERT_NE(first, 0u); + ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(lock_inode(), first); +} + +class CoreAIAcquisitionSaveFailureTest + : public CoreAIAcquisitionTest, + public ::testing::WithParamInterface {}; +TEST_P(CoreAIAcquisitionSaveFailureTest, PreservesBookmarkAndRetries) { + const bool prior = GetParam(); + BookmarkLoaderScope fake; + NSData* original = nil; + if (prior) { + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + original = saved(); + ASSERT_NE(original, nil); + EXPECT_FALSE(source_key_exists()); + fake.loader.restoreMiss = YES; + } + { + // Cold and rematerialized sources publish before the bookmark, so the + // bookmark rename is the second rename once the first success has + // cleaned the source. + StorageFaultScope fault(FaultPoint::Rename); + fault.match_index = 2; + EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_TRUE(original == nil ? saved() == nil : [saved() isEqual:original]); + EXPECT_TRUE([NSFileManager.defaultManager + fileExistsAtPath:fake.loader.lastSpecializeURL]); + const int calls = fake.loader->specializations; + fake.loader.restoreMiss = NO; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), calls + (prior ? 0 : 1)); + EXPECT_FALSE(source_key_exists()); +} +INSTANTIATE_TEST_SUITE_P( + Boundaries, + CoreAIAcquisitionSaveFailureTest, + ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return info.param ? "PriorBookmark" : "Cold"; + }); + +class CoreAIAcquisitionProcessTest : public CoreAIAcquisitionTest { + protected: + void expect_parallel(bool same_key) { + NSMutableDictionary* first = manifest_dict(); + NSMutableDictionary* second = [first mutableCopy]; + if (!same_key) + second[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity('2')}; + BookmarkChild child; + ASSERT_TRUE( + spawn_parallel_acquisition(child, root.url, first, second, same_key)); + ASSERT_TRUE(child.receive('C')); + ASSERT_TRUE(child.receive(1)) << "harness completed without timeouts"; + ASSERT_TRUE(child.receive(1)) + << "second acquisition blocked or entered as expected"; + ASSERT_TRUE(child.receive(1)) << "both acquisitions succeeded"; + ASSERT_TRUE(child.receive(same_key ? 1 : 2)) << "specializations"; + ASSERT_TRUE(child.receive(same_key ? 1 : 2)) << "named-data attempts"; + ASSERT_TRUE(child.receive(same_key ? 1 : 0)) << "restores"; + ASSERT_TRUE(child.receive(1)) << "no live fake SDK objects"; + ASSERT_TRUE(child.expect_exit(0)); + } +}; + +TEST_F(CoreAIAcquisitionProcessTest, ContendedSameKeySerializesSpecialization) { + expect_parallel(true); +} + +TEST_F(CoreAIAcquisitionProcessTest, IndependentKeysSpecializeInParallel) { + expect_parallel(false); +} + +TEST_F(CoreAIAcquisitionProcessTest, CrashBeforeSdkRetainsSourceAndRetries) { + BookmarkChild child; + ASSERT_TRUE(spawn_acquisition_crash( + child, root.url, FaultPoint::BeforeSDK, manifest_dict())); + ASSERT_TRUE(child.receive('C')); + ASSERT_TRUE(child.expect_exit(73)); + EXPECT_EQ(saved(), nil); + EXPECT_TRUE(source_key_exists()); + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(data.attempts.load(), 0); + EXPECT_FALSE(source_key_exists()); +} + +TEST_F(CoreAIAcquisitionProcessTest, CrashAfterSdkRetainsSourceAndRetries) { + BookmarkChild child; + ASSERT_TRUE(spawn_acquisition_crash( + child, root.url, FaultPoint::AfterSDK, manifest_dict())); + ASSERT_TRUE(child.receive('C')); + ASSERT_TRUE(child.expect_exit(73)); + EXPECT_EQ(saved(), nil); + auto staged = acquisition_staged_url(root.url, assets->bundle); + ASSERT_TRUE(staged.ok()); + EXPECT_TRUE( + [NSFileManager.defaultManager fileExistsAtPath:staged.get().path]); + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(data.attempts.load(), 0); + EXPECT_NE(saved(), nil); + EXPECT_FALSE(source_key_exists()); +} + +class CoreAILifecycleLogTest : public CoreAIAcquisitionTest { + protected: + void SetUp() override { + CoreAIAcquisitionTest::SetUp(); + if (!ET_LOG_ENABLED || LogLevel::Info < LogLevel::ET_MIN_LOG_LEVEL) { + GTEST_SKIP() << "Info logging is compiled out"; + } + } + + template + std::string capture(Operation operation, Error expected = Error::Ok) { + ::testing::internal::CaptureStderr(); + const auto error = operation(); + auto logs = ::testing::internal::GetCapturedStderr(); + EXPECT_EQ(error, expected); + return logs; + } + + void expect_order( + const std::string& logs, + std::initializer_list events) { + size_t position = 0; + for (const char* event : events) { + const auto found = logs.find(event, position); + ASSERT_NE(found, std::string::npos) << event << "\n" << logs; + position = found + 1; + } + EXPECT_NE(logs.find(key.UTF8String), std::string::npos); + } +}; + +TEST_F(CoreAILifecycleLogTest, ColdAcquisitionLogsValidatedMilestonesInOrder) { + BookmarkLoaderScope fake; + const auto logs = capture([&] { return acquire(fake.loader, &data); }); + expect_order( + logs, + { + "Core AI cache lookup:", + "Core AI cache miss: no bookmark", + "Core AI specialization started:", + "Core AI specialization succeeded:", + "Core AI cache bookmark saved durably:", + }); + EXPECT_EQ(logs.find("Core AI cache hit:"), std::string::npos); + EXPECT_TRUE([saved() isEqual:fake.loader.lastAcquired.bookmark]); +} + +TEST_F(CoreAILifecycleLogTest, WarmHitDoesNotLogSpecializationOrSave) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + const auto logs = capture([&] { return acquire(fake.loader, nullptr); }); + expect_order( + logs, + { + "Core AI cache lookup:", + "Core AI cache hit: restored model from bookmark", + }); + EXPECT_EQ(logs.find("Core AI cache miss:"), std::string::npos); + EXPECT_EQ(logs.find("Core AI specialization"), std::string::npos); + EXPECT_EQ(logs.find("Core AI cache bookmark saved"), std::string::npos); +} + +TEST_F(CoreAILifecycleLogTest, ConfirmedMissLogsRespecialization) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + fake.loader.restoreMiss = YES; + const auto logs = capture([&] { return acquire(fake.loader, &data); }); + expect_order( + logs, + { + "Core AI cache miss: SDK could not restore bookmark", + "Core AI specialization started:", + "Core AI specialization succeeded:", + "Core AI cache bookmark saved durably:", + }); + EXPECT_EQ(logs.find("Core AI cache hit:"), std::string::npos); +} + +TEST_F(CoreAILifecycleLogTest, RestoreErrorIsNotReportedAsHitOrMiss) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + fake.loader.restoreError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EIO + userInfo:nil]; + const auto logs = + capture([&] { return acquire(fake.loader, &data); }, Error::Internal); + EXPECT_EQ(logs.find("Core AI cache hit:"), std::string::npos); + EXPECT_EQ(logs.find("Core AI cache miss:"), std::string::npos); + EXPECT_EQ(logs.find("Core AI specialization"), std::string::npos); + EXPECT_EQ(logs.find("Core AI cache bookmark saved"), std::string::npos); +} + +TEST_F(CoreAILifecycleLogTest, SpecializationErrorDoesNotLogSuccessOrSave) { + BookmarkLoaderScope fake; + fake.loader.specializeError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EIO + userInfo:nil]; + const auto logs = + capture([&] { return acquire(fake.loader, &data); }, Error::Internal); + EXPECT_NE(logs.find("Core AI specialization started:"), std::string::npos); + EXPECT_EQ(logs.find("Core AI specialization succeeded:"), std::string::npos); + EXPECT_EQ(logs.find("Core AI cache bookmark saved"), std::string::npos); +} + +TEST_F(CoreAILifecycleLogTest, BeforeSdkFaultDoesNotLogSpecializationStart) { + BookmarkLoaderScope fake; + StorageFaultScope fault(FaultPoint::BeforeSDK); + const auto logs = + capture([&] { return acquire(fake.loader, &data); }, Error::AccessFailed); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(logs.find("Core AI specialization"), std::string::npos); + EXPECT_EQ(logs.find("Core AI cache bookmark saved"), std::string::npos); +} + +} // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_cache_entry_contention_test.mm b/backends/apple/coreai/runtime/test/coreai_cache_entry_contention_test.mm new file mode 100644 index 00000000000..cd0220a1b71 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_cache_entry_contention_test.mm @@ -0,0 +1,289 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include "coreai_cache_entry_fixture.h" +#include "coreai_fault_scope.h" +#include "coreai_filesystem_fixture.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +namespace { +constexpr auto kWaitTimeout = std::chrono::seconds(5); + +struct CacheEntryWaiter { + std::future attempting; + std::future>> result; +}; + +CacheEntryWaiter start_cache_entry_waiter( + NSString* path, + NSString* key, + bool existing_only = false) { + std::promise attempting; + std::promise>> result; + CacheEntryWaiter waiter{attempting.get_future(), result.get_future()}; + // Detached promise-backed work cannot block test teardown if acquisition + // stalls after a failed assertion. + std::thread([path, + key, + existing_only, + attempting = std::move(attempting), + result = std::move(result)]() mutable { + @autoreleasepool { + StorageFaultScope scope(FaultPoint::BeforeCacheEntryLock); + scope.armed = false; + auto* signal = &attempting; + scope.observe = ^(FaultPoint point) { + if (point == FaultPoint::BeforeCacheEntryLock) { + signal->set_value(); + } + }; + result.set_value( + existing_only ? acquire_existing_test_cache_entry(path, key) + : acquire_test_cache_entry(path, key)); + } + }).detach(); + return waiter; +} + +class CoreAIBookmarkContentionTest : public ::testing::Test { + protected: + TestDirectory root; + std::unique_ptr bundle; + NSString* key = nil; + + void SetUp() override { + ASSERT_NE(root.url, nil); + auto selected = selected_bundle(manifest_dict()); + ASSERT_TRUE(selected.ok()); + bundle = std::make_unique(selected.get()); + auto result = test_cache_key(*bundle); + ASSERT_TRUE(result.ok()); + key = result.get(); + } + + NSURL* lock_url() { + return test_cache_lock_url(root.url, key); + } +}; + +TEST_F( + CoreAIBookmarkContentionTest, + ThreadContentionSerializesSameKeyButNotDifferentStripes) { + auto holder = acquire_test_cache_entry(root.url.path, key); + ASSERT_TRUE(holder.ok()); + ASSERT_NE(holder.get(), nullptr); + auto waiter = start_cache_entry_waiter(root.url.path, key); + ASSERT_EQ( + waiter.attempting.wait_for(kWaitTimeout), std::future_status::ready); + EXPECT_EQ( + waiter.result.wait_for(std::chrono::milliseconds(100)), + std::future_status::timeout); + + auto independent_key = different_cache_key(*bundle, key, false); + ASSERT_TRUE(independent_key.ok()); + auto independent = + start_cache_entry_waiter(root.url.path, independent_key.get()); + ASSERT_EQ( + independent.result.wait_for(kWaitTimeout), std::future_status::ready); + auto other = independent.result.get(); + ASSERT_TRUE(other.ok()); + ASSERT_NE(other.get(), nullptr); + + holder.get().reset(); + ASSERT_EQ(waiter.result.wait_for(kWaitTimeout), std::future_status::ready); + auto acquired = waiter.result.get(); + ASSERT_TRUE(acquired.ok()); + EXPECT_NE(acquired.get(), nullptr); +} + +TEST_F(CoreAIBookmarkContentionTest, CollidingKeysSerializeAcrossThreads) { + auto collision = different_cache_key(*bundle, key, true); + ASSERT_TRUE(collision.ok()); + auto holder = acquire_test_cache_entry(root.url.path, key); + ASSERT_TRUE(holder.ok()); + auto waiter = start_cache_entry_waiter(root.url.path, collision.get()); + ASSERT_EQ( + waiter.attempting.wait_for(kWaitTimeout), std::future_status::ready); + EXPECT_EQ( + waiter.result.wait_for(std::chrono::milliseconds(100)), + std::future_status::timeout); + holder.get().reset(); + ASSERT_EQ(waiter.result.wait_for(kWaitTimeout), std::future_status::ready); + auto acquired = waiter.result.get(); + ASSERT_TRUE(acquired.ok()); + EXPECT_NE(acquired.get(), nullptr); +} + +TEST_F( + CoreAIBookmarkContentionTest, + MissingCollidingLookupDoesNotWaitOrCreateAnotherLock) { + auto collision = different_cache_key(*bundle, key, true); + ASSERT_TRUE(collision.ok()); + auto holder = acquire_test_cache_entry(root.url.path, key); + ASSERT_TRUE(holder.ok()); + auto waiter = start_cache_entry_waiter(root.url.path, collision.get(), true); + EXPECT_EQ(waiter.result.wait_for(kWaitTimeout), std::future_status::ready); + holder.get().reset(); + auto absent = waiter.result.get(); + ASSERT_TRUE(absent.ok()); + EXPECT_EQ(absent.get(), nullptr); + NSArray* names = [NSFileManager.defaultManager + contentsOfDirectoryAtPath:[root.url URLByAppendingPathComponent:@"locks"] + .path + error:nil]; + ASSERT_NE(names, nil); + EXPECT_EQ(names.count, 1u); +} + +TEST_F(CoreAIBookmarkContentionTest, ExistingWaiterRechecksEvictedEntry) { + auto holder = acquire_test_cache_entry(root.url.path, key); + ASSERT_TRUE(holder.ok()); + NSData* bytes = [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_EQ(holder.get()->write_bookmark(bytes), Error::Ok); + auto waiter = start_cache_entry_waiter(root.url.path, key, true); + ASSERT_EQ( + waiter.attempting.wait_for(kWaitTimeout), std::future_status::ready); + ASSERT_EQ( + waiter.result.wait_for(std::chrono::milliseconds(100)), + std::future_status::timeout); + NSURL* bookmark = test_bookmark_url(root.url, key); + ASSERT_EQ(unlink(bookmark.fileSystemRepresentation), 0); + holder.get().reset(); + ASSERT_EQ(waiter.result.wait_for(kWaitTimeout), std::future_status::ready); + auto absent = waiter.result.get(); + ASSERT_TRUE(absent.ok()); + EXPECT_EQ(absent.get(), nullptr); +} + +TEST_F( + CoreAIBookmarkContentionTest, + WaiterOpensBookmarkDirectoryAfterAcquiringLock) { + auto holder = acquire_test_cache_entry(root.url.path, key); + ASSERT_TRUE(holder.ok()); + auto waiter = start_cache_entry_waiter(root.url.path, key); + ASSERT_EQ( + waiter.attempting.wait_for(kWaitTimeout), std::future_status::ready); + ASSERT_EQ( + waiter.result.wait_for(std::chrono::milliseconds(100)), + std::future_status::timeout); + NSURL* bookmarks = [root.url URLByAppendingPathComponent:@"bookmarks"]; + ASSERT_EQ(rmdir(bookmarks.fileSystemRepresentation), 0); + ASSERT_EQ(mkdir(bookmarks.fileSystemRepresentation, 0700), 0); + holder.get().reset(); + ASSERT_EQ(waiter.result.wait_for(kWaitTimeout), std::future_status::ready); + auto acquired = waiter.result.get(); + ASSERT_TRUE(acquired.ok()); + ASSERT_NE(acquired.get(), nullptr); + NSData* bytes = [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_EQ(acquired.get()->write_bookmark(bytes), Error::Ok); + EXPECT_TRUE([[NSData dataWithContentsOfURL:test_bookmark_url(root.url, key)] + isEqual:bytes]); +} + +TEST_F( + CoreAIBookmarkContentionTest, + ProcessContentionAllowsIndependentKeyAndRecoversAfterHolderDeath) { + BookmarkChild holder; + ASSERT_TRUE(holder.spawn(@"--bookmark-hold", root.url.path, @"same")); + ASSERT_TRUE(holder.receive('L')); + NSURL* file = lock_url(); + struct stat before{}; + ASSERT_EQ(stat(file.fileSystemRepresentation, &before), 0); + BookmarkChild waiter; + ASSERT_TRUE(waiter.spawn(@"--bookmark-lock", root.url.path, @"same")); + ASSERT_TRUE(waiter.receive('B')); + ASSERT_TRUE(waiter.expect_blocked()); + BookmarkChild independent; + ASSERT_TRUE(independent.spawn(@"--bookmark-lock", root.url.path, @"other")); + ASSERT_TRUE(independent.receive('L')); + ASSERT_TRUE(independent.expect_exit(0)); + ASSERT_TRUE(holder.kill_and_reap()); + ASSERT_TRUE(waiter.receive('L')); + ASSERT_TRUE(waiter.expect_exit(0)); + struct stat after{}; + ASSERT_EQ(stat(file.fileSystemRepresentation, &after), 0); + EXPECT_EQ(before.st_dev, after.st_dev); + EXPECT_EQ(before.st_ino, after.st_ino); +} + +TEST_F( + CoreAIBookmarkContentionTest, + CollidingProcessWaiterRecoversAfterHolderDeath) { + BookmarkChild holder; + ASSERT_TRUE(holder.spawn(@"--bookmark-hold", root.url.path, @"same")); + ASSERT_TRUE(holder.receive('L')); + struct stat before{}; + ASSERT_EQ(lstat(lock_url().fileSystemRepresentation, &before), 0); + BookmarkChild waiter; + ASSERT_TRUE(waiter.spawn(@"--bookmark-lock", root.url.path, @"collision")); + ASSERT_TRUE(waiter.receive('B')); + ASSERT_TRUE(waiter.expect_blocked()); + ASSERT_TRUE(holder.kill_and_reap()); + ASSERT_TRUE(waiter.receive('L')); + ASSERT_TRUE(waiter.expect_exit(0)); + struct stat after{}; + ASSERT_EQ(lstat(lock_url().fileSystemRepresentation, &after), 0); + EXPECT_EQ(before.st_dev, after.st_dev); + EXPECT_EQ(before.st_ino, after.st_ino); +} + +TEST_F(CoreAIBookmarkContentionTest, LockInventoryIsBoundedAndStable) { + NSMutableDictionary* representatives = + [NSMutableDictionary dictionary]; + for (NSUInteger index = 0; + index < 100000 && representatives.count < kCacheEntryLockStripeCount; + ++index) { + auto candidate = test_cache_key( + *bundle, [NSString stringWithFormat:@"inventory-%lu", index]); + ASSERT_TRUE(candidate.ok()); + auto name = cache_entry_lock_name(candidate.get()); + ASSERT_TRUE(name.ok()); + representatives[name.get()] = candidate.get(); + } + ASSERT_EQ(representatives.count, kCacheEntryLockStripeCount); + + NSMutableDictionary* inodes = + [NSMutableDictionary dictionary]; + for (NSString* name in representatives) { + auto locked = + acquire_test_cache_entry(root.url.path, representatives[name]); + ASSERT_TRUE(locked.ok()); + NSURL* file = [[root.url URLByAppendingPathComponent:@"locks"] + URLByAppendingPathComponent:name]; + struct stat info {}; + ASSERT_EQ(lstat(file.fileSystemRepresentation, &info), 0); + ASSERT_TRUE(S_ISREG(info.st_mode)); + EXPECT_EQ(info.st_size, 0); + inodes[name] = @(info.st_ino); + } + NSURL* locks = [root.url URLByAppendingPathComponent:@"locks"]; + NSArray* children = + [NSFileManager.defaultManager contentsOfDirectoryAtPath:locks.path + error:nil]; + ASSERT_EQ(children.count, kCacheEntryLockStripeCount); + + for (NSString* name in representatives) { + auto locked = + acquire_test_cache_entry(root.url.path, representatives[name]); + ASSERT_TRUE(locked.ok()); + struct stat info {}; + NSURL* file = [locks URLByAppendingPathComponent:name]; + ASSERT_EQ(lstat(file.fileSystemRepresentation, &info), 0); + EXPECT_EQ(info.st_ino, inodes[name].longLongValue); + } +} + +} // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_cache_entry_fixture.h b/backends/apple/coreai/runtime/test/coreai_cache_entry_fixture.h new file mode 100644 index 00000000000..40dee635fdb --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_cache_entry_fixture.h @@ -0,0 +1,70 @@ +/* + * 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 +#include +#include +#include +#include "coreai_cache_entry.h" +#include "coreai_cache_key.h" +#include "coreai_manifest.h" +#include "coreai_manifest_fixture.h" + +namespace executorch::backends::coreai::testing { + +runtime::Result test_cache_directory( + NSString* path, + StorageCreationMode mode = StorageCreationMode::Create); +runtime::Result> acquire_test_cache_entry( + NSString* path, + NSString* key); +runtime::Result> +acquire_existing_test_cache_entry(NSString* path, NSString* key); +runtime::Result test_cache_key(const SelectedBundle& bundle); +runtime::Result test_cache_key( + const SelectedBundle& bundle, + NSString* specialization_identity); +runtime::Result different_cache_key( + const SelectedBundle& bundle, + NSString* key, + bool same_stripe); +NSURL* test_bookmark_url(NSURL* root, NSString* key); +NSURL* test_cache_lock_url(NSURL* root, NSString* key); +runtime::Result saved_bookmark(NSURL* root, NSString* key); +runtime::Result sized_bookmark(size_t size, uint32_t tag = 1); + +// Returns -1 unless argument is a nonnegative int in decimal. +int parse_child_fd(const char* argument); +bool send_child_byte(int fd, char byte); +void set_bookmark_test_executable(const char* path); +// Returns -1 for normal test invocations, otherwise a child-mode exit code. +int bookmark_child_mode(int argc, char** argv); + +class BookmarkChild { + public: + BookmarkChild() = default; + ~BookmarkChild(); + BookmarkChild(const BookmarkChild&) = delete; + BookmarkChild& operator=(const BookmarkChild&) = delete; + + ::testing::AssertionResult + spawn(NSString* mode, NSString* root, NSString* argument); + ::testing::AssertionResult receive(char expected); + ::testing::AssertionResult expect_blocked(); + ::testing::AssertionResult expect_exit(int expected); + ::testing::AssertionResult kill_and_reap(); + + private: + pid_t pid_ = -1; + int signal_fd_ = -1; + ::testing::AssertionResult reap(int& status); +}; + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_cache_entry_fixture.mm b/backends/apple/coreai/runtime/test/coreai_cache_entry_fixture.mm new file mode 100644 index 00000000000..f5d6b82a661 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_cache_entry_fixture.mm @@ -0,0 +1,442 @@ +/* + * 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. + */ + +#include "coreai_cache_entry_fixture.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "coreai_fault_scope.h" +#include "coreai_filesystem_fixture.h" +#include "coreai_manifest_fixture.h" + +extern char** environ; + +namespace executorch::backends::coreai::testing { +namespace { +using runtime::Error; +using runtime::Result; +using ::testing::AssertionFailure; +using ::testing::AssertionSuccess; +using Clock = std::chrono::steady_clock; +constexpr auto kChildTimeout = std::chrono::seconds(10); +constexpr auto kContentionObservation = std::chrono::milliseconds(100); +constexpr int kChildSignalFD = 3; +std::string executable; + +int remaining_ms(Clock::time_point deadline) { + const auto remaining = std::chrono::duration_cast( + deadline - Clock::now()); + return remaining.count() <= 0 ? 0 : static_cast(remaining.count()); +} + +bool send_child_byte_impl(int fd, char byte) { + const auto deadline = Clock::now() + kChildTimeout; + while (const int remaining = remaining_ms(deadline)) { + pollfd entry{fd, POLLOUT, 0}; + const int ready = poll(&entry, 1, remaining); + if (ready < 0 && errno == EINTR) { + continue; + } + if (ready <= 0 || !(entry.revents & POLLOUT)) { + return false; + } + const ssize_t written = write(fd, &byte, 1); + if (written == 1) { + return true; + } + if (written < 0 && (errno == EINTR || errno == EAGAIN)) { + continue; + } + return false; + } + return false; +} +} // namespace + +Result test_cache_directory( + NSString* path, + StorageCreationMode mode) { + auto directory = resolve_test_cache_directory(path, mode); + if (!directory.ok()) { + return directory.error(); + } + if (!directory->has_value()) { + return Error::NotFound; + } + return std::move(directory->value()); +} + +Result> acquire_test_cache_entry( + NSString* path, + NSString* key) { + auto directory = test_cache_directory(path); + if (!directory.ok()) { + return directory.error(); + } + return LockedCacheEntry::acquire(directory.get(), key); +} + +Result> acquire_existing_test_cache_entry( + NSString* path, + NSString* key) { + auto directory = + resolve_test_cache_directory(path, StorageCreationMode::OpenExisting); + if (!directory.ok()) { + return directory.error(); + } + if (!directory->has_value()) { + return std::unique_ptr(); + } + return LockedCacheEntry::acquire_existing(directory->value(), key); +} + +int parse_child_fd(const char* argument) { + char* end = nullptr; + errno = 0; + const long value = std::strtol(argument, &end, 10); + if (errno != 0 || end == argument || *end != '\0' || value < 0 || + value > INT_MAX) { + return -1; + } + return static_cast(value); +} + +bool send_child_byte(int fd, char byte) { + return send_child_byte_impl(fd, byte); +} + +Result test_cache_key(const SelectedBundle& bundle) { + return test_cache_key(bundle, @"default"); +} + +Result test_cache_key( + const SelectedBundle& bundle, + NSString* specialization_identity) { + return cache_variant_key( + bundle, @"macOS", @"arch_b", specialization_identity); +} + +Result different_cache_key( + const SelectedBundle& bundle, + NSString* key, + bool same_stripe) { + auto original_name = cache_entry_lock_name(key); + if (!original_name.ok()) { + return original_name.error(); + } + for (NSUInteger index = 0; index < 10000; ++index) { + NSString* identity = + [NSString stringWithFormat:@"specialization-%lu", (unsigned long)index]; + auto candidate = test_cache_key(bundle, identity); + if (!candidate.ok()) { + return candidate.error(); + } + if ([candidate.get() isEqual:key]) { + continue; + } + auto candidate_name = cache_entry_lock_name(candidate.get()); + if (!candidate_name.ok()) { + return candidate_name.error(); + } + if ([candidate_name.get() isEqual:original_name.get()] == same_stripe) { + return candidate.get(); + } + } + return Error::NotFound; +} + +NSURL* test_bookmark_url(NSURL* root, NSString* key) { + return [[root URLByAppendingPathComponent:@"bookmarks"] + URLByAppendingPathComponent:[key stringByAppendingString:@".bookmark"]]; +} + +NSURL* test_cache_lock_url(NSURL* root, NSString* key) { + auto name = cache_entry_lock_name(key); + EXPECT_TRUE(name.ok()); + if (!name.ok()) { + return nil; + } + return [[root URLByAppendingPathComponent:@"locks"] + URLByAppendingPathComponent:name.get()]; +} + +Result saved_bookmark(NSURL* root, NSString* key) { + auto locked = acquire_test_cache_entry(root.path, key); + if (!locked.ok()) { + return locked.error(); + } + return locked.get()->read_bookmark(); +} + +Result sized_bookmark(size_t size, uint32_t tag) { + if (size < sizeof(tag)) { + return Error::InvalidArgument; + } + NSMutableData* bytes = [NSMutableData dataWithLength:size]; + if (bytes == nil) { + return Error::MemoryAllocationFailed; + } + std::memcpy(bytes.mutableBytes, &tag, sizeof(tag)); + return [bytes copy]; +} + +void set_bookmark_test_executable(const char* path) { + executable = path == nullptr ? "" : path; +} + +int bookmark_child_mode(int argc, char** argv) { + if (argc < 2 || std::strncmp(argv[1], "--bookmark-", 11) != 0) { + return -1; + } + if (argc != 5) { + return 1; + } + const bool hold = std::strcmp(argv[1], "--bookmark-hold") == 0; + const bool lock = std::strcmp(argv[1], "--bookmark-lock") == 0; + const bool other = std::strcmp(argv[3], "other") == 0; + const bool same = std::strcmp(argv[3], "same") == 0; + const bool collision = std::strcmp(argv[3], "collision") == 0; + if ((!hold && !lock) || (!other && !same && !collision)) { + return 2; + } + const int parsed_fd = parse_child_fd(argv[4]); + if (parsed_fd < 0) { + return 3; + } + FileDescriptor signal_fd(parsed_fd); + if (fcntl(signal_fd.get(), F_GETFD) < 0) { + return 4; + } + signal(SIGPIPE, SIG_IGN); + NSString* root = [NSString stringWithUTF8String:argv[2]]; + if (root == nil) { + return 5; + } + auto bundle = selected_bundle(manifest_dict()); + if (!bundle.ok()) { + return 6; + } + auto key = test_cache_key(bundle.get()); + if (!key.ok()) { + return 7; + } + NSString* selected_key = key.get(); + if (other || collision) { + auto changed = different_cache_key(bundle.get(), key.get(), collision); + if (!changed.ok()) { + return 8; + } + selected_key = changed.get(); + } + StorageFaultScope scope(FaultPoint::BeforeCacheEntryLock); + scope.armed = false; + const int notify_fd = signal_fd.get(); + scope.observe = ^(FaultPoint point) { + if (lock && (same || collision) && + point == FaultPoint::BeforeCacheEntryLock && + !send_child_byte_impl(notify_fd, 'B')) { + _exit(12); + } + }; + auto locked = acquire_test_cache_entry(root, selected_key); + if (!locked.ok()) { + return 9; + } + if (!send_child_byte_impl(signal_fd.get(), 'L')) { + return 10; + } + if (close(signal_fd.release()) != 0) { + return 11; + } + if (hold) { + // The parent owns this PID and always kills and reaps it. + for (;;) { + pause(); + } + } + return 0; +} + +BookmarkChild::~BookmarkChild() { + if (pid_ > 0) { + const auto cleaned = kill_and_reap(); + if (!cleaned) { + ADD_FAILURE() << cleaned.message(); + } + } + if (signal_fd_ >= 0 && close(signal_fd_) != 0) { + ADD_FAILURE() << "Cannot close child signal descriptor: " << errno; + } +} + +::testing::AssertionResult +BookmarkChild::spawn(NSString* mode, NSString* root, NSString* argument) { + if (pid_ > 0 || signal_fd_ >= 0 || executable.empty() || root == nil) { + return AssertionFailure() << "Invalid bookmark child setup"; + } + int descriptors[2]; + if (pipe(descriptors) != 0) { + return AssertionFailure() << "pipe: " << errno; + } + FileDescriptor reader(descriptors[0]); + FileDescriptor writer(descriptors[1]); + for (int fd : descriptors) { + if (fcntl(fd, F_SETFD, FD_CLOEXEC) != 0 || + fcntl(fd, F_SETFL, O_NONBLOCK) != 0) { + return AssertionFailure() << "fcntl pipe: " << errno; + } + } + posix_spawn_file_actions_t actions; + int error = posix_spawn_file_actions_init(&actions); + if (error != 0) { + return AssertionFailure() << "spawn actions: " << error; + } + error = + posix_spawn_file_actions_adddup2(&actions, writer.get(), kChildSignalFD); + std::string descriptor = std::to_string(kChildSignalFD); + char* arguments[] = { + const_cast(executable.c_str()), + const_cast(mode.UTF8String), + const_cast(root.fileSystemRepresentation), + const_cast(argument.UTF8String), + descriptor.data(), + nullptr}; + if (error == 0) { + error = + posix_spawn(&pid_, arguments[0], &actions, nullptr, arguments, environ); + } + const int destroyed = posix_spawn_file_actions_destroy(&actions); + if (error != 0) { + pid_ = -1; + return AssertionFailure() << "posix_spawn: " << error; + } + signal_fd_ = reader.release(); + if (destroyed != 0) { + return AssertionFailure() << "destroy spawn actions: " << destroyed; + } + return AssertionSuccess(); +} + +::testing::AssertionResult BookmarkChild::receive(char expected) { + const auto deadline = Clock::now() + kChildTimeout; + while (const int remaining = remaining_ms(deadline)) { + pollfd entry{signal_fd_, POLLIN, 0}; + const int ready = poll(&entry, 1, remaining); + if (ready < 0 && errno == EINTR) { + continue; + } + if (ready < 0) { + return AssertionFailure() << "poll: " << errno; + } + if (ready == 0) { + break; + } + if (!(entry.revents & POLLIN)) { + return AssertionFailure() + << "Child closed signal pipe before " << expected; + } + char byte = 0; + const ssize_t count = read(signal_fd_, &byte, 1); + if (count < 0 && (errno == EINTR || errno == EAGAIN)) { + continue; + } + if (count != 1 || byte != expected) { + return AssertionFailure() + << "Expected child signal " << expected << ", got " << int(byte); + } + return AssertionSuccess(); + } + return AssertionFailure() + << "Timed out waiting for child signal " << expected; +} + +::testing::AssertionResult BookmarkChild::expect_blocked() { + const auto deadline = Clock::now() + kContentionObservation; + while (const int remaining = remaining_ms(deadline)) { + pollfd entry{signal_fd_, POLLIN, 0}; + const int ready = poll(&entry, 1, remaining); + if (ready < 0 && errno == EINTR) { + continue; + } + if (ready < 0) { + return AssertionFailure() << "poll: " << errno; + } + if (ready > 0) { + return AssertionFailure() << "Child completed while lock was held"; + } + return AssertionSuccess(); + } + return AssertionSuccess(); +} + +::testing::AssertionResult BookmarkChild::reap(int& status) { + if (pid_ <= 0) { + return AssertionFailure() << "No owned child to reap"; + } + const auto deadline = Clock::now() + kChildTimeout; + while (const int remaining = remaining_ms(deadline)) { + const pid_t result = waitpid(pid_, &status, WNOHANG); + if (result == pid_) { + pid_ = -1; + return AssertionSuccess(); + } + if (result < 0 && errno == EINTR) { + continue; + } + if (result < 0) { + return AssertionFailure() << "waitpid: " << errno; + } + poll(nullptr, 0, std::min(remaining, 10)); + } + return AssertionFailure() << "Timed out reaping child " << pid_; +} + +::testing::AssertionResult BookmarkChild::expect_exit(int expected) { + int status = 0; + auto result = reap(status); + if (!result) { + return result; + } + if (!WIFEXITED(status) || WEXITSTATUS(status) != expected) { + return AssertionFailure() + << "Expected child exit " << expected << ", wait status " << status; + } + return AssertionSuccess(); +} + +::testing::AssertionResult BookmarkChild::kill_and_reap() { + if (pid_ <= 0) { + return AssertionFailure() << "No owned child to kill"; + } + const int killed = kill(pid_, SIGKILL); + const int error = errno; + int status = 0; + auto result = reap(status); + if (!result) { + return result; + } + if (killed != 0 || !WIFSIGNALED(status) || WTERMSIG(status) != SIGKILL) { + return AssertionFailure() + << "Expected SIGKILL, kill errno " << (killed == 0 ? 0 : error) + << ", wait status " << status; + } + return AssertionSuccess(); +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_cache_entry_test.mm b/backends/apple/coreai/runtime/test/coreai_cache_entry_test.mm new file mode 100644 index 00000000000..b277344b27a --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_cache_entry_test.mm @@ -0,0 +1,521 @@ +/* + * 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. + */ + +#include +#include +#include +#include "coreai_cache_entry_fixture.h" +#include "coreai_fault_scope.h" +#include "coreai_filesystem_fixture.h" +#include "coreai_storage.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +namespace { +constexpr size_t kMiB = 1024 * 1024; + +class CoreAIBookmarkPathTest : public ::testing::Test { + protected: + TestDirectory root; + std::unique_ptr bundle; + NSString* key = nil; + + void SetUp() override { + ASSERT_NE(root.url, nil); + auto selected = selected_bundle(manifest_dict()); + ASSERT_TRUE(selected.ok()); + bundle = std::make_unique(selected.get()); + auto result = test_cache_key(*bundle); + ASSERT_TRUE(result.ok()); + key = result.get(); + } + + NSURL* lock_url() { + return test_cache_lock_url(root.url, key); + } +}; + +class CoreAIBookmarkIOTest : public CoreAIBookmarkPathTest { + protected: + std::unique_ptr locked; + NSURL* url = nil; + + void SetUp() override { + CoreAIBookmarkPathTest::SetUp(); + if (HasFatalFailure()) { + return; + } + auto result = acquire_test_cache_entry(root.url.path, key); + ASSERT_TRUE(result.ok()); + locked = std::move(result.get()); + ASSERT_NE(locked, nullptr); + ASSERT_EQ(locked->persistence(), CacheEntryPersistence::BookmarkCapable); + url = test_bookmark_url(root.url, key); + } +}; + +::testing::AssertionResult bookmark_bytes_equal( + const LockedCacheEntry& lock, + NSData* expected) { + auto result = lock.read_bookmark(); + if (!result.ok()) { + return ::testing::AssertionFailure() + << "Bookmark read error " << int(result.error()); + } + if (expected == nil ? result.get() != nil + : ![result.get() isEqual:expected]) { + return ::testing::AssertionFailure() << "Bookmark bytes differ"; + } + return ::testing::AssertionSuccess(); +} + +::testing::AssertionResult no_temporary_bookmarks(NSURL* root) { + NSURL* directory = [root URLByAppendingPathComponent:@"bookmarks"]; + FileDescriptor fd(open( + directory.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + if (fd.get() < 0) { + return ::testing::AssertionFailure() + << "open bookmark directory: " << errno; + } + auto children = storage_children(fd.get()); + if (!children.ok()) { + return ::testing::AssertionFailure() + << "Cannot enumerate bookmark directory"; + } + for (NSString* name in children.get()) { + if ([name hasPrefix:@".tmp-"]) { + return ::testing::AssertionFailure() << "Leaked " << name.UTF8String; + } + } + return ::testing::AssertionSuccess(); +} + +TEST_F(CoreAIBookmarkIOTest, MissingBookmarkIsAnEmptySuccessfulRead) { + EXPECT_TRUE(bookmark_bytes_equal(*locked, nil)); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:url.path]); +} + +TEST_F(CoreAIBookmarkIOTest, RoundTripsMaximumBookmarkAndReopens) { + auto bytes = sized_bookmark(8 * kMiB); + ASSERT_TRUE(bytes.ok()); + ASSERT_EQ(locked->write_bookmark(bytes.get()), Error::Ok); + EXPECT_TRUE(bookmark_bytes_equal(*locked, bytes.get())); + EXPECT_TRUE([[NSData dataWithContentsOfURL:url] isEqual:bytes.get()]); + locked.reset(); + auto saved = saved_bookmark(root.url, key); + ASSERT_TRUE(saved.ok()); + EXPECT_TRUE([saved.get() isEqual:bytes.get()]); +} + +TEST_F(CoreAIBookmarkIOTest, InvalidWritesPreserveExistingBookmark) { + auto maximum = sized_bookmark(8 * kMiB); + ASSERT_TRUE(maximum.ok()); + ASSERT_EQ(locked->write_bookmark(maximum.get()), Error::Ok); + auto oversized = sized_bookmark(8 * kMiB + 1); + ASSERT_TRUE(oversized.ok()); + for (NSData* invalid : {[NSData data], oversized.get()}) { + SCOPED_TRACE(invalid.length); + EXPECT_EQ(locked->write_bookmark(invalid), Error::InvalidExternalData); + EXPECT_TRUE([[NSData dataWithContentsOfURL:url] isEqual:maximum.get()]); + EXPECT_TRUE(bookmark_bytes_equal(*locked, maximum.get())); + EXPECT_TRUE(no_temporary_bookmarks(root.url)); + } +} + +TEST_F(CoreAIBookmarkIOTest, RejectsEmptyStoredBookmark) { + FileDescriptor file(open( + url.fileSystemRepresentation, + O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, + 0600)); + ASSERT_GE(file.get(), 0); + EXPECT_EQ(locked->read_bookmark().error(), Error::InvalidExternalData); +} + +TEST_F(CoreAIBookmarkIOTest, RejectsSymlinkWithoutReadingTarget) { + NSURL* target = [root.url URLByAppendingPathComponent:@"target"]; + NSData* bytes = [@"target bytes" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([bytes writeToURL:target atomically:NO]); + ASSERT_EQ( + symlink(target.fileSystemRepresentation, url.fileSystemRepresentation), + 0); + EXPECT_EQ(locked->read_bookmark().error(), Error::AccessFailed); + EXPECT_TRUE([[NSData dataWithContentsOfURL:target] isEqual:bytes]); +} + +TEST_F(CoreAIBookmarkIOTest, RejectsNonregularStoredBookmarks) { + ASSERT_EQ(mkdir(url.fileSystemRepresentation, 0700), 0); + EXPECT_EQ(locked->read_bookmark().error(), Error::InvalidExternalData); + ASSERT_EQ(rmdir(url.fileSystemRepresentation), 0); + ASSERT_EQ(mkfifo(url.fileSystemRepresentation, 0600), 0); + EXPECT_EQ(locked->read_bookmark().error(), Error::InvalidExternalData); +} + +TEST_F(CoreAIBookmarkIOTest, RejectsInaccessibleBookmark) { + if (geteuid() == 0) { + GTEST_SKIP() << "Root can bypass file read permissions"; + } + NSData* bytes = [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_EQ(locked->write_bookmark(bytes), Error::Ok); + ASSERT_EQ(chmod(url.fileSystemRepresentation, 0000), 0); + EXPECT_EQ(locked->read_bookmark().error(), Error::AccessFailed); + ASSERT_EQ(chmod(url.fileSystemRepresentation, 0600), 0); + EXPECT_TRUE(bookmark_bytes_equal(*locked, bytes)); +} + +TEST_F(CoreAIBookmarkIOTest, RejectsOversizedStoredBookmark) { + FileDescriptor file(open( + url.fileSystemRepresentation, + O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, + 0600)); + ASSERT_GE(file.get(), 0); + ASSERT_EQ(ftruncate(file.get(), 8 * kMiB + 1), 0); + ASSERT_EQ(close(file.release()), 0); + EXPECT_EQ(locked->read_bookmark().error(), Error::InvalidExternalData); + struct stat info {}; + ASSERT_EQ(lstat(url.fileSystemRepresentation, &info), 0); + EXPECT_EQ(info.st_size, static_cast(8 * kMiB + 1)); +} + +TEST_F(CoreAIBookmarkIOTest, ExistingLookupReadsWithoutPreparingSources) { + NSData* bytes = [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_EQ(locked->write_bookmark(bytes), Error::Ok); + locked.reset(); + auto existing = acquire_existing_test_cache_entry(root.url.path, key); + ASSERT_TRUE(existing.ok()); + ASSERT_NE(existing.get(), nullptr); + EXPECT_EQ( + existing.get()->persistence(), CacheEntryPersistence::BookmarkCapable); + EXPECT_TRUE(bookmark_bytes_equal(*existing.get(), bytes)); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:[root.url URLByAppendingPathComponent:@"sources"].path]); +} + +TEST_F(CoreAIBookmarkIOTest, PreparesSharedSourcesParentOnDemand) { + EXPECT_TRUE([locked->key() isEqual:key]); + NSString* expected = [root.url.URLByResolvingSymlinksInPath + URLByAppendingPathComponent:@"sources"] + .path; + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:expected]); + auto sources = locked->prepare_source_directory(); + ASSERT_TRUE(sources.ok()); + EXPECT_TRUE([sources->url.path isEqual:expected]); + struct stat info {}; + ASSERT_EQ(lstat(sources->url.fileSystemRepresentation, &info), 0); + EXPECT_TRUE(S_ISDIR(info.st_mode)); + auto repeated = locked->prepare_source_directory(); + ASSERT_TRUE(repeated.ok()); + EXPECT_TRUE([repeated->url.path isEqual:expected]); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:[expected stringByAppendingPathComponent:key]]); +} + +TEST_F(CoreAIBookmarkIOTest, CleanupDoesNotCreateMissingSources) { + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); + EXPECT_EQ(locked->remove_source(), Error::Ok); + EXPECT_EQ(fault.matches, 0); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:[root.url URLByAppendingPathComponent:@"sources"].path]); +} + +TEST_F(CoreAIBookmarkIOTest, CleanupRejectsUnexpectedTypesBeforeDeletingFiles) { + auto sources = locked->prepare_source_directory(); + ASSERT_TRUE(sources.ok()); + NSString* entry = [sources->url.path stringByAppendingPathComponent:key]; + ASSERT_EQ(mkdir(entry.fileSystemRepresentation, 0700), 0); + NSString* file = [entry stringByAppendingPathComponent:@"data"]; + NSData* bytes = [@"preserved" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([bytes writeToFile:file atomically:NO]); + NSString* pipe = [entry stringByAppendingPathComponent:@"pipe"]; + ASSERT_EQ(mkfifo(pipe.fileSystemRepresentation, 0600), 0); + + EXPECT_EQ(locked->remove_source(), Error::InvalidExternalData); + EXPECT_TRUE([[NSData dataWithContentsOfFile:file] isEqual:bytes]); + struct stat info; + ASSERT_EQ(lstat(pipe.fileSystemRepresentation, &info), 0); + EXPECT_TRUE(S_ISFIFO(info.st_mode)); +} + +TEST_F(CoreAIBookmarkIOTest, CleanupUsesLockedRootAfterPathReplacement) { + auto sources = locked->prepare_source_directory(); + ASSERT_TRUE(sources.ok()); + NSString* entry = [sources->url.path stringByAppendingPathComponent:key]; + ASSERT_EQ(mkdir(entry.fileSystemRepresentation, 0700), 0); + TestDirectory moved; + ASSERT_NE(moved.url, nil); + ASSERT_EQ( + rename( + root.url.fileSystemRepresentation, + moved.url.fileSystemRepresentation), + 0); + NSURL* replacement = [[root.url URLByAppendingPathComponent:@"sources"] + URLByAppendingPathComponent:key]; + ASSERT_TRUE([NSFileManager.defaultManager createDirectoryAtURL:replacement + withIntermediateDirectories:YES + attributes:nil + error:nil]); + + EXPECT_EQ(locked->remove_source(), Error::Ok); + NSURL* original = [[moved.url URLByAppendingPathComponent:@"sources"] + URLByAppendingPathComponent:key]; + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:original.path]); + EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:replacement.path]); +} + +TEST_F(CoreAIBookmarkIOTest, CreatesDirectoriesWithDefaultPermissions) { + NSURL* reference = [root.url URLByAppendingPathComponent:@"reference"]; + ASSERT_EQ(mkdir(reference.fileSystemRepresentation, 0777), 0); + struct stat expected; + ASSERT_EQ(stat(reference.fileSystemRepresentation, &expected), 0); + ASSERT_TRUE(locked->prepare_source_directory().ok()); + for (NSString* name in @[ @"locks", @"bookmarks", @"sources" ]) { + NSURL* directory = [root.url URLByAppendingPathComponent:name]; + struct stat actual; + ASSERT_EQ(stat(directory.fileSystemRepresentation, &actual), 0); + EXPECT_EQ(actual.st_mode & 0777, expected.st_mode & 0777); + } +} + +TEST_F(CoreAIBookmarkIOTest, ReusesDirectoriesWithoutCreation) { + ASSERT_TRUE(locked->prepare_source_directory().ok()); + ASSERT_EQ( + locked->write_bookmark( + [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]), + Error::Ok); + locked.reset(); + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); + auto created = acquire_test_cache_entry(root.url.path, key); + ASSERT_TRUE(created.ok()); + ASSERT_NE(created.get(), nullptr); + EXPECT_TRUE(created.get()->prepare_source_directory().ok()); + created.get().reset(); + auto existing = acquire_existing_test_cache_entry(root.url.path, key); + ASSERT_TRUE(existing.ok()); + ASSERT_NE(existing.get(), nullptr); + EXPECT_TRUE(existing.get()->prepare_source_directory().ok()); + EXPECT_EQ(fault.matches, 0); +} + +TEST_F(CoreAIBookmarkIOTest, RecreatesMissingSourcesParent) { + auto sources = locked->prepare_source_directory(); + ASSERT_TRUE(sources.ok()); + ASSERT_EQ(rmdir(sources->url.fileSystemRepresentation), 0); + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); + EXPECT_EQ(locked->prepare_source_directory().error(), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + auto recreated = locked->prepare_source_directory(); + ASSERT_TRUE(recreated.ok()); + EXPECT_TRUE([recreated->url isEqual:sources->url]); +} + +TEST_F(CoreAIBookmarkIOTest, AcceptsConcurrentSourcesCreation) { + NSURL* sources = [root.url URLByAppendingPathComponent:@"sources"]; + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); + fault.armed = false; + __block int creations = 0; + fault.observe = ^(FaultPoint point) { + if (point == FaultPoint::BeforeCreateDirectory) { + ++creations; + EXPECT_EQ(mkdir(sources.fileSystemRepresentation, 0700), 0); + } + }; + EXPECT_TRUE(locked->prepare_source_directory().ok()); + EXPECT_EQ(creations, 1); +} + +TEST_F(CoreAIBookmarkIOTest, RejectsInvalidSourcesParent) { + TestDirectory target; + ASSERT_NE(target.url, nil); + NSURL* sources = [root.url URLByAppendingPathComponent:@"sources"]; + ASSERT_EQ( + symlink( + target.url.fileSystemRepresentation, + sources.fileSystemRepresentation), + 0); + EXPECT_EQ(locked->prepare_source_directory().error(), Error::AccessFailed); + ASSERT_EQ(unlink(sources.fileSystemRepresentation), 0); + NSData* bytes = [@"not a directory" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([bytes writeToURL:sources atomically:NO]); + EXPECT_EQ(locked->prepare_source_directory().error(), Error::AccessFailed); + EXPECT_TRUE([[NSData dataWithContentsOfURL:sources] isEqual:bytes]); +} + +TEST_F(CoreAIBookmarkPathTest, AllowsTemporaryRootsAndRejectsInvalidKeys) { + auto allowed = acquire_test_cache_entry(root.url.path, key); + ASSERT_TRUE(allowed.ok()); + ASSERT_NE(allowed.get(), nullptr); + EXPECT_FALSE(acquire_test_cache_entry(root.url.path, @"../escape").ok()); +} + +TEST_F(CoreAIBookmarkPathTest, SnapshotsMutableRootAtAcquisition) { + TestDirectory other; + ASSERT_NE(other.url, nil); + NSMutableString* path = [root.url.path mutableCopy]; + auto locked = acquire_test_cache_entry(path, key); + ASSERT_TRUE(locked.ok()); + ASSERT_NE(locked.get(), nullptr); + [path setString:other.url.path]; + + auto sources = locked.get()->prepare_source_directory(); + ASSERT_TRUE(sources.ok()); + NSString* expected = [root.url.URLByResolvingSymlinksInPath + URLByAppendingPathComponent:@"sources"] + .path; + NSString* unrelated = [other.url URLByAppendingPathComponent:@"sources"].path; + EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:expected]); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:unrelated]); + EXPECT_TRUE([sources->url.path isEqual:expected]) + << "Returned sources path: " << sources->url.path.UTF8String + << "; descriptor-pinned sources path: " << expected.UTF8String; +} + +TEST_F(CoreAIBookmarkPathTest, ExcludesRootWithoutCreatingSources) { + auto locked = acquire_test_cache_entry(root.url.path, key); + ASSERT_TRUE(locked.ok()); + EXPECT_TRUE(backup_excluded(root.url)); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:[root.url URLByAppendingPathComponent:@"sources"].path]); +} + +TEST_F(CoreAIBookmarkPathTest, MissingLookupLeavesRootAndEntryUntouched) { + ASSERT_TRUE(set_backup_excluded(root.url, false)); + NSString* missing = [root.url URLByAppendingPathComponent:@"missing"].path; + for (NSString* path in @[ root.url.path, missing ]) { + SCOPED_TRACE(path.UTF8String); + auto resolved = + resolve_test_cache_directory(path, StorageCreationMode::OpenExisting); + ASSERT_TRUE(resolved.ok()); + const bool root_exists = [path isEqual:root.url.path]; + EXPECT_EQ(resolved->has_value(), root_exists); + if (resolved->has_value()) { + EXPECT_TRUE([resolved->value().path + isEqual:[NSURL fileURLWithPath:path] + .URLByResolvingSymlinksInPath.path]); + } + auto absent = acquire_existing_test_cache_entry(path, key); + ASSERT_TRUE(absent.ok()); + EXPECT_EQ(absent.get(), nullptr); + } + NSArray* contents = + [NSFileManager.defaultManager contentsOfDirectoryAtPath:root.url.path + error:nil]; + ASSERT_NE(contents, nil); + EXPECT_EQ(contents.count, 0u); + EXPECT_FALSE(backup_excluded(root.url)); +} + +TEST_F(CoreAIBookmarkPathTest, MissingLookupIgnoresOtherEntries) { + auto other_key = different_cache_key(*bundle, key, false); + ASSERT_TRUE(other_key.ok()); + auto other = acquire_test_cache_entry(root.url.path, other_key.get()); + ASSERT_TRUE(other.ok()); + NSData* bytes = [@"other" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_EQ(other.get()->write_bookmark(bytes), Error::Ok); + auto absent = acquire_existing_test_cache_entry(root.url.path, key); + ASSERT_TRUE(absent.ok()); + EXPECT_EQ(absent.get(), nullptr); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:lock_url().path]); + EXPECT_TRUE(bookmark_bytes_equal(*other.get(), bytes)); +} + +TEST_F(CoreAIBookmarkPathTest, SourceOnlyLookupHasExplicitPersistence) { + NSURL* staged = [[root.url URLByAppendingPathComponent:@"sources"] + URLByAppendingPathComponent:key]; + ASSERT_TRUE([NSFileManager.defaultManager createDirectoryAtURL:staged + withIntermediateDirectories:YES + attributes:nil + error:nil]); + auto existing = acquire_existing_test_cache_entry(root.url.path, key); + ASSERT_TRUE(existing.ok()); + ASSERT_NE(existing.get(), nullptr); + EXPECT_EQ(existing.get()->persistence(), CacheEntryPersistence::SourceOnly); + EXPECT_TRUE(bookmark_bytes_equal(*existing.get(), nil)); + NSData* bytes = [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + EXPECT_EQ(existing.get()->write_bookmark(bytes), Error::InvalidArgument); + EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:lock_url().path]); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:[root.url URLByAppendingPathComponent:@"bookmarks"] + .path]); +} + +TEST_F(CoreAIBookmarkPathTest, ResolvesAndCreatesRootWithBackupExclusion) { + ASSERT_TRUE(set_backup_excluded(root.url, false)); + NSString* path = [root.url URLByAppendingPathComponent:@"nested/cache"].path; + auto resolved = + resolve_test_cache_directory(path, StorageCreationMode::Create); + ASSERT_TRUE(resolved.ok()); + ASSERT_TRUE(resolved->has_value()); + EXPECT_TRUE([resolved->value().path + isEqual:[NSURL fileURLWithPath:path].URLByResolvingSymlinksInPath.path]); + EXPECT_TRUE(backup_excluded([NSURL fileURLWithPath:resolved->value().path])); + EXPECT_FALSE(backup_excluded(root.url)); + auto existing = + resolve_test_cache_directory(path, StorageCreationMode::OpenExisting); + ASSERT_TRUE(existing.ok()); + ASSERT_TRUE(existing->has_value()); + EXPECT_TRUE([existing->value().path + isEqual:[NSURL fileURLWithPath:path].URLByResolvingSymlinksInPath.path]); +} + +TEST_F(CoreAIBookmarkPathTest, RejectsInvalidResolvedRoots) { + for (StorageCreationMode mode : + {StorageCreationMode::OpenExisting, StorageCreationMode::Create}) { + for (NSString* path in + @[ @"relative", @"~", @"~/cache", @"~nobody/cache" ]) { + SCOPED_TRACE(path.UTF8String); + EXPECT_EQ( + resolve_test_cache_directory(path, mode).error(), + Error::InvalidArgument); + } + EXPECT_EQ( + resolve_test_cache_directory(nil, mode).error(), + Error::InvalidArgument); + } + NSURL* file = [root.url URLByAppendingPathComponent:@"file"]; + ASSERT_TRUE([[NSData data] writeToURL:file atomically:NO]); + EXPECT_EQ( + resolve_test_cache_directory(file.path, StorageCreationMode::OpenExisting) + .error(), + Error::AccessFailed); + EXPECT_EQ( + resolve_test_cache_directory( + [file.path stringByAppendingPathComponent:@"child"], + StorageCreationMode::OpenExisting) + .error(), + Error::AccessFailed); +} + +TEST_F(CoreAIBookmarkPathTest, RejectsInvalidStripeFiles) { + NSData* bytes = [@"untouched target" dataUsingEncoding:NSUTF8StringEncoding]; + NSURL* target = [root.url URLByAppendingPathComponent:@"target"]; + ASSERT_TRUE([bytes writeToURL:target atomically:NO]); + for (int kind = 0; kind < 3; ++kind) { + TestDirectory isolated; + ASSERT_NE(isolated.url, nil); + NSURL* locks = [isolated.url URLByAppendingPathComponent:@"locks"]; + ASSERT_EQ(mkdir(locks.fileSystemRepresentation, 0700), 0); + NSURL* stripe = test_cache_lock_url(isolated.url, key); + if (kind == 0) { + ASSERT_EQ( + symlink( + target.fileSystemRepresentation, stripe.fileSystemRepresentation), + 0); + } else if (kind == 1) { + ASSERT_EQ(mkdir(stripe.fileSystemRepresentation, 0700), 0); + } else { + ASSERT_EQ(mkfifo(stripe.fileSystemRepresentation, 0600), 0); + } + EXPECT_EQ( + acquire_test_cache_entry(isolated.url.path, key).error(), + Error::AccessFailed); + EXPECT_TRUE([[NSData dataWithContentsOfURL:target] isEqual:bytes]); + } +} + +} // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_cache_key_test.mm b/backends/apple/coreai/runtime/test/coreai_cache_key_test.mm new file mode 100644 index 00000000000..03fe45ca4a0 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_cache_key_test.mm @@ -0,0 +1,209 @@ +/* + * 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_cache_key.h" + +#include +#include +#include "coreai_cache_entry_fixture.h" +#include "coreai_manifest_fixture.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +namespace { + +NSString* key_for_device(NSString* device) { + NSString* model = [@"" stringByPaddingToLength:64 + withString:@"a" + startingAtIndex:0]; + NSString* options = [@"" stringByPaddingToLength:64 + withString:@"b" + startingAtIndex:0]; + return [NSString stringWithFormat:@"%@-%@-%@", model, device, options]; +} + +TEST(CoreAICacheKeyIdentityTest, IgnoresBindingsAndDeliveryLocation) { + auto baseline = selected_bundle(manifest_dict()); + ASSERT_TRUE(baseline.ok()); + auto expected = test_cache_key(baseline.get()); + ASSERT_TRUE(expected.ok()); + + auto dictionary = manifest_dict(); + dictionary[@"export_hash"] = @"cd"; + dictionary[@"path"] = @"cd/model.aimodel"; + dictionary[@"function"] = @"another_function"; + dictionary[@"input_names"] = @[ @"different_input" ]; + dictionary[@"output_names"] = @[ @"different_output" ]; + auto changed = selected_bundle(dictionary); + ASSERT_TRUE(changed.ok()); + auto actual = test_cache_key(changed.get()); + ASSERT_TRUE(actual.ok()); + EXPECT_TRUE([actual.get() isEqual:expected.get()]); +} + +TEST(CoreAICacheKeyIdentityTest, SeparatesExportedBundleIdentity) { + auto baseline = selected_bundle(manifest_dict()); + ASSERT_TRUE(baseline.ok()); + auto expected = test_cache_key(baseline.get()); + ASSERT_TRUE(expected.ok()); + + auto digest_changed = manifest_dict(); + digest_changed[@"bundle_digests"] = + @{@"model.aimodel" : fixture_identity('2')}; + auto reweighted = selected_bundle(digest_changed); + ASSERT_TRUE(reweighted.ok()); + + auto aot = selected_bundle(aot_manifest_dict()); + ASSERT_TRUE(aot.ok()); + for (const SelectedBundle* bundle : {&reweighted.get(), &aot.get()}) { + auto actual = test_cache_key(*bundle); + ASSERT_TRUE(actual.ok()); + EXPECT_FALSE([actual.get() isEqual:expected.get()]); + } +} + +TEST(CoreAICacheKeyIdentityTest, SeparatesBundleNameWithSameDigestAndMode) { + auto dictionary = aot_manifest_dict(); + dictionary[@"bundle_digests"] = @{ + @"model.arch_a.aimodelc" : fixture_identity('b'), + @"model.arch_b.aimodelc" : fixture_identity('b') + }; + auto manifest = parse_manifest(encode(dictionary)); + ASSERT_TRUE(manifest.ok()); + auto first = select_bundle(manifest.get(), @"arch_a", @"macOS"); + auto second = select_bundle(manifest.get(), @"arch_b", @"macOS"); + ASSERT_TRUE(first.ok()); + ASSERT_TRUE(second.ok()); + auto first_key = bundle_identity(first.get()); + auto second_key = bundle_identity(second.get()); + ASSERT_TRUE(first_key.ok()); + ASSERT_TRUE(second_key.ok()); + EXPECT_FALSE([first_key.get() isEqual:second_key.get()]); +} + +TEST(CoreAICacheKeyIdentityTest, RejectsUnknownBundleName) { + auto manifest = parse_manifest(encode(manifest_dict())); + ASSERT_TRUE(manifest.ok()); + EXPECT_EQ( + bundle_identity(manifest.get(), @"unknown.aimodel").error(), + Error::InvalidArgument); +} + +TEST(CoreAICacheKeyIdentityTest, SeparatesDeviceAndSpecializationIdentity) { + auto bundle = selected_bundle(manifest_dict()); + ASSERT_TRUE(bundle.ok()); + auto baseline = test_cache_key(bundle.get()); + auto platform = + cache_variant_key(bundle.get(), @"iOS", @"arch_b", @"default"); + auto architecture = + cache_variant_key(bundle.get(), @"macOS", @"arch_a", @"default"); + auto specialization = cache_variant_key( + bundle.get(), + @"macOS", + @"arch_b", + @"{\"version\":1,\"compute\":[\"cpu\"]}"); + ASSERT_TRUE(baseline.ok()); + for (auto* changed : {&platform, &architecture, &specialization}) { + ASSERT_TRUE(changed->ok()); + EXPECT_FALSE([changed->get() isEqual:baseline.get()]); + EXPECT_TRUE([cache_key_bundle_identity(changed->get()) + isEqual:cache_key_bundle_identity(baseline.get())]); + } +} + +TEST(CoreAICacheKeyIdentityTest, RejectsIncompletePlainIdentity) { + auto bundle = selected_bundle(manifest_dict()); + ASSERT_TRUE(bundle.ok()); + EXPECT_FALSE( + cache_variant_key(bundle.get(), nil, @"arch_b", @"default").ok()); + EXPECT_FALSE(cache_variant_key(bundle.get(), @"macOS", @"", @"default").ok()); + EXPECT_FALSE(cache_variant_key(bundle.get(), @"macOS", @"arch_b", nil).ok()); + EXPECT_FALSE( + cache_variant_key(bundle.get(), @"mac_OS", @"arch_b", @"default").ok()); +} + +TEST(CoreAICacheKeyFormatTest, AcceptsReadableDeviceComponents) { + for (NSString* device in + @[ @"macOS_h17p", @"iOS_arch_b", @"visionOS_h17p-high" ]) { + NSString* key = key_for_device(device); + EXPECT_TRUE(is_cache_key(key)); + EXPECT_TRUE( + [cache_key_bundle_identity(key) isEqual:[key substringToIndex:64]]); + } +} + +TEST(CoreAICacheKeyFormatTest, RejectsUnsafeOrIncompleteDevices) { + for (NSString* device in @[ + @"", + @"iOS", + @"iOS_", + @"_h17p", + @"iOS_../h17p", + @"iOS_h17p\\extra", + @"iOS_h17p extra", + @"iOS_h\u00e9", + @"iOS_h17p.bookmark" + ]) { + NSString* key = key_for_device(device); + EXPECT_FALSE(is_cache_key(key)); + EXPECT_EQ(cache_key_bundle_identity(key), nil); + } +} + +TEST(CoreAICacheKeyFormatTest, LeavesRoomForBookmarkExtension) { + constexpr NSUInteger max_device_length = NAME_MAX - 130 - 9; + NSString* device = [@"iOS_" stringByPaddingToLength:max_device_length + withString:@"a" + startingAtIndex:0]; + NSString* key = key_for_device(device); + ASSERT_EQ(key.length + 9, NAME_MAX); + EXPECT_TRUE(is_cache_key(key)); + EXPECT_FALSE( + is_cache_key(key_for_device([device stringByAppendingString:@"a"]))); +} + +TEST(CoreAICacheKeyFormatTest, RequiresExactHashBoundaries) { + NSString* key = key_for_device(@"iOS_h17p"); + EXPECT_FALSE(is_cache_key(nil)); + EXPECT_FALSE(is_cache_key(@"")); + EXPECT_FALSE(is_cache_key([key substringToIndex:64])); + EXPECT_FALSE(is_cache_key([key substringFromIndex:1])); + EXPECT_FALSE(is_cache_key([key stringByAppendingString:@"0"])); + EXPECT_FALSE(is_cache_key(key.uppercaseString)); + EXPECT_FALSE(is_cache_key([key stringByAppendingString:@".bookmark"])); + EXPECT_FALSE( + is_cache_key([key stringByReplacingCharactersInRange:NSMakeRange(64, 1) + withString:@"_"])); +} + +TEST(CoreAICacheKeyFormatTest, LockStripesStayWithinConfiguredPool) { + auto bundle = selected_bundle(manifest_dict()); + ASSERT_TRUE(bundle.ok()); + NSMutableSet* names = [NSMutableSet set]; + for (NSUInteger index = 0; index < 4096; ++index) { + auto key = test_cache_key( + bundle.get(), [NSString stringWithFormat:@"identity-%lu", index]); + ASSERT_TRUE(key.ok()); + auto name = cache_entry_lock_name(key.get()); + ASSERT_TRUE(name.ok()); + ASSERT_TRUE([name.get() hasPrefix:@"stripe-v1-"]); + ASSERT_TRUE([name.get() hasSuffix:@".lock"]); + NSString* encoded = + [name.get() substringWithRange:NSMakeRange(10, name.get().length - 15)]; + unsigned int stripe = 0; + ASSERT_TRUE([[NSScanner scannerWithString:encoded] scanHexInt:&stripe]); + EXPECT_LT(stripe, kCacheEntryLockStripeCount); + [names addObject:name.get()]; + } + EXPECT_GT(names.count, 1u); + EXPECT_LE(names.count, kCacheEntryLockStripeCount); +} + +} // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_cache_root_test.mm b/backends/apple/coreai/runtime/test/coreai_cache_root_test.mm index b6cf8141065..5fffbce9a6e 100644 --- a/backends/apple/coreai/runtime/test/coreai_cache_root_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_cache_root_test.mm @@ -18,6 +18,27 @@ using namespace executorch::backends::coreai; using namespace executorch::backends::coreai::testing; +TEST( + CoreAICacheRootTest, + ResolvesDefaultAndExplicitPathsWithoutCreatingStorage) { + auto default_path = cache_directory_path(nullptr); + ASSERT_TRUE(default_path.ok()); + EXPECT_TRUE(default_path.get().isAbsolutePath); + EXPECT_TRUE( + [default_path.get().lastPathComponent isEqual:@"executorch_coreai"]); + + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* missing = [storage.url URLByAppendingPathComponent:@"missing"]; + auto explicit_path = cache_directory_path(missing.path.UTF8String); + ASSERT_TRUE(explicit_path.ok()); + EXPECT_TRUE([explicit_path.get() isEqual:missing.path]); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:missing.path]); + for (const char* invalid : {"", "relative", "~", "~/cache", "\xff"}) { + EXPECT_EQ(cache_directory_path(invalid).error(), Error::InvalidArgument); + } +} + TEST(CoreAICacheRootTest, PreparesRootWithRootOnlyBackupPolicy) { NSFileManager* manager = NSFileManager.defaultManager; TestDirectory sandbox; @@ -30,29 +51,38 @@ ASSERT_TRUE(set_backup_excluded(sandbox.url, false)); ASSERT_TRUE(set_backup_excluded(parent, false)); NSURL* root = [parent URLByAppendingPathComponent:@"nested/models"]; - auto prepared = prepare_storage_root(root.path, StorageCreationMode::Create); - ASSERT_TRUE((prepared.ok() && [prepared.get() isEqualToString:root.path])); + auto prepared = + resolve_test_cache_directory(root.path, StorageCreationMode::Create); + ASSERT_TRUE(prepared.ok()); + ASSERT_TRUE(prepared->has_value()); + EXPECT_TRUE([prepared->value().path + isEqualToString:root.URLByResolvingSymlinksInPath.path]); ASSERT_TRUE((backup_excluded(root))); ASSERT_TRUE((!backup_excluded(parent) && !backup_excluded(sandbox.url))); ASSERT_TRUE((!backup_excluded(root.URLByDeletingLastPathComponent))); ASSERT_TRUE(set_backup_excluded(root, false)); - auto inspected = - prepare_storage_root(root.path, StorageCreationMode::OpenExisting); - ASSERT_TRUE((inspected.ok() && [inspected.get() isEqualToString:root.path])); + auto inspected = resolve_test_cache_directory( + root.path, StorageCreationMode::OpenExisting); + ASSERT_TRUE(inspected.ok()); + ASSERT_TRUE(inspected->has_value()); + EXPECT_TRUE([inspected->value().path + isEqualToString:root.URLByResolvingSymlinksInPath.path]); ASSERT_TRUE((!backup_excluded(root))); ASSERT_TRUE( - (prepare_storage_root(root.path, StorageCreationMode::Create).ok() && + (resolve_test_cache_directory(root.path, StorageCreationMode::Create) + .ok() && backup_excluded(root))); NSURL* missing = [sandbox.url URLByAppendingPathComponent:@"missing"]; - ASSERT_EQ( - prepare_storage_root(missing.path, StorageCreationMode::OpenExisting) - .error(), - Error::InvalidArgument); + auto absent = resolve_test_cache_directory( + missing.path, StorageCreationMode::OpenExisting); + ASSERT_TRUE(absent.ok()); + EXPECT_FALSE(absent->has_value()); EXPECT_FALSE([manager fileExistsAtPath:missing.path]); EXPECT_EQ( - prepare_storage_root(@"relative", StorageCreationMode::Create).error(), + resolve_test_cache_directory(@"relative", StorageCreationMode::Create) + .error(), Error::InvalidArgument); EXPECT_TRUE((!backup_excluded(parent) && !backup_excluded(sandbox.url))); } @@ -62,28 +92,33 @@ ASSERT_NE(storage.url, nil); NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; ASSERT_TRUE( - prepare_storage_root(root.path, StorageCreationMode::Create).ok()); + resolve_test_cache_directory(root.path, StorageCreationMode::Create) + .ok()); StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); EXPECT_TRUE( - prepare_storage_root(root.path, StorageCreationMode::Create).ok()); + resolve_test_cache_directory(root.path, StorageCreationMode::Create) + .ok()); EXPECT_TRUE( - prepare_storage_root(root.path, StorageCreationMode::OpenExisting).ok()); + resolve_test_cache_directory(root.path, StorageCreationMode::OpenExisting) + .ok()); EXPECT_EQ(fault.matches, 0); EXPECT_TRUE(backup_excluded(root)); NSURL* missing = [storage.url URLByAppendingPathComponent:@"missing"]; - EXPECT_EQ( - prepare_storage_root(missing.path, StorageCreationMode::OpenExisting) - .error(), - Error::InvalidArgument); + auto absent = resolve_test_cache_directory( + missing.path, StorageCreationMode::OpenExisting); + ASSERT_TRUE(absent.ok()); + EXPECT_FALSE(absent->has_value()); EXPECT_EQ(fault.matches, 0); EXPECT_EQ( - prepare_storage_root(missing.path, StorageCreationMode::Create).error(), + resolve_test_cache_directory(missing.path, StorageCreationMode::Create) + .error(), Error::AccessFailed); EXPECT_EQ(fault.hits, 1); EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:missing.path]); EXPECT_TRUE( - prepare_storage_root(missing.path, StorageCreationMode::Create).ok()); + resolve_test_cache_directory(missing.path, StorageCreationMode::Create) + .ok()); } TEST(CoreAICacheRootTest, CreatesDirectoriesWithDefaultPermissions) { @@ -95,7 +130,8 @@ ASSERT_EQ(stat(reference.fileSystemRepresentation, &expected), 0); NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; ASSERT_TRUE( - prepare_storage_root(root.path, StorageCreationMode::Create).ok()); + resolve_test_cache_directory(root.path, StorageCreationMode::Create) + .ok()); FileDescriptor directory( open(root.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_GE(directory.get(), 0); @@ -114,10 +150,12 @@ FileDescriptor directory( ASSERT_NE(storage.url, nil); NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; ASSERT_TRUE( - prepare_storage_root(root.path, StorageCreationMode::Create).ok()); + resolve_test_cache_directory(root.path, StorageCreationMode::Create) + .ok()); ASSERT_EQ(rmdir(root.fileSystemRepresentation), 0); ASSERT_TRUE( - prepare_storage_root(root.path, StorageCreationMode::Create).ok()); + resolve_test_cache_directory(root.path, StorageCreationMode::Create) + .ok()); EXPECT_TRUE(backup_excluded(root)); } @@ -126,7 +164,8 @@ FileDescriptor directory( ASSERT_NE(sandbox.url, nil); NSURL* root = [sandbox.url URLByAppendingPathComponent:@"cache"]; ASSERT_TRUE( - prepare_storage_root(root.path, StorageCreationMode::Create).ok()); + resolve_test_cache_directory(root.path, StorageCreationMode::Create) + .ok()); FileDescriptor directory( open_directory_at(AT_FDCWD, root.fileSystemRepresentation)); ASSERT_GE(directory.get(), 0); @@ -171,7 +210,7 @@ FileDescriptor directory( ASSERT_TRUE([before[NSFileProtectionKey] isEqual:NSFileProtectionComplete]); for (StorageCreationMode mode : {StorageCreationMode::OpenExisting, StorageCreationMode::Create}) { - ASSERT_TRUE(prepare_storage_root(root.path, mode).ok()); + ASSERT_TRUE(resolve_test_cache_directory(root.path, mode).ok()); NSDictionary* after = [manager attributesOfItemAtPath:root.path error:&error]; ASSERT_NE(after, nil) << error.description.UTF8String; @@ -179,6 +218,44 @@ FileDescriptor directory( } } +TEST(CoreAICacheRootTest, SupportsSymlinkAndAnchorsResolvedDirectory) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* first = [storage.url URLByAppendingPathComponent:@"first"]; + NSURL* second = [storage.url URLByAppendingPathComponent:@"second"]; + NSURL* configured = [storage.url URLByAppendingPathComponent:@"configured"]; + ASSERT_EQ(mkdir(first.fileSystemRepresentation, 0777), 0); + ASSERT_EQ(mkdir(second.fileSystemRepresentation, 0777), 0); + ASSERT_EQ( + symlink( + first.fileSystemRepresentation, configured.fileSystemRepresentation), + 0); + + auto resolved = resolve_test_cache_directory( + configured.path, StorageCreationMode::Create); + ASSERT_TRUE(resolved.ok()); + ASSERT_TRUE(resolved->has_value()); + struct stat target_info; + struct stat resolved_info; + ASSERT_EQ(stat(first.fileSystemRepresentation, &target_info), 0); + ASSERT_EQ(fstat(resolved->value().descriptor.get(), &resolved_info), 0); + EXPECT_EQ(resolved_info.st_ino, target_info.st_ino); + + ASSERT_EQ(unlink(configured.fileSystemRepresentation), 0); + ASSERT_EQ( + symlink( + second.fileSystemRepresentation, configured.fileSystemRepresentation), + 0); + ASSERT_EQ( + write_storage_file( + resolved->value().descriptor.get(), @"anchored", "x", 1), + Error::Ok); + EXPECT_TRUE([NSFileManager.defaultManager + fileExistsAtPath:[first URLByAppendingPathComponent:@"anchored"].path]); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:[second URLByAppendingPathComponent:@"anchored"].path]); +} + TEST(CoreAICacheRootTest, RejectsInvalidRootsConsistently) { TestDirectory storage; ASSERT_NE(storage.url, nil); @@ -188,19 +265,24 @@ FileDescriptor directory( for (StorageCreationMode mode : {StorageCreationMode::OpenExisting, StorageCreationMode::Create}) { SCOPED_TRACE(static_cast(mode)); - EXPECT_EQ(prepare_storage_root(nil, mode).error(), Error::InvalidArgument); + EXPECT_EQ( + resolve_test_cache_directory(nil, mode).error(), + Error::InvalidArgument); + for (NSString* path in + @[ @"", @"relative", @"~", @"~/cache", @"~nobody/cache" ]) { + SCOPED_TRACE(path.UTF8String); + EXPECT_EQ( + resolve_test_cache_directory(path, mode).error(), + Error::InvalidArgument); + } for (NSString* path in @[ - @"", - @"relative", - @"~", - @"~/cache", - @"~nobody/cache", regular.path, [regular.path stringByAppendingPathComponent:@"child"] ]) { - SCOPED_TRACE(path.UTF8String); EXPECT_EQ( - prepare_storage_root(path, mode).error(), Error::InvalidArgument); + resolve_test_cache_directory(path, mode).error(), + mode == StorageCreationMode::Create ? Error::InvalidArgument + : Error::AccessFailed); } } EXPECT_TRUE([[NSData dataWithContentsOfURL:regular] isEqual:bytes]); diff --git a/backends/apple/coreai/runtime/test/coreai_cache_test.mm b/backends/apple/coreai/runtime/test/coreai_cache_test.mm new file mode 100644 index 00000000000..a04fa9e0e96 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_cache_test.mm @@ -0,0 +1,841 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "coreai_cache.h" +#include "coreai_cache_entry_fixture.h" +#include "coreai_fake_loader.h" +#include "coreai_fault_scope.h" +#include "coreai_filesystem_fixture.h" +#include "coreai_pte_fixture.h" +#include "coreai_source_fixture.h" +#include "coreai_test_utils.h" + +namespace executorch::backends::coreai::testing { +namespace { + +using runtime::Error; +using SegmentType = runtime::DataLoader::SegmentInfo::Type; +using ::testing::AssertionFailure; +using ::testing::AssertionResult; +using ::testing::AssertionSuccess; + +class CoreAICacheTest : public ::testing::Test { + protected: + void SetUp() override { + runtime::runtime_init(); + } + + void drain_callbacks() { + auto drained = bridge.wait_for_callbacks(test_deadline()); + if (!drained) { + ADD_FAILURE() << drained.message(); + // Timed-out callbacks may still borrow this test's PTE or observations. + std::_Exit(EXIT_FAILURE); + } + } + + ScopedFakeBridgeState bridge; +}; + +struct CacheEntry { + NSString* key = nil; + NSData* bookmark = nil; + NSURL* stage = nil; +}; + +AssertionResult seed_cache( + NSURL* root, + NSDictionary* dict, + BookmarkFakeLoader* loader, + CacheEntry& output, + bool bookmark = true, + bool stage = true, + NSString* architecture = @"arch_b", + NSString* platform = @"macOS", + NSString* specialization = @"default") { + if (root == nil || loader == nil) { + return AssertionFailure() << "Cache seed requires a root and fake loader"; + } + auto parsed = parse_manifest(encode(dict)); + if (!parsed.ok()) { + return AssertionFailure() << "Cannot parse cache seed manifest"; + } + auto selected = select_bundle(parsed.get(), architecture, platform); + if (!selected.ok()) { + return AssertionFailure() << "Cannot select cache seed assets"; + } + auto key = + cache_variant_key(selected.get(), platform, architecture, specialization); + if (!key.ok()) { + return AssertionFailure() << "Cannot compute cache seed key"; + } + CacheEntry entry{ + key.get(), + nil, + [[root URLByAppendingPathComponent:@"sources"] + URLByAppendingPathComponent:key.get()]}; + auto lock = acquire_test_cache_entry(root.path, entry.key); + if (!lock.ok()) { + return AssertionFailure() << "Cannot lock cache seed bookmark"; + } + if (stage) { + TestData data; + if (selected->aot_compiled) { + aot_data(data, architecture); + } + auto staged = executorch::backends::coreai::prepare_source_bundle( + selected.get(), &data, *lock.get(), CorruptBundlePolicy::Reject); + if (data.requests != data.releases) { + return AssertionFailure() + << "Cache seed source buffers were not released"; + } + if (!staged.ok()) { + return AssertionFailure() << "Cannot stage cache seed source bundle"; + } + } + if (bookmark) { + entry.bookmark = + [NSUUID.UUID.UUIDString dataUsingEncoding:NSUTF8StringEncoding]; + BookmarkFakeModel* model = [[BookmarkFakeModel alloc] init]; + if (model == nil || entry.bookmark == nil) { + return AssertionFailure() << "Cannot create cache seed model/bookmark"; + } + model.owner = loader; + model.bookmark = entry.bookmark; + loader.models[entry.bookmark] = model; + if (lock.get()->write_bookmark(entry.bookmark) != Error::Ok) { + return AssertionFailure() << "Cannot write cache seed bookmark"; + } + } + output = entry; + return AssertionSuccess(); +} + +bool cache_exists(NSURL* url) { + return [NSFileManager.defaultManager fileExistsAtPath:url.path]; +} + +void cache_entry_removed(NSURL* root, const CacheEntry& entry) { + EXPECT_FALSE(cache_exists(test_bookmark_url(root, entry.key))); + EXPECT_FALSE(cache_exists(entry.stage)); + EXPECT_TRUE(cache_exists(test_cache_lock_url(root, entry.key))); +} + +NSMutableDictionary* cache_manifest(char identity) { + auto dict = manifest_dict(); + dict[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity(identity)}; + return dict; +} + +TEST_F(CoreAICacheTest, EmptyAndMissingRootsAreNoncreatingNoops) { + @autoreleasepool { + TestDirectory root; + ASSERT_NE(root.url, nil); + BookmarkLoaderScope fake; + ASSERT_NE(fake.loader, nil); + FBSyntheticPTE empty({{"forward", {}}}); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(root.url, snapshot)); + const int queries = bridge.state().architecture_queries; + bridge.state().available = false; + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(clear_cache_for_pte(empty, root.url.path.UTF8String), Error::Ok); + NSURL* missing = + [root.url URLByAppendingPathComponent:@"absent/nested/models"]; + EXPECT_EQ(clear_cache(missing.path.UTF8String), Error::Ok); + EXPECT_EQ(clear_cache_for_pte(empty, missing.path.UTF8String), Error::Ok); + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(clear_cache_for_pte(pte, missing.path.UTF8String), Error::Ok); + EXPECT_TRUE(pte.check_released()); + EXPECT_FALSE(cache_exists(missing)); + EXPECT_TRUE(snapshot_matches(root.url, snapshot)); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(bridge.state().architecture_queries.load(), queries); + EXPECT_TRUE(empty.check_released()); + } +} + +TEST_F(CoreAICacheTest, ValidatesExplicitRootsAndPtePaths) { + @autoreleasepool { + TestDirectory root; + ASSERT_NE(root.url, nil); + BookmarkLoaderScope fake; + ASSERT_NE(fake.loader, nil); + FBSyntheticPTE empty({{"forward", {}}}); + EXPECT_EQ(clear_cache("relative"), Error::InvalidArgument); + EXPECT_EQ(clear_cache_for_pte(empty, "relative"), Error::InvalidArgument); + EXPECT_EQ(clear_cache(""), Error::InvalidArgument); + EXPECT_EQ(clear_cache("\xff"), Error::InvalidArgument); + EXPECT_EQ( + clear_cache_for_pte( + static_cast(nullptr), root.url.path.UTF8String), + Error::InvalidArgument); + EXPECT_TRUE(empty.check_released()); + } +} + +TEST_F(CoreAICacheTest, SelectsAllMethodsDeduplicatesAndEvictsInKeyOrder) { + @autoreleasepool { + TestDirectory root; + TestDirectory isolated; + BookmarkLoaderScope fake; + auto source_dict = manifest_dict(); + auto second_dict = cache_manifest('2'); + auto aot_dict = aot_manifest_dict(); + CacheEntry source, second, aot, other_arch, unrelated, other_root; + ASSERT_TRUE(seed_cache(root.url, source_dict, fake.loader, source)); + ASSERT_TRUE(seed_cache(root.url, second_dict, fake.loader, second)); + ASSERT_TRUE(seed_cache(root.url, aot_dict, fake.loader, aot)); + ASSERT_TRUE(seed_cache( + root.url, aot_dict, fake.loader, other_arch, true, true, @"arch_a")); + ASSERT_TRUE( + seed_cache(root.url, cache_manifest('3'), fake.loader, unrelated)); + ASSERT_TRUE(seed_cache(isolated.url, source_dict, fake.loader, other_root)); + NSDictionary* isolation = nil; + NSDictionary* unrelated_stage = nil; + ASSERT_TRUE(asset_tree_snapshot(isolated.url, isolation)); + ASSERT_TRUE(asset_tree_snapshot(unrelated.stage, unrelated_stage)); + FBSyntheticPTE pte( + {{"empty", {}}, + {"first", + {{"UnregisteredBackend", + [@"not JSON" dataUsingEncoding:NSUTF8StringEncoding], + true}, + {"CoreAIBackend", encode(source_dict)}, + {"CoreAIBackend", encode(second_dict), true}}}, + {"last", + {{"CoreAIBackend", encode(source_dict), true}, + {"CoreAIBackend", encode(aot_manifest_dict()), true}}}}); + NSMutableArray* evictions = [NSMutableArray array]; + fake.loader.onEvict = ^(NSData* bookmark) { + [evictions addObject:bookmark]; + }; + const int queries = bridge.state().architecture_queries; + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + drain_callbacks(); + EXPECT_EQ(bridge.state().architecture_queries.load(), queries); + EXPECT_EQ(fake.loader->evictions.load(), 4); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(fake.loader->bindings.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + NSArray* ordered = + [@[ source.key, second.key, aot.key, other_arch.key ] + sortedArrayUsingSelector:@selector(compare:)]; + NSDictionary* bookmarks = @{ + source.key : source.bookmark, + second.key : second.bookmark, + aot.key : aot.bookmark, + other_arch.key : other_arch.bookmark + }; + ASSERT_EQ(evictions.count, ordered.count); + for (NSUInteger i = 0; i < ordered.count; ++i) { + SCOPED_TRACE(::testing::Message() << "sorted eviction index " << i); + EXPECT_TRUE([evictions[i] isEqual:bookmarks[ordered[i]]]); + } + cache_entry_removed(root.url, source); + cache_entry_removed(root.url, second); + cache_entry_removed(root.url, aot); + EXPECT_NE(fake.loader.models[other_root.bookmark], nil); + EXPECT_EQ(fake.loader.models[other_arch.bookmark], nil); + cache_entry_removed(root.url, other_arch); + EXPECT_TRUE(snapshot_matches(isolated.url, isolation)); + EXPECT_TRUE(snapshot_matches(unrelated.stage, unrelated_stage)); + EXPECT_EQ(pte.backend_indices, (std::vector{1, 2, 3})); + EXPECT_EQ( + pte.backend_descriptors, + (std::vector{ + "CoreAIBackend", "CoreAIBackend", "CoreAIBackend"})); + EXPECT_TRUE(pte.check_released()); + } +} + +TEST_F(CoreAICacheTest, FilePathOverloadClearsReferencedEntry) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + TestDirectory distribution; + ASSERT_NE(distribution.url, nil); + CacheEntry entry; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, entry)); + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); + NSURL* path = [distribution.url URLByAppendingPathComponent:@"model.pte"]; + ASSERT_TRUE([[NSData dataWithBytes:pte.bytes.data() + length:pte.bytes.size()] writeToURL:path + options:0 + error:nil]); + EXPECT_EQ( + clear_cache_for_pte(path.path.UTF8String, root.url.path.UTF8String), + Error::Ok); + drain_callbacks(); + EXPECT_EQ(fake.loader->evictions.load(), 1); + cache_entry_removed(root.url, entry); + } +} + +TEST_F(CoreAICacheTest, RejectsInvalidLaterDelegatesBeforeMutation) { + @autoreleasepool { + const char* rows[] = {"malformed manifest", "backend load failure"}; + for (int row = 0; row < 2; ++row) { + SCOPED_TRACE(rows[row]); + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry source; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, source)); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(root.url, snapshot)); + PteDelegateSpec later{"CoreAIBackend", encode(aot_manifest_dict()), true}; + if (row == 0) + later.data = [@"malformed" dataUsingEncoding:NSUTF8StringEncoding]; + FBSyntheticPTE pte( + {{"first", {{"CoreAIBackend", encode(manifest_dict())}}}, + {"later", {later}}}); + if (row == 1) { + pte.fault_call = 2; + pte.load_fault = FBSyntheticPTE::LoadFault::Error; + pte.load_error = Error::AccessFailed; + } + NSDictionary* models = [fake.loader.models copy]; + const auto error = clear_cache_for_pte(pte, root.url.path.UTF8String); + drain_callbacks(); + EXPECT_NE(error, Error::Ok); + if (row == 1) + EXPECT_EQ(error, Error::AccessFailed); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(fake.loader->bindings.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_TRUE(snapshot_matches(root.url, snapshot)); + EXPECT_TRUE([fake.loader.models isEqual:models]); + EXPECT_TRUE(pte.check_released()); + } + } +} + +TEST_F(CoreAICacheTest, ClearsWarmVariantsWithoutHostCompatibility) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + auto source = manifest_dict(); + auto aot = aot_manifest_dict(); + aot[@"platform"] = @"iOS"; + CacheEntry mac, ios, compiled; + ASSERT_TRUE(seed_cache( + root.url, source, fake.loader, mac, true, false, @"unknown_mac_arch")); + ASSERT_TRUE(seed_cache( + root.url, + source, + fake.loader, + ios, + true, + false, + @"unknown_ios_arch", + @"iOS")); + ASSERT_TRUE(seed_cache( + root.url, aot, fake.loader, compiled, true, false, @"arch_a", @"iOS")); + source[@"min_deployment_version"] = @"999999.0"; + aot[@"min_deployment_version"] = @"999999.0"; + bridge.state().device_architecture = nil; + const int queries = bridge.state().architecture_queries; + FBSyntheticPTE pte( + {{"forward", + {{"CoreAIBackend", encode(source)}, + {"CoreAIBackend", encode(aot), true}}}}); + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 3); + EXPECT_EQ(bridge.state().architecture_queries.load(), queries); + cache_entry_removed(root.url, mac); + cache_entry_removed(root.url, ios); + cache_entry_removed(root.url, compiled); + } +} + +TEST_F(CoreAICacheTest, ClearsAotBundles) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry first, second; + auto dict = aot_manifest_dict(); + ASSERT_TRUE(seed_cache( + root.url, + dict, + fake.loader, + first, + true, + false, + @"arch_a", + @"macOS", + @"default")); + ASSERT_TRUE(seed_cache( + root.url, + dict, + fake.loader, + second, + true, + false, + @"arch_b", + @"macOS", + @"default")); + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", encode(dict)}}}}); + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 2); + cache_entry_removed(root.url, first); + cache_entry_removed(root.url, second); + } +} + +TEST_F(CoreAICacheTest, ClearsOptionAndDeviceVariantsPreservingUnrelatedModel) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry base, specialized, device, unrelated; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, base)); + ASSERT_TRUE(seed_cache( + root.url, + manifest_dict(), + fake.loader, + specialized, + true, + true, + @"arch_b", + @"macOS", + @"tuned")); + ASSERT_TRUE(seed_cache( + root.url, + manifest_dict(), + fake.loader, + device, + true, + true, + @"other_arch")); + ASSERT_TRUE( + seed_cache(root.url, cache_manifest('3'), fake.loader, unrelated)); + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + drain_callbacks(); + EXPECT_EQ(fake.loader->evictions.load(), 3); + EXPECT_EQ(fake.loader->specializationQueries.load(), 0); + cache_entry_removed(root.url, base); + cache_entry_removed(root.url, specialized); + cache_entry_removed(root.url, device); + EXPECT_TRUE(cache_exists(test_bookmark_url(root.url, unrelated.key))); + EXPECT_TRUE(cache_exists(unrelated.stage)); + EXPECT_NE(fake.loader.models[unrelated.bookmark], nil); + EXPECT_EQ(fake.loader.models[base.bookmark], nil); + EXPECT_EQ(fake.loader.models[specialized.bookmark], nil); + EXPECT_EQ(fake.loader.models[device.bookmark], nil); + EXPECT_TRUE(pte.check_released()); + } +} + +TEST_F(CoreAICacheTest, ClearsSourceOnlyEntriesWithoutSdk) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry staged; + ASSERT_TRUE(seed_cache( + root.url, manifest_dict(), fake.loader, staged, false, true)); + bridge.state().available = false; + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 0); + cache_entry_removed(root.url, staged); + } +} + +TEST_F(CoreAICacheTest, MaintenanceUnlinksSourcesRootWithoutFollowingIt) { + @autoreleasepool { + TestDirectory root, outside; + BookmarkLoaderScope fake; + CacheEntry entry; + ASSERT_TRUE( + seed_cache(root.url, manifest_dict(), fake.loader, entry, true, false)); + NSURL* sentinel = [outside.url URLByAppendingPathComponent:@"sentinel"]; + NSData* contents = [@"untouched" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([contents writeToURL:sentinel options:0 error:nil]); + ASSERT_TRUE([NSFileManager.defaultManager + createDirectoryAtURL:[entry.stage URLByDeletingLastPathComponent] + withIntermediateDirectories:NO + attributes:nil + error:nil]); + ASSERT_EQ( + symlink( + outside.url.fileSystemRepresentation, + entry.stage.fileSystemRepresentation), + 0); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + EXPECT_TRUE([[NSData dataWithContentsOfURL:sentinel] isEqual:contents]); + EXPECT_EQ(fake.loader->evictions.load(), 1); + cache_entry_removed(root.url, entry); + } +} + +TEST_F(CoreAICacheTest, ProcessedBuffersRemainAliveThroughEviction) { + @autoreleasepool { + TestDirectory root; + CacheEntry source; + NSData* manifest = encode(manifest_dict()); + FBSyntheticPTE valid( + {{"forward", + {{"CoreAIBackend", manifest, true}, + {"CoreAIBackend", manifest, true}, + {"CoreAIBackend", manifest}}}}); + // Destroy the loader scope before the PTE captured by its callback. + BookmarkLoaderScope fake; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, source)); + __block int observations = 0; + __block bool program_alive = false; + __block int backend_releases = -1; + FBSyntheticPTE* observed = &valid; + fake.loader.onEvict = ^(NSData*) { + ++observations; + program_alive = observed->program_loads.load() == + observed->program_releases.load() + 1; + backend_releases = observed->backend_releases.load(); + }; + EXPECT_EQ(clear_cache_for_pte(valid, root.url.path.UTF8String), Error::Ok); + drain_callbacks(); + EXPECT_EQ(observations, 1); + EXPECT_TRUE(program_alive); + EXPECT_EQ(backend_releases, 0); + EXPECT_EQ(valid.backend_loads.load(), 2); + EXPECT_EQ(valid.backend_indices, (std::vector{0, 1})); + EXPECT_TRUE(valid.program_alive_at_backend_release); + EXPECT_TRUE(valid.check_released()); + ASSERT_FALSE(valid.release_order.empty()); + EXPECT_EQ(valid.release_order.back(), SegmentType::Program); + } +} + +TEST_F( + CoreAICacheTest, + InventoriesBookmarkAndSourceUnionPreservingUnownedFiles) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry both, bookmark_only, stage_only; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, both)); + ASSERT_TRUE(seed_cache( + root.url, + cache_manifest('2'), + fake.loader, + bookmark_only, + true, + false)); + ASSERT_TRUE(seed_cache( + root.url, cache_manifest('3'), fake.loader, stage_only, false, true)); + NSURL* temporary = + [root.url URLByAppendingPathComponent:@"sources/.pending-unkeyed"]; + ASSERT_TRUE([NSFileManager.defaultManager createDirectoryAtURL:temporary + withIntermediateDirectories:NO + attributes:nil + error:nil]); + NSURL* untouched = [temporary URLByAppendingPathComponent:@"partial"]; + ASSERT_TRUE([[@"do not infer ownership" + dataUsingEncoding:NSUTF8StringEncoding] writeToURL:untouched + options:0 + error:nil]); + NSURL* external_sentinel = + [root.url URLByAppendingPathComponent:@"caller.pte"]; + ASSERT_TRUE([[@"PTE" dataUsingEncoding:NSUTF8StringEncoding] + writeToURL:external_sentinel + options:0 + error:nil]); + NSDictionary* temp_snapshot = nil; + NSDictionary* locks = nil; + ASSERT_TRUE(asset_tree_snapshot(temporary, temp_snapshot)); + ASSERT_TRUE(asset_tree_snapshot( + [root.url URLByAppendingPathComponent:@"locks"], locks)); + const int queries = bridge.state().architecture_queries; + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 2); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(bridge.state().architecture_queries.load(), queries); + cache_entry_removed(root.url, both); + cache_entry_removed(root.url, bookmark_only); + cache_entry_removed(root.url, stage_only); + EXPECT_TRUE(snapshot_matches( + [root.url URLByAppendingPathComponent:@"locks"], locks)); + EXPECT_TRUE(snapshot_matches(temporary, temp_snapshot)); + EXPECT_TRUE(cache_exists(external_sentinel)); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 2); + } +} + +TEST_F(CoreAICacheTest, RemovesSourceOnlyEntriesWithoutSdkAvailability) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry stage_only; + ASSERT_TRUE(seed_cache( + root.url, manifest_dict(), fake.loader, stage_only, false, true)); + bridge.state().available = false; + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + cache_entry_removed(root.url, stage_only); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + } +} + +TEST_F(CoreAICacheTest, WarnsAndContinuesPastMissingOrCorruptBookmarks) { + for (bool pte_clear : {false, true}) { + for (int corruption = 0; corruption < 6; ++corruption) { + SCOPED_TRACE(pte_clear ? "PTE" : "root"); + SCOPED_TRACE(corruption); + @autoreleasepool { + TestDirectory root, outside; + NSURL* sentinel = [outside.url URLByAppendingPathComponent:@"sentinel"]; + NSData* contents = + [@"untouched" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([contents writeToURL:sentinel options:0 error:nil]); + BookmarkLoaderScope fake; + CacheEntry damaged, healthy; + ASSERT_TRUE( + seed_cache(root.url, manifest_dict(), fake.loader, damaged)); + ASSERT_TRUE( + seed_cache(root.url, cache_manifest('2'), fake.loader, healthy)); + if ([damaged.key compare:healthy.key] == NSOrderedDescending) { + std::swap(damaged, healthy); + } + NSURL* path = test_bookmark_url(root.url, damaged.key); + if (corruption == 0 || corruption >= 3) { + ASSERT_EQ(unlink(path.fileSystemRepresentation), 0); + if (corruption == 3) { + ASSERT_EQ(mkfifo(path.fileSystemRepresentation, 0600), 0); + } else if (corruption == 4) { + ASSERT_EQ( + symlink( + sentinel.fileSystemRepresentation, + path.fileSystemRepresentation), + 0); + } else if (corruption == 5) { + ASSERT_EQ(mkdir(path.fileSystemRepresentation, 0700), 0); + ASSERT_TRUE([contents + writeToURL:[path URLByAppendingPathComponent:@"invalid"] + options:0 + error:nil]); + } + } else { + FileDescriptor file( + open(path.fileSystemRepresentation, O_WRONLY | O_NOFOLLOW)); + ASSERT_GE(file.get(), 0); + ASSERT_EQ( + ftruncate(file.get(), corruption == 1 ? 0 : 8 * 1024 * 1024 + 1), + 0); + } + FBSyntheticPTE pte( + {{"forward", + {{"CoreAIBackend", encode(manifest_dict())}, + {"CoreAIBackend", encode(cache_manifest('2')), true}}}}); + auto clear = [&] { + return pte_clear ? clear_cache_for_pte(pte, root.url.path.UTF8String) + : clear_cache(root.url.path.UTF8String); + }; + ::testing::internal::CaptureStderr(); + const auto error = clear(); + const std::string logs = ::testing::internal::GetCapturedStderr(); + EXPECT_EQ(error, Error::Ok); + if constexpr ( + ET_LOG_ENABLED && + runtime::LogLevel::Info >= runtime::LogLevel::ET_MIN_LOG_LEVEL) { + EXPECT_NE(logs.find(damaged.key.UTF8String), std::string::npos); + EXPECT_NE( + logs.find( + corruption == 0 ? "has a missing bookmark" + : "has a corrupt bookmark"), + std::string::npos); + EXPECT_NE( + logs.find("SDK eviction could not be confirmed"), + std::string::npos); + EXPECT_EQ( + logs.find( + std::string( + static_cast(damaged.bookmark.bytes), + damaged.bookmark.length)), + std::string::npos); + } + cache_entry_removed(root.url, damaged); + cache_entry_removed(root.url, healthy); + EXPECT_TRUE([[NSData dataWithContentsOfURL:sentinel] isEqual:contents]); + EXPECT_EQ(fake.loader->evictions.load(), 1); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_NE(fake.loader.models[damaged.bookmark], nil); + EXPECT_EQ(fake.loader.models[healthy.bookmark], nil); + ::testing::internal::CaptureStderr(); + const auto repeated = clear(); + const std::string repeated_logs = + ::testing::internal::GetCapturedStderr(); + EXPECT_EQ(repeated, Error::Ok); + EXPECT_TRUE(repeated_logs.empty()); + EXPECT_EQ(fake.loader->evictions.load(), 1); + EXPECT_TRUE(pte.check_released()); + } + } + } +} + +TEST_F(CoreAICacheTest, CorruptBookmarkCleanupDoesNotRequireSdkAndCanRetry) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry damaged; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, damaged)); + NSURL* path = test_bookmark_url(root.url, damaged.key); + ASSERT_TRUE([NSData.data writeToURL:path options:0 error:nil]); + bridge.state().available = false; + { + StorageFaultScope fault(FaultPoint::Remove); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + EXPECT_TRUE(cache_exists(path)); + EXPECT_TRUE(cache_exists(damaged.stage)); + } + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + cache_entry_removed(root.url, damaged); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_NE(fake.loader.models[damaged.bookmark], nil); + } +} + +TEST_F(CoreAICacheTest, NilLoaderFactoryPreservesEntriesAndAllowsRetry) { + @autoreleasepool { + ASSERT_TRUE(bridge.state().available); + ASSERT_FALSE(bridge.state().fail_loader_factory); + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry entry; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, entry)); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(root.url, snapshot)); + bridge.state().fail_loader_factory = true; + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Internal); + bridge.state().fail_loader_factory = false; + EXPECT_TRUE(snapshot_matches(root.url, snapshot)); + EXPECT_NE(fake.loader.models[entry.bookmark], nil); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(fake.loader->bindings.load(), 0); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + cache_entry_removed(root.url, entry); + } +} + +TEST_F(CoreAICacheTest, UnavailableSdkPreservesTrackedEntriesAndAllowsRetry) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry entry; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, entry)); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(root.url, snapshot)); + bridge.state().available = false; + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::NotSupported); + bridge.state().available = true; + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_TRUE(snapshot_matches(root.url, snapshot)); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + cache_entry_removed(root.url, entry); + } +} + +TEST_F(CoreAICacheTest, SourceRemovalFailuresKeepBookmarkForRetry) { + @autoreleasepool { + const struct { + int match_index; + bool bookmark_removal; + const char* name; + } rows[] = { + {2, false, "staged-tree removal"}, {6, true, "bookmark removal"}}; + for (const auto& row : rows) { + SCOPED_TRACE(row.name); + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry entry; + ASSERT_TRUE( + seed_cache(root.url, aot_manifest_dict(), fake.loader, entry)); + { + StorageFaultScope fault(FaultPoint::Remove); + fault.match_index = row.match_index; + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + drain_callbacks(); + auto saved = saved_bookmark(root.url, entry.key); + ASSERT_TRUE(saved.ok()); + EXPECT_TRUE([saved.get() isEqual:entry.bookmark]); + EXPECT_EQ(cache_exists(entry.stage), !row.bookmark_removal); + if (fake.loader.models[entry.bookmark] == nil) { + fake.loader.evictError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:ENOENT + userInfo:nil]; + } + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + cache_entry_removed(root.url, entry); + } + } +} + +TEST_F(CoreAICacheTest, IndependentKeysContinueInSortedOrderAfterFailures) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry first, second, third; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, first)); + ASSERT_TRUE(seed_cache(root.url, cache_manifest('2'), fake.loader, second)); + ASSERT_TRUE(seed_cache(root.url, cache_manifest('3'), fake.loader, third)); + NSArray* ordered = [@[ first.key, second.key, third.key ] + sortedArrayUsingSelector:@selector(compare:)]; + NSDictionary* bookmarks = @{ + first.key : first.bookmark, + second.key : second.bookmark, + third.key : third.bookmark + }; + NSMutableArray* evicted = [NSMutableArray array]; + BookmarkFakeLoader* loader = fake.loader; + fake.loader.onEvict = ^(NSData* bookmark) { + [evicted addObject:bookmark]; + loader.evictError = [bookmark isEqual:bookmarks[ordered[1]]] + ? [NSError errorWithDomain:NSPOSIXErrorDomain code:EBUSY userInfo:nil] + : nil; + }; + StorageFaultScope fault(FaultPoint::BeforeEvict); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::AccessFailed); + drain_callbacks(); + EXPECT_EQ(fault.hits, 1); + ASSERT_EQ(evicted.count, 2u); + EXPECT_TRUE([evicted[0] isEqual:bookmarks[ordered[1]]]); + EXPECT_TRUE([evicted[1] isEqual:bookmarks[ordered[2]]]); + EXPECT_TRUE(cache_exists(test_bookmark_url(root.url, ordered[0]))); + EXPECT_TRUE(cache_exists(test_bookmark_url(root.url, ordered[1]))); + EXPECT_FALSE(cache_exists(test_bookmark_url(root.url, ordered[2]))); + } +} + +} // namespace +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_delegate_fixture.h b/backends/apple/coreai/runtime/test/coreai_delegate_fixture.h new file mode 100644 index 00000000000..7d9a9b64fa1 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_delegate_fixture.h @@ -0,0 +1,123 @@ +/* + * 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 +#include +#include +#include + +#import "coreai_fake_loader.h" +#include "coreai_source_fixture.h" + +namespace executorch::backends::coreai::testing { + +class DelegateFixture final { + public: + DelegateFixture( + const runtime::BackendInterface* backend, + ScopedFakeBridgeState& bridge, + NSDictionary* manifest = manifest_dict()); + ~DelegateFixture(); + DelegateFixture(const DelegateFixture&) = delete; + DelegateFixture& operator=(const DelegateFixture&) = delete; + + runtime::Error init(NSString* root = nil); + runtime::Error init_with_options( + runtime::Span options, + runtime::MemoryAllocator* allocator = nullptr); + void reset(); + runtime::Error execute(runtime::Span args); + + TestData data; + TestDirectory storage; + int freed = 0; + NSData* json; + runtime::FreeableBuffer processed; + runtime::DelegateHandle* handle = nullptr; + + private: + const runtime::BackendInterface* backend_; + ScopedFakeBridgeState& bridge_; +}; + +struct FloatTensors final { + float input[2] = {3, 7}; + float output[2] = {-1, -1}; + executorch::aten::TensorImpl::SizesType input_size[1] = {2}; + executorch::aten::TensorImpl::SizesType output_size[1] = {2}; + executorch::aten::TensorImpl::DimOrderType order[1] = {0}; + executorch::aten::TensorImpl::StridesType input_stride[1] = {1}; + executorch::aten::TensorImpl::StridesType output_stride[1] = {1}; + executorch::aten::TensorImpl input_impl{ + executorch::aten::ScalarType::Float, + 1, + input_size, + input, + order, + input_stride}; + executorch::aten::TensorImpl output_impl{ + executorch::aten::ScalarType::Float, + 1, + output_size, + output, + order, + output_stride, + runtime::TensorShapeDynamism::DYNAMIC_BOUND}; + runtime::EValue in{executorch::aten::Tensor(&input_impl)}; + runtime::EValue out{executorch::aten::Tensor(&output_impl)}; + runtime::EValue* args[2] = {&in, &out}; +}; + +NSMutableDictionary* multiple_io_manifest(); + +template +struct RangeTensors final { + static constexpr executorch::aten::ScalarType dtype = std::is_same_v + ? executorch::aten::ScalarType::Float + : executorch::aten::ScalarType::Half; + T input[12] = {}; + T output[12] = {}; + executorch::aten::TensorImpl::SizesType input_shape[2] = {4, 2}; + executorch::aten::TensorImpl::SizesType output_shape[2] = {4, 2}; + executorch::aten::TensorImpl::DimOrderType order[2] = {0, 1}; + executorch::aten::TensorImpl::StridesType input_strides[2] = {2, 1}; + executorch::aten::TensorImpl::StridesType output_strides[2] = {2, 1}; + executorch::aten::TensorImpl input_impl{ + dtype, + 2, + input_shape, + input + 2, + order, + input_strides, + runtime::TensorShapeDynamism::DYNAMIC_BOUND}; + executorch::aten::TensorImpl output_impl{ + dtype, + 2, + output_shape, + output + 2, + order, + output_strides, + runtime::TensorShapeDynamism::DYNAMIC_BOUND}; + runtime::EValue in{executorch::aten::Tensor(&input_impl)}; + runtime::EValue out{executorch::aten::Tensor(&output_impl)}; + runtime::EValue* args[2] = {&in, &out}; + + RangeTensors() { + for (size_t i = 0; i < 12; ++i) { + input[i] = static_cast(i + 1); + } + fill_output(); + } + + void fill_output() { + std::fill_n(output, 12, static_cast(99)); + } +}; + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_delegate_fixture.mm b/backends/apple/coreai/runtime/test/coreai_delegate_fixture.mm new file mode 100644 index 00000000000..3d450ee4d24 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_delegate_fixture.mm @@ -0,0 +1,87 @@ +/* + * 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_delegate_fixture.h" +#include "coreai_test_utils.h" + +namespace executorch::backends::coreai::testing { +using namespace executorch::runtime; + +DelegateFixture::DelegateFixture( + const BackendInterface* backend, + ScopedFakeBridgeState& bridge, + NSDictionary* manifest) + : json(encode(manifest)), + processed( + json.bytes, + json.length, + [](void* context, void*, size_t) { ++*static_cast(context); }, + &freed), + backend_(backend), + bridge_(bridge) {} + +DelegateFixture::~DelegateFixture() { + reset(); + EXPECT_TRUE(bridge_.wait_for_callbacks(test_deadline())); + processed.Free(); + if (json != nil) { + EXPECT_EQ(freed, 1); + } + EXPECT_EQ(data.requests.load(), data.releases.load()); +} + +Error DelegateFixture::init(NSString* root) { + if (storage.url == nil || json == nil) { + return Error::Internal; + } + BackendOptions<1> options; + const Error error = options.set_option( + "coreai_cache_dir", (root ?: storage.url.path).UTF8String); + if (error != Error::Ok) { + return error; + } + return init_with_options({options.view().data(), options.view().size()}); +} + +Error DelegateFixture::init_with_options( + Span options, + MemoryAllocator* allocator) { + reset(); + BackendInitContext context(allocator, nullptr, "forward", &data, options); + auto result = backend_->init(context, &processed, {}); + if (result.ok()) { + handle = result.get(); + } + auto drained = bridge_.wait_for_callbacks(test_deadline()); + if (!drained) { + ADD_FAILURE() << drained.message(); + reset(); + return Error::Internal; + } + return result.ok() ? Error::Ok : result.error(); +} + +void DelegateFixture::reset() { + if (handle != nullptr) { + backend_->destroy(handle); + } + handle = nullptr; +} + +Error DelegateFixture::execute(Span args) { + BackendExecutionContext context; + return backend_->execute(context, handle, args); +} + +NSMutableDictionary* multiple_io_manifest() { + NSMutableDictionary* manifest = manifest_dict(); + manifest[@"input_names"] = @[ @"input_1", @"input_0" ]; + manifest[@"output_names"] = @[ @"output_1", @"output_0" ]; + return manifest; +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm new file mode 100644 index 00000000000..1321bed50f0 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm @@ -0,0 +1,912 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "coreai_cache.h" +#include "coreai_cache_entry_fixture.h" +#import "coreai_delegate_fixture.h" +#include "coreai_test_utils.h" + +namespace executorch::backends::coreai::testing { +namespace { +using namespace executorch::runtime; +using executorch::aten::ScalarType; +using executorch::aten::Tensor; +using executorch::aten::TensorImpl; + +class CoreAIDelegateTest : public ::testing::Test { + protected: + void SetUp() override { + backend = get_backend_class("CoreAIBackend"); + ASSERT_NE(backend, nullptr); + } + ScopedFakeBridgeState bridge; + const BackendInterface* backend = nullptr; +}; + +// All preflight rows must fail before reading payloads or creating storage. +void expect_preflight_rejection( + const BackendInterface* backend, + ScopedFakeBridgeState& bridge, + NSDictionary* manifest, + Error expected) { + DelegateFixture fixture(backend, bridge, manifest); + ASSERT_NE(fixture.storage.url, nil); + NSURL* absent = + [fixture.storage.url URLByAppendingPathComponent:@"uncreated/models"]; + NSString* previous_bundle = bridge.state().last_bundle; + EXPECT_EQ(fixture.init(absent.path), expected); + EXPECT_EQ(fixture.data.attempts.load(), 0); + EXPECT_EQ(fixture.data.metadata_requests.load(), 0); + EXPECT_EQ(fixture.freed, 0); + EXPECT_EQ(bridge.state().last_bundle, previous_bundle); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:absent.path]); + NSError* error = nil; + NSArray* children = [NSFileManager.defaultManager + contentsOfDirectoryAtPath:fixture.storage.url.path + error:&error]; + EXPECT_NE(children, nil); + EXPECT_EQ(error, nil); + EXPECT_EQ(children.count, 0u); +} + +TEST_F(CoreAIDelegateTest, AvailabilityRejectsWithoutTakingManifestOwnership) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + bridge.state().available = false; + EXPECT_FALSE(backend->is_available()); + EXPECT_EQ(fixture.init(), Error::NotSupported); + EXPECT_EQ(fixture.freed, 0); + EXPECT_EQ(fixture.data.attempts.load(), 0); + EXPECT_EQ(bridge.state().sessions.load(), 0); + bridge.state().available = true; + EXPECT_TRUE(backend->is_available()); + ASSERT_EQ(fixture.init(), Error::Ok); + EXPECT_EQ(fixture.freed, 1); + EXPECT_EQ(bridge.state().sessions.load(), 1); + } +} + +TEST_F( + CoreAIDelegateTest, + LoaderFactoryFailureLeavesManifestAndStorageUntouched) { + @autoreleasepool { + bridge.state().fail_loader_factory = true; + expect_preflight_rejection( + backend, bridge, manifest_dict(), Error::Internal); + EXPECT_EQ(bridge.state().loaders.load(), 0); + EXPECT_EQ(bridge.state().sessions.load(), 0); + } +} + +TEST_F( + CoreAIDelegateTest, + PreflightRejectsBeforePayloadReadsOrStorageCreation) { + @autoreleasepool { + auto missing_function = aot_manifest_dict(); + [missing_function removeObjectForKey:@"function"]; + auto numeric_platform = aot_manifest_dict(); + numeric_platform[@"platform"] = @1; + auto escaping_arch = aot_manifest_dict(); + NSMutableDictionary* archs = [escaping_arch[@"archs"] mutableCopy]; + archs[@"arch_a"] = @"../model.arch_a.aimodelc"; + escaping_arch[@"archs"] = archs; + auto escaping_file = aot_manifest_dict(); + escaping_file[@"files"] = @{@"model.arch_a.aimodelc/../escape" : @1}; + auto ios = aot_manifest_dict(); + ios[@"platform"] = @"iOS"; + const struct { + const char* name; + NSMutableDictionary* manifest; + NSString* device_architecture; + Error expected; + } rows[] = { + {"missing function", + missing_function, + @"arch_b", + Error::InvalidProgram}, + {"non-string platform", + numeric_platform, + @"arch_b", + Error::InvalidProgram}, + {"escaping architecture path", + escaping_arch, + @"arch_b", + Error::InvalidProgram}, + {"escaping file entry", + escaping_file, + @"arch_b", + Error::InvalidProgram}, + {"incompatible platform", + ios, + @"arch_b", + Error::DelegateInvalidCompatibility}, + {"missing device architecture", + aot_manifest_dict(), + nil, + Error::DelegateInvalidCompatibility}, + }; + for (const auto& row : rows) { + SCOPED_TRACE(row.name); + bridge.state().device_architecture = row.device_architecture; + ASSERT_NO_FATAL_FAILURE(expect_preflight_rejection( + backend, bridge, row.manifest, row.expected)); + } + bridge.state().device_architecture = @"arch_b"; + } +} + +TEST_F(CoreAIDelegateTest, PreflightRejectsFutureDeploymentForSourceAndAot) { + @autoreleasepool { + for (bool aot : {false, true}) { + SCOPED_TRACE(aot ? "AOT" : "source"); + auto dict = aot ? aot_manifest_dict() : manifest_dict(); + dict[@"min_deployment_version"] = @"999999.0"; + expect_preflight_rejection( + backend, bridge, dict, Error::DelegateInvalidCompatibility); + } + } +} + +TEST_F(CoreAIDelegateTest, AotInitPassesOnlySelectedArchitectureBundleToSdk) { + @autoreleasepool { + NSString* arch = @"arch_a"; + DelegateFixture fixture(backend, bridge, aot_manifest_dict()); + aot_data(fixture.data, arch); + bridge.state().device_architecture = arch; + auto parsed = parse_manifest(fixture.json); + ASSERT_TRUE(parsed.ok()); + auto selected = select_bundle(parsed.get(), arch, @"macOS"); + ASSERT_TRUE(selected.ok()); + ASSERT_EQ(fixture.init(), Error::Ok); + auto key = cache_variant_key(selected.get(), @"macOS", arch, @"default"); + ASSERT_TRUE(key.ok()); + NSString* expected = [fixture.storage.url.path + stringByAppendingPathComponent: + [NSString stringWithFormat:@"sources/%@/%@", + key.get(), + selected->path.lastPathComponent]]; + EXPECT_TRUE([bridge.state().last_bundle isEqualToString:expected]); + fixture.reset(); + EXPECT_EQ(bridge.state().sessions.load(), 0); + EXPECT_EQ(fixture.data.attempts.load(), 2); + EXPECT_EQ(fixture.data.releases.load(), fixture.data.requests.load()); + } +} + +TEST_F( + CoreAIDelegateTest, + InitBindingFailurePreservesManifestAndAllowsWarmRetry) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + bridge.state().fail_load = true; + EXPECT_EQ(fixture.init(), Error::InvalidProgram); + EXPECT_EQ(fixture.freed, 0); + EXPECT_EQ(bridge.state().sessions.load(), 0); + ASSERT_NE(bridge.state().last_bundle, nil); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:bridge.state().last_bundle]); + const int attempts = fixture.data.attempts.load(); + bridge.state().fail_load = false; + ASSERT_EQ(fixture.init(), Error::Ok); + EXPECT_EQ(fixture.freed, 1); + EXPECT_EQ(bridge.state().sessions.load(), 1); + EXPECT_EQ(fixture.data.attempts.load(), attempts); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:bridge.state().last_bundle]); + EXPECT_TRUE([bridge.state().last_bundle + hasPrefix:[fixture.storage.url.URLByResolvingSymlinksInPath.path + stringByAppendingString:@"/"]]); + } + EXPECT_EQ(bridge.state().sessions.load(), 0); +} + +TEST_F(CoreAIDelegateTest, RetainedBindingCompletionDoesNotPinSession) { + @autoreleasepool { + bridge.state().retain_binding_completion = true; + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + EXPECT_EQ(bridge.state().sessions.load(), 1); + fixture.reset(); + EXPECT_EQ(bridge.state().sessions.load(), 0); + EXPECT_EQ(bridge.state().prepared_models.load(), 0); + } +} + +TEST_F(CoreAIDelegateTest, OrderedIoBindingUsesManifestFunctionAndNames) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + EXPECT_EQ(bridge.state().binding_mismatches.load(), 0); + FloatTensors tensors; + ASSERT_EQ(fixture.execute({tensors.args, 2}), Error::Ok); + EXPECT_EQ(tensors.output[0], 3); + EXPECT_EQ(tensors.output[1], 7); + EXPECT_EQ(bridge.state().last_input_bytes.load(), tensors.input); + EXPECT_EQ(bridge.state().executions.load(), 1); + EXPECT_EQ(bridge.state().output_copies.load(), 1); + EXPECT_EQ(bridge.state().last_output_destination.load(), tensors.output); + } +} + +TEST_F(CoreAIDelegateTest, AdjacentInputAndOutputStorageCopiesDirectly) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + float storage[] = {3, 7, -1, -1}; + tensors.input_impl.set_data(storage); + tensors.output_impl.set_data(storage + 2); + ASSERT_EQ(fixture.execute({tensors.args, 2}), Error::Ok); + EXPECT_EQ(storage[2], 3); + EXPECT_EQ(storage[3], 7); + EXPECT_EQ(bridge.state().output_copies.load(), 1); + EXPECT_EQ(bridge.state().last_output_destination.load(), storage + 2); + } +} + +TEST_F(CoreAIDelegateTest, PartialInputOutputOverlapStagesBeforeWriting) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + for (bool destination_first : {false, true}) { + SCOPED_TRACE(destination_first); + FloatTensors tensors; + float storage[] = {3, 7, 11}; + float* source = storage + (destination_first ? 1 : 0); + float* destination = storage + (destination_first ? 0 : 1); + const float expected[] = {source[0], source[1]}; + tensors.input_impl.set_data(source); + tensors.output_impl.set_data(destination); + const int copies = bridge.state().output_copies.load(); + ASSERT_EQ(fixture.execute({tensors.args, 2}), Error::Ok); + EXPECT_EQ(destination[0], expected[0]); + EXPECT_EQ(destination[1], expected[1]); + EXPECT_EQ(bridge.state().output_copies.load(), copies + 1); + EXPECT_NE(bridge.state().last_output_destination.load(), destination); + } + } +} + +TEST_F(CoreAIDelegateTest, CrossOutputInputAliasesPreserveEverySource) { + @autoreleasepool { + bridge.state().multiple_io = true; + DelegateFixture fixture(backend, bridge, multiple_io_manifest()); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors first, second; + second.input[0] = 11; + second.input[1] = 13; + first.output_impl.set_data(second.input); + second.output_impl.set_data(first.input); + EValue* args[] = {&first.in, &second.in, &first.out, &second.out}; + ASSERT_EQ(fixture.execute({args, 4}), Error::Ok); + EXPECT_EQ(first.input[0], 11); + EXPECT_EQ(first.input[1], 13); + EXPECT_EQ(second.input[0], 3); + EXPECT_EQ(second.input[1], 7); + EXPECT_EQ(bridge.state().output_copies.load(), 2); + EXPECT_NE(bridge.state().last_output_destination.load(), first.input); + } +} + +TEST_F(CoreAIDelegateTest, OverlappingOutputDestinationsPreserveWriteOrder) { + @autoreleasepool { + bridge.state().multiple_io = true; + DelegateFixture fixture(backend, bridge, multiple_io_manifest()); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors first, second; + second.input[0] = 11; + second.input[1] = 13; + float storage[] = {-1, -1, -1}; + first.output_impl.set_data(storage); + second.output_impl.set_data(storage + 1); + EValue* args[] = {&first.in, &second.in, &first.out, &second.out}; + ASSERT_EQ(fixture.execute({args, 4}), Error::Ok); + EXPECT_EQ(storage[0], 3); + EXPECT_EQ(storage[1], 11); + EXPECT_EQ(storage[2], 13); + EXPECT_EQ(bridge.state().output_copies.load(), 2); + EXPECT_NE(bridge.state().last_output_destination.load(), storage + 1); + } +} + +TEST_F(CoreAIDelegateTest, MultipleInputsAndOutputsPreserveManifestOrder) { + @autoreleasepool { + bridge.state().multiple_io = true; + DelegateFixture fixture(backend, bridge, multiple_io_manifest()); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors first, second; + second.input[0] = 11; + second.input[1] = 13; + EValue* args[] = {&first.in, &second.in, &first.out, &second.out}; + ASSERT_EQ(fixture.execute({args, 4}), Error::Ok); + EXPECT_EQ(first.output[0], 3); + EXPECT_EQ(first.output[1], 7); + EXPECT_EQ(second.output[0], 11); + EXPECT_EQ(second.output[1], 13); + EXPECT_EQ(bridge.state().binding_mismatches.load(), 0); + } +} + +TEST_F(CoreAIDelegateTest, LaterOutputResizeFailureLeavesResultBytesUntouched) { + @autoreleasepool { + bridge.state().multiple_io = true; + DelegateFixture fixture(backend, bridge, multiple_io_manifest()); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors first, second; + float larger[] = {11, 13, 17}; + TensorImpl::SizesType one[] = {1}, three[] = {3}; + TensorImpl first_input( + ScalarType::Float, + 1, + one, + first.input, + first.order, + first.input_stride); + TensorImpl second_input( + ScalarType::Float, 1, three, larger, second.order, second.input_stride); + EValue first_in{Tensor(&first_input)}, second_in{Tensor(&second_input)}; + EValue* args[] = {&first_in, &second_in, &first.out, &second.out}; + EXPECT_EQ(fixture.execute({args, 4}), Error::InvalidArgument); + EXPECT_EQ(bridge.state().output_copies.load(), 0); + EXPECT_EQ(first.out.toTensor().size(0), 1); + EXPECT_EQ(second.out.toTensor().size(0), 2); + for (size_t i = 0; i < 2; ++i) { + EXPECT_EQ(first.output[i], -1); + EXPECT_EQ(second.output[i], -1); + } + } +} + +TEST_F(CoreAIDelegateTest, ConflictingSharedOutputShapesDoNotCopyBytes) { + @autoreleasepool { + bridge.state().multiple_io = true; + DelegateFixture fixture(backend, bridge, multiple_io_manifest()); + ASSERT_EQ(fixture.init(), Error::Ok); + const struct { + const char* name; + TensorImpl::SizesType shape[2]; + } cases[] = { + {"different byte counts", {2, 2}}, + {"equal byte counts", {2, 1}}, + }; + for (const auto& row : cases) { + SCOPED_TRACE(row.name); + for (bool separate_value : {false, true}) { + SCOPED_TRACE(separate_value); + RangeTensors first, second; + const TensorImpl::SizesType first_shape[] = {1, 2}; + ASSERT_EQ(resize_tensor(first.in.toTensor(), first_shape), Error::Ok); + ASSERT_EQ(resize_tensor(second.in.toTensor(), row.shape), Error::Ok); + EValue alias{Tensor(&first.output_impl)}; + EValue* args[] = { + &first.in, + &second.in, + &first.out, + separate_value ? &alias : &first.out}; + EXPECT_EQ(fixture.execute({args, 4}), Error::InvalidArgument); + EXPECT_EQ(bridge.state().output_copies.load(), 0); + for (float value : first.output) { + EXPECT_EQ(value, 99); + } + } + } + } +} + +TEST_F(CoreAIDelegateTest, MatchingSharedOutputShapesPreserveCopyOrder) { + @autoreleasepool { + bridge.state().multiple_io = true; + DelegateFixture fixture(backend, bridge, multiple_io_manifest()); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors first, second; + second.input[0] = 11; + second.input[1] = 13; + EValue alias{Tensor(&first.output_impl)}; + EValue* args[] = {&first.in, &second.in, &first.out, &alias}; + ASSERT_EQ(fixture.execute({args, 4}), Error::Ok); + EXPECT_EQ(bridge.state().output_copies.load(), 2); + EXPECT_EQ(first.out.toTensor().size(0), 2); + EXPECT_EQ(first.output[0], 11); + EXPECT_EQ(first.output[1], 13); + } +} + +TEST_F(CoreAIDelegateTest, WrongOutputCountDoesNotWriteCallerStorage) { + @autoreleasepool { + bridge.state().multiple_io = true; + bridge.state().bad_output = true; + DelegateFixture fixture(backend, bridge, multiple_io_manifest()); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors first, second; + EValue* args[] = {&first.in, &second.in, &first.out, &second.out}; + EXPECT_EQ(fixture.execute({args, 4}), Error::Internal); + EXPECT_EQ(bridge.state().output_copies.load(), 0); + EXPECT_EQ(first.out.toTensor().size(0), 2); + EXPECT_EQ(second.out.toTensor().size(0), 2); + for (size_t i = 0; i < 2; ++i) { + EXPECT_EQ(first.output[i], -1); + EXPECT_EQ(second.output[i], -1); + } + } +} + +TEST_F(CoreAIDelegateTest, ArgumentValidationRejectsBeforeSdkExecution) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + EXPECT_EQ(fixture.execute({tensors.args, 1}), Error::InvalidArgument); + EValue scalar(int64_t(3)); + EValue* args[] = {&scalar, &tensors.out}; + EXPECT_EQ(fixture.execute({args, 2}), Error::InvalidArgument); + EXPECT_EQ(bridge.state().executions.load(), 0); + } +} + +// No assertion may leave a caller thread blocked behind its input-read gate. +class BorrowedInputWorker final { + public: + BorrowedInputWorker( + DelegateFixture& fixture, + FloatTensors& tensors, + FakeBridgeState& state) + : fixture_(fixture), tensors_(tensors), state_(state) { + state_.input_ready = dispatch_semaphore_create(0); + state_.allow_input_read = dispatch_semaphore_create(0); + } + ~BorrowedInputWorker() { + release_and_join(); + state_.input_ready = nil; + state_.allow_input_read = nil; + } + BorrowedInputWorker(const BorrowedInputWorker&) = delete; + BorrowedInputWorker& operator=(const BorrowedInputWorker&) = delete; + int start() { + const int error = pthread_create(&thread_, nullptr, run, this); + started_ = error == 0; + return error; + } + long wait_ready() { + return dispatch_semaphore_wait(state_.input_ready, test_deadline(5)); + } + long wait_finished(int64_t timeout) { + return dispatch_semaphore_wait( + finished_, dispatch_time(DISPATCH_TIME_NOW, timeout)); + } + void release() { + dispatch_semaphore_signal(state_.allow_input_read); + } + void release_and_join() { + release(); + if (started_) { + const long returned = dispatch_semaphore_wait(returned_, test_deadline()); + if (returned != 0) { + ADD_FAILURE() + << "Delegate worker did not return after releasing its input gate"; + // A stuck worker still borrows this stack; unwinding would be unsafe. + std::_Exit(EXIT_FAILURE); + } + EXPECT_EQ(pthread_join(thread_, nullptr), 0); + started_ = false; + } + } + Error status = Error::Internal; + + private: + static void* run(void* context) { + auto& worker = *static_cast(context); + @autoreleasepool { + worker.status = worker.fixture_.execute({worker.tensors_.args, 2}); + } + dispatch_semaphore_signal(worker.finished_); + dispatch_semaphore_signal(worker.returned_); + return nullptr; + } + DelegateFixture& fixture_; + FloatTensors& tensors_; + FakeBridgeState& state_; + dispatch_semaphore_t finished_ = dispatch_semaphore_create(0); + dispatch_semaphore_t returned_ = dispatch_semaphore_create(0); + pthread_t thread_{}; + bool started_ = false; +}; + +TEST_F(CoreAIDelegateTest, BorrowedInputStorageStaysLiveUntilSuccessOrError) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + for (bool failure : {false, true}) { + SCOPED_TRACE(failure ? "runtime error" : "success"); + bridge.state().fail_execute = failure; + bridge.state().last_input_bytes = nullptr; + const auto first_byte = + reinterpret_cast(tensors.input)[0]; + bridge.state().last_input_first_byte = + static_cast(first_byte ^ 0xff); + const int executions = bridge.state().executions.load(); + BorrowedInputWorker worker(fixture, tensors, bridge.state()); + const int started = worker.start(); + if (started != 0) { + worker.release_and_join(); + FAIL() << "pthread_create failed: " << started; + } + const long ready = worker.wait_ready(); + const void* observed = bridge.state().last_input_bytes.load(); + const long premature = worker.wait_finished(50 * NSEC_PER_MSEC); + worker.release(); + const long finished = worker.wait_finished(5 * NSEC_PER_SEC); + worker.release_and_join(); + ASSERT_TRUE(bridge.wait_for_callbacks(test_deadline())); + EXPECT_EQ(ready, 0); + EXPECT_EQ(observed, tensors.input); + EXPECT_NE(premature, 0); + EXPECT_EQ(finished, 0); + EXPECT_EQ(worker.status, failure ? Error::Internal : Error::Ok); + EXPECT_EQ(bridge.state().executions.load(), executions + 1); + EXPECT_EQ( + bridge.state().last_input_first_byte.load(), + reinterpret_cast(tensors.input)[0]); + } + } +} + +TEST_F(CoreAIDelegateTest, RetainedExecutionCompletionReleasesResults) { + @autoreleasepool { + bridge.state().retain_execution_completion = true; + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + ASSERT_EQ(fixture.execute({tensors.args, 2}), Error::Ok); + EXPECT_EQ(bridge.state().output_owners.load(), 0); + bridge.state().fail_execute = true; + EXPECT_EQ(fixture.execute({tensors.args, 2}), Error::Internal); + EXPECT_EQ(bridge.state().output_owners.load(), 0); + } +} + +TEST_F(CoreAIDelegateTest, ExecuteMapsRuntimeErrors) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + bridge.state().fail_execute = true; + EXPECT_EQ(fixture.execute({tensors.args, 2}), Error::Internal); + } +} + +TEST_F(CoreAIDelegateTest, MalformedOutputDoesNotWriteCallerStorage) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + bridge.state().bad_output = true; + EXPECT_EQ(fixture.execute({tensors.args, 2}), Error::Internal); + EXPECT_EQ(tensors.output[0], -1); + EXPECT_EQ(tensors.output[1], -1); + } +} + +TEST_F(CoreAIDelegateTest, NoncontiguousInputIsNotSupported) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + tensors.input_stride[0] = 2; + EXPECT_EQ(fixture.execute({tensors.args, 2}), Error::NotSupported); + EXPECT_EQ(bridge.state().executions.load(), 0); + } +} + +TEST_F(CoreAIDelegateTest, NonemptyOutputRequiresStorage) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + tensors.output_impl.set_data(nullptr); + EXPECT_EQ(fixture.execute({tensors.args, 2}), Error::InvalidArgument); + } +} + +TEST_F(CoreAIDelegateTest, OutputResizesDownAndBackWithinCapacity) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + TensorImpl::SizesType one[] = {1}; + TensorImpl smaller( + ScalarType::Float, + 1, + one, + tensors.input, + tensors.order, + tensors.input_stride); + EValue small{Tensor(&smaller)}; + EValue* args[] = {&small, &tensors.out}; + ASSERT_EQ(fixture.execute({args, 2}), Error::Ok); + EXPECT_EQ(tensors.out.toTensor().size(0), 1); + ASSERT_EQ(fixture.execute({tensors.args, 2}), Error::Ok); + EXPECT_EQ(tensors.out.toTensor().size(0), 2); + } +} + +TEST_F(CoreAIDelegateTest, OutputResizeFailurePreservesStorage) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + TensorImpl::SizesType capacity[] = {1}; + TensorImpl output( + ScalarType::Float, + 1, + capacity, + tensors.output, + tensors.order, + tensors.output_stride, + TensorShapeDynamism::DYNAMIC_BOUND); + EValue too_small{Tensor(&output)}; + EValue* args[] = {&tensors.in, &too_small}; + tensors.output[0] = -2; + EXPECT_EQ(fixture.execute({args, 2}), Error::InvalidArgument); + EXPECT_EQ(tensors.output[0], -2); + } +} + +TEST_F(CoreAIDelegateTest, StaticOutputRejectsShapeChangeWithinCapacity) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + TensorImpl::SizesType shape[] = {1}; + TensorImpl input( + ScalarType::Float, + 1, + shape, + tensors.input, + tensors.order, + tensors.input_stride); + TensorImpl output( + ScalarType::Float, + 1, + tensors.output_size, + tensors.output, + tensors.order, + tensors.output_stride, + TensorShapeDynamism::STATIC); + EValue in{Tensor(&input)}, out{Tensor(&output)}; + EValue* args[] = {&in, &out}; + EXPECT_EQ(fixture.execute({args, 2}), Error::InvalidArgument); + EXPECT_EQ(out.toTensor().size(0), 2); + EXPECT_EQ(tensors.output[0], -1); + EXPECT_EQ(tensors.output[1], -1); + } +} + +TEST_F(CoreAIDelegateTest, HalfDtypePreservesBitsAndRejectsMixedOutputDtype) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + uint16_t input[] = {0x3c00, 0x4000}, output[] = {0, 0}; + TensorImpl input_impl( + ScalarType::Half, + 1, + tensors.input_size, + input, + tensors.order, + tensors.input_stride); + TensorImpl output_impl( + ScalarType::Half, + 1, + tensors.output_size, + output, + tensors.order, + tensors.output_stride); + EValue in{Tensor(&input_impl)}, out{Tensor(&output_impl)}; + EValue* args[] = {&in, &out}; + ASSERT_EQ(fixture.execute({args, 2}), Error::Ok); + EXPECT_EQ(output[0], input[0]); + EXPECT_EQ(output[1], input[1]); + EXPECT_EQ(bridge.state().last_input_bytes.load(), input); + EValue* mixed[] = {&in, &tensors.out}; + EXPECT_EQ(fixture.execute({mixed, 2}), Error::InvalidArgument); + } +} + +TEST_F(CoreAIDelegateTest, ScalarTensorsExecuteAndRejectOutputRankMismatch) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + TensorImpl input(ScalarType::Float, 0, nullptr, tensors.input); + TensorImpl output(ScalarType::Float, 0, nullptr, tensors.output); + EValue in{Tensor(&input)}, out{Tensor(&output)}; + EValue* args[] = {&in, &out}; + ASSERT_EQ(fixture.execute({args, 2}), Error::Ok); + EXPECT_EQ(tensors.output[0], tensors.input[0]); + EXPECT_EQ(bridge.state().last_input_bytes.load(), tensors.input); + EValue* wrong_rank[] = {&tensors.in, &out}; + EXPECT_EQ(fixture.execute({wrong_rank, 2}), Error::InvalidArgument); + } +} + +TEST_F(CoreAIDelegateTest, EmptyVectorAndMatrixBorrowNullStorage) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + for (int rank : {1, 2}) { + SCOPED_TRACE(rank); + TensorImpl::SizesType shape[] = {2, 0}; + TensorImpl::DimOrderType order[] = {0, 1}; + TensorImpl::StridesType strides[] = {1, 1}; + auto* sizes = rank == 1 ? shape + 1 : shape; + TensorImpl input(ScalarType::Float, rank, sizes, nullptr, order, strides); + TensorImpl output( + ScalarType::Float, rank, sizes, nullptr, order, strides); + EValue in{Tensor(&input)}, out{Tensor(&output)}; + EValue* args[] = {&in, &out}; + const int executions = bridge.state().executions.load(); + bridge.state().last_input_bytes = &input; + bridge.state().last_input_byte_count = 1; + ASSERT_EQ(fixture.execute({args, 2}), Error::Ok); + EXPECT_EQ(bridge.state().executions.load(), executions + 1); + EXPECT_EQ(bridge.state().last_input_bytes.load(), nullptr); + EXPECT_EQ(bridge.state().last_input_byte_count.load(), 0); + EXPECT_EQ(bridge.state().output_copies.load(), 0); + } + } +} + +TEST_F(CoreAIDelegateTest, BorrowedStorageMayAliasOutputAfterExecution) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + tensors.output_impl.set_data(tensors.input); + ASSERT_EQ(fixture.execute({tensors.args, 2}), Error::Ok); + EXPECT_EQ(tensors.input[0], 3); + EXPECT_EQ(tensors.input[1], 7); + } +} + +TEST_F(CoreAIDelegateTest, BorrowedStorageUsesCurrentInputPointer) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + FloatTensors tensors; + ASSERT_EQ(fixture.execute({tensors.args, 2}), Error::Ok); + float next_input[] = {11, 13}; + tensors.input_impl.set_data(next_input); + ASSERT_EQ(fixture.execute({tensors.args, 2}), Error::Ok); + EXPECT_EQ(bridge.state().last_input_bytes.load(), next_input); + EXPECT_EQ(tensors.output[0], 11); + EXPECT_EQ(tensors.output[1], 13); + } +} + +TEST_F(CoreAIDelegateTest, DynamicShapesBorrowCurrentSubrange) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + ASSERT_EQ(fixture.init(), Error::Ok); + RangeTensors tensors; + for (TensorImpl::SizesType rows : {3, 0}) { + SCOPED_TRACE(rows); + TensorImpl::SizesType shape[] = {rows, 2}; + ASSERT_EQ(resize_tensor(tensors.in.toTensor(), {shape, 2}), Error::Ok); + tensors.fill_output(); + ASSERT_EQ(fixture.execute({tensors.args, 2}), Error::Ok); + const size_t elements = static_cast(rows) * 2; + EXPECT_EQ(bridge.state().last_input_bytes.load(), tensors.input + 2); + EXPECT_EQ( + bridge.state().last_input_byte_count.load(), + elements * sizeof(float)); + EXPECT_EQ(tensors.out.toTensor().size(0), rows); + EXPECT_EQ(tensors.out.toTensor().const_data_ptr(), tensors.output + 2); + for (size_t i = 0; i < 12; ++i) { + SCOPED_TRACE(i); + EXPECT_EQ( + tensors.output[i], + i >= 2 && i < 2 + elements ? tensors.input[i] : 99.0f); + } + } + } +} + +TEST_F(CoreAIDelegateTest, StorageOptionsRejectWrongTypeAndRelativePaths) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + { + SCOPED_TRACE("integer cache directory"); + BackendOption wrong_type{"coreai_cache_dir", 1}; + EXPECT_EQ( + fixture.init_with_options({&wrong_type, 1}), Error::InvalidArgument); + } + { + SCOPED_TRACE("relative cache directory"); + BackendOptions<1> relative; + ASSERT_EQ(relative.set_option("coreai_cache_dir", "relative"), Error::Ok); + EXPECT_EQ( + fixture.init_with_options( + {relative.view().data(), relative.view().size()}), + Error::InvalidArgument); + } + EXPECT_EQ(fixture.freed, 0); + } +} + +TEST_F(CoreAIDelegateTest, ConfiguredSymlinkCacheDirectoryLoadsAndClears) { + @autoreleasepool { + DelegateFixture fixture(backend, bridge); + NSURL* target = + [fixture.storage.url URLByAppendingPathComponent:@"cache-target"]; + NSURL* configured = + [fixture.storage.url URLByAppendingPathComponent:@"cache-link"]; + ASSERT_EQ(mkdir(target.fileSystemRepresentation, 0777), 0); + ASSERT_EQ( + symlink( + target.fileSystemRepresentation, + configured.fileSystemRepresentation), + 0); + ASSERT_EQ(fixture.init(configured.path), Error::Ok); + EXPECT_TRUE([bridge.state().last_bundle containsString:@"/cache-target/"]); + EXPECT_FALSE([bridge.state().last_bundle containsString:@"/cache-link/"]); + fixture.reset(); + EXPECT_EQ(clear_cache(configured.path.UTF8String), Error::Ok); + } +} + +TEST_F( + CoreAIDelegateTest, + LifetimeDestroyAndWarmReloadIsolateRootsWithoutSource) { + @autoreleasepool { + DelegateFixture first(backend, bridge); + ASSERT_EQ(first.init(), Error::Ok); + EXPECT_EQ(first.freed, 1); + EXPECT_EQ(bridge.state().sessions.load(), 1); + NSString* first_bundle = bridge.state().last_bundle; + ASSERT_NE(first_bundle, nil); + first.reset(); + EXPECT_EQ(bridge.state().sessions.load(), 0); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:first_bundle]); + backend->destroy(nullptr); + + DelegateFixture reloaded(backend, bridge); + ASSERT_EQ(reloaded.init(first.storage.url.path), Error::Ok); + EXPECT_EQ(reloaded.data.attempts.load(), 0); + EXPECT_TRUE([bridge.state().last_bundle isEqualToString:first_bundle]); + DelegateFixture other(backend, bridge); + NSURL* nested = + [other.storage.url URLByAppendingPathComponent:@"nested/models"]; + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:nested.path]); + ASSERT_EQ(other.init(nested.path), Error::Ok); + EXPECT_EQ(bridge.state().sessions.load(), 2); + EXPECT_TRUE([bridge.state().last_bundle + hasPrefix:[other.storage.url.URLByResolvingSymlinksInPath.path + stringByAppendingString:@"/"]]); + EXPECT_FALSE([bridge.state().last_bundle isEqualToString:first_bundle]); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:first_bundle]); + other.reset(); + reloaded.reset(); + EXPECT_EQ(bridge.state().sessions.load(), 0); + } + EXPECT_EQ(bridge.state().sessions.load(), 0); + EXPECT_EQ(bridge.state().prepared_models.load(), 0); + EXPECT_EQ(bridge.state().loaders.load(), 0); +} + +} // namespace +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_eviction_test.mm b/backends/apple/coreai/runtime/test/coreai_eviction_test.mm new file mode 100644 index 00000000000..6a45af50a37 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_eviction_test.mm @@ -0,0 +1,470 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include "coreai_acquisition_fixture.h" +#include "coreai_cache_entry_fixture.h" +#include "coreai_fake_loader.h" +#include "coreai_fault_scope.h" +#include "coreai_load_coordinator.h" +#include "coreai_source_fixture.h" +#include "coreai_test_utils.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +namespace { +NSError* bookmark_test_error(NSInteger code = EIO) { + return [NSError errorWithDomain:NSPOSIXErrorDomain code:code userInfo:nil]; +} + +// These cases never retain a delegate handle beyond initialization. +Error initialize_and_destroy( + NSDictionary* dict, + const NamedDataMap* data, + NSString* root) { + if (root == nil) + return Error::InvalidArgument; + BackendOptions<1> options; + auto error = options.set_option("coreai_cache_dir", root.UTF8String); + if (error != Error::Ok) + return error; + auto* backend = get_backend_class("CoreAIBackend"); + if (backend == nullptr) + return Error::Internal; + BackendInitContext context( + nullptr, + nullptr, + "forward", + data, + {options.view().data(), options.view().size()}); + NSData* json = encode(dict); + if (json == nil) + return Error::InvalidArgument; + FreeableBuffer processed(json.bytes, json.length, nullptr); + auto handle = backend->init(context, &processed, {}); + if (!handle.ok()) + return handle.error(); + backend->destroy(handle.get()); + return Error::Ok; +} + +class CoreAIEvictionTest : public ::testing::Test { + protected: + ScopedFakeBridgeState bridge; + void SetUp() override { + runtime_init(); + } +}; + +// Release the callback before waiting; never unwind storage borrowed by a hang. +struct ReleaseAndJoin final { + dispatch_semaphore_t gate; + dispatch_semaphore_t eviction_done; + dispatch_semaphore_t loading_done; + ScopedFakeBridgeState& bridge; + pthread_t eviction{}; + pthread_t loading{}; + bool eviction_started = false; + bool loading_started = false; + bool released = false; + bool joined = false; + long eviction_wait = -1; + long loading_wait = -1; + + void release() { + if (!released) { + dispatch_semaphore_signal(gate); + released = true; + } + } + void join() { + release(); + if (joined) + return; + const auto completion_deadline = test_deadline(); + if (eviction_started) + eviction_wait = + dispatch_semaphore_wait(eviction_done, completion_deadline); + if (loading_started) + loading_wait = dispatch_semaphore_wait(loading_done, completion_deadline); + if ((eviction_started && eviction_wait != 0) || + (loading_started && loading_wait != 0)) { + ADD_FAILURE() + << "Eviction/load workers did not return after releasing the gate"; + std::_Exit(EXIT_FAILURE); + } + const int eviction_join = + eviction_started ? pthread_join(eviction, nullptr) : 0; + const int loading_join = + loading_started ? pthread_join(loading, nullptr) : 0; + if (eviction_join != 0 || loading_join != 0) { + ADD_FAILURE() << "Cannot join eviction/load workers: " << eviction_join + << ", " << loading_join; + std::_Exit(EXIT_FAILURE); + } + auto drained = bridge.wait_for_callbacks(test_deadline()); + if (!drained) { + ADD_FAILURE() << drained.message(); + std::_Exit(EXIT_FAILURE); + } + joined = true; + } + ~ReleaseAndJoin() { + join(); + } +}; +} // namespace + +@interface PinBookmarkLoader : BookmarkFakeLoader +@property(nonatomic, weak) BookmarkFakeModel* temporary; +@property(nonatomic, copy) ETCoreAIAcquisitionCompletion retainedCompletion; +@end + +@implementation PinBookmarkLoader +- (void)specializeModelAtURL:(NSURL*)url + completion:(ETCoreAIAcquisitionCompletion)completion { + ++specializations; + self.retainedCompletion = completion; + @autoreleasepool { + BookmarkFakeModel* model = [[BookmarkFakeModel alloc] init]; + model.owner = self; + model.bookmark = + [@"pinned-bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + self.temporary = model; + completion(ETCoreAIAcquisitionStatusHit, model, nil); + model = nil; + } +} +- (void)restoreModelFromBookmark:(NSData*)bookmark + completion:(ETCoreAIAcquisitionCompletion)completion { + ++restores; + self.retainedCompletion = completion; + @autoreleasepool { + BookmarkFakeModel* model = self.temporary; + completion( + model == nil ? ETCoreAIAcquisitionStatusMiss + : ETCoreAIAcquisitionStatusHit, + model, + nil); + model = nil; + } +} +- (void)evictModelWithBookmark:(NSData*)bookmark + completion:(void (^)(NSError*))completion { + ++evictions; + bool busy; + @autoreleasepool { + busy = self.temporary != nil; + } + completion(busy ? bookmark_test_error(EBUSY) : nil); +} +- (void)clearRetainedState { + self.retainedCompletion = nil; + [super clearRetainedState]; +} +@end + +namespace { +TEST_F(CoreAIEvictionTest, ConfirmedMissRecoveryEvictsOnlyCurrentBookmark) { + @autoreleasepool { + TestDirectory root; + ASSERT_NE(root.url, nil); + BookmarkLoaderScope fake; + ASSERT_NE(fake.loader, nil); + TestData data; + auto selected = selected_bundle(manifest_dict()); + ASSERT_TRUE(selected.ok()); + auto key_result = test_cache_key(selected.get()); + ASSERT_TRUE(key_result.ok()); + NSString* key = key_result.get(); + NSMutableDictionary* dict = manifest_dict(); + ASSERT_EQ(initialize_and_destroy(dict, &data, root.url.path), Error::Ok); + auto original_result = saved_bookmark(root.url, key); + ASSERT_TRUE(original_result.ok()); + NSData* original = original_result.get(); + ASSERT_NE(original, nil); + NSString* source_path = fake.loader.lastSpecializeURL; + ASSERT_NE(source_path, nil); + NSURL* source_url = [NSURL fileURLWithPath:source_path]; + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + + fake.loader.restoreMiss = YES; + fake.loader.specializeError = bookmark_test_error(); + EXPECT_NE(initialize_and_destroy(dict, &data, root.url.path), Error::Ok); + auto after_failure = saved_bookmark(root.url, key); + ASSERT_TRUE(after_failure.ok()); + EXPECT_TRUE([after_failure.get() isEqual:original]); + EXPECT_TRUE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + fake.loader.specializeError = nil; + ASSERT_EQ(initialize_and_destroy(dict, &data, root.url.path), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 3); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_TRUE([fake.loader.lastSpecializeURL isEqual:source_path]); + auto saved_result = saved_bookmark(root.url, key); + ASSERT_TRUE(saved_result.ok()); + NSData* saved = saved_result.get(); + ASSERT_NE(saved, nil); + EXPECT_FALSE([saved isEqual:original]); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + ASSERT_EQ( + clear_test_cache_entry(root.url.path, key, fake.loader), Error::Ok); + EXPECT_NE(fake.loader.models[original], nil); + EXPECT_EQ(fake.loader.models[saved], nil); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + fake.loader.restoreMiss = NO; + fake.loader.emptyBookmark = YES; + EXPECT_NE(initialize_and_destroy(dict, &data, root.url.path), Error::Ok); + auto empty = saved_bookmark(root.url, key); + ASSERT_TRUE(empty.ok()); + EXPECT_EQ(empty.get(), nil); + fake.loader.emptyBookmark = NO; + ASSERT_EQ(initialize_and_destroy(dict, &data, root.url.path), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 5); + EXPECT_EQ(data.requests.load(), data.releases.load()); + } +} + +TEST_F( + CoreAIEvictionTest, + BusyAndUnknownEvictionRetainBookmarkUntilConfirmedAbsent) { + @autoreleasepool { + TestDirectory root; + ASSERT_NE(root.url, nil); + BookmarkLoaderScope fake; + ASSERT_NE(fake.loader, nil); + TestData data; + auto selected = selected_bundle(manifest_dict()); + ASSERT_TRUE(selected.ok()); + auto key_result = test_cache_key(selected.get()); + ASSERT_TRUE(key_result.ok()); + NSString* key = key_result.get(); + EXPECT_EQ( + clear_test_cache_entry(root.url.path, key, fake.loader), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 0); + // Retain the source bundle after publication to exercise clearing's failure + // paths. + { + StorageFaultScope fault(FaultPoint::Remove); + ASSERT_EQ( + initialize_and_destroy(manifest_dict(), &data, root.url.path), + Error::Ok); + EXPECT_EQ(fault.hits, 1); + } + auto saved_result = saved_bookmark(root.url, key); + ASSERT_TRUE(saved_result.ok()); + NSData* saved = saved_result.get(); + ASSERT_NE(saved, nil); + NSURL* source_url = [[[root.url URLByAppendingPathComponent:@"sources"] + URLByAppendingPathComponent:key] + URLByAppendingPathComponent:selected->path.lastPathComponent]; + EXPECT_TRUE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + for (bool probe_error : {false, true}) { + SCOPED_TRACE( + probe_error ? "unknown restore failure" : "restore hit/busy"); + fake.loader.evictError = bookmark_test_error(EBUSY); + fake.loader.restoreError = probe_error ? bookmark_test_error() : nil; + EXPECT_NE( + clear_test_cache_entry(root.url.path, key, fake.loader), Error::Ok); + auto retained = saved_bookmark(root.url, key); + ASSERT_TRUE(retained.ok()); + EXPECT_TRUE([retained.get() isEqual:saved]); + EXPECT_TRUE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + } + fake.loader.restoreError = nil; + fake.loader.evictError = nil; + { + StorageFaultScope fault(FaultPoint::AfterEvict); + EXPECT_NE( + clear_test_cache_entry(root.url.path, key, fake.loader), Error::Ok); + EXPECT_EQ(fault.hits, 1); + } + auto retained = saved_bookmark(root.url, key); + ASSERT_TRUE(retained.ok()); + EXPECT_TRUE([retained.get() isEqual:saved]); + fake.loader.evictError = bookmark_test_error(ENOENT); + ASSERT_EQ( + clear_test_cache_entry(root.url.path, key, fake.loader), Error::Ok); + auto absent = saved_bookmark(root.url, key); + ASSERT_TRUE(absent.ok()); + EXPECT_EQ(absent.get(), nil); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + fake.loader.evictError = nil; + ASSERT_EQ( + initialize_and_destroy(manifest_dict(), &data, root.url.path), + Error::Ok); + { + StorageFaultScope fault(FaultPoint::Remove); + EXPECT_NE( + clear_test_cache_entry(root.url.path, key, fake.loader), Error::Ok); + EXPECT_EQ(fault.hits, 1); + auto before_remove = saved_bookmark(root.url, key); + ASSERT_TRUE(before_remove.ok()); + EXPECT_NE(before_remove.get(), nil); + } + ASSERT_EQ( + clear_test_cache_entry(root.url.path, key, fake.loader), Error::Ok); + auto removed = saved_bookmark(root.url, key); + ASSERT_TRUE(removed.ok()); + EXPECT_EQ(removed.get(), nil); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + EXPECT_EQ(data.requests.load(), data.releases.load()); + } +} + +TEST_F( + CoreAIEvictionTest, + PreparedPinsBlockEvictionWithoutRetainedCompletionLeak) { + @autoreleasepool { + TestDirectory root; + ASSERT_NE(root.url, nil); + TestData data; + auto selected = selected_bundle(manifest_dict()); + ASSERT_TRUE(selected.ok()); + auto key_result = test_cache_key(selected.get()); + ASSERT_TRUE(key_result.ok()); + NSString* key = key_result.get(); + PinBookmarkLoader* loader = [[PinBookmarkLoader alloc] init]; + ASSERT_NE(loader, nil); + // Scope teardown clears retained completions even after fatal test + // failures. + bridge.state().bookmark_loader = loader; + SelectedHostBundle assets{selected.get(), @"macOS", @"arch_b"}; + @autoreleasepool { + auto pinned = acquire_test_model(assets, &data, root.url.path, loader); + ASSERT_TRUE(pinned.ok()); + EXPECT_NE(loader.temporary, nil); + EXPECT_NE(clear_test_cache_entry(root.url.path, key, loader), Error::Ok); + auto saved = saved_bookmark(root.url, key); + ASSERT_TRUE(saved.ok()); + EXPECT_NE(saved.get(), nil); + } + EXPECT_EQ(loader.temporary, nil); + EXPECT_NE(loader.retainedCompletion, nil); + ASSERT_EQ(clear_test_cache_entry(root.url.path, key, loader), Error::Ok); + auto removed = saved_bookmark(root.url, key); + ASSERT_TRUE(removed.ok()); + EXPECT_EQ(removed.get(), nil); + EXPECT_EQ(loader->evictions.load(), 2); + NSURL* source_url = [[[root.url URLByAppendingPathComponent:@"sources"] + URLByAppendingPathComponent:key] + URLByAppendingPathComponent:selected->path.lastPathComponent]; + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + [loader clearRetainedState]; + EXPECT_EQ(data.requests.load(), data.releases.load()); + } +} + +TEST_F(CoreAIEvictionTest, LoadContendsOnActualKeyLockUntilEvictionFinishes) { + @autoreleasepool { + TestDirectory root; + ASSERT_NE(root.url, nil); + BookmarkLoaderScope fake; + ASSERT_NE(fake.loader, nil); + TestData data; + auto selected = selected_bundle(manifest_dict()); + ASSERT_TRUE(selected.ok()); + auto key_result = test_cache_key(selected.get()); + ASSERT_TRUE(key_result.ok()); + NSString* key = key_result.get(); + ASSERT_EQ( + initialize_and_destroy(manifest_dict(), &data, root.url.path), + Error::Ok); + dispatch_semaphore_t deleting = dispatch_semaphore_create(0); + dispatch_semaphore_t proceed = dispatch_semaphore_create(0); + dispatch_semaphore_t loading_observed = dispatch_semaphore_create(0); + dispatch_semaphore_t eviction_done = dispatch_semaphore_create(0); + dispatch_semaphore_t loading_done = dispatch_semaphore_create(0); + __block long gate_wait = -1; + fake.loader.onEvict = ^(NSData*) { + dispatch_semaphore_signal(deleting); + gate_wait = dispatch_semaphore_wait(proceed, test_deadline()); + }; + Error eviction_error = Error::Internal; + bool load_ok = false; + auto evict = [&] { + @autoreleasepool { + eviction_error = + clear_test_cache_entry(root.url.path, key, fake.loader); + } + dispatch_semaphore_signal(eviction_done); + }; + SelectedHostBundle assets{selected.get(), @"macOS", @"arch_b"}; + auto load = [&] { + @autoreleasepool { + auto acquired = + acquire_test_model(assets, &data, root.url.path, fake.loader); + load_ok = acquired.ok() && acquired.get() != nil; + } + dispatch_semaphore_signal(loading_observed); + dispatch_semaphore_signal(loading_done); + }; + ReleaseAndJoin cleanup{proceed, eviction_done, loading_done, bridge}; + const int eviction_start = pthread_create( + &cleanup.eviction, + nullptr, + [](void* context) -> void* { + (*static_cast(context))(); + return nullptr; + }, + &evict); + cleanup.eviction_started = eviction_start == 0; + const long entered = cleanup.eviction_started + ? dispatch_semaphore_wait(deleting, test_deadline()) + : -1; + int loading_start = -1; + long premature_completion = -1; + if (entered == 0) { + loading_start = pthread_create( + &cleanup.loading, + nullptr, + [](void* context) -> void* { + (*static_cast(context))(); + return nullptr; + }, + &load); + cleanup.loading_started = loading_start == 0; + if (cleanup.loading_started) { + premature_completion = dispatch_semaphore_wait( + loading_observed, + dispatch_time(DISPATCH_TIME_NOW, 50 * NSEC_PER_MSEC)); + } + } + cleanup.join(); + fake.loader.onEvict = nil; + + EXPECT_EQ(eviction_start, 0); + EXPECT_EQ(loading_start, 0); + EXPECT_EQ(entered, 0); + EXPECT_NE(premature_completion, 0); + EXPECT_EQ(gate_wait, 0); + EXPECT_EQ(cleanup.eviction_wait, 0); + EXPECT_EQ(cleanup.loading_wait, 0); + EXPECT_EQ(eviction_error, Error::Ok); + EXPECT_TRUE(load_ok); + EXPECT_EQ(fake.loader->specializations.load(), 2); + auto saved = saved_bookmark(root.url, key); + ASSERT_TRUE(saved.ok()); + EXPECT_NE(saved.get(), nil); + EXPECT_EQ(data.requests.load(), data.releases.load()); + } +} +} // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.h b/backends/apple/coreai/runtime/test/coreai_fake_loader.h new file mode 100644 index 00000000000..5fb287ed79a --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.h @@ -0,0 +1,128 @@ +/* + * 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 +#include +#import "ETCoreAIBridge.h" + +@class CoreAIFakeBridgeContext; +@class BookmarkFakeLoader; + +namespace executorch::backends::coreai::testing { + +// Configure on the test thread before starting workers; join before changing +// it. +struct FakeBridgeState { + bool available = true; + bool fail_loader_factory = false; + NSString* device_architecture = @"arch_b"; + id bookmark_loader = nil; + bool fail_load = false; + bool fail_execute = false; + bool bad_output = false; + bool multiple_io = false; + bool retain_binding_completion = false; + bool retain_execution_completion = false; + NSString* last_bundle = nil; + dispatch_semaphore_t input_ready = nil; + dispatch_semaphore_t allow_input_read = nil; + std::atomic architecture_queries{0}; + std::atomic executions{0}; + std::atomic output_copies{0}; + std::atomic output_owners{0}; + std::atomic last_output_destination{nullptr}; + std::atomic sessions{0}; + std::atomic prepared_models{0}; + std::atomic loaders{0}; + std::atomic last_input_bytes{nullptr}; + std::atomic last_input_byte_count{0}; + std::atomic last_input_first_byte{0}; + // Worker-side contract failures are checked by the scope on the test thread. + std::atomic binding_mismatches{0}; + std::atomic missing_bundles{0}; + std::atomic input_wait_timeouts{0}; +}; + +// Process-visible, not thread-local. Scopes and configuration changes are +// serial. +class ScopedFakeBridgeState final { + public: + ScopedFakeBridgeState(); + ~ScopedFakeBridgeState(); + FakeBridgeState& state(); + // Release input gates and join caller threads before draining callbacks. + ::testing::AssertionResult wait_for_callbacks(dispatch_time_t deadline); + ScopedFakeBridgeState(const ScopedFakeBridgeState&) = delete; + ScopedFakeBridgeState& operator=(const ScopedFakeBridgeState&) = delete; + + private: + CoreAIFakeBridgeContext* context_; + CoreAIFakeBridgeContext* previous_; +}; + +} // namespace executorch::backends::coreai::testing + +@interface FakeSession : NSObject +@property(nonatomic, strong) id preparedPin; +@end + +@interface FakePreparedModel : NSObject +@property(nonatomic, copy) NSData* bookmark; +@end + +@interface FakeLoader : NSObject +@end + +@interface BookmarkFakeModel : NSObject +@property(nonatomic, weak) BookmarkFakeLoader* owner; +@property(nonatomic, copy) NSData* bookmark; +@end + +@interface BookmarkFakeLoader : NSObject { + @public + std::atomic restores; + std::atomic specializations; + std::atomic bindings; + std::atomic evictions; + std::atomic specializationQueries; +} +@property(nonatomic, strong) + NSMutableDictionary* models; +@property(nonatomic, strong) NSError* restoreError; +@property(nonatomic, strong) NSError* specializeError; +@property(nonatomic, strong) NSError* bindError; +@property(nonatomic, strong) NSError* evictError; +@property(nonatomic) BOOL restoreMiss; +@property(nonatomic) BOOL emptyBookmark; +@property(nonatomic, copy) void (^onSpecialize)(NSURL*); +@property(nonatomic, copy) void (^onEvict)(NSData*); +@property(nonatomic, strong) BookmarkFakeModel* lastAcquired; +@property(nonatomic, strong) BookmarkFakeModel* lastBound; +@property(nonatomic, copy) NSString* lastSpecializeURL; +// Call only after workers/callbacks finish and caller-owned handles are +// released. +- (void)clearRetainedState; +@end + +namespace executorch::backends::coreai::testing { + +struct BookmarkLoaderScope final { + BookmarkFakeLoader* loader = nil; + BookmarkLoaderScope(); + ~BookmarkLoaderScope(); + BookmarkLoaderScope(const BookmarkLoaderScope&) = delete; + BookmarkLoaderScope& operator=(const BookmarkLoaderScope&) = delete; + + private: + CoreAIFakeBridgeContext* context_; + id previous_; +}; + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm new file mode 100644 index 00000000000..5eefda12f9d --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm @@ -0,0 +1,545 @@ +/* + * 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_fake_loader.h" +#include +#include "coreai_test_utils.h" + +using executorch::backends::coreai::testing::FakeBridgeState; +using executorch::backends::coreai::testing::test_deadline; + +// ARC keeps failed asynchronous work from outliving its state storage. +@interface CoreAIFakeBridgeContext : NSObject { + @public + FakeBridgeState state; +} +@property(nonatomic, strong) dispatch_group_t callbacks; +@property(nonatomic, copy) void (^retainedBindingCompletion) + (id, NSError*); +@property(nonatomic, copy) void (^retainedExecutionCompletion) + (NSArray*, NSError*); +@end + +@implementation CoreAIFakeBridgeContext +- (instancetype)init { + if ((self = [super init])) { + self.callbacks = dispatch_group_create(); + } + return self; +} +@end + +namespace { +CoreAIFakeBridgeContext* current_context = nil; + +NSError* fake_error(ETCoreAIErrorCode code) { + return [NSError errorWithDomain:ETCoreAIErrorDomain code:code userInfo:nil]; +} + +bool valid_binding( + CoreAIFakeBridgeContext* context, + NSString* function, + NSArray* inputs, + NSArray* outputs) { + // Reversed names make accidental sorting observable in multi-I/O tests. + NSArray* expected_inputs = context->state.multiple_io + ? @[ @"input_1", @"input_0" ] + : @[ @"input_0" ]; + NSArray* expected_outputs = context->state.multiple_io + ? @[ @"output_1", @"output_0" ] + : @[ @"output_0" ]; + const bool valid = [function isEqualToString:@"main"] && + [inputs isEqualToArray:expected_inputs] && + [outputs isEqualToArray:expected_outputs]; + if (!valid) { + ++context->state.binding_mismatches; + } + return valid; +} + +// Release callback captures before publishing that all fake work has drained. +void run_async(CoreAIFakeBridgeContext* context, void (^work)()) { + dispatch_group_t group = context.callbacks; + __block void (^task)() = [work copy]; + dispatch_group_enter(group); + dispatch_async(dispatch_get_global_queue(QOS_CLASS_DEFAULT, 0), ^{ + @autoreleasepool { + task(); + task = nil; + } + dispatch_group_leave(group); + }); +} + +::testing::AssertionResult wait_for_callbacks( + CoreAIFakeBridgeContext* context, + dispatch_time_t deadline) { + if (dispatch_group_wait(context.callbacks, deadline) != 0) { + return ::testing::AssertionFailure() << "Fake SDK callbacks did not finish"; + } + return ::testing::AssertionSuccess(); +} +} // namespace + +@interface FakeOutputOwner : NSObject +@property(nonatomic, strong) CoreAIFakeBridgeContext* context; +- (instancetype)initWithContext:(CoreAIFakeBridgeContext*)context; +@end + +@implementation FakeOutputOwner +- (instancetype)initWithContext:(CoreAIFakeBridgeContext*)context { + if ((self = [super init])) { + self.context = context; + ++context->state.output_owners; + } + return self; +} +- (void)dealloc { + --_context->state.output_owners; +} +@end + +@interface FakeSession () +@property(nonatomic, strong) CoreAIFakeBridgeContext* context; +- (instancetype)initWithContext:(CoreAIFakeBridgeContext*)context; +@end + +@implementation FakeSession +- (instancetype)init { + return [self initWithContext:current_context]; +} +- (instancetype)initWithContext:(CoreAIFakeBridgeContext*)context { + if (context == nil) { + return nil; + } + if ((self = [super init])) { + self.context = context; + ++context->state.sessions; + } + return self; +} +- (void)dealloc { + if (_context != nil) { + --_context->state.sessions; + } +} +- (void)executeInputs:(NSArray*)inputs + completion: + (void (^)(NSArray*, NSError*))completion { + CoreAIFakeBridgeContext* context = self.context; + ++context->state.executions; + const bool should_fail = context->state.fail_execute; + const bool malformed = context->state.bad_output; + dispatch_semaphore_t ready = context->state.input_ready; + dispatch_semaphore_t proceed = context->state.allow_input_read; + if (context->state.retain_execution_completion) { + context.retainedExecutionCompletion = completion; + } + run_async(context, ^{ + ETCoreAIInputTensor* input = inputs[0]; + context->state.last_input_bytes = input.bytes; + context->state.last_input_byte_count = input.byteCount; + if (ready != nil) { + dispatch_semaphore_signal(ready); + if (proceed == nil || + dispatch_semaphore_wait(proceed, test_deadline()) != 0) { + ++context->state.input_wait_timeouts; + completion(nil, fake_error(ETCoreAIErrorRuntime)); + return; + } + } + if (input.byteCount > 0) { + context->state.last_input_first_byte = + static_cast(input.bytes)[0]; + } + if (should_fail) { + completion(nil, fake_error(ETCoreAIErrorRuntime)); + return; + } + if (malformed) { + completion( + @[ [[ETCoreAITensor alloc] initWithData:[NSData data] + shape:@[ @2 ] + scalarType:ETCoreAIScalarTypeFloat32] ], + nil); + return; + } + NSMutableArray* outputs = + [NSMutableArray arrayWithCapacity:inputs.count]; + for (ETCoreAIInputTensor* value in inputs) { + FakeOutputOwner* owner = + [[FakeOutputOwner alloc] initWithContext:context]; + [outputs addObject:[[ETCoreAITensor alloc] + initWithByteCount:value.byteCount + shape:value.shape + scalarType:value.scalarType + copyBytes:^(void* destination) { + (void)owner; + ++context->state.output_copies; + context->state.last_output_destination = + destination; + std::memcpy( + destination, + value.bytes, + value.byteCount); + }]]; + } + completion(outputs, nil); + }); +} +@end + +@interface FakePreparedModel () +@property(nonatomic, strong) CoreAIFakeBridgeContext* context; +- (instancetype)initWithContext:(CoreAIFakeBridgeContext*)context; +@end + +@implementation FakePreparedModel +- (instancetype)init { + return [self initWithContext:current_context]; +} +- (instancetype)initWithContext:(CoreAIFakeBridgeContext*)context { + if (context == nil) { + return nil; + } + if ((self = [super init])) { + self.context = context; + ++context->state.prepared_models; + } + return self; +} +- (void)dealloc { + if (_context != nil) { + --_context->state.prepared_models; + } +} +- (NSData*)copyBookmarkData { + return self.bookmark; +} +- (void)loadFunctionNamed:(NSString*)functionName + inputNames:(NSArray*)inputNames + outputNames:(NSArray*)outputNames + completion:(void (^)(id, NSError*))completion { + CoreAIFakeBridgeContext* context = self.context; + const bool valid = + valid_binding(context, functionName, inputNames, outputNames); + const bool should_fail = context->state.fail_load; + if (context->state.retain_binding_completion) { + context.retainedBindingCompletion = completion; + } + run_async(context, ^{ + if (!valid || should_fail) { + completion(nil, fake_error(ETCoreAIErrorInvalidModel)); + return; + } + FakeSession* session = [[FakeSession alloc] initWithContext:context]; + session.preparedPin = self; + completion(session, nil); + }); +} +@end + +@interface FakeLoader () +@property(nonatomic, strong) CoreAIFakeBridgeContext* context; +@end + +@implementation FakeLoader +- (instancetype)init { + if (current_context == nil) { + return nil; + } + if ((self = [super init])) { + self.context = current_context; + ++_context->state.loaders; + } + return self; +} +- (void)dealloc { + if (_context != nil) { + --_context->state.loaders; + } +} +- (NSString*)specializationIdentifier { + return @"default"; +} +- (void)restoreModelFromBookmark:(NSData*)bookmark + completion:(ETCoreAIAcquisitionCompletion)completion { + FakePreparedModel* model = + [[FakePreparedModel alloc] initWithContext:self.context]; + model.bookmark = bookmark; + _context->state.last_bundle = + [[NSString alloc] initWithData:bookmark encoding:NSUTF8StringEncoding]; + completion(ETCoreAIAcquisitionStatusHit, model, nil); +} +- (void)specializeModelAtURL:(NSURL*)url + completion:(ETCoreAIAcquisitionCompletion)completion { + if (![NSFileManager.defaultManager fileExistsAtPath:url.path]) { + ++_context->state.missing_bundles; + completion( + ETCoreAIAcquisitionStatusError, + nil, + fake_error(ETCoreAIErrorInvalidModel)); + return; + } + _context->state.last_bundle = url.path; + FakePreparedModel* model = + [[FakePreparedModel alloc] initWithContext:self.context]; + model.bookmark = [url.path dataUsingEncoding:NSUTF8StringEncoding]; + completion(ETCoreAIAcquisitionStatusHit, model, nil); +} +- (void)evictModelWithBookmark:(NSData*)bookmark + completion:(void (^)(NSError*))completion { + completion(nil); +} +@end + +@interface BookmarkFakeModel () +@property(nonatomic, strong) CoreAIFakeBridgeContext* context; +- (instancetype)initWithContext:(CoreAIFakeBridgeContext*)context; +@end + +@interface BookmarkFakeLoader () +@property(nonatomic, strong) CoreAIFakeBridgeContext* context; +@end + +@implementation BookmarkFakeModel +- (instancetype)init { + return [self initWithContext:current_context]; +} +- (instancetype)initWithContext:(CoreAIFakeBridgeContext*)context { + if (context == nil) { + return nil; + } + if ((self = [super init])) { + self.context = context; + ++context->state.prepared_models; + } + return self; +} +- (void)dealloc { + if (_context != nil) { + --_context->state.prepared_models; + } +} +- (NSData*)copyBookmarkData { + return self.owner.emptyBookmark ? [NSData data] : self.bookmark; +} +- (void)loadFunctionNamed:(NSString*)functionName + inputNames:(NSArray*)inputNames + outputNames:(NSArray*)outputNames + completion:(void (^)(id, NSError*))completion { + CoreAIFakeBridgeContext* context = self.context; + BookmarkFakeLoader* owner = self.owner; + if (!valid_binding(context, functionName, inputNames, outputNames) || + owner == nil) { + completion(nil, fake_error(ETCoreAIErrorInvalidModel)); + return; + } + ++owner->bindings; + owner.lastBound = self; + run_async(context, ^{ + if (owner.bindError != nil) { + completion(nil, owner.bindError); + return; + } + FakeSession* session = [[FakeSession alloc] initWithContext:context]; + session.preparedPin = self; + completion(session, nil); + }); +} +@end + +@implementation BookmarkFakeLoader +- (instancetype)init { + if (current_context == nil) { + return nil; + } + if ((self = [super init])) { + self.context = current_context; + self.models = [NSMutableDictionary dictionary]; + restores = 0; + specializations = 0; + bindings = 0; + evictions = 0; + specializationQueries = 0; + ++_context->state.loaders; + } + return self; +} +- (void)dealloc { + if (_context != nil) { + --_context->state.loaders; + } +} +- (NSString*)specializationIdentifier { + ++specializationQueries; + return @"default"; +} +- (void)restoreModelFromBookmark:(NSData*)bookmark + completion:(ETCoreAIAcquisitionCompletion)completion { + ++restores; + run_async(self.context, ^{ + @synchronized(self) { + if (self.restoreError != nil) { + completion(ETCoreAIAcquisitionStatusError, nil, self.restoreError); + return; + } + BookmarkFakeModel* model = self.restoreMiss ? nil : self.models[bookmark]; + self.lastAcquired = model; + completion( + model == nil ? ETCoreAIAcquisitionStatusMiss + : ETCoreAIAcquisitionStatusHit, + model, + nil); + } + }); +} +- (void)specializeModelAtURL:(NSURL*)url + completion:(ETCoreAIAcquisitionCompletion)completion { + ++specializations; + run_async(self.context, ^{ + if (self.onSpecialize != nil) { + self.onSpecialize(url); + } + @synchronized(self) { + self.lastSpecializeURL = url.path; + if (self.specializeError != nil) { + completion(ETCoreAIAcquisitionStatusError, nil, self.specializeError); + return; + } + if (![NSFileManager.defaultManager fileExistsAtPath:url.path]) { + ++self.context->state.missing_bundles; + completion( + ETCoreAIAcquisitionStatusError, + nil, + fake_error(ETCoreAIErrorInvalidModel)); + return; + } + BookmarkFakeModel* model = + [[BookmarkFakeModel alloc] initWithContext:self.context]; + model.owner = self; + model.bookmark = + [NSUUID.UUID.UUIDString dataUsingEncoding:NSUTF8StringEncoding]; + self.models[model.bookmark] = model; + self.lastAcquired = model; + completion(ETCoreAIAcquisitionStatusHit, model, nil); + } + }); +} +- (void)evictModelWithBookmark:(NSData*)bookmark + completion:(void (^)(NSError*))completion { + ++evictions; + run_async(self.context, ^{ + if (self.onEvict != nil) { + self.onEvict(bookmark); + } + @synchronized(self) { + if (self.evictError == nil) { + [self.models removeObjectForKey:bookmark]; + } + completion(self.evictError); + } + }); +} +- (void)clearRetainedState { + self.onSpecialize = nil; + self.onEvict = nil; + self.lastAcquired = nil; + self.lastBound = nil; + [self.models removeAllObjects]; +} +@end + +namespace executorch::backends::coreai::testing { + +ScopedFakeBridgeState::ScopedFakeBridgeState() + : context_([[CoreAIFakeBridgeContext alloc] init]), + previous_(current_context) { + current_context = context_; +} + +ScopedFakeBridgeState::~ScopedFakeBridgeState() { + auto drained = wait_for_callbacks(test_deadline()); + EXPECT_TRUE(drained); + if (drained) { + id loader = context_->state.bookmark_loader; + if ([loader isKindOfClass:BookmarkFakeLoader.class]) { + [(BookmarkFakeLoader*)loader clearRetainedState]; + } + context_->state.bookmark_loader = nil; + loader = nil; + context_.retainedBindingCompletion = nil; + context_.retainedExecutionCompletion = nil; + context_->state.input_ready = nil; + context_->state.allow_input_read = nil; + EXPECT_EQ(context_->state.sessions.load(), 0); + EXPECT_EQ(context_->state.prepared_models.load(), 0); + EXPECT_EQ(context_->state.loaders.load(), 0); + EXPECT_EQ(context_->state.output_owners.load(), 0); + EXPECT_EQ(context_->state.binding_mismatches.load(), 0); + EXPECT_EQ(context_->state.missing_bundles.load(), 0); + EXPECT_EQ(context_->state.input_wait_timeouts.load(), 0); + } + EXPECT_EQ(current_context, context_); + current_context = previous_; +} + +FakeBridgeState& ScopedFakeBridgeState::state() { + return context_->state; +} + +::testing::AssertionResult ScopedFakeBridgeState::wait_for_callbacks( + dispatch_time_t deadline) { + return ::wait_for_callbacks(context_, deadline); +} + +BookmarkLoaderScope::BookmarkLoaderScope() + : context_(current_context), + previous_(context_ == nil ? nil : context_->state.bookmark_loader) { + EXPECT_NE(context_, nil) + << "BookmarkLoaderScope requires ScopedFakeBridgeState"; + if (context_ != nil) { + loader = [[BookmarkFakeLoader alloc] init]; + context_->state.bookmark_loader = loader; + } +} + +BookmarkLoaderScope::~BookmarkLoaderScope() { + if (context_ == nil) { + return; + } + auto drained = ::wait_for_callbacks(context_, test_deadline()); + EXPECT_TRUE(drained); + if (drained) { + [loader clearRetainedState]; + } + EXPECT_EQ(context_->state.bookmark_loader, loader); + context_->state.bookmark_loader = previous_; + loader = nil; +} + +} // namespace executorch::backends::coreai::testing + +BOOL ETCoreAIIsAvailable(void) { + return current_context != nil && current_context->state.available; +} + +NSString* ETCoreAIDeviceArchitectureName(void) { + if (current_context == nil) { + return nil; + } + ++current_context->state.architecture_queries; + return current_context->state.device_architecture; +} + +id ETCoreAICreateModelLoader(void) { + if (current_context == nil || current_context->state.fail_loader_factory) { + return nil; + } + return current_context->state.bookmark_loader ?: [[FakeLoader alloc] init]; +} diff --git a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h index 625f1f48a8a..7379eff4267 100644 --- a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h +++ b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h @@ -10,6 +10,9 @@ #import #include +#include +#include +#include "coreai_cache_root.h" namespace executorch::backends::coreai::testing { @@ -23,5 +26,25 @@ struct TestDirectory { ::testing::AssertionResult backup_excluded(NSURL* url); ::testing::AssertionResult set_backup_excluded(NSURL* url, bool excluded); +::testing::AssertionResult asset_tree_snapshot( + NSURL* root, + NSDictionary* __strong& result); +::testing::AssertionResult snapshot_matches( + NSURL* root, + NSDictionary* expected); +NSArray* pending_directories(NSURL* root); + +// Lets one test body cover both cache directory entry points. +inline runtime::Result> +resolve_test_cache_directory(NSString* path, StorageCreationMode mode) { + if (mode == StorageCreationMode::OpenExisting) { + return open_cache_directory(path); + } + auto created = create_cache_directory(path); + if (!created.ok()) { + return created.error(); + } + return std::optional(std::move(created.get())); +} } // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm index 6c60dddd6c5..1e5a7bdcb9a 100644 --- a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm +++ b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm @@ -7,6 +7,7 @@ */ #include "coreai_filesystem_fixture.h" +#include #include namespace executorch::backends::coreai::testing { @@ -68,4 +69,101 @@ << "Backup flag did not settle: " << url.path.UTF8String; } +::testing::AssertionResult asset_tree_snapshot( + NSURL* root, + NSDictionary* __strong& result) { + result = nil; + if (root == nil) + return ::testing::AssertionFailure() << "Missing snapshot root"; + NSFileManager* manager = NSFileManager.defaultManager; + NSError* error = nil; + NSArray* children = [manager subpathsOfDirectoryAtPath:root.path + error:&error]; + if (children == nil || error != nil) { + return ::testing::AssertionFailure() + << "Cannot list " << root.path.UTF8String; + } + NSMutableDictionary* snapshot = [NSMutableDictionary dictionary]; + for (NSString* path in [@[ @"" ] arrayByAddingObjectsFromArray:children]) { + NSURL* url = + path.length == 0 ? root : [root URLByAppendingPathComponent:path]; + struct stat info; + if (lstat(url.fileSystemRepresentation, &info) != 0) { + return ::testing::AssertionFailure() + << "Cannot stat " << url.path.UTF8String; + } + NSMutableDictionary* item = [@{ + @"inode" : @(info.st_ino), + @"device" : @(info.st_dev), + @"mode" : @(info.st_mode), + @"size" : @(info.st_size), + @"mtime_seconds" : @(info.st_mtimespec.tv_sec), + @"mtime_nanos" : @(info.st_mtimespec.tv_nsec) + } mutableCopy]; + error = nil; + if (S_ISREG(info.st_mode)) { + NSData* bytes = [NSData dataWithContentsOfURL:url options:0 error:&error]; + if (bytes == nil || error != nil) { + return ::testing::AssertionFailure() + << "Cannot read " << url.path.UTF8String; + } + item[@"bytes"] = bytes; + } else if (S_ISLNK(info.st_mode)) { + NSString* destination = [manager destinationOfSymbolicLinkAtPath:url.path + error:&error]; + if (destination == nil || error != nil) { + return ::testing::AssertionFailure() + << "Cannot read link " << url.path.UTF8String; + } + item[@"link"] = destination; + } + snapshot[path] = item; + } + result = snapshot; + return ::testing::AssertionSuccess(); +} + +::testing::AssertionResult snapshot_matches( + NSURL* root, + NSDictionary* expected) { + if (expected == nil) + return ::testing::AssertionFailure() << "Missing expected snapshot"; + NSDictionary* actual = nil; + auto read = asset_tree_snapshot(root, actual); + if (!read) + return read; + return [actual isEqual:expected] ? ::testing::AssertionSuccess() + : ::testing::AssertionFailure() + << "Filesystem snapshot changed: " << root.path.UTF8String; +} + +NSArray* pending_directories(NSURL* root) { + if (root == nil) { + ADD_FAILURE() << "Missing cache root"; + return nil; + } + NSURL* cache_sources = [root URLByAppendingPathComponent:@"sources" + isDirectory:YES]; + if ([NSFileManager.defaultManager fileExistsAtPath:cache_sources.path]) { + root = cache_sources; + } + NSError* error = nil; + NSArray* children = + [NSFileManager.defaultManager contentsOfDirectoryAtURL:root + includingPropertiesForKeys:nil + options:0 + error:&error]; + if (children == nil || error != nil) { + ADD_FAILURE() << "Cannot list pending directories: " + << root.path.UTF8String; + return nil; + } + NSMutableArray* pending = [NSMutableArray array]; + for (NSURL* child in children) { + if ([child.lastPathComponent hasPrefix:@".pending-"]) + [pending addObject:child]; + } + return pending; +} + } // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_host_test.mm b/backends/apple/coreai/runtime/test/coreai_host_test.mm index 5c857dfddf4..3702c812e0d 100644 --- a/backends/apple/coreai/runtime/test/coreai_host_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_host_test.mm @@ -9,10 +9,23 @@ #import #include #include +#include "coreai_acquisition_fixture.h" +#include "coreai_cache_entry_fixture.h" int main(int argc, char** argv) { executorch::runtime::runtime_init(); @autoreleasepool { + executorch::backends::coreai::testing::set_bookmark_test_executable( + argv[0]); + const int acquisition_result = + executorch::backends::coreai::testing::acquisition_child_mode( + argc, argv); + if (acquisition_result >= 0) + return acquisition_result; + const int child_result = + executorch::backends::coreai::testing::bookmark_child_mode(argc, argv); + if (child_result >= 0) + return child_result; ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); } diff --git a/backends/apple/coreai/runtime/test/coreai_manifest_fixture.h b/backends/apple/coreai/runtime/test/coreai_manifest_fixture.h index 68223c980a9..bf9af55c0bb 100644 --- a/backends/apple/coreai/runtime/test/coreai_manifest_fixture.h +++ b/backends/apple/coreai/runtime/test/coreai_manifest_fixture.h @@ -23,7 +23,7 @@ inline NSString* fixture_identity(char digit) { inline NSMutableDictionary* manifest_dict() { return [@{ @"packaging" : @"inline", - @"hash" : @"ab", + @"export_hash" : @"ab", @"path" : @"ab/model.aimodel", @"function" : @"main", @"input_names" : @[ @"input_0" ], @@ -60,4 +60,19 @@ inline NSData* encode(NSDictionary* dict) { return [NSJSONSerialization dataWithJSONObject:dict options:0 error:nil]; } +inline runtime::Result selected_bundle( + NSDictionary* dict, + NSString* architecture = @"arch_b", + NSString* platform = @"macOS") { + NSData* encoded = encode(dict); + if (encoded == nil) { + return runtime::Error::InvalidExternalData; + } + auto manifest = parse_manifest(encoded); + if (!manifest.ok()) { + return manifest.error(); + } + return select_bundle(manifest.get(), architecture, platform); +} + } // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_manifest_test.mm b/backends/apple/coreai/runtime/test/coreai_manifest_test.mm index cd7d2d22f65..2e0fa5393e7 100644 --- a/backends/apple/coreai/runtime/test/coreai_manifest_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_manifest_test.mm @@ -170,7 +170,7 @@ auto parsed = parse_manifest(encode(dict)); ASSERT_TRUE(parsed.ok()); EXPECT_EQ(parsed->aot_compiled, GetParam()); - EXPECT_TRUE([parsed->model_hash isEqual:dict[@"hash"]]); + EXPECT_TRUE([parsed->export_hash isEqual:dict[@"export_hash"]]); EXPECT_TRUE([parsed->function isEqual:dict[@"function"]]); EXPECT_TRUE([parsed->inputs isEqual:dict[@"input_names"]]); EXPECT_TRUE([parsed->outputs isEqual:dict[@"output_names"]]); @@ -246,12 +246,12 @@ TEST_P(CoreAIManifestMetadataTest, RejectsInvalidHashes) { auto missing = manifest(); - [missing removeObjectForKey:@"hash"]; + [missing removeObjectForKey:@"export_hash"]; EXPECT_EQ(parse_manifest(encode(missing)).error(), Error::InvalidProgram); for (id hash in @[ @"", @"g", @"AB", @"../x", @"a/b", @1, NSNull.null ]) { SCOPED_TRACE([hash description].UTF8String); auto dict = manifest(); - dict[@"hash"] = hash; + dict[@"export_hash"] = hash; if (GetParam()) { dict[@"archs"] = @{ @"arch_a" : @@ -391,12 +391,12 @@ [parsed->architectures copy]; for (NSString* arch in @[ @"arch_a", @"arch_b" ]) { SCOPED_TRACE(arch.UTF8String); - auto selected = select_assets(parsed.get(), arch, @"macOS"); + auto selected = select_bundle(parsed.get(), arch, @"macOS"); ASSERT_TRUE(selected.ok()); EXPECT_TRUE([selected->path isEqualToString:dict[@"archs"][arch]]); EXPECT_EQ(selected->files.count, 2); NSString* bundle = selected->path.lastPathComponent; - EXPECT_TRUE([selected->model_hash isEqual:parsed->model_hash]); + EXPECT_TRUE([selected->export_hash isEqual:parsed->export_hash]); EXPECT_TRUE([selected->function isEqual:parsed->function]); EXPECT_TRUE([selected->inputs isEqual:parsed->inputs]); EXPECT_TRUE([selected->outputs isEqual:parsed->outputs]); @@ -420,10 +420,10 @@ ASSERT_TRUE(parsed.ok()); ASSERT_FALSE(parsed->aot_compiled); // Selection ignores the architecture and platform for source bundles. - auto selected = select_assets(parsed.get(), @"anything", @"elsewhere"); + auto selected = select_bundle(parsed.get(), @"anything", @"elsewhere"); ASSERT_TRUE(selected.ok()); EXPECT_TRUE([selected->path isEqualToString:dict[@"path"]]); - EXPECT_TRUE([selected->model_hash isEqual:parsed->model_hash]); + EXPECT_TRUE([selected->export_hash isEqual:parsed->export_hash]); EXPECT_TRUE([selected->function isEqual:parsed->function]); EXPECT_TRUE([selected->inputs isEqual:parsed->inputs]); EXPECT_TRUE([selected->outputs isEqual:parsed->outputs]); @@ -441,7 +441,7 @@ ASSERT_TRUE(parsed.ok()); [bytes setLength:0]; [original removeAllObjects]; - auto selected = select_assets(parsed.get(), @"arch_b", @"macOS"); + auto selected = select_bundle(parsed.get(), @"arch_b", @"macOS"); ASSERT_TRUE(selected.ok()); EXPECT_TRUE([selected->path isEqual:@"ab/model.arch_b.aimodelc"]); EXPECT_EQ(selected->files.count, 2); @@ -470,30 +470,31 @@ } TEST(CoreAIManifestTest, ValidationHelpersMatchParserRules) { - EXPECT_FALSE(nonempty_string(nil)); - EXPECT_FALSE(nonempty_string(@"")); - EXPECT_TRUE(nonempty_string(@"ab")); + EXPECT_FALSE(is_nonempty_string(nil)); + EXPECT_FALSE(is_nonempty_string(@"")); + EXPECT_TRUE(is_nonempty_string(@"ab")); const unichar characters[] = {'a', 0, 'b'}; NSString* embedded_nul = [NSString stringWithCharacters:characters length:3]; ASSERT_EQ(embedded_nul.length, 3); - EXPECT_FALSE(nonempty_string(embedded_nul)); - EXPECT_FALSE(relative_path(embedded_nul)); - EXPECT_FALSE(lowercase_hex(embedded_nul)); - EXPECT_FALSE(nonempty_string(@1)); - EXPECT_FALSE(lowercase_hex((id) @1, 64)); - EXPECT_FALSE(lowercase_hex((id)NSNull.null, 64)); - EXPECT_FALSE(relative_path(@"/ab/model.aimodel")); - EXPECT_FALSE(relative_path(@"ab/../model.aimodel")); - EXPECT_FALSE(relative_path(@"a\\b")); - EXPECT_FALSE(relative_path(@"a//b")); - EXPECT_TRUE(relative_path(@"ab/model.aimodel")); - EXPECT_TRUE(lowercase_hex(@"ab")); - EXPECT_FALSE(lowercase_hex(@"")); - EXPECT_FALSE(lowercase_hex(@"AB")); - EXPECT_FALSE(lowercase_hex(@"ag")); - EXPECT_TRUE(lowercase_hex(fixture_identity('a'), 64)); - EXPECT_FALSE(lowercase_hex([fixture_identity('a') substringFromIndex:1], 64)); - EXPECT_FALSE(lowercase_hex(fixture_identity('a').uppercaseString, 64)); + EXPECT_FALSE(is_nonempty_string(embedded_nul)); + EXPECT_FALSE(is_relative_path(embedded_nul)); + EXPECT_FALSE(is_lowercase_hex(embedded_nul)); + EXPECT_FALSE(is_nonempty_string(@1)); + EXPECT_FALSE(is_lowercase_hex((id) @1, 64)); + EXPECT_FALSE(is_lowercase_hex((id)NSNull.null, 64)); + EXPECT_FALSE(is_relative_path(@"/ab/model.aimodel")); + EXPECT_FALSE(is_relative_path(@"ab/../model.aimodel")); + EXPECT_FALSE(is_relative_path(@"a\\b")); + EXPECT_FALSE(is_relative_path(@"a//b")); + EXPECT_TRUE(is_relative_path(@"ab/model.aimodel")); + EXPECT_TRUE(is_lowercase_hex(@"ab")); + EXPECT_FALSE(is_lowercase_hex(@"")); + EXPECT_FALSE(is_lowercase_hex(@"AB")); + EXPECT_FALSE(is_lowercase_hex(@"ag")); + EXPECT_TRUE(is_lowercase_hex(fixture_identity('a'), 64)); + EXPECT_FALSE( + is_lowercase_hex([fixture_identity('a') substringFromIndex:1], 64)); + EXPECT_FALSE(is_lowercase_hex(fixture_identity('a').uppercaseString, 64)); } TEST(CoreAIManifestTest, RejectsMismatchedAotPlatformAndArchitecture) { @@ -501,19 +502,19 @@ ios[@"platform"] = @"iOS"; auto parsed_ios = parse_manifest(encode(ios)); ASSERT_TRUE(parsed_ios.ok()); - EXPECT_TRUE(select_assets(parsed_ios.get(), @"arch_b", @"iOS").ok()); + EXPECT_TRUE(select_bundle(parsed_ios.get(), @"arch_b", @"iOS").ok()); auto parsed = parse_manifest(encode(aot_manifest_dict())); ASSERT_TRUE(parsed.ok()); EXPECT_EQ( - select_assets(parsed.get(), @"arch_b", @"iOS").error(), + select_bundle(parsed.get(), @"arch_b", @"iOS").error(), Error::DelegateInvalidCompatibility); EXPECT_EQ( - select_assets(parsed.get(), @"arch_b", nil).error(), + select_bundle(parsed.get(), @"arch_b", nil).error(), Error::DelegateInvalidCompatibility); for (id arch in @[ @"missing", @"ARCH_B", @"", NSNull.null ]) { SCOPED_TRACE([arch description].UTF8String); EXPECT_EQ( - select_assets(parsed.get(), arch == NSNull.null ? nil : arch, @"macOS") + select_bundle(parsed.get(), arch == NSNull.null ? nil : arch, @"macOS") .error(), Error::DelegateInvalidCompatibility); } diff --git a/backends/apple/coreai/runtime/test/coreai_pte_export_test.cpp b/backends/apple/coreai/runtime/test/coreai_pte_export_test.cpp new file mode 100644 index 00000000000..b78362f989a --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_pte_export_test.cpp @@ -0,0 +1,55 @@ +/* + * 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. + */ + +#include "coreai_pte.h" + +#include +#include + +#include +#include + +namespace { +using executorch::extension::FileDataLoader; +using executorch::runtime::DataLoader; + +bool matches_export(const char* pte_path, const char* manifest_path) { + auto loader = FileDataLoader::from(pte_path); + auto manifest_loader = FileDataLoader::from(manifest_path); + if (!loader.ok() || !manifest_loader.ok()) { + return false; + } + auto result = executorch::backends::coreai::inspect_coreai_pte(loader.get()); + if (!result.ok() || result->processed.size() != 1) { + return false; + } + auto size = manifest_loader->size(); + if (!size.ok() || size.get() == 0) { + return false; + } + auto expected = manifest_loader->load( + 0, + size.get(), + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + return expected.ok() && result->processed[0].size() == size.get() && + std::memcmp(result->processed[0].data(), expected->data(), size.get()) == + 0; +} +} // namespace + +int main(int argc, char** argv) { + executorch::runtime::runtime_init(); + if (argc != 3 || !matches_export(argv[1], argv[2])) { + std::fputs( + "PTE inspection did not reproduce the exported Core AI manifest\n", + stderr); + return 1; + } + std::puts("PTE inspection matches the exported Core AI manifest"); + return 0; +} diff --git a/backends/apple/coreai/runtime/test/coreai_pte_fixture.h b/backends/apple/coreai/runtime/test/coreai_pte_fixture.h new file mode 100644 index 00000000000..6274ce64ecb --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_pte_fixture.h @@ -0,0 +1,76 @@ +/* + * 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 + +#include +#include +#include + +#include +#include +#include +#include + +namespace executorch::backends::coreai::testing { + +struct PteDelegateSpec { + const char* backend; + NSData* data; + bool segmented = false; + int index_override = -1; + bool missing_processed = false; + std::optional backend_length = std::nullopt; +}; + +struct PteMethodSpec { + const char* name; + std::vector delegates; +}; + +class FBSyntheticPTE final : public executorch::runtime::DataLoader { + public: + enum class LoadFault { None, Error, Misaligned, Short, Oversized }; + struct Request { + size_t offset; + size_t size; + SegmentInfo::Type type; + size_t index; + std::string descriptor; + const void* data = nullptr; + }; + + std::vector bytes; + mutable std::vector requests; + mutable std::vector release_order; + mutable bool program_alive_at_backend_release = true; + executorch::runtime::Error size_error = executorch::runtime::Error::Ok; + executorch::runtime::Error load_error = + executorch::runtime::Error::NotSupported; + int fault_call = -1; + LoadFault load_fault = LoadFault::None; + mutable std::atomic program_loads{0}; + mutable std::atomic program_releases{0}; + mutable std::atomic backend_loads{0}; + mutable std::atomic backend_releases{0}; + mutable std::vector backend_indices; + mutable std::vector backend_descriptors; + + explicit FBSyntheticPTE( + const std::vector& methods, + bool extended_header = true); + + executorch::runtime::Result size() const override; + executorch::runtime::Result + load(size_t offset, size_t size, const SegmentInfo& info) const override; + ::testing::AssertionResult check_released() const; +}; + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_pte_fixture.mm b/backends/apple/coreai/runtime/test/coreai_pte_fixture.mm new file mode 100644 index 00000000000..cf463cfaac0 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_pte_fixture.mm @@ -0,0 +1,254 @@ +/* + * 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. + */ + +#include "coreai_pte_fixture.h" + +#include +#include + +#include +#include +#include + +namespace executorch::backends::coreai::testing { +namespace fb = executorch_flatbuffer; +using executorch::runtime::Error; +using executorch::runtime::ExtendedHeader; +using executorch::runtime::FreeableBuffer; +using executorch::runtime::Result; + +FBSyntheticPTE::FBSyntheticPTE( + const std::vector& methods, + bool extended_header) { + flatbuffers::FlatBufferBuilder builder; + builder.ForceDefaults(true); + // Malformed-field tests must mutate only the selected table's vtable. + builder.DedupVtables(false); + std::vector> plans; + std::vector> inline_data; + std::vector> segments; + std::vector segment_bytes; + auto append_segment = [&](NSData* data) { + while (segment_bytes.size() % 16 != 0) { + segment_bytes.push_back(0); + } + const size_t index = segments.size(); + segments.push_back( + fb::CreateDataSegment(builder, segment_bytes.size(), data.length)); + const auto* begin = static_cast(data.bytes); + if (data.length != 0) { + segment_bytes.insert(segment_bytes.end(), begin, begin + data.length); + } + return static_cast(index); + }; + for (const auto& method : methods) { + std::vector> delegates; + for (const auto& delegate : method.delegates) { + uint32_t index; + auto location = fb::DataLocation::INLINE; + if (delegate.segmented) { + index = append_segment(delegate.data); + location = fb::DataLocation::SEGMENT; + } else { + index = static_cast(inline_data.size()); + inline_data.push_back(fb::CreateBackendDelegateInlineData( + builder, + builder.CreateVector( + static_cast(delegate.data.bytes), + delegate.data.length))); + } + if (delegate.index_override >= 0) { + index = static_cast(delegate.index_override); + } + auto reference = delegate.missing_processed + ? flatbuffers::Offset{} + : fb::CreateBackendDelegateDataReference(builder, location, index); + auto name = delegate.backend == nullptr + ? flatbuffers::Offset{} + : builder.CreateString( + delegate.backend, + delegate.backend_length.value_or( + std::strlen(delegate.backend))); + delegates.push_back(fb::CreateBackendDelegate(builder, name, reference)); + } + auto name = method.name == nullptr + ? flatbuffers::Offset{} + : builder.CreateString(method.name); + auto values = + builder.CreateVector(std::vector>{}); + auto inputs = builder.CreateVector(std::vector{}); + auto outputs = builder.CreateVector(std::vector{}); + auto backends = builder.CreateVector(delegates); + auto sizes = builder.CreateVector(std::vector{0}); + fb::ExecutionPlanBuilder plan(builder); + plan.add_name(name); + plan.add_values(values); + plan.add_inputs(inputs); + plan.add_outputs(outputs); + plan.add_delegates(backends); + plan.add_non_const_buffer_sizes(sizes); + plans.push_back(plan.Finish()); + } + // Delegate metadata access must not request named asset payloads. + uint32_t asset_segment = append_segment( + [@"unread named asset" dataUsingEncoding:NSUTF8StringEncoding]); + auto named = fb::CreateNamedData( + builder, + builder.CreateString("coreai/ab/model.aimodel/graph.bin"), + asset_segment); + auto names = builder.CreateVector( + std::vector>{named}); + auto constants = fb::CreateSubsegmentOffsets( + builder, 0, builder.CreateVector(std::vector{0})); + auto plan_vector = builder.CreateVector(plans); + auto inline_vector = builder.CreateVector(inline_data); + auto segment_vector = builder.CreateVector(segments); + fb::ProgramBuilder program(builder); + program.add_version(executorch::runtime::Program::kMaxSupportedSchemaVersion); + program.add_execution_plan(plan_vector); + program.add_backend_delegate_data(inline_vector); + program.add_segments(segment_vector); + program.add_constant_segment(constants); + program.add_named_data(names); + fb::FinishProgramBuffer(builder, program.Finish()); + if (!extended_header) { + bytes.assign( + builder.GetBufferPointer(), + builder.GetBufferPointer() + builder.GetSize()); + return; + } + // Preserve flatbuffer alignment when inserting the extended header. + constexpr size_t extra = 32; + const size_t program_size = builder.GetSize() + extra; + const size_t segment_base = (program_size + 15) & ~size_t(15); + bytes.resize(segment_base + segment_bytes.size()); + std::memcpy( + bytes.data(), builder.GetBufferPointer(), ExtendedHeader::kHeaderOffset); + std::memcpy( + bytes.data() + ExtendedHeader::kHeaderOffset + extra, + builder.GetBufferPointer() + ExtendedHeader::kHeaderOffset, + builder.GetSize() - ExtendedHeader::kHeaderOffset); + auto root_offset = + flatbuffers::ReadScalar(bytes.data()); + flatbuffers::WriteScalar( + bytes.data(), root_offset + extra); + std::memcpy( + bytes.data() + 8, ExtendedHeader::kMagic, ExtendedHeader::kMagicSize); + flatbuffers::WriteScalar(bytes.data() + 12, extra); + flatbuffers::WriteScalar(bytes.data() + 16, program_size); + flatbuffers::WriteScalar(bytes.data() + 24, segment_base); + flatbuffers::WriteScalar(bytes.data() + 32, segment_bytes.size()); + std::memcpy( + bytes.data() + segment_base, segment_bytes.data(), segment_bytes.size()); +} + +Result FBSyntheticPTE::size() const { + if (size_error != Error::Ok) { + return size_error; + } + return bytes.size(); +} + +Result FBSyntheticPTE::load( + size_t offset, + size_t size, + const SegmentInfo& info) const { + const size_t call = requests.size(); + requests.push_back( + {offset, + size, + info.segment_type, + info.segment_index, + info.descriptor == nullptr ? "" : info.descriptor}); + const bool fault = static_cast(call) == fault_call; + if (fault && load_fault == LoadFault::Error) { + return load_error; + } + if (offset > bytes.size() || size > bytes.size() - offset) { + return Error::AccessFailed; + } + if (info.segment_type != SegmentInfo::Type::Program && + info.segment_type != SegmentInfo::Type::Backend) { + ADD_FAILURE() << "Unexpected segment type: " + << static_cast(info.segment_type); + return Error::InvalidArgument; + } + const bool backend = info.segment_type == SegmentInfo::Type::Backend; + if (backend) { + if (info.descriptor == nullptr) { + ADD_FAILURE() << "Backend segment has no descriptor"; + return Error::InvalidArgument; + } + backend_indices.push_back(info.segment_index); + backend_descriptors.emplace_back(info.descriptor); + } else if (info.descriptor != nullptr) { + ADD_FAILURE() << "Program segment has a backend descriptor"; + return Error::InvalidArgument; + } + auto* loads = backend ? &backend_loads : &program_loads; + const bool misaligned = fault && load_fault == LoadFault::Misaligned; + size_t loaded_size = size; + if (fault && load_fault == LoadFault::Short && size != 0) { + --loaded_size; + } else if (fault && load_fault == LoadFault::Oversized) { + ++loaded_size; + } + void* allocation = std::malloc(std::max(loaded_size, size_t(1)) + misaligned); + if (allocation == nullptr) { + ADD_FAILURE() << "Could not allocate " << size << " bytes"; + return Error::MemoryAllocationFailed; + } + auto* data = static_cast(allocation) + misaligned; + if (loaded_size != 0) { + std::memcpy(data, bytes.data() + offset, std::min(size, loaded_size)); + if (loaded_size > size) { + data[size] = 0; + } + } + struct Allocation { + const FBSyntheticPTE* loader; + void* base; + bool backend; + }; + auto* context = new Allocation{this, allocation, backend}; + requests.back().data = data; + ++*loads; + return FreeableBuffer( + data, + loaded_size, + [](void* opaque, void*, size_t) { + auto* context = static_cast(opaque); + const auto* loader = context->loader; + if (context->backend) { + loader->program_alive_at_backend_release &= + loader->program_loads > loader->program_releases; + ++loader->backend_releases; + } else { + ++loader->program_releases; + } + loader->release_order.push_back( + context->backend ? SegmentInfo::Type::Backend + : SegmentInfo::Type::Program); + std::free(context->base); + delete context; + }, + context); +} + +::testing::AssertionResult FBSyntheticPTE::check_released() const { + if (program_loads == program_releases && backend_loads == backend_releases) { + return ::testing::AssertionSuccess(); + } + return ::testing::AssertionFailure() + << "Program loads/releases: " << program_loads.load() << "/" + << program_releases.load() + << "; backend loads/releases: " << backend_loads.load() << "/" + << backend_releases.load(); +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_pte_test.mm b/backends/apple/coreai/runtime/test/coreai_pte_test.mm new file mode 100644 index 00000000000..b9839952fce --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_pte_test.mm @@ -0,0 +1,538 @@ +/* + * 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. + */ + +#include "coreai_pte_fixture.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace executorch::backends::coreai::testing { +namespace { +namespace fb = executorch_flatbuffer; +using executorch::runtime::DataLoader; +using executorch::runtime::Error; +using executorch::runtime::ExtendedHeader; +using executorch::runtime::FreeableBuffer; +using executorch::runtime::Program; +using SegmentType = DataLoader::SegmentInfo::Type; +using LoadFault = FBSyntheticPTE::LoadFault; + +class CoreAIPteTest : public ::testing::Test { + protected: + void SetUp() override { + executorch::runtime::runtime_init(); + } +}; + +NSData* payload(const char* value = "abc") { + return [NSData dataWithBytes:value length:std::strlen(value)]; +} + +template +void set_field(const void* table, flatbuffers::voffset_t field, T value) { + ASSERT_NE(table, nullptr); + auto* fields = static_cast(const_cast(table)); + ASSERT_TRUE(fields->SetField(field, value)); +} + +void clear_field(const void* table, flatbuffers::voffset_t field) { + ASSERT_NE(table, nullptr); + auto* fields = static_cast(table); + ASSERT_NE(fields->GetOptionalFieldOffset(field), 0); + auto* vtable = const_cast(fields->GetVTable()); + flatbuffers::WriteScalar(vtable + field, 0); +} + +void expect_payload(const FreeableBuffer& buffer, NSData* expected) { + ASSERT_EQ(buffer.size(), expected.length); + ASSERT_NE(buffer.data(), nullptr); + EXPECT_EQ(std::memcmp(buffer.data(), expected.bytes, expected.length), 0); +} + +void expect_invalid(FBSyntheticPTE& pte, Error error = Error::InvalidProgram) { + { + auto result = inspect_coreai_pte(pte); + ASSERT_FALSE(result.ok()); + EXPECT_EQ(result.error(), error); + } + EXPECT_TRUE(pte.check_released()); +} + +TEST_F(CoreAIPteTest, BorrowsInlineAndOwnsSegmentsAcrossMethods) { + NSData* first = payload("abc"); + NSData* second = payload("xyz"); + FBSyntheticPTE pte( + {{"forward", + {{"CoreAIBackend", first}, {"UnregisteredBackend", first, true}}}, + {"other", {{"CoreAIBackend", second, true}}}}); + { + auto result = inspect_coreai_pte(pte); + ASSERT_TRUE(result.ok()); + ASSERT_EQ(result->processed.size(), 2); + ASSERT_NO_FATAL_FAILURE(expect_payload(result->processed[0], first)); + ASSERT_NO_FATAL_FAILURE(expect_payload(result->processed[1], second)); + const auto* program = fb::GetProgram(result->program_storage.data()); + const auto* inline_data = + program->backend_delegate_data()->Get(0)->data()->data(); + EXPECT_EQ(result->processed[0].data(), inline_data); + EXPECT_NE(result->program_storage.data(), pte.bytes.data()); + ASSERT_EQ(pte.requests.size(), 3); + EXPECT_EQ(result->program_storage.data(), pte.requests[1].data); + EXPECT_EQ(result->processed[1].data(), pte.requests[2].data); + EXPECT_EQ(pte.backend_indices, (std::vector{1})); + EXPECT_EQ( + pte.backend_descriptors, (std::vector{"CoreAIBackend"})); + EXPECT_EQ(pte.program_loads, 2); + EXPECT_EQ(pte.program_releases, 1); + EXPECT_EQ(pte.backend_releases, 0); + const int releases = pte.program_releases; + result->processed[0].Free(); + EXPECT_EQ(pte.program_releases, releases); + EXPECT_EQ(std::memcmp(inline_data, first.bytes, first.length), 0); + FreeableBuffer moved(std::move(result->processed[1])); + result->processed[1].Free(); + EXPECT_EQ(pte.backend_releases, 0); + ASSERT_NO_FATAL_FAILURE(expect_payload(moved, second)); + moved.Free(); + moved.Free(); + EXPECT_EQ(pte.backend_releases, 1); + } + EXPECT_TRUE(pte.program_alive_at_backend_release); + EXPECT_TRUE(pte.check_released()); +} + +TEST_F(CoreAIPteTest, RoutesNonidentityGlobalReferencesInMethodOrder) { + for (bool segmented : {false, true}) { + SCOPED_TRACE(segmented ? "segmented" : "inline"); + auto row = [&] { + NSData* first = payload("abc"); + NSData* second = payload("xyz"); + ASSERT_EQ(first.length, second.length); + ASSERT_NE(std::memcmp(first.bytes, second.bytes, first.length), 0); + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", payload("---"), segmented, 2}}}, + {"other", + {{"CoreAIBackend", second, segmented, 1}, + {"UnregisteredBackend", first, segmented}}}}); + const auto* source = fb::GetProgram(pte.bytes.data()); + const auto header = + ExtendedHeader::Parse(pte.bytes.data(), pte.bytes.size()); + ASSERT_TRUE(header.ok()); + { + auto result = inspect_coreai_pte(pte); + ASSERT_TRUE(result.ok()); + ASSERT_EQ(result->processed.size(), 2); + ASSERT_NO_FATAL_FAILURE(expect_payload(result->processed[0], first)); + ASSERT_NO_FATAL_FAILURE(expect_payload(result->processed[1], second)); + ASSERT_GE(pte.requests.size(), 2); + for (size_t i = 0; i < 2; ++i) { + EXPECT_EQ(pte.requests[i].offset, 0); + EXPECT_EQ(pte.requests[i].type, SegmentType::Program); + EXPECT_EQ(pte.requests[i].index, 0); + EXPECT_TRUE(pte.requests[i].descriptor.empty()); + } + EXPECT_EQ(pte.requests[0].size, ExtendedHeader::kNumHeadBytes); + EXPECT_EQ(pte.requests[1].size, header->program_size); + if (segmented) { + ASSERT_EQ(pte.requests.size(), 4); + EXPECT_EQ(pte.backend_indices, (std::vector{2, 1})); + for (size_t i = 0; i < 2; ++i) { + const size_t index = i == 0 ? 2 : 1; + const auto* segment = source->segments()->Get(index); + const auto& request = pte.requests[i + 2]; + EXPECT_EQ(request.type, SegmentType::Backend); + EXPECT_EQ(request.index, index); + EXPECT_EQ(request.descriptor, "CoreAIBackend"); + EXPECT_EQ( + request.offset, + header->segment_base_offset + segment->offset()); + EXPECT_EQ(request.size, segment->size()); + EXPECT_EQ(request.data, result->processed[i].data()); + } + } else { + ASSERT_EQ(pte.requests.size(), 2); + const auto* program = fb::GetProgram(result->program_storage.data()); + EXPECT_EQ( + result->processed[0].data(), + program->backend_delegate_data()->Get(2)->data()->data()); + EXPECT_EQ( + result->processed[1].data(), + program->backend_delegate_data()->Get(1)->data()->data()); + } + } + EXPECT_TRUE(pte.check_released()); + EXPECT_TRUE(pte.program_alive_at_backend_release); + }; + ASSERT_NO_FATAL_FAILURE(row()); + } +} + +TEST_F(CoreAIPteTest, MatchesCompleteBackendIdentifier) { + const char embedded[] = "CoreAIBackend\0suffix"; + PteDelegateSpec nul_id{embedded, payload(), true}; + nul_id.backend_length = sizeof(embedded) - 1; + nul_id.missing_processed = true; + FBSyntheticPTE pte( + {{"forward", + {{"CoreAIBackendSuffix", payload(), true, 999}, + {"CoreAIBacken", payload(), true, 999}, + {"", payload(), true, 999}, + {nullptr, payload(), true, 999, true}, + nul_id, + {"CoreAIBackend", payload("yes")}}}}); + { + auto result = inspect_coreai_pte(pte); + ASSERT_TRUE(result.ok()); + ASSERT_EQ(result->processed.size(), 1); + ASSERT_NO_FATAL_FAILURE( + expect_payload(result->processed[0], payload("yes"))); + EXPECT_EQ(pte.backend_loads, 0); + EXPECT_EQ(pte.requests.size(), 2); + } + EXPECT_TRUE(pte.check_released()); +} + +TEST_F(CoreAIPteTest, AcceptsEmptyAndUnselectedPrograms) { + for (bool other_backend : {false, true}) { + SCOPED_TRACE(other_backend ? "no CoreAI delegate" : "no plans"); + std::vector methods; + if (other_backend) { + methods.push_back( + {"forward", {{"UnregisteredBackend", payload(), true, 999}}}); + } + FBSyntheticPTE pte(methods); + { + auto result = inspect_coreai_pte(pte); + ASSERT_TRUE(result.ok()); + EXPECT_TRUE(result->processed.empty()); + EXPECT_EQ(pte.backend_loads, 0); + } + EXPECT_TRUE(pte.check_released()); + } +} + +TEST_F(CoreAIPteTest, AcceptsHeaderlessAndExtendedPrograms) { + for (bool extended : {false, true}) { + SCOPED_TRACE(extended ? "extended" : "headerless"); + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", payload(), extended}}}}, extended); + { + auto result = inspect_coreai_pte(pte); + ASSERT_TRUE(result.ok()); + ASSERT_EQ(result->processed.size(), 1); + ASSERT_NO_FATAL_FAILURE(expect_payload(result->processed[0], payload())); + ASSERT_GE(pte.requests.size(), 2); + EXPECT_EQ(pte.requests[0].size, ExtendedHeader::kNumHeadBytes); + if (!extended) + EXPECT_EQ(pte.requests[1].size, pte.bytes.size()); + } + EXPECT_TRUE(pte.check_released()); + } +} + +TEST_F(CoreAIPteTest, AcceptsUnusedMethodMetadata) { + for (auto field : + {fb::ExecutionPlan::VT_NAME, + fb::ExecutionPlan::VT_NON_CONST_BUFFER_SIZES, + fb::ExecutionPlan::VT_INPUTS, + fb::ExecutionPlan::VT_OUTPUTS}) { + SCOPED_TRACE(field); + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", payload()}}}, + {"other", {{"CoreAIBackend", payload("xyz")}}}}); + const auto* plan = + fb::GetProgram(pte.bytes.data())->execution_plan()->Get(0); + ASSERT_NO_FATAL_FAILURE(clear_field(plan, field)); + { + auto program = + Program::load(&pte, Program::Verification::InternalConsistency); + ASSERT_TRUE(program.ok()); + auto result = inspect_coreai_pte(pte); + ASSERT_TRUE(result.ok()); + ASSERT_EQ(result->processed.size(), 2); + ASSERT_NO_FATAL_FAILURE(expect_payload(result->processed[0], payload())); + ASSERT_NO_FATAL_FAILURE( + expect_payload(result->processed[1], payload("xyz"))); + } + EXPECT_TRUE(pte.check_released()); + } +} + +TEST_F(CoreAIPteTest, AcceptsDuplicateMethodNames) { + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", payload()}}}, + {"forward", {{"CoreAIBackend", payload("xyz")}}}}); + { + auto program = + Program::load(&pte, Program::Verification::InternalConsistency); + ASSERT_TRUE(program.ok()); + auto result = inspect_coreai_pte(pte); + ASSERT_TRUE(result.ok()); + ASSERT_EQ(result->processed.size(), 2); + ASSERT_NO_FATAL_FAILURE(expect_payload(result->processed[0], payload())); + ASSERT_NO_FATAL_FAILURE( + expect_payload(result->processed[1], payload("xyz"))); + } + EXPECT_TRUE(pte.check_released()); +} + +TEST_F(CoreAIPteTest, AcceptsEmptyProcessedBuffers) { + for (bool segmented : {false, true}) { + SCOPED_TRACE(segmented); + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", payload(""), segmented}}}}); + { + auto result = inspect_coreai_pte(pte); + ASSERT_TRUE(result.ok()); + ASSERT_EQ(result->processed.size(), 1); + EXPECT_EQ(result->processed[0].size(), 0); + } + EXPECT_TRUE(pte.check_released()); + } +} + +TEST_F(CoreAIPteTest, RejectsMissingAndOutOfRangeCoreAIReferences) { + for (bool segmented : {false, true}) { + for (bool missing : {false, true}) { + SCOPED_TRACE(segmented); + SCOPED_TRACE(missing); + PteDelegateSpec delegate{"CoreAIBackend", payload(), segmented}; + delegate.missing_processed = missing; + delegate.index_override = 999; + FBSyntheticPTE pte({{"forward", {delegate}}}); + ASSERT_NO_FATAL_FAILURE(expect_invalid( + pte, missing ? Error::InvalidProgram : Error::NotFound)); + EXPECT_EQ(pte.backend_loads, 0); + EXPECT_EQ(pte.requests.size(), 2); + } + } +} + +TEST_F(CoreAIPteTest, RejectsSegmentsWithoutExtendedHeader) { + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", payload(), true}}}}, false); + ASSERT_NO_FATAL_FAILURE(expect_invalid(pte, Error::NotFound)); + EXPECT_EQ(pte.backend_loads, 0); + EXPECT_EQ(pte.requests.size(), 2); +} + +TEST_F(CoreAIPteTest, RejectsInvalidExtendedHeaderBounds) { + FBSyntheticPTE baseline({{"forward", {{"CoreAIBackend", payload(), true}}}}); + const auto header = + ExtendedHeader::Parse(baseline.bytes.data(), baseline.bytes.size()); + ASSERT_TRUE(header.ok()); + struct Case { + const char* name; + size_t offset; + uint64_t value; + }; + for (const auto& test : + {Case{"short program", 16, ExtendedHeader::kNumHeadBytes - 1}, + Case{"program past EOF", 16, baseline.bytes.size() + 1}, + Case{"segments overlap program", 24, header->program_size - 1}, + Case{"segment base past EOF", 24, baseline.bytes.size() + 1}, + Case{"segment bytes without a base", 24, 0}, + Case{"segment bytes past EOF", 32, header->segment_data_size + 1}, + Case{ + "segment extent overflow", + 32, + std::numeric_limits::max()}}) { + SCOPED_TRACE(test.name); + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", payload(), true}}}}); + flatbuffers::WriteScalar( + pte.bytes.data() + test.offset, test.value); + ASSERT_NO_FATAL_FAILURE(expect_invalid(pte)); + EXPECT_EQ(pte.requests.size(), 1); + EXPECT_EQ(pte.backend_loads, 0); + } +} + +TEST_F(CoreAIPteTest, RejectsInvalidSelectedSegmentBounds) { + for (int fault = 0; fault < 4; ++fault) { + SCOPED_TRACE(fault); + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", payload(), true}}}}); + const auto header = + ExtendedHeader::Parse(pte.bytes.data(), pte.bytes.size()); + ASSERT_TRUE(header.ok()); + const auto* segment = fb::GetProgram(pte.bytes.data())->segments()->Get(0); + if (fault == 0) { + ASSERT_NO_FATAL_FAILURE(set_field( + segment, + fb::DataSegment::VT_OFFSET, + std::numeric_limits::max())); + } else if (fault == 1) { + ASSERT_NO_FATAL_FAILURE(set_field( + segment, + fb::DataSegment::VT_SIZE, + std::numeric_limits::max())); + } else if (fault == 2) { + ASSERT_NO_FATAL_FAILURE(set_field( + segment, fb::DataSegment::VT_OFFSET, header->segment_data_size)); + } else { + flatbuffers::WriteScalar( + pte.bytes.data() + 32, segment->size() - 1); + } + ASSERT_NO_FATAL_FAILURE(expect_invalid(pte)); + EXPECT_EQ(pte.requests.size(), 2); + EXPECT_EQ(pte.backend_loads, 0); + } +} + +TEST_F(CoreAIPteTest, RejectsShortLoadsFromEveryPhase) { + for (int phase = 0; phase < 3; ++phase) { + SCOPED_TRACE(phase); + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", payload(), true}}}}); + pte.fault_call = phase; + pte.load_fault = LoadFault::Short; + ASSERT_NO_FATAL_FAILURE(expect_invalid(pte)); + EXPECT_EQ(pte.requests.size(), phase + 1); + EXPECT_TRUE(pte.program_alive_at_backend_release); + } +} + +TEST_F(CoreAIPteTest, RejectsMisalignedProgramImages) { + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", payload()}}}}); + pte.fault_call = 1; + pte.load_fault = LoadFault::Misaligned; + ASSERT_NO_FATAL_FAILURE(expect_invalid(pte, Error::InvalidArgument)); + EXPECT_EQ(pte.requests.size(), 2); + EXPECT_EQ(pte.backend_loads, 0); +} + +TEST_F(CoreAIPteTest, RejectsMalformedFlatbufferAndUnsupportedVersion) { + for (bool version : {false, true}) { + SCOPED_TRACE(version ? "unsupported version" : "bad identifier"); + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", payload()}}}}, false); + if (version) { + ASSERT_NO_FATAL_FAILURE(set_field( + fb::GetProgram(pte.bytes.data()), + fb::Program::VT_VERSION, + Program::kMaxSupportedSchemaVersion + 1)); + } else { + pte.bytes[4] = 'X'; + } + ASSERT_NO_FATAL_FAILURE(expect_invalid(pte)); + EXPECT_EQ(pte.backend_loads, 0); + } +} + +TEST_F(CoreAIPteTest, RejectsSemanticallyInvalidVerifiedProgram) { + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", payload()}}}}); + const auto* program = fb::GetProgram(pte.bytes.data()); + ASSERT_NO_FATAL_FAILURE(clear_field( + program->execution_plan()->Get(0), fb::ExecutionPlan::VT_VALUES)); + flatbuffers::Verifier verifier(pte.bytes.data(), pte.bytes.size()); + ASSERT_TRUE(fb::VerifyProgramBuffer(verifier)); + EXPECT_EQ(runtime::validate_program(program), Error::InvalidProgram); + ASSERT_NO_FATAL_FAILURE(expect_invalid(pte)); + EXPECT_EQ(pte.backend_loads, 0); +} + +TEST_F(CoreAIPteTest, RejectsUnknownProcessedDataLocation) { + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", payload()}}}}); + const auto* reference = fb::GetProgram(pte.bytes.data()) + ->execution_plan() + ->Get(0) + ->delegates() + ->Get(0) + ->processed(); + ASSERT_NO_FATAL_FAILURE(set_field( + reference, fb::BackendDelegateDataReference::VT_LOCATION, 255)); + flatbuffers::Verifier verifier(pte.bytes.data(), pte.bytes.size()); + ASSERT_TRUE(fb::VerifyProgramBuffer(verifier)); + ASSERT_NO_FATAL_FAILURE(expect_invalid(pte)); + EXPECT_EQ(pte.backend_loads, 0); +} + +TEST_F(CoreAIPteTest, AcceptsLegacyHeaderWithUnusedSegmentBase) { + for (uint64_t base : {uint64_t(1), std::numeric_limits::max()}) { + SCOPED_TRACE(base); + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", payload()}}}}); + flatbuffers::WriteScalar(pte.bytes.data() + 24, base); + flatbuffers::WriteScalar(pte.bytes.data() + 32, 0); + { + auto program = + Program::load(&pte, Program::Verification::InternalConsistency); + ASSERT_TRUE(program.ok()); + auto result = inspect_coreai_pte(pte); + ASSERT_TRUE(result.ok()); + ASSERT_EQ(result->processed.size(), 1); + ASSERT_NO_FATAL_FAILURE(expect_payload(result->processed[0], payload())); + EXPECT_EQ(pte.backend_loads, 0); + } + EXPECT_TRUE(pte.check_released()); + } +} + +TEST_F(CoreAIPteTest, ChecksSelectedLegacySegmentBaseBeforeSubtracting) { + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", payload(), true}}}}); + flatbuffers::WriteScalar( + pte.bytes.data() + 24, std::numeric_limits::max()); + flatbuffers::WriteScalar(pte.bytes.data() + 32, 0); + ASSERT_NO_FATAL_FAILURE(expect_invalid(pte)); + EXPECT_EQ(pte.requests.size(), 2); + EXPECT_EQ(pte.backend_loads, 0); +} + +TEST_F(CoreAIPteTest, RejectsSegmentsWithZeroBase) { + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", payload(), true}}}}); + flatbuffers::WriteScalar(pte.bytes.data() + 24, 0); + flatbuffers::WriteScalar(pte.bytes.data() + 32, 0); + ASSERT_NO_FATAL_FAILURE(expect_invalid(pte, Error::NotFound)); + EXPECT_EQ(pte.backend_loads, 0); +} + +TEST_F(CoreAIPteTest, PropagatesLoadErrorsFromEveryPhase) { + const char* phases[] = { + "size", "prefix", "retained program", "selected segment"}; + for (int phase = 0; phase < 4; ++phase) { + SCOPED_TRACE(phases[phase]); + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", payload(), true}}}}); + if (phase == 0) { + pte.size_error = Error::AccessFailed; + } else { + pte.fault_call = phase - 1; + pte.load_fault = LoadFault::Error; + pte.load_error = Error::AccessFailed; + } + ASSERT_NO_FATAL_FAILURE(expect_invalid(pte, Error::AccessFailed)); + EXPECT_EQ(pte.requests.size(), phase); + } +} + +TEST_F(CoreAIPteTest, ReleasesEarlierBuffersOnLateFailure) { + for (auto fault : + {LoadFault::Error, LoadFault::Short, LoadFault::Oversized}) { + SCOPED_TRACE(static_cast(fault)); + FBSyntheticPTE pte( + {{"forward", + {{"CoreAIBackend", payload()}, + {"CoreAIBackend", payload("one"), true}, + {"CoreAIBackend", payload("two"), true}}}}); + pte.fault_call = 3; + pte.load_fault = fault; + const bool load_error = fault == LoadFault::Error; + ASSERT_NO_FATAL_FAILURE(expect_invalid( + pte, load_error ? Error::NotSupported : Error::InvalidProgram)); + EXPECT_EQ(pte.backend_loads, load_error ? 1 : 2); + EXPECT_TRUE(pte.program_alive_at_backend_release); + ASSERT_FALSE(pte.release_order.empty()); + EXPECT_EQ(pte.release_order.back(), SegmentType::Program); + } +} + +} // namespace +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_source_fixture.h b/backends/apple/coreai/runtime/test/coreai_source_fixture.h new file mode 100644 index 00000000000..a5377e4f93d --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_fixture.h @@ -0,0 +1,56 @@ +/* + * 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 +#include +#include +#include "coreai_cache_entry_fixture.h" +#include "coreai_filesystem_fixture.h" +#include "coreai_manifest_fixture.h" +#include "coreai_source_bundle.h" + +namespace executorch::backends::coreai::testing { + +class TestData final : public runtime::NamedDataMap { + public: + mutable std::atomic releases{0}; + mutable std::atomic requests{0}; + mutable std::atomic attempts{0}; + mutable std::atomic metadata_requests{0}; + int fail_at = 0; + runtime::Error failure_error = runtime::Error::AccessFailed; + std::string required_prefix; + std::map files{ + {"coreai/ab/model.aimodel/graph.bin", std::string("model data", 11)}}; + runtime::Result get_tensor_layout( + std::string_view) const override; + runtime::Result get_data( + std::string_view key) const override; + runtime::Error load_data_into(std::string_view, void*, size_t) const override; + runtime::Result get_num_keys() const override; + runtime::Result get_key(uint32_t i) const override; +}; + +void aot_data(TestData& data, NSString* arch); +NSString* source_key(const SelectedBundle& bundle); +NSURL* sources_root(NSURL* cache_root); +NSURL* source_entry(NSURL* cache_root, const SelectedBundle& bundle); + +class CoreAISourceTest : public ::testing::Test { + protected: + void check_concurrent_source(); + runtime::Result prepare_source_bundle( + const SelectedBundle& bundle, + const runtime::NamedDataMap* data, + NSString* cache_root, + CorruptBundlePolicy corrupt_policy = CorruptBundlePolicy::Reject); +}; + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_source_fixture.mm b/backends/apple/coreai/runtime/test/coreai_source_fixture.mm new file mode 100644 index 00000000000..c06cef50cf2 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_fixture.mm @@ -0,0 +1,112 @@ +/* + * 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. + */ + +#include "coreai_source_fixture.h" + +#include +#include "coreai_cache_entry.h" +#include "coreai_cache_key.h" + +namespace executorch::backends::coreai::testing { +using runtime::Error; +using runtime::Result; + +Result TestData::get_tensor_layout( + std::string_view) const { + return Error::NotSupported; +} + +Result TestData::get_data(std::string_view key) const { + if (!required_prefix.empty() && + key.substr(0, required_prefix.size()) != required_prefix) { + ADD_FAILURE() << "NamedDataMap key " << key + << " is outside selected prefix " << required_prefix; + return Error::InvalidArgument; + } + if (++attempts == fail_at) { + return failure_error; + } + auto found = files.find(std::string(key)); + if (found == files.end()) { + return Error::NotFound; + } + ++requests; + return runtime::FreeableBuffer( + found->second.data(), + found->second.size(), + [](void* context, void*, size_t) { + ++static_cast(context)->releases; + }, + const_cast(this)); +} + +Error TestData::load_data_into(std::string_view, void*, size_t) const { + return Error::NotSupported; +} + +Result TestData::get_num_keys() const { + ++metadata_requests; + return static_cast(files.size()); +} + +Result TestData::get_key(uint32_t i) const { + ++metadata_requests; + if (i >= files.size()) { + return Error::NotFound; + } + auto found = files.begin(); + std::advance(found, i); + return found->first.c_str(); +} + +void aot_data(TestData& data, NSString* arch) { + data.required_prefix = + [NSString stringWithFormat:@"coreai/ab/model.%@.aimodelc/", arch] + .UTF8String; + data.files.clear(); + data.files[data.required_prefix + "graph.bin"] = + std::string("compiled\0graph", 14); + data.files[data.required_prefix + "nested/weights.bin"] = "weights"; +} + +NSString* source_key(const SelectedBundle& bundle) { + auto key = cache_variant_key(bundle, @"macOS", @"arch_b", @"default"); + EXPECT_TRUE(key.ok()); + return key.ok() ? key.get() : nil; +} + +NSURL* sources_root(NSURL* cache_root) { + return [cache_root URLByAppendingPathComponent:@"sources" isDirectory:YES]; +} + +NSURL* source_entry(NSURL* cache_root, const SelectedBundle& bundle) { + return + [sources_root(cache_root) URLByAppendingPathComponent:source_key(bundle) + isDirectory:YES]; +} + +Result CoreAISourceTest::prepare_source_bundle( + const SelectedBundle& bundle, + const runtime::NamedDataMap* data, + NSString* cache_root, + CorruptBundlePolicy corrupt_policy) { + auto key = cache_variant_key(bundle, @"macOS", @"arch_b", @"default"); + if (!key.ok()) { + ADD_FAILURE() << "Cannot compute source fixture key"; + return key.error(); + } + auto lock = acquire_test_cache_entry(cache_root, key.get()); + if (!lock.ok()) { + ADD_FAILURE() << "Cannot acquire source fixture lock"; + return lock.error(); + } + return executorch::backends::coreai::prepare_source_bundle( + bundle, data, *lock.get(), corrupt_policy); +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_source_integrity_test.mm b/backends/apple/coreai/runtime/test/coreai_source_integrity_test.mm new file mode 100644 index 00000000000..264042988e7 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_integrity_test.mm @@ -0,0 +1,304 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include "coreai_cache_entry.h" +#include "coreai_cache_key.h" +#include "coreai_fault_scope.h" +#include "coreai_source_fixture.h" +namespace executorch::backends::coreai::testing { +using runtime::Error; +using runtime::Result; +TEST_F(CoreAISourceTest, SourcesRemainAnchoredWhenConfiguredSymlinkChanges) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* target = [storage.url URLByAppendingPathComponent:@"target"]; + NSURL* replacement = [storage.url URLByAppendingPathComponent:@"replacement"]; + NSURL* configured = [storage.url URLByAppendingPathComponent:@"configured"]; + ASSERT_EQ(mkdir(target.fileSystemRepresentation, 0777), 0); + ASSERT_EQ(mkdir(replacement.fileSystemRepresentation, 0777), 0); + ASSERT_EQ( + symlink( + target.fileSystemRepresentation, configured.fileSystemRepresentation), + 0); + auto bundle = selected_bundle(manifest_dict()); + ASSERT_TRUE(bundle.ok()); + auto key = test_cache_key(bundle.get()); + ASSERT_TRUE(key.ok()); + auto directory = test_cache_directory(configured.path); + ASSERT_TRUE(directory.ok()); + NSString* canonical_target = [directory->path copy]; + auto locked = LockedCacheEntry::acquire(directory.get(), key.get()); + ASSERT_TRUE(locked.ok()); + + ASSERT_EQ(unlink(configured.fileSystemRepresentation), 0); + ASSERT_EQ( + symlink( + replacement.fileSystemRepresentation, + configured.fileSystemRepresentation), + 0); + TestData data; + auto prepared = executorch::backends::coreai::prepare_source_bundle( + bundle.get(), &data, *locked.get()); + ASSERT_TRUE(prepared.ok()); + EXPECT_TRUE([prepared.get().path hasPrefix:canonical_target]); + EXPECT_FALSE([prepared.get().path + hasPrefix:replacement.URLByResolvingSymlinksInPath.path]); + EXPECT_TRUE( + [NSFileManager.defaultManager fileExistsAtPath:prepared.get().path]); +} + +TEST_F(CoreAISourceTest, RejectsMalformedStoredBundlesWithoutReplacement) { + for (bool aot : {false, true}) { + SCOPED_TRACE(aot ? "aot" : "source"); + for (NSString* damage in @[ + @"missing-file", + @"extra-file", + @"wrong-size", + @"symlink-file", + @"file-at-key", + @"extra-directory" + ]) { + SCOPED_TRACE(damage.UTF8String); + TestDirectory storage; + TestDirectory target; + ASSERT_NE(storage.url, nil); + ASSERT_NE(target.url, nil); + TestData data; + if (aot) + aot_data(data, @"arch_b"); + auto manifest = + selected_bundle(aot ? aot_manifest_dict() : manifest_dict()); + ASSERT_TRUE(manifest.ok()); + auto source = + prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(source.ok()); + NSURL* entry = source.get().URLByDeletingLastPathComponent; + NSURL* graph = [source.get() URLByAppendingPathComponent:@"graph.bin"]; + NSData* bytes = [NSData dataWithContentsOfURL:graph]; + ASSERT_NE(bytes, nil); + NSFileManager* manager = NSFileManager.defaultManager; + if ([damage isEqual:@"missing-file"]) { + ASSERT_TRUE([manager removeItemAtURL:graph error:nil]); + } else if ([damage isEqual:@"extra-file"]) { + ASSERT_TRUE([[NSData data] + writeToURL:[source.get() URLByAppendingPathComponent:@"extra"] + atomically:NO]); + } else if ([damage isEqual:@"wrong-size"]) { + ASSERT_TRUE([[NSData data] writeToURL:graph atomically:NO]); + } else if ([damage isEqual:@"symlink-file"]) { + NSURL* destination = + [target.url URLByAppendingPathComponent:@"graph.bin"]; + ASSERT_TRUE([bytes writeToURL:destination atomically:NO]); + ASSERT_TRUE([manager removeItemAtURL:graph error:nil]); + ASSERT_TRUE([manager createSymbolicLinkAtPath:graph.path + withDestinationPath:destination.path + error:nil]); + } else if ([damage isEqual:@"file-at-key"]) { + ASSERT_TRUE([manager removeItemAtURL:entry error:nil]); + ASSERT_TRUE([bytes writeToURL:entry atomically:NO]); + } else { + ASSERT_TRUE([manager + createDirectoryAtURL: + [source.get() URLByAppendingPathComponent:@"extra"] + withIntermediateDirectories:NO + attributes:nil + error:nil]); + } + NSDictionary* snapshot = nil; + NSDictionary* target_snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, snapshot)); + ASSERT_TRUE(asset_tree_snapshot(target.url, target_snapshot)); + const int attempts = data.attempts; + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path) + .error(), + Error::InvalidExternalData); + EXPECT_EQ(data.attempts, attempts); + EXPECT_TRUE(snapshot_matches(storage.url, snapshot)); + EXPECT_TRUE(snapshot_matches(target.url, target_snapshot)); + EXPECT_EQ(pending_directories(storage.url).count, 0u); + } + } +} + +TEST_F(CoreAISourceTest, ExplicitReplacementRepairsMalformedStoredBundles) { + for (NSString* damage in @[ + @"missing-file", + @"extra-file", + @"wrong-size", + @"symlink-file", + @"fifo-file", + @"file-at-key", + @"symlink-at-key", + @"extra-directory" + ]) { + SCOPED_TRACE(damage.UTF8String); + TestDirectory storage; + TestDirectory target; + ASSERT_NE(storage.url, nil); + ASSERT_NE(target.url, nil); + TestData data; + auto bundle = selected_bundle(manifest_dict()); + ASSERT_TRUE(bundle.ok()); + auto source = prepare_source_bundle(bundle.get(), &data, storage.url.path); + ASSERT_TRUE(source.ok()); + NSURL* entry = source.get().URLByDeletingLastPathComponent; + NSURL* graph = [source.get() URLByAppendingPathComponent:@"graph.bin"]; + NSFileManager* manager = NSFileManager.defaultManager; + if ([damage isEqual:@"file-at-key"] || [damage isEqual:@"symlink-at-key"]) { + ASSERT_TRUE([manager removeItemAtURL:entry error:nil]); + if ([damage isEqual:@"file-at-key"]) { + ASSERT_TRUE([[NSData data] writeToURL:entry atomically:NO]); + } else { + ASSERT_EQ( + symlink( + target.url.fileSystemRepresentation, + entry.fileSystemRepresentation), + 0); + } + } else if ([damage isEqual:@"extra-file"]) { + ASSERT_TRUE([[NSData data] + writeToURL:[source.get() URLByAppendingPathComponent:@"extra"] + atomically:NO]); + } else if ([damage isEqual:@"extra-directory"]) { + ASSERT_TRUE([manager + createDirectoryAtURL:[source.get() + URLByAppendingPathComponent:@"extra"] + withIntermediateDirectories:NO + attributes:nil + error:nil]); + } else { + ASSERT_TRUE([manager removeItemAtURL:graph error:nil]); + if ([damage isEqual:@"wrong-size"]) { + ASSERT_TRUE([[NSData data] writeToURL:graph atomically:NO]); + } else if ([damage isEqual:@"symlink-file"]) { + NSURL* sentinel = [target.url URLByAppendingPathComponent:@"sentinel"]; + ASSERT_TRUE([@"external" writeToURL:sentinel + atomically:NO + encoding:NSUTF8StringEncoding + error:nil]); + ASSERT_EQ( + symlink( + sentinel.fileSystemRepresentation, + graph.fileSystemRepresentation), + 0); + } else if ([damage isEqual:@"fifo-file"]) { + ASSERT_EQ(mkfifo(graph.fileSystemRepresentation, 0600), 0); + } + } + NSDictionary* target_snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(target.url, target_snapshot)); + auto repaired = prepare_source_bundle( + bundle.get(), &data, storage.url.path, CorruptBundlePolicy::Replace); + ASSERT_TRUE(repaired.ok()); + EXPECT_TRUE([repaired.get() isEqual:source.get()]); + EXPECT_TRUE(([[NSData + dataWithContentsOfURL:[repaired.get() + URLByAppendingPathComponent:@"graph.bin"]] + isEqual:[NSData dataWithBytes:"model data" length:11]])); + EXPECT_EQ(data.attempts, 2); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(pending_directories(storage.url).count, 0u); + EXPECT_TRUE(snapshot_matches(target.url, target_snapshot)); + } +} + +TEST_F(CoreAISourceTest, PublishedDirectoriesUseDefaultPermissions) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + NSURL* reference = [storage.url URLByAppendingPathComponent:@"reference"]; + ASSERT_EQ(mkdir(reference.fileSystemRepresentation, 0777), 0); + struct stat expected; + ASSERT_EQ(stat(reference.fileSystemRepresentation, &expected), 0); + TestData data; + auto manifest = selected_bundle(manifest_dict()); + ASSERT_TRUE(manifest.ok()); + auto source = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(source.ok()); + for (NSURL* directory in + @[ source.get().URLByDeletingLastPathComponent, source.get() ]) { + SCOPED_TRACE(directory.path.UTF8String); + struct stat actual; + ASSERT_EQ(stat(directory.fileSystemRepresentation, &actual), 0); + EXPECT_EQ(actual.st_mode & 0777, expected.st_mode & 0777); + } +} + +TEST_F(CoreAISourceTest, PermissionErrorsRemainDistinctFromCorruptEntries) { + if (geteuid() == 0) { + GTEST_SKIP() << "Root can bypass directory read permissions"; + } + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto manifest = selected_bundle(manifest_dict()); + ASSERT_TRUE(manifest.ok()); + auto source = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(source.ok()); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, snapshot)); + for (NSURL* directory in + @[ source.get().URLByDeletingLastPathComponent, source.get() ]) { + SCOPED_TRACE(directory.path.UTF8String); + struct stat original; + ASSERT_EQ(stat(directory.fileSystemRepresentation, &original), 0); + ASSERT_EQ(chmod(directory.fileSystemRepresentation, 0000), 0); + auto denied = + prepare_source_bundle(manifest.get(), &data, storage.url.path); + EXPECT_EQ( + chmod(directory.fileSystemRepresentation, original.st_mode & 07777), 0); + EXPECT_EQ(denied.error(), Error::AccessFailed); + EXPECT_TRUE(snapshot_matches(storage.url, snapshot)); + EXPECT_EQ(data.attempts, 1); + EXPECT_TRUE( + prepare_source_bundle(manifest.get(), nullptr, storage.url.path).ok()); + } +} + +TEST_F(CoreAISourceTest, PreservesNamedDataErrorsWithoutRetrying) { + for (Error error : {Error::AccessFailed, Error::MemoryAllocationFailed}) { + SCOPED_TRACE(static_cast(error)); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + data.fail_at = 1; + data.failure_error = error; + auto manifest = selected_bundle(manifest_dict()); + ASSERT_TRUE(manifest.ok()); + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path).error(), + error); + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, 0); + EXPECT_EQ(data.releases, 0); + EXPECT_EQ(pending_directories(storage.url).count, 0u); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:source_entry(storage.url, manifest.get()).path]); + } +} + +TEST_F(CoreAISourceTest, MissingNamedDataIsInvalidExternalData) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + data.files.clear(); + auto manifest = selected_bundle(manifest_dict()); + ASSERT_TRUE(manifest.ok()); + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path).error(), + Error::InvalidExternalData); + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, 0); + EXPECT_EQ(pending_directories(storage.url).count, 0u); +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_source_publication_test.mm b/backends/apple/coreai/runtime/test/coreai_source_publication_test.mm new file mode 100644 index 00000000000..5179f875826 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_publication_test.mm @@ -0,0 +1,295 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include "coreai_cache_entry.h" +#include "coreai_cache_key.h" +#include "coreai_fault_scope.h" +#include "coreai_source_fixture.h" +namespace executorch::backends::coreai::testing { +using runtime::Error; +using runtime::Result; +TEST_F(CoreAISourceTest, RejectsMismatchedBufferSizesAndReleasesPayload) { + for (size_t size : {0u, 10u, 12u}) { + SCOPED_TRACE(size); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + data.files["coreai/ab/model.aimodel/graph.bin"] = std::string(size, 'x'); + auto manifest = selected_bundle(manifest_dict()); + ASSERT_TRUE(manifest.ok()); + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path).error(), + Error::InvalidExternalData); + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, 1); + EXPECT_EQ(data.releases, 1); + EXPECT_EQ(pending_directories(storage.url).count, 0u); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:source_entry(storage.url, manifest.get()).path]); + } +} + +TEST_F(CoreAISourceTest, FullyFlushesSourceBeforeAndAfterPublication) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + aot_data(data, @"arch_b"); + auto manifest = selected_bundle(aot_manifest_dict()); + ASSERT_TRUE(manifest.ok()); + NSURL* final = source_entry(storage.url, manifest.get()); + NSURL* root = sources_root(storage.url); + const NSUInteger file_count = manifest->files.count; + __block int synced_files = 0; + __block int full_syncs = 0; + __block int publications = 0; + __block bool ordered = true; + StorageFaultScope fault(FaultPoint::FullSync); + fault.armed = false; + fault.observe = ^(FaultPoint operation) { + if (operation == FaultPoint::SyncFile) + ++synced_files; + if (operation == FaultPoint::FullSync) { + ++full_syncs; + ordered &= synced_files == file_count; + const bool published = + [NSFileManager.defaultManager fileExistsAtPath:final.path]; + ordered &= full_syncs == 1 ? !published : published; + ordered &= pending_directories(root).count == (full_syncs == 1 ? 1u : 0u); + } + if (operation == FaultPoint::Rename) { + ++publications; + ordered &= full_syncs == 1; + } + }; + auto source = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(source.ok()); + EXPECT_TRUE(ordered); + EXPECT_EQ(full_syncs, 2); + EXPECT_EQ(publications, 1); + EXPECT_EQ(data.attempts, file_count); + EXPECT_EQ(data.requests, data.releases); +} + +TEST_F(CoreAISourceTest, PublicationFailuresCleanPendingWithoutAutomaticRetry) { + for (auto operation : + {FaultPoint::CreatePendingDirectory, + FaultPoint::OpenPendingDirectory, + FaultPoint::Write, + FaultPoint::SyncFile, + FaultPoint::SyncDirectory, + FaultPoint::FullSync, + FaultPoint::Rename}) { + SCOPED_TRACE(static_cast(operation)); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto manifest = selected_bundle(manifest_dict()); + ASSERT_TRUE(manifest.ok()); + NSString* key = source_key(manifest.get()); + auto lock = acquire_test_cache_entry(storage.url.path, key); + ASSERT_TRUE(lock.ok()); + NSURL* final = source_entry(storage.url, manifest.get()); + { + StorageFaultScope fault(operation); + EXPECT_EQ( + executorch::backends::coreai::prepare_source_bundle( + manifest.get(), &data, *lock.get()) + .error(), + Error::AccessFailed); + EXPECT_EQ(fault.matches, 1); + EXPECT_EQ(fault.hits, 1); + } + const int attempts = operation == FaultPoint::CreatePendingDirectory || + operation == FaultPoint::OpenPendingDirectory || + operation == FaultPoint::SyncDirectory + ? 0 + : 1; + EXPECT_EQ(data.attempts, attempts); + EXPECT_EQ(data.requests, data.releases); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:final.path]); + EXPECT_EQ(pending_directories(storage.url).count, 0u); + auto retried = executorch::backends::coreai::prepare_source_bundle( + manifest.get(), &data, *lock.get()); + ASSERT_TRUE(retried.ok()); + EXPECT_EQ(data.attempts, attempts + 1); + } +} + +TEST_F(CoreAISourceTest, RetriesInterruptedPendingAndPublicationCalls) { + for (auto operation : + {FaultPoint::CreatePendingDirectory, + FaultPoint::OpenPendingDirectory, + FaultPoint::SyncDirectory, + FaultPoint::Rename}) { + SCOPED_TRACE(static_cast(operation)); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto manifest = selected_bundle(manifest_dict()); + ASSERT_TRUE(manifest.ok()); + NSString* key = source_key(manifest.get()); + auto lock = acquire_test_cache_entry(storage.url.path, key); + ASSERT_TRUE(lock.ok()); + StorageFaultScope fault(operation); + fault.error = EINTR; + auto source = executorch::backends::coreai::prepare_source_bundle( + manifest.get(), &data, *lock.get()); + ASSERT_TRUE(source.ok()); + EXPECT_EQ(fault.hits, 1); + EXPECT_GE(fault.matches, 2); + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(pending_directories(storage.url).count, 0u); + } +} + +TEST_F(CoreAISourceTest, PostRenameFlushFailureLeavesCompleteSourceVisible) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto manifest = selected_bundle(manifest_dict()); + ASSERT_TRUE(manifest.ok()); + NSURL* final = source_entry(storage.url, manifest.get()); + { + StorageFaultScope fault(FaultPoint::FullSync); + fault.match_index = 2; + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path).error(), + Error::AccessFailed); + EXPECT_EQ(fault.matches, 2); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(pending_directories(storage.url).count, 0u); + NSURL* graph = [final URLByAppendingPathComponent:@"model.aimodel/graph.bin"]; + EXPECT_TRUE(([[NSData dataWithContentsOfURL:graph] + isEqual:[NSData dataWithBytes:"model data" length:11]])); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, snapshot)); + { + StorageFaultScope fault(FaultPoint::FullSync); + EXPECT_EQ( + prepare_source_bundle(manifest.get(), nullptr, storage.url.path) + .error(), + Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + EXPECT_TRUE(snapshot_matches(storage.url, snapshot)); + } + StorageFaultScope flush(FaultPoint::FullSync); + flush.error = EINTR; + auto reused = + prepare_source_bundle(manifest.get(), nullptr, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_EQ(flush.hits, 1); + EXPECT_EQ(flush.matches, 2); + EXPECT_TRUE(snapshot_matches(storage.url, snapshot)); + EXPECT_EQ(data.attempts, 1); +} + +TEST_F(CoreAISourceTest, IncompleteConcurrentWinnerFailsWithoutReplacement) { + TestDirectory storage; + TestDirectory winner_storage; + ASSERT_NE(storage.url, nil); + ASSERT_NE(winner_storage.url, nil); + TestData data; + TestData winner_data; + auto manifest = selected_bundle(manifest_dict()); + ASSERT_TRUE(manifest.ok()); + auto winner = prepare_source_bundle( + manifest.get(), &winner_data, winner_storage.url.path); + ASSERT_TRUE(winner.ok()); + ASSERT_TRUE([NSFileManager.defaultManager + removeItemAtURL:[winner.get() URLByAppendingPathComponent:@"graph.bin"] + error:nil]); + NSURL* winner_entry = winner.get().URLByDeletingLastPathComponent; + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(winner_entry, snapshot)); + NSURL* root = sources_root(storage.url); + NSURL* final = source_entry(storage.url, manifest.get()); + __block int publications = 0; + __block bool callback_ok = true; + { + StorageFaultScope fault(FaultPoint::Rename); + fault.armed = false; + fault.observe = ^(FaultPoint operation) { + if (operation != FaultPoint::Rename) + return; + ++publications; + callback_ok &= pending_directories(root).count == 1; + callback_ok &= [NSFileManager.defaultManager moveItemAtURL:winner_entry + toURL:final + error:nil]; + }; + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path).error(), + Error::InvalidExternalData); + EXPECT_TRUE(callback_ok); + EXPECT_EQ(publications, 1); + } + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(pending_directories(storage.url).count, 0u); + EXPECT_TRUE(snapshot_matches(final, snapshot)); +} + +TEST_F( + CoreAISourceTest, + ExclusivePublicationValidatesWinnerAndRemovesLosingPending) { + TestDirectory storage; + TestDirectory winner_storage; + ASSERT_NE(storage.url, nil); + ASSERT_NE(winner_storage.url, nil); + TestData data; + TestData winner_data; + auto manifest = selected_bundle(manifest_dict()); + ASSERT_TRUE(manifest.ok()); + auto winner = prepare_source_bundle( + manifest.get(), &winner_data, winner_storage.url.path); + ASSERT_TRUE(winner.ok()); + NSURL* winner_entry = winner.get().URLByDeletingLastPathComponent; + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(winner_entry, snapshot)); + NSURL* root = sources_root(storage.url); + NSURL* final = source_entry(storage.url, manifest.get()); + __block int publications = 0; + __block bool callback_ok = true; + { + StorageFaultScope fault(FaultPoint::FullSync); + fault.match_index = 2; + fault.error = EINTR; + fault.observe = ^(FaultPoint operation) { + if (operation != FaultPoint::Rename) + return; + ++publications; + callback_ok &= pending_directories(root).count == 1; + callback_ok &= [NSFileManager.defaultManager moveItemAtURL:winner_entry + toURL:final + error:nil]; + }; + auto result = + prepare_source_bundle(manifest.get(), &data, storage.url.path); + EXPECT_TRUE(callback_ok); + ASSERT_TRUE(result.ok()); + EXPECT_TRUE([result.get() + isEqual:[final URLByAppendingPathComponent:manifest->bundle_name]]); + EXPECT_EQ(publications, 1); + EXPECT_EQ(fault.hits, 1); + EXPECT_EQ(fault.matches, 3); + } + EXPECT_EQ(data.attempts, manifest->files.count); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(pending_directories(storage.url).count, 0u); + EXPECT_TRUE(snapshot_matches(final, snapshot)); +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_source_test.mm b/backends/apple/coreai/runtime/test/coreai_source_test.mm new file mode 100644 index 00000000000..67ec531c6a0 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_test.mm @@ -0,0 +1,428 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include "coreai_cache_entry.h" +#include "coreai_cache_key.h" +#include "coreai_fault_scope.h" +#include "coreai_source_fixture.h" +#include "coreai_test_utils.h" + +namespace executorch::backends::coreai::testing { +using runtime::Error; +using runtime::Result; + +TEST_F(CoreAISourceTest, UnicodePathsMaterializeAndReuse) { + NSString* composed = [NSString stringWithFormat:@"caf%C", 0x00e9]; + for (NSString* name in + @[ composed, composed.decomposedStringWithCanonicalMapping ]) { + SCOPED_TRACE(name.UTF8String); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + data.files.clear(); + NSString* relative = [NSString stringWithFormat:@"%@/%@.bin", name, name]; + NSString* file = [@"model.aimodel" stringByAppendingPathComponent:relative]; + auto dict = manifest_dict(); + dict[@"files"] = @{file : @7}; + data.files[[NSString stringWithFormat:@"coreai/ab/%@", file].UTF8String] = + "unicode"; + auto manifest = selected_bundle(dict); + ASSERT_TRUE(manifest.ok()); + auto source = + prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(source.ok()); + NSData* bytes = [NSData + dataWithContentsOfURL:[source.get() + URLByAppendingPathComponent:relative]]; + EXPECT_TRUE([bytes isEqual:[NSData dataWithBytes:"unicode" length:7]]); + auto reused = + prepare_source_bundle(manifest.get(), nullptr, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:source.get()]); + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, data.releases); + } +} + +TEST_F(CoreAISourceTest, RejectsAliasedPathsBeforeReadingAssets) { + NSString* composed = [NSString stringWithFormat:@"caf%C", 0x00e9]; + NSString* decomposed = composed.decomposedStringWithCanonicalMapping; + for (NSArray* pair in @[ + @[ @"A.bin", @"a.bin" ], + @[ @"Dir/x", @"dir/y" ], + @[ @"Dir", @"dir/x" ], + @[ composed, decomposed ], + @[ + [composed stringByAppendingPathComponent:@"x"], + [decomposed stringByAppendingPathComponent:@"y"] + ] + ]) { + SCOPED_TRACE(pair.description.UTF8String); + TestData data; + auto dictionary = manifest_dict(); + dictionary[@"files"] = @{ + [@"model.aimodel" stringByAppendingPathComponent:pair[0]] : @1, + [@"model.aimodel" stringByAppendingPathComponent:pair[1]] : @1 + }; + EXPECT_EQ( + parse_manifest(encode(dictionary)).error(), Error::InvalidProgram); + EXPECT_EQ(data.attempts, 0); + } +} + +TEST_F(CoreAISourceTest, AotMaterializesOnlySelectedArchitecture) { + auto parsed = parse_manifest(encode(aot_manifest_dict())); + ASSERT_TRUE(parsed.ok()); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + for (NSString* arch in @[ @"arch_a", @"arch_b" ]) { + SCOPED_TRACE(arch.UTF8String); + TestData data; + aot_data(data, arch); + auto selected = select_bundle(parsed.get(), arch, @"macOS"); + ASSERT_TRUE(selected.ok()); + auto first = prepare_source_bundle(selected.get(), &data, storage.url.path); + ASSERT_TRUE(first.ok()); + EXPECT_EQ(data.attempts, 2); + EXPECT_EQ(data.requests, 2); + NSURL* url = first.get(); + EXPECT_TRUE([url.lastPathComponent + isEqualToString:selected->path.lastPathComponent]); + NSData* contents = [NSData + dataWithContentsOfURL:[url URLByAppendingPathComponent:@"graph.bin"]]; + ASSERT_TRUE([contents isEqual:[NSData dataWithBytes:"compiled\0graph" + length:14]]); + + NSString* other_arch = + [arch isEqualToString:@"arch_a"] ? @"arch_b" : @"arch_a"; + NSString* other_key = [NSString + stringWithFormat:@"coreai/ab/model.%@.aimodelc/graph.bin", other_arch]; + data.files[other_key.UTF8String] = "must not affect identity"; + auto reused = + prepare_source_bundle(selected.get(), &data, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:url]); + EXPECT_EQ(data.attempts, 2); + + auto unselected_changed = aot_manifest_dict(); + NSMutableDictionary* digests = + [unselected_changed[@"bundle_digests"] mutableCopy]; + digests[[NSString stringWithFormat:@"model.%@.aimodelc", other_arch]] = + fixture_identity('f'); + unselected_changed[@"bundle_digests"] = digests; + NSMutableDictionary* files = [unselected_changed[@"files"] mutableCopy]; + files[[NSString stringWithFormat:@"model.%@.aimodelc/extra", other_arch]] = + @0; + unselected_changed[@"files"] = files; + auto changed_parsed = parse_manifest(encode(unselected_changed)); + ASSERT_TRUE(changed_parsed.ok()); + auto unchanged = select_bundle(changed_parsed.get(), arch, @"macOS"); + ASSERT_TRUE(unchanged.ok()); + auto unchanged_source = + prepare_source_bundle(unchanged.get(), &data, storage.url.path); + ASSERT_TRUE(unchanged_source.ok()); + EXPECT_TRUE([unchanged_source.get() isEqual:url]); + EXPECT_EQ(data.attempts, 2); + EXPECT_EQ(data.releases, data.requests); + } +} + +void CoreAISourceTest::check_concurrent_source() { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto selected = selected_bundle(manifest_dict()); + ASSERT_TRUE(selected.ok()); + for (int pass = 0; pass < 2; ++pass) { + SCOPED_TRACE(pass == 0 ? "cold" : "warm"); + using FixturePointer = decltype(this); + struct Worker { + FixturePointer fixture = nullptr; + const SelectedBundle* manifest = nullptr; + TestData* data = nullptr; + NSString* root = nil; + dispatch_semaphore_t start = nullptr; + bool started = false; + Error error = Error::Internal; + std::string path; + } workers[4]; + pthread_t threads[4]; + dispatch_semaphore_t start = dispatch_semaphore_create(0); + int launch_error = 0; + size_t launched = 0; + for (; launched < 4; ++launched) { + auto& worker = workers[launched]; + worker.fixture = this; + worker.manifest = &selected.get(); + worker.data = &data; + worker.root = storage.url.path; + worker.start = start; + launch_error = pthread_create( + &threads[launched], + nullptr, + [](void* context) -> void* { + auto& worker = *static_cast(context); + @autoreleasepool { + worker.started = + dispatch_semaphore_wait(worker.start, test_deadline()) == 0; + if (!worker.started) + return nullptr; + auto assets = worker.fixture->prepare_source_bundle( + *worker.manifest, worker.data, worker.root); + worker.error = assets.error(); + if (assets.ok()) + worker.path = assets.get().path.UTF8String; + } + return nullptr; + }, + &worker); + if (launch_error != 0) + break; + } + for (size_t i = 0; i < launched; ++i) + dispatch_semaphore_signal(start); + int join_errors[4] = {}; + for (size_t i = 0; i < launched; ++i) + join_errors[i] = pthread_join(threads[i], nullptr); + ASSERT_EQ(launch_error, 0); + for (size_t i = 0; i < 4; ++i) { + SCOPED_TRACE(i); + ASSERT_EQ(join_errors[i], 0); + ASSERT_TRUE(workers[i].started); + ASSERT_EQ(workers[i].error, Error::Ok); + EXPECT_EQ(workers[i].path, workers[0].path); + } + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, 1); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(data.metadata_requests, 0); + EXPECT_EQ( + [NSFileManager.defaultManager + contentsOfDirectoryAtPath:sources_root(storage.url).path + error:nil] + .count, + 1u); + } +} + +TEST_F(CoreAISourceTest, PerKeyLockSerializesColdAndWarmPreparation) { + check_concurrent_source(); +} + +TEST_F(CoreAISourceTest, MissingSourceRequiresNamedData) { + auto bundle = selected_bundle(manifest_dict()); + ASSERT_TRUE(bundle.ok()); + TestData data; + TestDirectory storage; + ASSERT_NE(storage.url, nil); + EXPECT_EQ( + prepare_source_bundle(bundle.get(), nullptr, storage.url.path).error(), + Error::InvalidProgram); + auto recovered = prepare_source_bundle(bundle.get(), &data, storage.url.path); + ASSERT_TRUE(recovered.ok()); + EXPECT_EQ( + [NSData + dataWithContentsOfURL:[recovered.get() + URLByAppendingPathComponent:@"graph.bin"]] + .length, + 11u); + EXPECT_EQ(data.requests, data.releases); +} + +TEST_F(CoreAISourceTest, PublishedSourceOutlivesResult) { + TestData data; + TestDirectory storage; + ASSERT_NE(storage.url, nil); + auto manifest = selected_bundle(manifest_dict()); + ASSERT_TRUE(manifest.ok()); + NSURL* first_url = nil; + { + auto assets = + prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(assets.ok()); + EXPECT_EQ(data.releases, data.requests); + first_url = assets.get(); + EXPECT_EQ( + [NSData + dataWithContentsOfURL:[first_url + URLByAppendingPathComponent:@"graph.bin"]] + .length, + 11u); + } + EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:first_url.path]); + auto reused = + prepare_source_bundle(manifest.get(), nullptr, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:first_url]); +} + +TEST_F(CoreAISourceTest, PersistentIdentityTracksDigestNotFileOrder) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto dict = manifest_dict(); + dict[@"files"] = @{ + @"model.aimodel/graph.bin" : @11, + @"model.aimodel/nested/weights.bin" : @7 + }; + dict[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity('2')}; + data.files["coreai/ab/model.aimodel/nested/weights.bin"] = "weights"; + auto manifest = selected_bundle(dict); + ASSERT_TRUE(manifest.ok()); + auto first = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(first.ok()); + NSURL* first_url = first.get(); + dict[@"files"] = @{ + @"model.aimodel/nested/weights.bin" : @7, + @"model.aimodel/graph.bin" : @11 + }; + auto reordered_manifest = selected_bundle(dict); + ASSERT_TRUE(reordered_manifest.ok()); + auto reordered = + prepare_source_bundle(reordered_manifest.get(), &data, storage.url.path); + ASSERT_TRUE(reordered.ok()); + EXPECT_TRUE([reordered.get() isEqual:first_url]); + data.files["coreai/ab/model.aimodel/nested/weights.bin"] = "changed"; + dict[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity('3')}; + auto changed_manifest = selected_bundle(dict); + ASSERT_TRUE(changed_manifest.ok()); + auto changed = + prepare_source_bundle(changed_manifest.get(), &data, storage.url.path); + ASSERT_TRUE(changed.ok()); + EXPECT_FALSE([changed.get() isEqual:first_url]); + EXPECT_TRUE( + [[NSData dataWithContentsOfURL:[first_url URLByAppendingPathComponent: + @"nested/weights.bin"]] + isEqual:[@"weights" dataUsingEncoding:NSUTF8StringEncoding]]); + EXPECT_EQ(data.releases, data.requests); +} + +TEST_F(CoreAISourceTest, PartialMaterializationCleansPendingAndCanRetry) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto dict = manifest_dict(); + dict[@"files"] = @{ + @"model.aimodel/graph.bin" : @11, + @"model.aimodel/nested/weights.bin" : @7 + }; + dict[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity('2')}; + data.files["coreai/ab/model.aimodel/nested/weights.bin"] = "weights"; + auto manifest = selected_bundle(dict); + ASSERT_TRUE(manifest.ok()); + data.fail_at = 2; + EXPECT_EQ( + prepare_source_bundle(manifest.get(), &data, storage.url.path).error(), + Error::AccessFailed); + EXPECT_EQ(data.attempts, 2); + EXPECT_EQ( + [NSFileManager.defaultManager + contentsOfDirectoryAtPath:sources_root(storage.url).path + error:nil] + .count, + 0u); + EXPECT_EQ(data.releases, data.requests); + data.fail_at = 0; + EXPECT_TRUE( + prepare_source_bundle(manifest.get(), &data, storage.url.path).ok()); + EXPECT_EQ( + [NSFileManager.defaultManager + contentsOfDirectoryAtPath:sources_root(storage.url).path + error:nil] + .count, + 1u); +} + +TEST_F( + CoreAISourceTest, + NestedEmptyAndLargeFilesMaterializeAndReuseWithoutNamedData) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto dict = manifest_dict(); + constexpr size_t large_size = 2 * 1024 * 1024 + 17; + dict[@"files"] = @{ + @"model.aimodel/graph.bin" : @11, + @"model.aimodel/nested/empty" : @0, + @"model.aimodel/nested/deeper/large.bin" : @(large_size) + }; + data.files["coreai/ab/model.aimodel/nested/empty"] = ""; + data.files["coreai/ab/model.aimodel/nested/deeper/large.bin"] = + std::string(large_size, 'x'); + auto manifest = selected_bundle(dict); + ASSERT_TRUE(manifest.ok()); + auto first = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(first.ok()); + EXPECT_EQ(data.attempts, 3); + EXPECT_EQ(data.requests, data.releases); + NSData* empty = [NSData + dataWithContentsOfURL:[first.get() + URLByAppendingPathComponent:@"nested/empty"]]; + ASSERT_NE(empty, nil); + EXPECT_EQ(empty.length, 0u); + NSData* large = [NSData + dataWithContentsOfURL: + [first.get() URLByAppendingPathComponent:@"nested/deeper/large.bin"]]; + ASSERT_NE(large, nil); + EXPECT_EQ(large.length, large_size); + + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, snapshot)); + const int attempts = data.attempts; + auto reused = + prepare_source_bundle(manifest.get(), nullptr, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:first.get()]); + EXPECT_EQ(data.attempts, attempts); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(data.metadata_requests, 0); + EXPECT_TRUE(snapshot_matches(storage.url, snapshot)); +} + +TEST_F( + CoreAISourceTest, + ExplicitSourceKeyReusesSourceWithoutWritesOrRepublication) { + TestDirectory directory; + ASSERT_NE(directory.url, nil); + TestData data; + auto bundle = selected_bundle(manifest_dict()); + ASSERT_TRUE(bundle.ok()); + auto key = cache_variant_key( + bundle.get(), @"macOS", @"arch_b", @"explicit-test-options"); + ASSERT_TRUE(key.ok()); + auto lock = acquire_test_cache_entry(directory.url.path, key.get()); + ASSERT_TRUE(lock.ok()); + auto first = executorch::backends::coreai::prepare_source_bundle( + bundle.get(), &data, *lock.get()); + ASSERT_TRUE(first.ok()); + EXPECT_EQ(data.attempts, 1); + NSString* expected_path = [directory.url.URLByResolvingSymlinksInPath.path + stringByAppendingPathComponent: + [NSString stringWithFormat:@"sources/%@/model.aimodel", key.get()]]; + EXPECT_TRUE([first.get().path isEqualToString:expected_path]); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(directory.url, snapshot)); + const int attempts = data.attempts; + const int requests = data.requests; + data.files.clear(); + auto reused = executorch::backends::coreai::prepare_source_bundle( + bundle.get(), &data, *lock.get()); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:first.get()]); + EXPECT_EQ(data.attempts, attempts); + EXPECT_EQ(data.requests, requests); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(data.metadata_requests, 0); + EXPECT_TRUE(snapshot_matches(directory.url, snapshot)); +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_storage_test.mm b/backends/apple/coreai/runtime/test/coreai_storage_test.mm index 266dfbd34bf..e5f135724f0 100644 --- a/backends/apple/coreai/runtime/test/coreai_storage_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_storage_test.mm @@ -10,6 +10,7 @@ #include #include #include +#include #include #include "coreai_fault_scope.h" #include "coreai_file.h" @@ -35,12 +36,38 @@ FileDescriptor root( ASSERT_EQ(lstat_at(root.get(), "directory", &info), 0); EXPECT_TRUE(S_ISDIR(info.st_mode)); EXPECT_EQ(open_directory_at(root.get(), "link"), -1); + EXPECT_EQ( + open_storage_directory( + root.get(), @"link", StorageCreationMode::OpenExisting) + .error(), + Error::AccessFailed); EXPECT_EQ(open_directory_at(root.get(), "missing"), -1); EXPECT_EQ(errno, ENOENT); EXPECT_EQ(lstat_at(root.get(), "missing", &info), -1); EXPECT_EQ(errno, ENOENT); } +TEST(CoreAIStorageTest, FileDescriptorMovesTransferSingleOwnership) { + FileDescriptor source(open("/dev/null", O_RDONLY | O_CLOEXEC)); + ASSERT_GE(source.get(), 0); + const int source_fd = source.get(); + FileDescriptor moved(std::move(source)); + EXPECT_EQ(source.get(), -1); + EXPECT_EQ(moved.get(), source_fd); + + FileDescriptor destination(open("/dev/null", O_RDONLY | O_CLOEXEC)); + ASSERT_GE(destination.get(), 0); + const int replaced_fd = destination.get(); + destination = std::move(moved); + EXPECT_EQ(moved.get(), -1); + EXPECT_EQ(destination.get(), source_fd); + EXPECT_EQ(fcntl(replaced_fd, F_GETFD), -1); + EXPECT_EQ(errno, EBADF); + + destination = std::move(destination); + EXPECT_EQ(destination.get(), source_fd); +} + TEST(CoreAIStorageTest, OpensDirectoriesAndReportsCreation) { TestDirectory storage; ASSERT_NE(storage.url, nil); @@ -51,16 +78,17 @@ FileDescriptor root(open( auto missing = open_storage_directory( root.get(), @"child", StorageCreationMode::OpenExisting); ASSERT_TRUE(missing.ok()); - EXPECT_EQ(missing.get().fd, -1); - EXPECT_FALSE(missing.get().created); + EXPECT_FALSE(missing->descriptor.has_value()); + EXPECT_FALSE(missing->created); StorageFaultScope sync_fault(FaultPoint::SyncDirectory); auto created = open_storage_directory(root.get(), @"child", StorageCreationMode::Create); ASSERT_TRUE(created.ok()); - FileDescriptor first(created.get().fd); + ASSERT_TRUE(created->descriptor.has_value()); + FileDescriptor first = std::move(*created->descriptor); ASSERT_GE(first.get(), 0); - EXPECT_TRUE(created.get().created); + EXPECT_TRUE(created->created); EXPECT_NE(fcntl(first.get(), F_GETFD) & FD_CLOEXEC, 0); EXPECT_EQ(sync_fault.matches, 0); struct stat original; @@ -72,9 +100,10 @@ FileDescriptor root(open( {StorageCreationMode::OpenExisting, StorageCreationMode::Create}) { auto reopened = open_storage_directory(root.get(), @"child", mode); ASSERT_TRUE(reopened.ok()); - FileDescriptor next(reopened.get().fd); + ASSERT_TRUE(reopened->descriptor.has_value()); + FileDescriptor next = std::move(*reopened->descriptor); ASSERT_GE(next.get(), 0); - EXPECT_FALSE(reopened.get().created); + EXPECT_FALSE(reopened->created); struct stat current; ASSERT_EQ(fstat(next.get(), ¤t), 0); EXPECT_EQ(current.st_ino, original.st_ino); @@ -99,9 +128,10 @@ FileDescriptor root(open( auto opened = open_storage_directory(root.get(), @"child", StorageCreationMode::Create); ASSERT_TRUE(opened.ok()); - FileDescriptor directory(opened.get().fd); + ASSERT_TRUE(opened->descriptor.has_value()); + FileDescriptor directory = std::move(*opened->descriptor); EXPECT_GE(directory.get(), 0); - EXPECT_FALSE(opened.get().created); + EXPECT_FALSE(opened->created); } TEST(CoreAIStorageTest, RejectsInvalidDirectoryComponents) { @@ -398,6 +428,162 @@ FileDescriptor root(open( EXPECT_EQ(lseek(created.get(), 0, SEEK_CUR), 4); } +TEST(CoreAIStorageTest, StrictRemovalPreflightsBeforeMutation) { + TestDirectory storage; + TestDirectory target; + ASSERT_NE(storage.url, nil); + ASSERT_NE(target.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + ASSERT_EQ(mkdirat(root.get(), "entry", 0777), 0); + FileDescriptor entry(open_directory_at(root.get(), "entry")); + ASSERT_GE(entry.get(), 0); + ASSERT_EQ(write_storage_file(entry.get(), @"a/file", "x", 1), Error::Ok); + ASSERT_EQ( + symlinkat(target.url.fileSystemRepresentation, entry.get(), "z-link"), 0); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot( + [storage.url URLByAppendingPathComponent:@"entry"], snapshot)); + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Strict), + Error::InvalidExternalData); + EXPECT_TRUE(snapshot_matches( + [storage.url URLByAppendingPathComponent:@"entry"], snapshot)); +} + +TEST(CoreAIStorageTest, StrictRemovalRevalidatesDuringRemoval) { + TestDirectory storage; + TestDirectory target; + ASSERT_NE(storage.url, nil); + ASSERT_NE(target.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + ASSERT_EQ(mkdirat(root.get(), "entry", 0777), 0); + FileDescriptor entry(open_directory_at(root.get(), "entry")); + ASSERT_GE(entry.get(), 0); + ASSERT_EQ(write_storage_file(entry.get(), @"a", "x", 1), Error::Ok); + ASSERT_EQ(write_storage_file(entry.get(), @"z", "x", 1), Error::Ok); + __block int removals = 0; + const int entry_fd = entry.get(); + NSString* target_path = target.url.path; + StorageFaultScope fault(FaultPoint::Remove); + fault.armed = false; + fault.observe = ^(FaultPoint operation) { + if (operation != FaultPoint::Remove || ++removals != 1) { + return; + } + EXPECT_EQ(unlinkat(entry_fd, "z", 0), 0); + EXPECT_EQ( + symlinkat(target_path.fileSystemRepresentation, entry_fd, "z"), 0); + }; + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Strict), + Error::InvalidExternalData); + EXPECT_EQ(removals, 1); + struct stat info; + EXPECT_EQ(lstat_at(entry.get(), "a", &info), -1); + ASSERT_EQ(lstat_at(entry.get(), "z", &info), 0); + EXPECT_TRUE(S_ISLNK(info.st_mode)); + EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:target.url.path]); +} + +TEST(CoreAIStorageTest, MaintenanceRemovalDoesNotFollowMalformedEntries) { + TestDirectory storage; + TestDirectory target; + ASSERT_NE(storage.url, nil); + ASSERT_NE(target.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + ASSERT_EQ(mkdirat(root.get(), "entry", 0777), 0); + FileDescriptor entry(open_directory_at(root.get(), "entry")); + ASSERT_GE(entry.get(), 0); + ASSERT_EQ(write_storage_file(entry.get(), @"nested/file", "x", 1), Error::Ok); + ASSERT_EQ( + symlinkat(target.url.fileSystemRepresentation, entry.get(), "link"), 0); + ASSERT_EQ(mkfifoat(entry.get(), "fifo", 0600), 0); + NSDictionary* target_snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(target.url, target_snapshot)); + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Maintenance), + Error::Ok); + struct stat info; + EXPECT_EQ(lstat_at(root.get(), "entry", &info), -1); + EXPECT_EQ(errno, ENOENT); + EXPECT_TRUE(snapshot_matches(target.url, target_snapshot)); + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Maintenance), + Error::Ok); +} + +TEST(CoreAIStorageTest, EntriesMissingDuringRemovalAreAlreadyRemoved) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + ASSERT_EQ(mkdirat(root.get(), "entry", 0777), 0); + FileDescriptor entry(open_directory_at(root.get(), "entry")); + ASSERT_GE(entry.get(), 0); + ASSERT_EQ(write_storage_file(entry.get(), @"a", "x", 1), Error::Ok); + ASSERT_EQ(write_storage_file(entry.get(), @"z", "x", 1), Error::Ok); + const int entry_fd = entry.get(); + __block int removals = 0; + StorageFaultScope fault(FaultPoint::Remove); + fault.armed = false; + fault.observe = ^(FaultPoint operation) { + if (operation == FaultPoint::Remove && ++removals == 1) { + EXPECT_EQ(unlinkat(entry_fd, "z", 0), 0); + } + }; + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Strict), + Error::Ok); + EXPECT_EQ(removals, 2); + struct stat info; + EXPECT_EQ(lstat_at(root.get(), "entry", &info), -1); + EXPECT_EQ(errno, ENOENT); +} + +TEST(CoreAIStorageTest, RemovalFailuresCanBeRetried) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + FileDescriptor root(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + ASSERT_EQ(mkdirat(root.get(), "entry", 0777), 0); + FileDescriptor entry(open_directory_at(root.get(), "entry")); + ASSERT_GE(entry.get(), 0); + ASSERT_EQ(write_storage_file(entry.get(), @"nested/file", "x", 1), Error::Ok); + { + StorageFaultScope fault(FaultPoint::Remove); + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Strict), + Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_EQ( + remove_storage_directory( + root.get(), @"entry", StorageRemovalPolicy::Strict), + Error::Ok); + struct stat info; + EXPECT_EQ(lstat_at(root.get(), "entry", &info), -1); + EXPECT_EQ(errno, ENOENT); +} + TEST(CoreAIStorageTest, RetriesEintr) { int calls = 0; const ssize_t count = retry_eintr([&]() -> ssize_t { diff --git a/backends/apple/coreai/runtime/test/coreai_test_utils.h b/backends/apple/coreai/runtime/test/coreai_test_utils.h new file mode 100644 index 00000000000..88c4bf35ec0 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_test_utils.h @@ -0,0 +1,20 @@ +/* + * 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 +#include + +namespace executorch::backends::coreai::testing { + +inline dispatch_time_t test_deadline(uint64_t seconds = 10) { + return dispatch_time(DISPATCH_TIME_NOW, seconds * NSEC_PER_SEC); +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/export_bridge_fixture.py b/backends/apple/coreai/runtime/test/export_bridge_fixture.py new file mode 100644 index 00000000000..2d92aa8025d --- /dev/null +++ b/backends/apple/coreai/runtime/test/export_bridge_fixture.py @@ -0,0 +1,125 @@ +# 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 argparse +import json +import shutil +import uuid +from pathlib import Path + +import torch + +from executorch.backends.apple.coreai import ( + CoreAIPartitioner, + get_default_compile_config, + get_default_passes, +) +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.lowered_backend_module import get_lowered_backend_modules + + +class Arithmetic(torch.nn.Module): + """Distinct outputs make reordered bindings observable.""" + + def forward(self, left, right): + """Return the elementwise sum and difference, in that order.""" + return left + right, left - right + + +class WeightedArithmetic(torch.nn.Module): + """External weights expose any dependency on removed source files.""" + + def __init__(self, first: float, second: float): + """Build a weighted model with distinguishable output offsets.""" + super().__init__() + self.first = first + self.second = second + self.network = torch.nn.Sequential( + torch.nn.Linear(2, 4096), torch.nn.ReLU(), torch.nn.Linear(4096, 2) + ) + + def forward(self, left, right): + """Apply the shared network to two distinct input combinations.""" + return ( + self.network(left + right) + self.first, + self.network(left - right) + self.second, + ) + + +def _export_model( + model, inputs, output_dir: Path, *, dynamic_shapes=None, include_pte=False +) -> dict: + exported = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes) + lowered = to_edge_transform_and_lower( + exported, + transform_passes=get_default_passes(), + partitioner=[CoreAIPartitioner()], + compile_config=get_default_compile_config(), + ) + delegates = get_lowered_backend_modules(lowered.exported_program().graph_module) + assert len(delegates) == 1, "Expected the fixture to lower to one Core AI delegate" + manifest_bytes = bytes(delegates[0].processed_bytes) + manifest = json.loads(manifest_bytes) + store = delegates[0].named_data_store_output + assert store is not None + if output_dir.exists(): + shutil.rmtree(output_dir) + for relative_path in manifest["files"]: + entry = store.pte_data[f"coreai/{manifest['export_hash']}/{relative_path}"] + destination = output_dir / relative_path + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(store.buffers[entry.buffer_index]) + (output_dir / "manifest.json").write_bytes(manifest_bytes) + if include_pte: + (output_dir / "model.pte").write_bytes(lowered.to_executorch().buffer) + return manifest + + +def export_fixture( + output_dir: Path, *, dynamic: bool = False, include_pte: bool = False +) -> None: + """Export into a dedicated fixture directory, replacing its contents.""" + size = torch.export.Dim("size", min=1, max=8) + _export_model( + Arithmetic().eval(), + (torch.ones(2), torch.ones(2)), + output_dir, + dynamic_shapes=({0: size}, {0: size}) if dynamic else None, + include_pte=include_pte, + ) + + +def export_weighted_fixture(output_dir: Path) -> None: + """Export isolated SDK cache contents and their expected output values.""" + # Unique model contents keep cleanup from evicting another test's SDK entry. + seed = uuid.uuid4().int + model = WeightedArithmetic( + (seed & 0xFFFFFF) / 2**24, ((seed >> 24) & 0xFFFFFF) / 2**24 + ).eval() + inputs = (torch.tensor([5.0, 9.0]), torch.tensor([1.0, 2.0])) + manifest = _export_model(model, inputs, output_dir) + with torch.no_grad(): + expected = [tensor.tolist() for tensor in model(*inputs)] + (output_dir / "fixture.json").write_text( + json.dumps( + { + "function": manifest["function"], + "inputs": manifest["input_names"], + "outputs": manifest["output_names"], + "expected": expected, + } + ) + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "output_dir", type=Path, help="Dedicated fixture directory, replaced on export" + ) + output_dir = parser.parse_args().output_dir + export_fixture(output_dir, include_pte=True) + export_fixture(output_dir / "dynamic", dynamic=True) diff --git a/backends/apple/coreai/runtime/test/run_sdk_source_lifetime.py b/backends/apple/coreai/runtime/test/run_sdk_source_lifetime.py new file mode 100644 index 00000000000..dc6c044b3aa --- /dev/null +++ b/backends/apple/coreai/runtime/test/run_sdk_source_lifetime.py @@ -0,0 +1,114 @@ +# 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 argparse +import shutil +import subprocess +import tempfile +from pathlib import Path + + +def run(command, expected_status=0): + """Run a probe command, asserting its exit status.""" + result = subprocess.run(command, capture_output=True, text=True, timeout=60) + if result.returncode != expected_status: + raise RuntimeError( + f"{command} returned {result.returncode}, expected {expected_status}\n" + f"{result.stdout}\n{result.stderr}" + ) + return result.stdout.strip() + + +def has_device_architecture(architecture: str) -> bool: + """Report whether the SDK identified a device architecture.""" + normalized = architecture.strip().lower() + return bool(normalized and normalized != "unknown") + + +def _export_weighted_fixture(directory: Path) -> None: + from export_bridge_fixture import export_weighted_fixture + + export_weighted_fixture(directory) + + +def compile_source(source: Path, compiled: Path, architecture: str) -> None: + """Compile a portable source to an AOT model.""" + run( + [ + "xcrun", + "coreai-build", + "compile", + str(source), + "--output", + str(compiled), + "--platform", + "macOS", + "--architecture", + architecture, + "--min-deployment-version", + "27.0", + ] + ) + shutil.rmtree(source) + + +def run_variant( + probe: Path, kind: str, source: Path, fixture: Path, bookmark: Path +) -> None: + """Exercise specialize, restore, and evict for one source variant.""" + arguments = [str(source), str(fixture), str(bookmark)] + try: + print(run([str(probe), "specialize", *arguments]), flush=True) + if source.exists() or not bookmark.is_file(): + raise AssertionError("Expected a saved bookmark and no original source") + for _ in range(2): + print(run([str(probe), "restore", *arguments]), flush=True) + except BaseException: + try: + if bookmark.is_file(): + run([str(probe), "evict", *arguments]) + except Exception as cleanup: + print(f"cleanup evict failed: {cleanup}", flush=True) + raise + if bookmark.is_file(): + run([str(probe), "evict", *arguments]) + run([str(probe), "restore", *arguments], expected_status=3) + print(f"{kind}: persistent cache lifetime and eviction control passed", flush=True) + + +def main(probe: Path) -> None: + """Run portable and compiled lifetime probes against a test binary.""" + architecture = run([str(probe), "architecture"]) + with tempfile.TemporaryDirectory(prefix="coreai-lifetime-") as temporary: + root = Path(temporary) + for kind in ("portable", "compiled"): + # AOT compile rejects unknown architectures. + if kind == "compiled" and not has_device_architecture(architecture): + print( + f"compiled: skipped (unsupported architecture {architecture!r})", + flush=True, + ) + continue + directory = root / kind + _export_weighted_fixture(directory) + source = directory / "model.aimodel" + if kind == "compiled": + compiled = directory / "model.aimodelc" + compile_source(source, compiled, architecture) + source = compiled + run_variant( + probe, + kind, + source, + directory / "fixture.json", + directory / "saved.bookmark", + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("probe", type=Path) + main(parser.parse_args().probe.resolve()) diff --git a/backends/apple/coreai/runtime/test/runtime_smoke.cpp b/backends/apple/coreai/runtime/test/runtime_smoke.cpp new file mode 100644 index 00000000000..a136dceb1df --- /dev/null +++ b/backends/apple/coreai/runtime/test/runtime_smoke.cpp @@ -0,0 +1,48 @@ +/* + * 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. + */ + +#include +#include +#include + +#include + +namespace { +using executorch::runtime::DataLoader; +using executorch::runtime::Error; +using RootClear = Error (*)(const char*); +using FileClear = Error (*)(const char*, const char*); +using LoaderClear = Error (*)(DataLoader&, const char*); + +// Keep every overload in the link check without touching cache storage. +RootClear volatile clear_root = &executorch::backends::coreai::clear_cache; +FileClear volatile clear_file = + &executorch::backends::coreai::clear_cache_for_pte; +LoaderClear volatile clear_loader = + &executorch::backends::coreai::clear_cache_for_pte; +} // namespace + +int main() { + executorch::runtime::runtime_init(); + if (clear_root == nullptr || clear_file == nullptr || + clear_loader == nullptr) { + std::fputs("Core AI cache APIs were not linked\n", stderr); + return 1; + } + const auto* backend = executorch::runtime::get_backend_class("CoreAIBackend"); + if (backend == nullptr) { + std::fputs("CoreAIBackend was not retained by the linker\n", stderr); + return 1; + } + if (!backend->is_available()) { + std::fputs("CoreAIBackend requires OS 27 or newer\n", stderr); + return 1; + } + std::puts("CoreAIBackend is registered and available"); + return 0; +} diff --git a/backends/apple/coreai/runtime/test/test_export_bridge_fixture.py b/backends/apple/coreai/runtime/test/test_export_bridge_fixture.py new file mode 100644 index 00000000000..c1f4b34a812 --- /dev/null +++ b/backends/apple/coreai/runtime/test/test_export_bridge_fixture.py @@ -0,0 +1,72 @@ +# 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 json +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from executorch.backends.apple.coreai.runtime.test import ( + export_bridge_fixture as fixture, +) + + +@pytest.mark.parametrize("dynamic", [False, True]) +@pytest.mark.parametrize("include_pte", [False, True]) +def test_export_replaces_stale_fixture_files(tmp_path, dynamic, include_pte): + """Re-exports must not retain stale assets or freeze dynamic extents.""" + output = tmp_path / "fixture" + bundle = output / "model.aimodel" + bundle.mkdir(parents=True) + (bundle / "obsolete.bin").write_bytes(b"obsolete") + (output / "stale.json").write_text("{}") + manifest = {"export_hash": "ab", "files": {"model.aimodel/graph.bin": 3}} + manifest_bytes = json.dumps(manifest).encode() + delegate = SimpleNamespace( + processed_bytes=manifest_bytes, + named_data_store_output=SimpleNamespace( + pte_data={ + "coreai/ab/model.aimodel/graph.bin": SimpleNamespace(buffer_index=0) + }, + buffers=[b"abc"], + ), + ) + with ( + patch.object(fixture, "to_edge_transform_and_lower") as lower, + patch.object(fixture, "get_lowered_backend_modules", return_value=[delegate]), + patch.object(fixture, "get_default_passes", return_value=[]), + patch.object(fixture, "get_default_compile_config"), + patch.object(fixture, "CoreAIPartitioner"), + patch.object( + fixture.torch.export, "export", wraps=fixture.torch.export.export + ) as export, + ): + lower.return_value.to_executorch.return_value.buffer = b"serialized-pte" + fixture.export_fixture(output, dynamic=dynamic, include_pte=include_pte) + if include_pte: + lower.return_value.to_executorch.assert_called_once_with() + assert (output / "model.pte").read_bytes() == b"serialized-pte" + else: + lower.return_value.to_executorch.assert_not_called() + + assert (output / "manifest.json").read_bytes() == manifest_bytes + assert (bundle / "graph.bin").read_bytes() == b"abc" + assert sorted( + str(p.relative_to(output)) for p in output.rglob("*") if p.is_file() + ) == [ + "manifest.json", + "model.aimodel/graph.bin", + ] + ( + ["model.pte"] if include_pte else [] + ) + shapes = export.call_args.kwargs["dynamic_shapes"] + if dynamic: + assert shapes[0][0] is shapes[1][0] + assert shapes[0][0].min == 1 + assert shapes[0][0].max == 8 + else: + assert shapes is None diff --git a/backends/apple/coreai/runtime/test/test_run_sdk_source_lifetime.py b/backends/apple/coreai/runtime/test/test_run_sdk_source_lifetime.py new file mode 100644 index 00000000000..548ffe1f9c6 --- /dev/null +++ b/backends/apple/coreai/runtime/test/test_run_sdk_source_lifetime.py @@ -0,0 +1,97 @@ +# 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. + +from pathlib import Path +from unittest.mock import patch + +import pytest + +from executorch.backends.apple.coreai.runtime.test import ( + run_sdk_source_lifetime as driver, +) + + +@pytest.mark.parametrize("architecture", ["unknown", "Unknown", "", " "]) +def test_compiled_variant_skipped_without_supported_architecture( + tmp_path, capsys, architecture +): + """AOT compile rejects unknown architectures, so only portable runs.""" + calls = [] + with ( + patch.object(driver, "run", return_value=architecture), + patch.object( + driver, + "_export_weighted_fixture", + side_effect=lambda d: d.mkdir(parents=True), + ), + patch.object( + driver, "compile_source", side_effect=lambda *args: calls.append("compile") + ), + patch.object( + driver, + "run_variant", + side_effect=lambda *args: calls.append(args[1]), + ), + patch.object(driver.tempfile, "TemporaryDirectory") as temporary, + ): + temporary.return_value.__enter__.return_value = str(tmp_path) + driver.main(Path("probe")) + + assert calls == ["portable"] + assert "compiled: skipped" in capsys.readouterr().out + + +def test_compiled_variant_runs_with_known_architecture(tmp_path): + """Physical known-arch Macs still exercise the AOT path.""" + calls = [] + with ( + patch.object(driver, "run", return_value="h16c"), + patch.object( + driver, + "_export_weighted_fixture", + side_effect=lambda d: d.mkdir(parents=True), + ), + patch.object( + driver, "compile_source", side_effect=lambda *args: calls.append("compile") + ), + patch.object( + driver, + "run_variant", + side_effect=lambda *args: calls.append(args[1]), + ), + patch.object(driver.tempfile, "TemporaryDirectory") as temporary, + ): + temporary.return_value.__enter__.return_value = str(tmp_path) + driver.main(Path("probe")) + + assert calls == ["portable", "compile", "compiled"] + + +@pytest.mark.parametrize( + "failure", [RuntimeError("probe failed"), KeyboardInterrupt("interrupted")] +) +def test_cleanup_failure_preserves_original_error(tmp_path, capsys, failure): + """A failing evict must not mask the probe failure it cleans up after.""" + + def fake_run(command, expected_status=0): + if "specialize" in command: + raise failure + raise RuntimeError("evict failed") + + bookmark = tmp_path / "saved.bookmark" + bookmark.write_bytes(b"bookmark") + with patch.object(driver, "run", side_effect=fake_run): + with pytest.raises(type(failure)) as raised: + driver.run_variant( + Path("probe"), + "portable", + tmp_path / "model.aimodel", + tmp_path / "fixture.json", + bookmark, + ) + + assert raised.value is failure + assert "cleanup evict failed: evict failed" in capsys.readouterr().out diff --git a/backends/apple/coreai/test/test_preprocess.py b/backends/apple/coreai/test/test_preprocess.py index 6570dac278b..5e5bd4b162f 100644 --- a/backends/apple/coreai/test/test_preprocess.py +++ b/backends/apple/coreai/test/test_preprocess.py @@ -144,7 +144,7 @@ def test_inline_embeds_files_in_nds(self): out = result.data_store_output self.assertIsNotNone(out) for rel in manifest["files"]: - self.assertIn(f"coreai/{manifest['hash']}/{rel}", out.pte_data) + self.assertIn(f"coreai/{manifest['export_hash']}/{rel}", out.pte_data) self.assertEqual(out.external_data, {}) # nothing external for inline def test_inline_keys_keep_the_bundle_directory(self): @@ -188,7 +188,7 @@ def test_sizes_and_digests_cover_every_delivered_file(self, _build): result = CoreAIBackend.preprocess(_edge_program(), specs) manifest = json.loads(result.processed_bytes) out = result.data_store_output - prefix = f"coreai/{manifest['hash']}/" + prefix = f"coreai/{manifest['export_hash']}/" embedded = { key[len(prefix) :]: len(out.buffers[entry.buffer_index]) for key, entry in out.pte_data.items() @@ -324,7 +324,7 @@ def test_archs_map_points_at_embedded_bundles(self, _build): result = CoreAIBackend.preprocess(_edge_program(), [_aot_spec(config)]) manifest = json.loads(result.processed_bytes) for arch, rel in manifest["archs"].items(): - self.assertEqual(rel, f"{manifest['hash']}/model.{arch}.aimodelc") + self.assertEqual(rel, f"{manifest['export_hash']}/model.{arch}.aimodelc") self.assertIn(f"coreai/{rel}/model.mil", result.data_store_output.pte_data) @mock.patch(_MOCK_BUILD, side_effect=_fake_run_coreai_build) @@ -348,7 +348,7 @@ def test_two_delegates_get_separate_asset_keys(self, _build): ) lbms = get_lowered_backend_modules(lowered.exported_program().graph_module) self.assertEqual(len(lbms), 2) - hashes = {json.loads(bytes(lbm.processed_bytes))["hash"] for lbm in lbms} + hashes = {json.loads(bytes(lbm.processed_bytes))["export_hash"] for lbm in lbms} self.assertEqual(len(hashes), 2) diff --git a/tools/cmake/executorch-config.cmake b/tools/cmake/executorch-config.cmake index 3d3b650f33e..6e468b1da0c 100644 --- a/tools/cmake/executorch-config.cmake +++ b/tools/cmake/executorch-config.cmake @@ -105,6 +105,7 @@ set(optional_lib_list coreml_util coreml_inmemoryfs coremldelegate + coreaidelegate mlxdelegate mlx metal_backend diff --git a/tools/cmake/preset/default.cmake b/tools/cmake/preset/default.cmake index 16e2b574d0e..cf3a335dd84 100644 --- a/tools/cmake/preset/default.cmake +++ b/tools/cmake/preset/default.cmake @@ -24,7 +24,7 @@ define_overridable_option( EXECUTORCH_BUILD_COREML "Build the Core ML backend" BOOL OFF ) define_overridable_option( - EXECUTORCH_BUILD_COREAI "Build the Core AI backend" BOOL OFF + EXECUTORCH_BUILD_COREAI "Build the Core AI backend (Apple SDK 27+)" BOOL OFF ) define_overridable_option( EXECUTORCH_FLATBUFFERS_MAX_ALIGNMENT @@ -440,6 +440,10 @@ check_required_options_on( EXECUTORCH_BUILD_EXTENSION_LLM ) +check_required_options_on( + IF_ON EXECUTORCH_BUILD_COREAI REQUIRES EXECUTORCH_BUILD_EXTENSION_DATA_LOADER +) + check_required_options_on( IF_ON EXECUTORCH_BUILD_EXTENSION_MODULE REQUIRES EXECUTORCH_BUILD_EXTENSION_DATA_LOADER EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR