diff --git a/.github/workflows/coreai.yml b/.github/workflows/coreai.yml index f09b5972d05..33367a13819 100644 --- a/.github/workflows/coreai.yml +++ b/.github/workflows/coreai.yml @@ -10,7 +10,11 @@ on: pull_request: paths: - .github/workflows/coreai.yml + - .clang-format + - CMakeLists.txt - backends/apple/coreai/** + - runtime/** + - tools/cmake/** workflow_dispatch: concurrency: @@ -59,10 +63,36 @@ jobs: echo "::group::Install ExecuTorch and Core AI requirements" ${CONDA_RUN} python install_executorch.py --editable > /dev/null - ${CONDA_RUN} pip install -r backends/apple/coreai/requirements.txt + ${CONDA_RUN} pip install -r backends/apple/coreai/requirements.txt "cmake>=3.31,<4" ninja clang-format==18.1.3 ${CONDA_RUN} pip list echo "::endgroup::" + echo "::group::Check Core AI formatting" + ${CONDA_RUN} python - <<'PY' + from pathlib import Path + import subprocess + import yaml + + root = yaml.safe_load(Path(".clang-format").read_text()) + coreai = yaml.safe_load(Path("backends/apple/coreai/.clang-format").read_text()) + root.pop("Language") + assert coreai == root, "Sync the Core AI clang-format configuration with the root style (without Language)" + sources = sorted(Path("backends/apple/coreai").rglob("*.mm")) + subprocess.run(["clang-format", "--dry-run", "--Werror", *map(str, sources)], check=True) + PY + echo "::endgroup::" + echo "::group::Run Core AI export tests" ${CONDA_RUN} python -m pytest backends/apple/coreai echo "::endgroup::" + + echo "::group::Build and run Core AI runtime tests" + ${CONDA_RUN} cmake -S . -B cmake-out-coreai -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 + ${CONDA_RUN} cmake --build cmake-out-coreai --target backends/apple/coreai/all -j$(sysctl -n hw.ncpu) + ${CONDA_RUN} ctest --test-dir cmake-out-coreai/backends/apple/coreai --output-on-failure --no-tests=error + echo "::endgroup::" diff --git a/CMakeLists.txt b/CMakeLists.txt index 7ae92a457b5..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 @@ -843,6 +852,11 @@ if(EXECUTORCH_BUILD_COREML) list(APPEND _executorch_backends coremldelegate) endif() +if(EXECUTORCH_BUILD_COREAI) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/apple/coreai) + list(APPEND _executorch_backends coreaidelegate) +endif() + if(EXECUTORCH_BUILD_MLX) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/mlx) list(APPEND _executorch_backends mlxdelegate) diff --git a/backends/apple/coreai/.clang-format b/backends/apple/coreai/.clang-format new file mode 100644 index 00000000000..8bd9536477b --- /dev/null +++ b/backends/apple/coreai/.clang-format @@ -0,0 +1,244 @@ +# Root style without its Cpp-only Language key, so Objective-C headers format too. +--- +AccessModifierOffset: -1 +AlignAfterOpenBracket: AlwaysBreak +AlignArrayOfStructures: None +AlignConsecutiveAssignments: + Enabled: false + AcrossEmptyLines: false + AcrossComments: false + AlignCompound: false + AlignFunctionPointers: false + PadOperators: true +AlignConsecutiveBitFields: + Enabled: false + AcrossEmptyLines: false + AcrossComments: false + AlignCompound: false + AlignFunctionPointers: false + PadOperators: true +AlignConsecutiveDeclarations: + Enabled: false + AcrossEmptyLines: false + AcrossComments: false + AlignCompound: false + AlignFunctionPointers: false + PadOperators: true +AlignConsecutiveMacros: + Enabled: false + AcrossEmptyLines: false + AcrossComments: false + AlignCompound: false + AlignFunctionPointers: false + PadOperators: true +AlignConsecutiveShortCaseStatements: + Enabled: false + AcrossEmptyLines: false + AcrossComments: false + AlignCaseColons: false +AlignEscapedNewlines: Left +AlignOperands: DontAlign +AlignTrailingComments: + Kind: Never + OverEmptyLines: 0 +AllowAllArgumentsOnNextLine: true +AllowAllParametersOfDeclarationOnNextLine: false +AllowBreakBeforeNoexceptSpecifier: Never +AllowShortBlocksOnASingleLine: Never +AllowShortCaseLabelsOnASingleLine: false +AllowShortCompoundRequirementOnASingleLine: true +AllowShortEnumsOnASingleLine: true +AllowShortFunctionsOnASingleLine: Empty +AllowShortIfStatementsOnASingleLine: Never +AllowShortLambdasOnASingleLine: All +AllowShortLoopsOnASingleLine: false +AlwaysBreakAfterDefinitionReturnType: None +AlwaysBreakAfterReturnType: None +AlwaysBreakBeforeMultilineStrings: true +AlwaysBreakTemplateDeclarations: Yes +AttributeMacros: + - __capability +BinPackArguments: false +BinPackParameters: false +BitFieldColonSpacing: Both +BraceWrapping: + AfterCaseLabel: false + AfterClass: false + AfterControlStatement: Never + AfterEnum: false + AfterExternBlock: false + AfterFunction: false + AfterNamespace: false + AfterObjCDeclaration: false + AfterStruct: false + AfterUnion: false + BeforeCatch: false + BeforeElse: false + BeforeLambdaBody: false + BeforeWhile: false + IndentBraces: false + SplitEmptyFunction: true + SplitEmptyRecord: true + SplitEmptyNamespace: true +BreakAdjacentStringLiterals: true +BreakAfterAttributes: Leave +BreakAfterJavaFieldAnnotations: false +BreakArrays: true +BreakBeforeBinaryOperators: None +BreakBeforeConceptDeclarations: Always +BreakBeforeBraces: Attach +BreakBeforeInlineASMColon: OnlyMultiline +BreakBeforeTernaryOperators: true +BreakConstructorInitializers: BeforeColon +BreakInheritanceList: BeforeColon +BreakStringLiterals: false +ColumnLimit: 80 +CommentPragmas: '^ IWYU pragma:' +CompactNamespaces: false +ConstructorInitializerIndentWidth: 4 +ContinuationIndentWidth: 4 +Cpp11BracedListStyle: true +DerivePointerAlignment: false +DisableFormat: false +EmptyLineAfterAccessModifier: Never +EmptyLineBeforeAccessModifier: LogicalBlock +ExperimentalAutoDetectBinPacking: false +FixNamespaceComments: true +ForEachMacros: + - FOR_EACH + - FOR_EACH_R + - FOR_EACH_RANGE +IfMacros: + - KJ_IF_MAYBE +IncludeBlocks: Preserve +IncludeCategories: + - Regex: '^<.*\.h(pp)?>' + Priority: 1 + SortPriority: 0 + CaseSensitive: false + - Regex: '^<.*' + Priority: 2 + SortPriority: 0 + CaseSensitive: false + - Regex: '.*' + Priority: 3 + SortPriority: 0 + CaseSensitive: false +IncludeIsMainRegex: '(Test)?$' +IncludeIsMainSourceRegex: '' +IndentAccessModifiers: false +IndentCaseBlocks: false +IndentCaseLabels: true +IndentExternBlock: AfterExternBlock +IndentGotoLabels: true +IndentPPDirectives: None +IndentRequiresClause: true +IndentWidth: 2 +IndentWrappedFunctionNames: false +InsertBraces: false +InsertNewlineAtEOF: false +InsertTrailingCommas: None +IntegerLiteralSeparator: + Binary: 0 + BinaryMinDigits: 0 + Decimal: 0 + DecimalMinDigits: 0 + Hex: 0 + HexMinDigits: 0 +JavaScriptQuotes: Leave +JavaScriptWrapImports: true +KeepEmptyLinesAtTheStartOfBlocks: false +KeepEmptyLinesAtEOF: false +LambdaBodyIndentation: Signature +LineEnding: DeriveLF +MacroBlockBegin: '' +MacroBlockEnd: '' +MaxEmptyLinesToKeep: 1 +NamespaceIndentation: None +ObjCBinPackProtocolList: Auto +ObjCBlockIndentWidth: 2 +ObjCBreakBeforeNestedBlockParam: true +ObjCSpaceAfterProperty: false +ObjCSpaceBeforeProtocolList: false +PackConstructorInitializers: NextLine +PenaltyBreakAssignment: 2 +PenaltyBreakBeforeFirstCallParameter: 1 +PenaltyBreakComment: 300 +PenaltyBreakFirstLessLess: 120 +PenaltyBreakOpenParenthesis: 0 +PenaltyBreakScopeResolution: 500 +PenaltyBreakString: 1000 +PenaltyBreakTemplateDeclaration: 10 +PenaltyExcessCharacter: 1000000 +PenaltyIndentedWhitespace: 0 +PenaltyReturnTypeOnItsOwnLine: 200 +PointerAlignment: Left +PPIndentWidth: -1 +QualifierAlignment: Leave +ReferenceAlignment: Pointer +ReflowComments: true +RemoveBracesLLVM: false +RemoveParentheses: Leave +RemoveSemicolon: false +RequiresClausePosition: OwnLine +RequiresExpressionIndentation: OuterScope +SeparateDefinitionBlocks: Leave +ShortNamespaceLines: 1 +SkipMacroDefinitionBody: false +SortIncludes: CaseSensitive +SortJavaStaticImport: Before +SortUsingDeclarations: LexicographicNumeric +SpaceAfterCStyleCast: false +SpaceAfterLogicalNot: false +SpaceAfterTemplateKeyword: true +SpaceAroundPointerQualifiers: Default +SpaceBeforeAssignmentOperators: true +SpaceBeforeCaseColon: false +SpaceBeforeCpp11BracedList: false +SpaceBeforeCtorInitializerColon: true +SpaceBeforeInheritanceColon: true +SpaceBeforeJsonColon: false +SpaceBeforeParens: ControlStatements +SpaceBeforeParensOptions: + AfterControlStatements: true + AfterForeachMacros: true + AfterFunctionDefinitionName: false + AfterFunctionDeclarationName: false + AfterIfMacros: true + AfterOverloadedOperator: false + AfterPlacementOperator: true + AfterRequiresInClause: false + AfterRequiresInExpression: false + BeforeNonEmptyParentheses: false +SpaceBeforeRangeBasedForLoopColon: true +SpaceBeforeSquareBrackets: false +SpaceInEmptyBlock: false +SpacesBeforeTrailingComments: 1 +SpacesInAngles: Never +SpacesInContainerLiterals: true +SpacesInLineCommentPrefix: + Minimum: 1 + Maximum: -1 +SpacesInParens: Never +SpacesInParensOptions: + InCStyleCasts: false + InConditionalStatements: false + InEmptyParentheses: false + Other: false +SpacesInSquareBrackets: false +Standard: Latest +StatementAttributeLikeMacros: + - Q_EMIT +StatementMacros: + - Q_UNUSED + - QT_REQUIRE_VERSION +TabWidth: 8 +UseTab: Never +VerilogBreakBetweenInstancePorts: true +WhitespaceSensitiveMacros: + - BOOST_PP_STRINGIZE + - CF_SWIFT_NAME + - NS_SWIFT_NAME + - PP_STRINGIZE + - STRINGIZE +... diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt new file mode 100644 index 00000000000..f3c269aa5be --- /dev/null +++ b/backends/apple/coreai/CMakeLists.txt @@ -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. + +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_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 + ) + set_target_properties( + ${target} PROPERTIES OBJCXX_STANDARD 17 OBJCXX_STANDARD_REQUIRED YES + ) + target_compile_options( + ${target} + PRIVATE "$<$:-fobjc-arc;-fno-exceptions;-fno-rtti>" + ) +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 + 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_FAULT_INJECTION=1) + find_library(COREAI_FOUNDATION_FRAMEWORK Foundation REQUIRED) + target_link_libraries( + 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 + "${PYTHON_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 2e6ce4d61b4..e76c209991e 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -23,3 +23,435 @@ python -m pytest backends/apple/coreai `CoreAIAOTCompileTest` needs `xcrun coreai-build` from the Metal Toolchain (`xcodebuild -downloadComponent MetalToolchain`) and skips without it. + +## Runtime + +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. `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 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 +optional `architectures` list. For example, `AOTCompileConfig(platform="iOS", +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`. + +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. + +### Cache directory + +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. + +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. + +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, 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: + +```bash +cmake -S . -B -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 --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 new file mode 100644 index 00000000000..f2c3f696202 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache_root.h @@ -0,0 +1,39 @@ +/* + * 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 "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_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 new file mode 100644 index 00000000000..abc039bed0f --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache_root.mm @@ -0,0 +1,194 @@ +/* + * 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_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(NSString* path) { + NSURL* directory = [NSURL fileURLWithPath:path isDirectory:YES]; + NSNumber* excluded = nil; + NSError* error = nil; + ET_CHECK_OR_RETURN_ERROR( + [directory 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( + [directory setResourceValue:@YES + forKey:NSURLIsExcludedFromBackupKey + error:&error], + AccessFailed, + "Cannot exclude Core AI cache directory from backup: %s", + error.localizedDescription.UTF8String); + return Error::Ok; +} + +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; +} + +// 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( + 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; +} + +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( + is_cache_directory_path(path), + InvalidArgument, + "Core AI cache directory requires an absolute path"); + struct stat info; + if (retry_eintr([&] { return stat(path.fileSystemRepresentation, &info); }) != + 0) { + if (errno == ENOTDIR) { + return not_a_directory(mode); + } + ET_CHECK_OR_RETURN_ERROR( + errno == ENOENT, + AccessFailed, + "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(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 new file mode 100644 index 00000000000..28c8cd0fbbb --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_fault_injection.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 +#include +#include "coreai_file.h" + +namespace executorch::backends::coreai { + +enum class FaultPoint { + Write, + SyncFile, + SyncDirectory, + FullSync, + Rename, + Remove, + CreatePendingDirectory, + OpenPendingDirectory, + BeforeCreateDirectory, + BeforeCacheEntryLock, + BeforeSDK, + AfterSDK, + BeforeEvict, + AfterEvict, +}; + +#if defined(COREAI_FAULT_INJECTION) && COREAI_FAULT_INJECTION +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 + +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_fault_injection.mm b/backends/apple/coreai/runtime/coreai_fault_injection.mm new file mode 100644 index 00000000000..84d20e05cc8 --- /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_FAULT_INJECTION) && COREAI_FAULT_INJECTION +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 new file mode 100644 index 00000000000..473c8d5bb7d --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_file.h @@ -0,0 +1,77 @@ +/* + * 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 + +namespace executorch::backends::coreai { + +class FileDescriptor { + public: + 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_; + } + 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; +} + +// 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_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 new file mode 100644 index 00000000000..9627761c0ce --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_manifest.h @@ -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. + */ + +#pragma once + +#import +#include + +namespace executorch::backends::coreai { + +class Manifest final { + public: + NSString* const export_hash; + NSString* const function; + NSArray* const inputs; + NSArray* const outputs; + NSDictionary* const files; + NSDictionary* const bundle_digests; + const bool aot_compiled; + // Source bundles do not constrain the platform or architecture. + NSString* const platform; + NSDictionary* const architectures; + const NSOperatingSystemVersion minimum_version; + + Manifest(const Manifest&) = default; + Manifest(Manifest&&) = default; + Manifest& operator=(const Manifest&) = delete; + Manifest& operator=(Manifest&&) = delete; + + private: + Manifest( + NSString* export_hash, + NSString* function, + NSArray* inputs, + NSArray* outputs, + NSDictionary* files, + NSDictionary* bundle_digests, + bool aot_compiled, + NSString* platform, + NSDictionary* architectures, + NSOperatingSystemVersion minimum_version); + + friend runtime::Result parse_manifest(NSData* data); +}; + +class SelectedBundle final { + public: + NSString* const export_hash; + NSString* const path; + NSString* const bundle_name; + NSString* const digest; + NSString* const function; + NSArray* const inputs; + NSArray* const outputs; + NSDictionary* const files; + const bool aot_compiled; + + SelectedBundle(const SelectedBundle&) = default; + SelectedBundle(SelectedBundle&&) = default; + SelectedBundle& operator=(const SelectedBundle&) = delete; + SelectedBundle& operator=(SelectedBundle&&) = delete; + + private: + SelectedBundle( + const Manifest& manifest, + NSString* path, + NSDictionary* files); + + friend runtime::Result select_bundle( + const Manifest& manifest, + NSString* device_architecture, + NSString* platform); +}; + +runtime::Result parse_manifest(NSData* data); +runtime::Result parse_manifest(const void* data, size_t size); +runtime::Result select_bundle( + const Manifest& manifest, + NSString* device_architecture, + NSString* platform); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_manifest.mm b/backends/apple/coreai/runtime/coreai_manifest.mm new file mode 100644 index 00000000000..af139c8f09d --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_manifest.mm @@ -0,0 +1,326 @@ +/* + * 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_manifest.h" +#include +#include +#import "coreai_validation.h" + +namespace executorch::backends::coreai { + +Manifest::Manifest( + NSString* export_hash, + NSString* function, + NSArray* inputs, + NSArray* outputs, + NSDictionary* files, + NSDictionary* bundle_digests, + bool aot_compiled, + NSString* platform, + NSDictionary* architectures, + NSOperatingSystemVersion minimum_version) + : export_hash(export_hash), + function(function), + inputs(inputs), + outputs(outputs), + files(files), + bundle_digests(bundle_digests), + aot_compiled(aot_compiled), + platform(platform), + architectures(architectures), + minimum_version(minimum_version) {} + +SelectedBundle::SelectedBundle( + const Manifest& manifest, + NSString* path, + NSDictionary* files) + : export_hash(manifest.export_hash), + path(path), + bundle_name(path.lastPathComponent), + digest(manifest.bundle_digests[path.lastPathComponent]), + function(manifest.function), + inputs(manifest.inputs), + outputs(manifest.outputs), + files([files copy]), + aot_compiled(manifest.aot_compiled) {} + +namespace { +using runtime::Error; +using runtime::Result; + +bool names_array(id value) { + if (![value isKindOfClass:NSArray.class]) { + return false; + } + NSMutableSet* seen = [NSMutableSet set]; + for (id name in value) { + if (!is_nonempty_string(name) || [seen containsObject:name]) { + return false; + } + [seen addObject:name]; + } + return true; +} + +bool file_size(id value) { + return [value isKindOfClass:NSNumber.class] && + CFGetTypeID((__bridge CFTypeRef)value) != CFBooleanGetTypeID() && + !CFNumberIsFloatType((__bridge CFNumberRef)value) && + [value compare:@0] != NSOrderedAscending && + [value compare:@(std::numeric_limits::max())] != + NSOrderedDescending && + [value isEqualToNumber:@([value unsignedLongLongValue])]; +} + +Error validate_files( + NSDictionary* filenames, + NSSet* bundles) { + ET_CHECK_OR_RETURN_ERROR( + [filenames isKindOfClass:NSDictionary.class] && filenames.count > 0, + InvalidProgram, + "Core AI manifest requires asset file sizes"); + NSSet* files = [NSSet setWithArray:filenames.allKeys]; + NSMutableSet* populated = [NSMutableSet set]; + NSMutableDictionary* spellings = + [NSMutableDictionary dictionary]; + NSLocale* locale = [NSLocale localeWithLocaleIdentifier:@"en_US_POSIX"]; + for (NSString* file in filenames) { + ET_CHECK_OR_RETURN_ERROR( + is_relative_path(file) && file_size(filenames[file]) && + [bundles containsObject:[file componentsSeparatedByString:@"/"] + .firstObject] && + [file containsString:@"/"] && + [file canBeConvertedToEncoding:NSUTF8StringEncoding], + InvalidProgram, + "Invalid Core AI asset path or size"); + [populated addObject:[file componentsSeparatedByString:@"/"].firstObject]; + // Aliased parent directories also merge otherwise distinct file paths. + NSString* path = file; + while (path.length > 0) { + NSString* identity = + [path stringByFoldingWithOptions:NSCaseInsensitiveSearch + locale:locale] + .decomposedStringWithCanonicalMapping; + NSString* previous = spellings[identity]; + ET_CHECK_OR_RETURN_ERROR( + previous == nil || [previous isEqualToString:path], + InvalidProgram, + "Case or Unicode alias in Core AI asset paths"); + spellings[identity] = path; + path = path.stringByDeletingLastPathComponent; + ET_CHECK_OR_RETURN_ERROR( + ![files containsObject:path], + InvalidProgram, + "Conflicting Core AI asset paths"); + } + } + ET_CHECK_OR_RETURN_ERROR( + [populated isEqualToSet:bundles], + InvalidProgram, + "Core AI manifest requires files for every declared bundle"); + return Error::Ok; +} + +bool parse_version(id value, NSOperatingSystemVersion& version) { + // A null floor means the exporter used the SDK's save_asset default. + if (value == NSNull.null) { + version = {}; + return true; + } + if (!is_nonempty_string(value)) { + return false; + } + NSArray* parts = [value componentsSeparatedByString:@"."]; + if (parts.count > 3) { + return false; + } + NSInteger numbers[3] = {0, 0, 0}; + NSCharacterSet* invalid = [[NSCharacterSet + characterSetWithCharactersInString:@"0123456789"] invertedSet]; + for (NSUInteger i = 0; i < parts.count; ++i) { + NSString* part = parts[i]; + if (part.length == 0 || part.length > 6 || + [part rangeOfCharacterFromSet:invalid].location != NSNotFound) { + return false; + } + numbers[i] = part.integerValue; + } + version = {numbers[0], numbers[1], numbers[2]}; + return numbers[0] > 0; +} + +} // namespace + +Result parse_manifest(NSData* data) { + ET_CHECK_OR_RETURN_ERROR( + data != nil, InvalidProgram, "Missing Core AI manifest JSON"); + NSError* error = nil; + id object = [NSJSONSerialization JSONObjectWithData:data + options:0 + error:&error]; + ET_CHECK_OR_RETURN_ERROR( + [object isKindOfClass:NSDictionary.class], + InvalidProgram, + "Invalid Core AI manifest JSON"); + NSDictionary* dict = object; + NSString* packaging = dict[@"packaging"]; + ET_CHECK_OR_RETURN_ERROR( + [packaging isKindOfClass:NSString.class], + InvalidProgram, + "Missing Core AI packaging"); + ET_CHECK_OR_RETURN_ERROR( + [packaging isEqualToString:@"inline"] || + [packaging isEqualToString:@"aot_compiled_inline"], + NotSupported, + "Unsupported Core AI packaging; expected .aimodel or AOT .aimodelc"); + const bool aot_compiled = [packaging hasPrefix:@"aot_compiled_"]; + 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( + 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( + is_nonempty_string(platform), + InvalidProgram, + "Core AI AOT manifest requires a target platform"); + ET_CHECK_OR_RETURN_ERROR( + [architectures isKindOfClass:NSDictionary.class] && + architectures.count > 0, + InvalidProgram, + "Core AI AOT manifest requires a nonempty architecture map"); + for (id arch in architectures) { + ET_CHECK_OR_RETURN_ERROR( + 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( + is_relative_path(path) && + [path isEqualToString:[export_hash + stringByAppendingPathComponent:bundle]], + InvalidProgram, + "Invalid Core AI AOT bundle path for architecture %s", + [arch UTF8String]); + [bundles addObject:bundle]; + } + } else { + source_path = dict[@"path"]; + ET_CHECK_OR_RETURN_ERROR( + is_relative_path(source_path) && + [source_path + isEqualToString:[export_hash + stringByAppendingString:@"/model.aimodel"]], + InvalidProgram, + "Invalid Core AI source bundle path"); + [bundles addObject:@"model.aimodel"]; + } + ET_CHECK_OR_RETURN_ERROR( + is_nonempty_string(function) && names_array(inputs) && + names_array(outputs), + InvalidProgram, + "Core AI manifest requires ordered input/output names"); + NSOperatingSystemVersion minimum_version{}; + ET_CHECK_OR_RETURN_ERROR( + parse_version(dict[@"min_deployment_version"], minimum_version), + InvalidProgram, + "Invalid Core AI minimum deployment version"); + NSDictionary* files = dict[@"files"]; + ET_CHECK_OK_OR_RETURN_ERROR(validate_files(files, bundles)); + NSDictionary* bundle_digests = dict[@"bundle_digests"]; + ET_CHECK_OR_RETURN_ERROR( + [bundle_digests isKindOfClass:NSDictionary.class] && + [[NSSet setWithArray:bundle_digests.allKeys] isEqualToSet:bundles], + InvalidProgram, + "Core AI bundle digests must exactly cover declared bundles"); + for (NSString* bundle in bundles) { + id digest = bundle_digests[bundle]; + ET_CHECK_OR_RETURN_ERROR( + is_lowercase_hex(digest, 64), + InvalidProgram, + "Invalid Core AI exported bundle digest"); + } + return Manifest( + export_hash, + function, + inputs, + outputs, + files, + bundle_digests, + aot_compiled, + platform, + architectures, + minimum_version); +} + +Result parse_manifest(const void* data, size_t size) { + ET_CHECK_OR_RETURN_ERROR( + data != nullptr, InvalidProgram, "Missing Core AI manifest JSON"); + return parse_manifest([NSData dataWithBytes:data length:size]); +} + +Result select_bundle( + const Manifest& manifest, + NSString* device_architecture, + NSString* platform) { + if (!manifest.aot_compiled) { + return SelectedBundle( + manifest, + [manifest.export_hash stringByAppendingString:@"/model.aimodel"], + manifest.files); + } + ET_CHECK_OR_RETURN_ERROR( + [manifest.platform isEqualToString:platform], + DelegateInvalidCompatibility, + "Core AI AOT target platform %s does not match runtime platform %s; " + "re-export the model for this platform", + manifest.platform.UTF8String, + platform == nil ? "nil" : platform.UTF8String); + NSArray* architectures = [manifest.architectures.allKeys + sortedArrayUsingComparator:^NSComparisonResult(NSString* a, NSString* b) { + return [a compare:b options:NSLiteralSearch]; + }]; + NSString* available = [architectures componentsJoinedByString:@", "]; + ET_CHECK_OR_RETURN_ERROR( + 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", + device_architecture == nil ? "nil" : device_architecture.UTF8String, + available.UTF8String); + NSString* path = manifest.architectures[device_architecture]; + ET_CHECK_OR_RETURN_ERROR( + path != nil, + DelegateInvalidCompatibility, + "Core AI AOT has no bundle for device architecture %s; available " + "architectures: [%s]; re-export the model for this device architecture", + device_architecture.UTF8String, + available.UTF8String); + NSString* prefix = [path.lastPathComponent stringByAppendingString:@"/"]; + NSMutableDictionary* selected_files = + [NSMutableDictionary dictionary]; + for (NSString* file in manifest.files) { + if ([file hasPrefix:prefix]) { + selected_files[file] = manifest.files[file]; + } + } + return SelectedBundle(manifest, path, selected_files); +} + +} // namespace executorch::backends::coreai 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 new file mode 100644 index 00000000000..b61699b6ff0 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -0,0 +1,82 @@ +/* + * 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_file.h" + +namespace executorch::backends::coreai { + +enum class StorageCreationMode { + OpenExisting, + Create, +}; + +runtime::Error sync_storage_directory(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. 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. +// 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( + int root_fd, + NSString* relative_path, + const void* bytes, + size_t size); +// 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, + 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 new file mode 100644 index 00000000000..585152356d7 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -0,0 +1,450 @@ +/* + * 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_fault_injection.h" +#include "coreai_file.h" +#include "coreai_validation.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::coreai { +namespace { +using runtime::Error; +using runtime::Result; + +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, + const void* data, + size_t size, + StorageSyncMode sync_mode) { + 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( + is_storage_component(component), + InvalidArgument, + "Invalid storage path"); + } + 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) { + auto opened = open_storage_directory( + parent.get(), components[i], StorageCreationMode::Create); + if (!opened.ok()) { + return opened.error(); + } + FileDescriptor next = std::move(*opened->descriptor); + ET_CHECK_OK_OR_RETURN_ERROR(sync_storage_directory(parent.get())); + parent = std::move(next); + } + 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 result = [&]() -> Error { + const auto* bytes = static_cast(data); + while (size > 0) { + 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( + count > 0, AccessFailed, "Cannot write Core AI storage file"); + bytes += count; + size -= count; + } + ET_CHECK_OR_RETURN_ERROR( + 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( + ::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); + } + } + 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( + int parent_fd, + NSString* name, + StorageCreationMode mode) { + ET_CHECK_OR_RETURN_ERROR( + is_storage_component(name), + InvalidArgument, + "Invalid storage directory name"); + const auto open_directory = [&] { + 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( + fault_point(FaultPoint::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{std::nullopt, false}; + } + ET_CHECK_OR_RETURN_ERROR( + fd >= 0, AccessFailed, "Cannot open Core AI storage directory"); + return StorageDirectory{FileDescriptor(fd), created}; +} + +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); })); + 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"); + copy.release(); + 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) { + if (policy == StorageNamePolicy::SkipInvalid) { + 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:)]; +} + +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( + 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 FileDescriptor(fd); +} + +Error sync_storage_directory(int fd) { + ET_CHECK_OR_RETURN_ERROR( + 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, + const void* data, + size_t size) { + return write_file( + root_fd, relative_path, data, size, StorageSyncMode::Normal); +} + +Error publish_storage_data(int root_fd, NSString* name, NSData* data) { + ET_CHECK_OR_RETURN_ERROR( + is_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, StorageSyncMode::Full); + if (result == Error::Ok && fault_aware_retry(FaultPoint::Rename, [&] { + return renameat( + root_fd, + temporary.fileSystemRepresentation, + root_fd, + name.fileSystemRepresentation); + }) != 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( + full_fsync(root_fd), + 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); + } + return result; +} + +} // namespace executorch::backends::coreai 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 new file mode 100644 index 00000000000..14f5e9fd63e --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_validation.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 + +namespace executorch::backends::coreai { + +// Nonempty text without embedded NULs. +bool is_nonempty_string(id value); +// Relative bundle path without escapes, backslashes, or empty components. +bool is_relative_path(NSString* path); +// Single path component for an fd-relative open: no slashes, ".", or "..". +// 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 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 new file mode 100644 index 00000000000..4b635772561 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_validation.mm @@ -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. + */ + +#import "coreai_validation.h" + +namespace executorch::backends::coreai { +namespace { + +NSCharacterSet* invalid_hex_characters() { + static NSCharacterSet* invalid = [[NSCharacterSet + characterSetWithCharactersInString:@"0123456789abcdef"] invertedSet]; + return invalid; +} + +} // namespace + +bool is_nonempty_string(id value) { + return [value isKindOfClass:NSString.class] && [value length] > 0 && + [value rangeOfString:@"\0"].location == NSNotFound; +} + +bool is_relative_path(NSString* path) { + if (!is_nonempty_string(path) || path.isAbsolutePath || + [path containsString:@"\\"]) { + return false; + } + for (NSString* component in [path componentsSeparatedByString:@"/"]) { + if (component.length == 0 || [component isEqualToString:@"."] || + [component isEqualToString:@".."]) { + return false; + } + } + return true; +} + +bool is_storage_component(NSString* name) { + return is_nonempty_string(name) && ![name isEqualToString:@"."] && + ![name isEqualToString:@".."] && ![name containsString:@"/"]; +} + +bool is_lowercase_hex(NSString* value) { + return is_nonempty_string(value) && + [value rangeOfCharacterFromSet:invalid_hex_characters()].location == + NSNotFound; +} + +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..d70bd29d0b7 --- /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 new file mode 100644 index 00000000000..e31e6d84378 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_cache_root_test.mm @@ -0,0 +1,294 @@ +/* + * 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, + 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; + 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 = + 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 = 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( + (resolve_test_cache_directory(root.path, StorageCreationMode::Create) + .ok() && + backup_excluded(root))); + + NSURL* missing = [sandbox.url URLByAppendingPathComponent:@"missing"]; + 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( + resolve_test_cache_directory(@"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( + resolve_test_cache_directory(root.path, StorageCreationMode::Create) + .ok()); + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); + EXPECT_TRUE( + resolve_test_cache_directory(root.path, StorageCreationMode::Create) + .ok()); + EXPECT_TRUE( + 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"]; + 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( + 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( + resolve_test_cache_directory(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( + 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); + 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( + resolve_test_cache_directory(root.path, StorageCreationMode::Create) + .ok()); + ASSERT_EQ(rmdir(root.fileSystemRepresentation), 0); + ASSERT_TRUE( + resolve_test_cache_directory(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( + resolve_test_cache_directory(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; + NSString* protection = before[NSFileProtectionKey]; + ASSERT_NE(protection, nil); + if (![protection isEqual:NSFileProtectionComplete]) { + GTEST_SKIP() << "Filesystem did not apply NSFileProtectionComplete; got " + << protection.UTF8String; + } + for (StorageCreationMode mode : + {StorageCreationMode::OpenExisting, StorageCreationMode::Create}) { + 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; + EXPECT_TRUE([after[NSFileProtectionKey] isEqual:NSFileProtectionComplete]); + } +} + +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); + 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( + 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 @[ + regular.path, + [regular.path stringByAppendingPathComponent:@"child"] + ]) { + EXPECT_EQ( + 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..e6ffd8ed795 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_cache_test.mm @@ -0,0 +1,840 @@ +/* + * 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_fault_scope.h b/backends/apple/coreai/runtime/test/coreai_fault_scope.h new file mode 100644 index 00000000000..f60092cd91a --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_fault_scope.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 + +#import +#include +#include "coreai_fault_injection.h" + +namespace executorch::backends::coreai::testing { + +struct StorageFaultScope { + FaultPoint operation; + int error = EIO; + int match_index = 1; + int matches = 0; + int hits = 0; + bool armed = true; + void (^observe)(FaultPoint) = nil; + + explicit StorageFaultScope(FaultPoint op) + : operation(op), previous_(current_), previous_callback_(fault_callback) { + current_ = this; + fault_callback = inject; + } + ~StorageFaultScope() { + 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_; + FaultCallback previous_callback_; + + static int inject(FaultPoint 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..7379eff4267 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h @@ -0,0 +1,50 @@ +/* + * 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::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); +::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 new file mode 100644 index 00000000000..1e5a7bdcb9a --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm @@ -0,0 +1,169 @@ +/* + * 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 +#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; +} + +::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 new file mode 100644 index 00000000000..3702c812e0d --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_host_test.mm @@ -0,0 +1,32 @@ +/* + * 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 +#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 new file mode 100644 index 00000000000..bf9af55c0bb --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_manifest_fixture.h @@ -0,0 +1,78 @@ +/* + * 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 "coreai_manifest.h" + +namespace executorch::backends::coreai::testing { + +// Synthetic identities keep host tests independent of the exporter hash. +inline NSString* fixture_identity(char digit) { + EXPECT_TRUE((digit >= '0' && digit <= '9') || (digit >= 'a' && digit <= 'f')); + return [NSString stringWithUTF8String:std::string(64, digit).c_str()]; +} + +inline NSMutableDictionary* manifest_dict() { + return [@{ + @"packaging" : @"inline", + @"export_hash" : @"ab", + @"path" : @"ab/model.aimodel", + @"function" : @"main", + @"input_names" : @[ @"input_0" ], + @"output_names" : @[ @"output_0" ], + @"files" : @{@"model.aimodel/graph.bin" : @11}, + @"bundle_digests" : @{@"model.aimodel" : fixture_identity('1')}, + @"min_deployment_version" : @"26.0" + } mutableCopy]; +} + +inline NSMutableDictionary* aot_manifest_dict() { + auto dict = manifest_dict(); + [dict removeObjectForKey:@"path"]; + dict[@"packaging"] = @"aot_compiled_inline"; + dict[@"platform"] = @"macOS"; + dict[@"archs"] = @{ + @"arch_b" : @"ab/model.arch_b.aimodelc", + @"arch_a" : @"ab/model.arch_a.aimodelc" + }; + dict[@"files"] = @{ + @"model.arch_b.aimodelc/nested/weights.bin" : @7, + @"model.arch_a.aimodelc/graph.bin" : @14, + @"model.arch_b.aimodelc/graph.bin" : @14, + @"model.arch_a.aimodelc/nested/weights.bin" : @7 + }; + dict[@"bundle_digests"] = @{ + @"model.arch_a.aimodelc" : fixture_identity('a'), + @"model.arch_b.aimodelc" : fixture_identity('b') + }; + return dict; +} + +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 new file mode 100644 index 00000000000..2e0fa5393e7 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_manifest_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 "coreai_manifest_fixture.h" +#include "coreai_validation.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +TEST(CoreAIManifestTest, AcceptsInlineAndRejectsInvalidJson) { + EXPECT_TRUE(parse_manifest(encode(manifest_dict())).ok()); + EXPECT_EQ(parse_manifest(nil).error(), Error::InvalidProgram); + EXPECT_FALSE( + parse_manifest([@"[]" dataUsingEncoding:NSUTF8StringEncoding]).ok()); + EXPECT_FALSE( + parse_manifest([@"garbage" dataUsingEncoding:NSUTF8StringEncoding]).ok()); +} + +TEST(CoreAIManifestTest, RejectsUnknownPackaging) { + auto dict = manifest_dict(); + dict[@"packaging"] = @"unknown"; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::NotSupported); +} + +TEST(CoreAIManifestTest, RejectsUnsafeSourcePaths) { + for (NSString* path in @[ + @"../model.aimodel", + @"/ab/model.aimodel", + @"ab/../model.aimodel", + @"xx/model.aimodel" + ]) { + SCOPED_TRACE(path.UTF8String); + auto dict = manifest_dict(); + dict[@"path"] = path; + EXPECT_FALSE(parse_manifest(encode(dict)).ok()); + } +} + +TEST(CoreAIManifestTest, RejectsInvalidFileMaps) { + for (id files in @[ + @[], + @{}, + NSNull.null, + @"files", + @{@"model.aimodel/../escape" : @1}, + @{@"model.aimodel/bad\\name" : @1}, + @{@"/tmp/file" : @1}, + @{@"model.aimodel/one" : @1, + @"model.aimodel/one/two" : @2} + ]) { + SCOPED_TRACE([files description].UTF8String); + auto dict = manifest_dict(); + dict[@"files"] = files; + EXPECT_FALSE(parse_manifest(encode(dict)).ok()); + } +} + +TEST(CoreAIManifestTest, RejectsInvalidInputNames) { + for (id names in + @[ @[ @"same", @"same" ], @[ @1 ], @[ @"" ], @"input", NSNull.null ]) { + SCOPED_TRACE([names description].UTF8String); + auto dict = manifest_dict(); + dict[@"input_names"] = names; + EXPECT_FALSE(parse_manifest(encode(dict)).ok()); + } +} + +TEST(CoreAIManifestTest, ValidatesDeploymentFloor) { + for (NSString* version in @[ @"", @"27.x", @"27.0.0.0" ]) { + SCOPED_TRACE(version.UTF8String); + auto dict = manifest_dict(); + dict[@"min_deployment_version"] = version; + EXPECT_FALSE(parse_manifest(encode(dict)).ok()); + } + auto dict = manifest_dict(); + dict[@"min_deployment_version"] = NSNull.null; + EXPECT_TRUE(parse_manifest(encode(dict)).ok()); +} + +TEST(CoreAIManifestTest, ParsesDeploymentVersionComponents) { + const struct { + id value; + NSOperatingSystemVersion expected; + } rows[] = { + {@"26", {26, 0, 0}}, + {@"26.4", {26, 4, 0}}, + {@"26.4.2", {26, 4, 2}}, + {NSNull.null, {0, 0, 0}}}; + for (const auto& row : rows) { + SCOPED_TRACE([row.value description].UTF8String); + auto dict = manifest_dict(); + dict[@"min_deployment_version"] = row.value; + auto parsed = parse_manifest(encode(dict)); + ASSERT_TRUE(parsed.ok()); + EXPECT_EQ(parsed->minimum_version.majorVersion, row.expected.majorVersion); + EXPECT_EQ(parsed->minimum_version.minorVersion, row.expected.minorVersion); + EXPECT_EQ(parsed->minimum_version.patchVersion, row.expected.patchVersion); + } +} + +TEST(CoreAIManifestTest, RequiresMetadataMaps) { + auto dict = manifest_dict(); + ASSERT_TRUE(parse_manifest(encode(dict)).ok()); + for (NSString* key in @[ @"files", @"bundle_digests" ]) { + SCOPED_TRACE(key.UTF8String); + NSMutableDictionary* invalid = [dict mutableCopy]; + [invalid removeObjectForKey:key]; + EXPECT_EQ(parse_manifest(encode(invalid)).error(), Error::InvalidProgram); + invalid[key] = @[]; + EXPECT_EQ(parse_manifest(encode(invalid)).error(), Error::InvalidProgram); + } +} + +TEST(CoreAIManifestTest, RejectsInvalidFileSizes) { + for (id size in @[ @(-1), @0.5, @(std::numeric_limits::max()) ]) { + SCOPED_TRACE([size description].UTF8String); + auto dict = manifest_dict(); + dict[@"files"] = @{@"model.aimodel/graph.bin" : size}; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + +TEST(CoreAIManifestTest, AcceptsIntegralFileSizes) { + for (NSNumber* size in @[ @0, @11, @(std::numeric_limits::max()) ]) { + SCOPED_TRACE(size.description.UTF8String); + auto dict = manifest_dict(); + dict[@"files"] = @{@"model.aimodel/graph.bin" : size}; + EXPECT_TRUE(parse_manifest(encode(dict)).ok()); + } +} + +TEST(CoreAIManifestTest, RequiresExactLowercaseDigests) { + for (NSString* digest in @[ + [fixture_identity('a') substringFromIndex:1], + fixture_identity('a').uppercaseString, + [@"g" stringByPaddingToLength:64 withString:@"g" startingAtIndex:0] + ]) { + SCOPED_TRACE(digest.UTF8String); + auto dict = manifest_dict(); + dict[@"bundle_digests"] = @{@"model.aimodel" : digest}; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + +class CoreAIManifestMetadataTest : public ::testing::TestWithParam { + protected: + NSMutableDictionary* manifest() { + return GetParam() ? aot_manifest_dict() : manifest_dict(); + } + NSString* file() { + return GetParam() ? @"model.arch_b.aimodelc/graph.bin" + : @"model.aimodel/graph.bin"; + } +}; + +TEST_P(CoreAIManifestMetadataTest, ReturnsManifestFieldsAndNameOrder) { + auto dict = manifest(); + dict[@"function"] = @"predict"; + dict[@"input_names"] = @[ @"input_1", @"input_0" ]; + dict[@"output_names"] = @[ @"output_2", @"output_0", @"output_1" ]; + dict[@"min_deployment_version"] = @"26.4.2"; + auto parsed = parse_manifest(encode(dict)); + ASSERT_TRUE(parsed.ok()); + EXPECT_EQ(parsed->aot_compiled, GetParam()); + 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"]]); + EXPECT_TRUE([parsed->files isEqual:dict[@"files"]]); + EXPECT_TRUE([parsed->bundle_digests isEqual:dict[@"bundle_digests"]]); + EXPECT_EQ(parsed->minimum_version.majorVersion, 26); + EXPECT_EQ(parsed->minimum_version.minorVersion, 4); + EXPECT_EQ(parsed->minimum_version.patchVersion, 2); + if (GetParam()) { + EXPECT_TRUE([parsed->architectures isEqual:dict[@"archs"]]); + EXPECT_TRUE([parsed->platform isEqual:dict[@"platform"]]); + } else { + EXPECT_EQ(parsed->architectures, nil); + EXPECT_EQ(parsed->platform, nil); + } +} + +TEST_P(CoreAIManifestMetadataTest, RejectsFilesystemAliases) { + NSString* composed = [NSString stringWithFormat:@"caf%C", 0x00e9]; + NSString* decomposed = composed.decomposedStringWithCanonicalMapping; + for (NSArray* pair in @[ + @[ @"One", @"one" ], + @[ @"One", @"one/two" ], + @[ @"Dir/a", @"dir/b" ], + @[ composed, decomposed ], + @[ composed, [decomposed stringByAppendingPathComponent:@"child"] ], + @[ + [composed stringByAppendingPathComponent:@"a"], + [decomposed stringByAppendingPathComponent:@"b"] + ], + @[ composed.uppercaseString, decomposed ] + ]) { + SCOPED_TRACE(pair.description.UTF8String); + auto dict = manifest(); + NSMutableDictionary* files = [dict[@"files"] mutableCopy]; + NSString* bundle = file().stringByDeletingLastPathComponent; + files[[bundle stringByAppendingPathComponent:pair[0]]] = @1; + files[[bundle stringByAppendingPathComponent:pair[1]]] = @2; + dict[@"files"] = files; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + +TEST_P(CoreAIManifestMetadataTest, PreservesDistinctUnicodePaths) { + auto dict = manifest(); + NSMutableDictionary* files = [dict[@"files"] mutableCopy]; + NSString* bundle = file().stringByDeletingLastPathComponent; + NSString* composed = [NSString stringWithFormat:@"caf%C", 0x00e9]; + for (NSString* name in + @[ composed, @"same/one", @"same/two", @"other/ONE" ]) { + files[[bundle stringByAppendingPathComponent:name]] = @1; + } + dict[@"files"] = files; + auto parsed = parse_manifest(encode(dict)); + ASSERT_TRUE(parsed.ok()); + EXPECT_TRUE([parsed->files isEqual:files]); +} + +TEST_P(CoreAIManifestMetadataTest, RequiresExactBundleCoverage) { + for (bool extra : {false, true}) { + SCOPED_TRACE(extra ? "extra digest" : "missing digest"); + auto dict = manifest(); + NSMutableDictionary* digests = [dict[@"bundle_digests"] mutableCopy]; + if (extra) { + digests[@"unexpected.aimodel"] = fixture_identity('e'); + } else { + [digests removeObjectForKey:file().stringByDeletingLastPathComponent]; + } + dict[@"bundle_digests"] = digests; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + +TEST_P(CoreAIManifestMetadataTest, RejectsInvalidHashes) { + auto missing = manifest(); + [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[@"export_hash"] = hash; + if (GetParam()) { + dict[@"archs"] = @{ + @"arch_a" : + [NSString stringWithFormat:@"%@/model.arch_a.aimodelc", hash], + @"arch_b" : + [NSString stringWithFormat:@"%@/model.arch_b.aimodelc", hash] + }; + } else { + dict[@"path"] = [NSString stringWithFormat:@"%@/model.aimodel", hash]; + } + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + +TEST_P(CoreAIManifestMetadataTest, RejectsFilesForUndeclaredBundles) { + auto dict = manifest(); + NSMutableDictionary* files = [dict[@"files"] mutableCopy]; + files[@"model.undeclared.aimodelc/graph.bin"] = @1; + dict[@"files"] = files; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); +} + +TEST_P(CoreAIManifestMetadataTest, RequiresFilesForEveryDeclaredBundle) { + auto valid = manifest(); + for (NSString* bundle in valid[@"bundle_digests"]) { + SCOPED_TRACE(bundle.UTF8String); + NSMutableDictionary* dict = [valid mutableCopy]; + NSMutableDictionary* files = [dict[@"files"] mutableCopy]; + NSString* prefix = [bundle stringByAppendingString:@"/"]; + for (NSString* file in dict[@"files"]) { + if ([file hasPrefix:prefix]) { + [files removeObjectForKey:file]; + } + } + dict[@"files"] = files; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + +TEST(CoreAIManifestTest, RequiresAotPlatform) { + auto dict = aot_manifest_dict(); + [dict removeObjectForKey:@"platform"]; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + for (id platform in @[ @"", @1, @[], @{}, NSNull.null ]) { + SCOPED_TRACE([platform description].UTF8String); + dict[@"platform"] = platform; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + +TEST(CoreAIManifestTest, RequiresNonemptyAotArchitectureMap) { + auto dict = aot_manifest_dict(); + [dict removeObjectForKey:@"archs"]; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + for (id archs in @[ @{}, @[], @"arch_a", @1, NSNull.null ]) { + SCOPED_TRACE([archs description].UTF8String); + dict[@"archs"] = archs; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + +TEST(CoreAIManifestTest, RejectsUnsafeAotArchitectureNames) { + for (NSString* arch in @[ @"", @".", @"..", @"../arch_a", @"a/b", @"a\\b" ]) { + SCOPED_TRACE(arch.UTF8String); + auto dict = aot_manifest_dict(); + NSString* bundle = [NSString stringWithFormat:@"model.%@.aimodelc", arch]; + dict[@"archs"] = @{arch : [@"ab/" stringByAppendingString:bundle]}; + dict[@"files"] = @{[bundle stringByAppendingString:@"/graph.bin"] : @1}; + dict[@"bundle_digests"] = @{bundle : fixture_identity('a')}; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + +TEST(CoreAIManifestTest, RejectsAliasedAotBundles) { + NSString* composed = [NSString stringWithFormat:@"arch_%C", 0x00e9]; + for (NSArray* pair in @[ + @[ @"arch_a", @"ARCH_A" ], + @[ composed, composed.decomposedStringWithCanonicalMapping ] + ]) { + SCOPED_TRACE(pair.description.UTF8String); + auto dict = aot_manifest_dict(); + NSMutableDictionary* archs = [NSMutableDictionary dictionary]; + NSMutableDictionary* files = [NSMutableDictionary dictionary]; + NSMutableDictionary* digests = [NSMutableDictionary dictionary]; + for (NSString* arch in pair) { + NSString* bundle = [NSString stringWithFormat:@"model.%@.aimodelc", arch]; + archs[arch] = [@"ab" stringByAppendingPathComponent:bundle]; + files[[bundle stringByAppendingPathComponent:@"graph.bin"]] = @1; + digests[bundle] = fixture_identity('a'); + } + ASSERT_EQ(archs.count, 2); + dict[@"archs"] = archs; + dict[@"files"] = files; + dict[@"bundle_digests"] = digests; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + +TEST(CoreAIManifestTest, RejectsInvalidAotBundlePaths) { + for (id path in @[ + @"", + @"../model.arch_a.aimodelc", + @"/ab/model.arch_a.aimodelc", + @"ab/../model.arch_a.aimodelc", + @"xx/model.arch_a.aimodelc", + @"ab/model.arch_b.aimodelc", + @"ab/model.aimodel", + @1, + NSNull.null + ]) { + SCOPED_TRACE([path description].UTF8String); + auto dict = aot_manifest_dict(); + NSMutableDictionary* archs = [dict[@"archs"] mutableCopy]; + archs[@"arch_a"] = path; + dict[@"archs"] = archs; + EXPECT_EQ(parse_manifest(encode(dict)).error(), Error::InvalidProgram); + } +} + +INSTANTIATE_TEST_SUITE_P( + Delivery, + CoreAIManifestMetadataTest, + ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return info.param ? "Aot" : "Source"; + }); + +TEST(CoreAIManifestTest, SelectsAotArchitectureWithoutMutatingManifest) { + auto dict = aot_manifest_dict(); + auto parsed = parse_manifest(encode(dict)); + ASSERT_TRUE(parsed.ok()); + ASSERT_TRUE(parsed->aot_compiled); + NSDictionary* files_snapshot = [parsed->files copy]; + NSDictionary* digests_snapshot = + [parsed->bundle_digests copy]; + NSDictionary* architectures_snapshot = + [parsed->architectures copy]; + for (NSString* arch in @[ @"arch_a", @"arch_b" ]) { + SCOPED_TRACE(arch.UTF8String); + 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->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]); + EXPECT_EQ(selected->aot_compiled, parsed->aot_compiled); + EXPECT_TRUE([selected->bundle_name isEqual:bundle]); + EXPECT_TRUE([selected->digest isEqual:dict[@"bundle_digests"][bundle]]); + EXPECT_FALSE([selected->files isKindOfClass:NSMutableDictionary.class]); + EXPECT_TRUE(([selected->files isEqual:@{ + [bundle stringByAppendingPathComponent:@"graph.bin"] : @14, + [bundle stringByAppendingPathComponent:@"nested/weights.bin"] : @7 + }])); + EXPECT_TRUE([parsed->files isEqual:files_snapshot]); + EXPECT_TRUE([parsed->bundle_digests isEqual:digests_snapshot]); + EXPECT_TRUE([parsed->architectures isEqual:architectures_snapshot]); + } +} + +TEST(CoreAIManifestTest, SelectsSourceBundleUnchanged) { + auto dict = manifest_dict(); + auto parsed = parse_manifest(encode(dict)); + ASSERT_TRUE(parsed.ok()); + ASSERT_FALSE(parsed->aot_compiled); + // Selection ignores the architecture and platform for source bundles. + auto selected = select_bundle(parsed.get(), @"anything", @"elsewhere"); + ASSERT_TRUE(selected.ok()); + EXPECT_TRUE([selected->path isEqualToString:dict[@"path"]]); + 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]); + EXPECT_TRUE([selected->files isEqual:parsed->files]); + EXPECT_TRUE([selected->bundle_name isEqual:@"model.aimodel"]); + EXPECT_TRUE( + [selected->digest isEqual:parsed->bundle_digests[@"model.aimodel"]]); + EXPECT_EQ(selected->aot_compiled, parsed->aot_compiled); +} + +TEST(CoreAIManifestTest, ParsesRawBytesAndKeepsIndependentMetadata) { + auto original = aot_manifest_dict(); + NSMutableData* bytes = [encode(original) mutableCopy]; + auto parsed = parse_manifest(bytes.bytes, bytes.length); + ASSERT_TRUE(parsed.ok()); + [bytes setLength:0]; + [original removeAllObjects]; + 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); + EXPECT_EQ(parsed->bundle_digests.count, 2); + EXPECT_FALSE([parsed->files isKindOfClass:NSMutableDictionary.class]); + EXPECT_FALSE([parsed->inputs isKindOfClass:NSMutableArray.class]); + EXPECT_FALSE([parsed->architectures isKindOfClass:NSMutableDictionary.class]); + EXPECT_EQ(parse_manifest(nullptr, 0).error(), Error::InvalidProgram); + EXPECT_EQ(parse_manifest("", 0).error(), Error::InvalidProgram); + EXPECT_EQ(parse_manifest("garbage", 7).error(), Error::InvalidProgram); +} + +TEST(CoreAIManifestTest, ManifestTypesAreNotDefaultConstructibleOrAssignable) { + static_assert(std::is_copy_constructible_v); + static_assert(std::is_move_constructible_v); + static_assert(std::is_copy_constructible_v); + static_assert(std::is_move_constructible_v); + static_assert(!std::is_convertible_v); + static_assert(!std::is_default_constructible_v); + static_assert(!std::is_copy_assignable_v); + static_assert(!std::is_move_assignable_v); + static_assert(!std::is_default_constructible_v); + static_assert(!std::is_copy_assignable_v); + static_assert(!std::is_move_assignable_v); + SUCCEED(); +} + +TEST(CoreAIManifestTest, ValidationHelpersMatchParserRules) { + 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(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) { + auto ios = aot_manifest_dict(); + ios[@"platform"] = @"iOS"; + auto parsed_ios = parse_manifest(encode(ios)); + ASSERT_TRUE(parsed_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_bundle(parsed.get(), @"arch_b", @"iOS").error(), + Error::DelegateInvalidCompatibility); + EXPECT_EQ( + 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_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 new file mode 100644 index 00000000000..e5f135724f0 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_storage_test.mm @@ -0,0 +1,712 @@ +/* + * 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 "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, 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_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); + 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", StorageCreationMode::OpenExisting); + ASSERT_TRUE(missing.ok()); + 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()); + ASSERT_TRUE(created->descriptor.has_value()); + FileDescriptor first = std::move(*created->descriptor); + ASSERT_GE(first.get(), 0); + EXPECT_TRUE(created->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(FaultPoint::BeforeCreateDirectory); + for (StorageCreationMode mode : + {StorageCreationMode::OpenExisting, StorageCreationMode::Create}) { + auto reopened = open_storage_directory(root.get(), @"child", mode); + ASSERT_TRUE(reopened.ok()); + ASSERT_TRUE(reopened->descriptor.has_value()); + FileDescriptor next = std::move(*reopened->descriptor); + ASSERT_GE(next.get(), 0); + EXPECT_FALSE(reopened->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(FaultPoint::BeforeCreateDirectory); + fault.armed = false; + fault.observe = ^(FaultPoint operation) { + if (operation == FaultPoint::BeforeCreateDirectory) + EXPECT_EQ(mkdir(child.fileSystemRepresentation, 0777), 0); + }; + auto opened = + open_storage_directory(root.get(), @"child", StorageCreationMode::Create); + ASSERT_TRUE(opened.ok()); + ASSERT_TRUE(opened->descriptor.has_value()); + FileDescriptor directory = std::move(*opened->descriptor); + EXPECT_GE(directory.get(), 0); + EXPECT_FALSE(opened->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, StorageCreationMode::OpenExisting) + .error(), + Error::InvalidArgument); + for (NSString* name in @[ @"", @".", @"..", @"nested/child" ]) { + EXPECT_EQ( + open_storage_directory( + root.get(), name, StorageCreationMode::OpenExisting) + .error(), + Error::InvalidArgument); + } +} + +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]; + } + 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); + } +} + +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 (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); + } + } +} + +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, 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(FaultPoint::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(FaultPoint::BeforeCreateDirectory); + fault.armed = false; + __block int creations = 0; + fault.observe = ^(FaultPoint operation) { + if (operation == FaultPoint::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(FaultPoint::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 { + FaultPoint operation; + int match_index; + } rows[] = { + {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)); + 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(FaultPoint::SyncDirectory); + fault.armed = false; + fault.observe = ^(FaultPoint 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{ + FaultPoint::BeforeCreateDirectory, + FaultPoint::SyncDirectory, + FaultPoint::BeforeCreateDirectory, + FaultPoint::SyncDirectory, + FaultPoint::Write, + FaultPoint::SyncFile, + FaultPoint::SyncDirectory})); +} + +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, 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 { + 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); + for (FaultPoint operation : + {FaultPoint::Write, + FaultPoint::SyncFile, + FaultPoint::SyncDirectory, + FaultPoint::FullSync, + FaultPoint::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())); + } + + 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(FaultPoint::FullSync); + fault.armed = false; + 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{ + FaultPoint::Write, + FaultPoint::SyncFile, + FaultPoint::FullSync, + FaultPoint::SyncDirectory, + FaultPoint::Rename, + FaultPoint::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(FaultPoint::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 (FaultPoint operation : + {FaultPoint::Write, + FaultPoint::SyncFile, + FaultPoint::SyncDirectory, + FaultPoint::FullSync, + FaultPoint::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())); + } +} 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 039c5ce3712..cf3a335dd84 100644 --- a/tools/cmake/preset/default.cmake +++ b/tools/cmake/preset/default.cmake @@ -23,6 +23,9 @@ define_overridable_option( define_overridable_option( EXECUTORCH_BUILD_COREML "Build the Core ML backend" BOOL OFF ) +define_overridable_option( + EXECUTORCH_BUILD_COREAI "Build the Core AI backend (Apple SDK 27+)" BOOL OFF +) define_overridable_option( EXECUTORCH_FLATBUFFERS_MAX_ALIGNMENT "Exir lets users set the alignment of tensor data embedded in the flatbuffer, and some users need an alignment larger than the default, which is typically 32." @@ -437,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