From 011fb302c8fa17e94c747480f091b54968952bf6 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Sun, 4 Oct 2026 19:50:51 -0700 Subject: [PATCH] [CoreAI] Backend-local PTE inspection Summary: Add a private reader that extracts the Core AI processed buffers from a PTE without constructing a `Program` or `Method`, so cache maintenance can find a program's Core AI models without loading it. - `runtime/coreai_pte.{h,cpp}`: `inspect_coreai_pte(DataLoader&)` returns the `CoreAIBackend` processed buffers in method/delegate order. - It always runs the flatbuffer structural verifier, the supported schema version check and the runtime's existing semantic validator, even when `ET_ENABLE_PROGRAM_VERIFICATION=0`. It also checks method metadata, unique method names, exact backend IDs, data ranges, alignment and CPU accessibility. - It reuses the existing extended-header parser and verifier rather than duplicating program parsing. No shared runtime code changes. - Inline buffers borrow the retained program image. Only external segments referenced by `CoreAIBackend` delegates are loaded, and they keep their original `FreeableBuffer` callbacks. Unrelated external payloads and named data are not loaded. - `CoreAIPteData` owns the program image and buffers. It is move-only, and the loader must outlive it. - Tests: a synthetic PTE builder (`runtime/test/coreai_pte_fixture.{h,mm}`). The suite covers inline and segmented delegates, method-order routing, exact backend-ID matching, header formats, load-error propagation and buffer release on late failure. The PTE is treated as trusted exporter output, so malformed-PTE robustness is limited to the verifier and version checks. - README: "Backend-local PTE Inspection". No changes outside `backends/apple/coreai`. Test Plan: ``` # Local: configure and build only cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_OSX_DEPLOYMENT_TARGET=27.0 -DEXECUTORCH_BUILD_COREAI=ON \ -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON cmake --build build --target backends/apple/coreai/all # CI (macOS 27 runner, .github/workflows/coreai.yml) ctest --test-dir build/backends/apple/coreai --output-on-failure --no-tests=error ``` CTest runs `coreai_host_test` (adding the PTE inspection suite), `coreai_swift_bridge_test` and `coreai_runtime_smoke`. --- backends/apple/coreai/CMakeLists.txt | 10 +- backends/apple/coreai/README.md | 20 ++ backends/apple/coreai/runtime/coreai_pte.cpp | 222 +++++++++++++++ backends/apple/coreai/runtime/coreai_pte.h | 37 +++ .../runtime/test/coreai_program_data_test.mm | 266 ++++++++++++++++++ .../coreai/runtime/test/coreai_pte_fixture.h | 77 +++++ .../coreai/runtime/test/coreai_pte_fixture.mm | 216 ++++++++++++++ 7 files changed, 844 insertions(+), 4 deletions(-) create mode 100644 backends/apple/coreai/runtime/coreai_pte.cpp create mode 100644 backends/apple/coreai/runtime/coreai_pte.h create mode 100644 backends/apple/coreai/runtime/test/coreai_program_data_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_pte_fixture.h create mode 100644 backends/apple/coreai/runtime/test/coreai_pte_fixture.mm diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index e3100ff0147..1ab52fdaf91 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -24,7 +24,7 @@ enable_language(OBJCXX) set(_coreai_runtime_sources runtime/coreai_backend.mm runtime/coreai_assets.mm runtime/coreai_storage.mm runtime/coreai_bookmarks.mm - runtime/coreai_load_coordinator.mm + runtime/coreai_load_coordinator.mm runtime/coreai_pte.cpp ) # Keep ARC and C++ exception settings away from the private Swift module. @@ -59,6 +59,8 @@ if(EXECUTORCH_BUILD_TESTS) runtime/test/coreai_acquisition_test.mm runtime/test/coreai_acquisition_fixture.mm runtime/test/coreai_delegate_test.mm + runtime/test/coreai_program_data_test.mm + runtime/test/coreai_pte_fixture.mm runtime/ETCoreAITensor.mm ${_coreai_runtime_sources} ) @@ -72,8 +74,8 @@ if(EXECUTORCH_BUILD_TESTS) set(_coreai_host_runtime executorch) endif() target_link_libraries( - coreai_host_test PRIVATE ${_coreai_host_runtime} GTest::gtest - ${COREAI_FOUNDATION_FRAMEWORK} + coreai_host_test PRIVATE ${_coreai_host_runtime} 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) @@ -306,7 +308,7 @@ 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} - executorch::coreai_dependencies + program_schema executorch::coreai_dependencies ) executorch_target_link_options_shared_lib(coreaidelegate) if(EXECUTORCH_BUILD_SHARED) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index f6a33e79623..ae69f87dcf8 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -241,6 +241,26 @@ registered with CTest; otherwise it is excluded from the default build. On OS 27 running it 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 returns Core AI processed +buffers in method/delegate order without loading a Program or Method, registering +backends, or calling the SDK. It always verifies the program structure, supported +schema version, and existing runtime semantic invariants, even when generic +runtime verification is disabled (`ET_ENABLE_PROGRAM_VERIFICATION=0`). It also +validates method metadata, unique method names, and selected processed-data +references. + +The reader first reads a 64-byte prefix when the input is long enough, then +loads and verifies the program region, including inline constants and unrelated +inline blobs. Only external processed-data segments selected by exact +`CoreAIBackend` IDs are requested; unselected external payloads, the external +constants segment and external named assets are not loaded. Inline buffers +borrow the retained program image. Selected segment buffers retain their +original DataLoader callbacks. Keep the returned `CoreAIPteData` owner and loader +alive while using the buffers, and keep the loader's data stable throughout +inspection and use. The reader does not derive cache keys or mutate cache state. + ## Host Tests `EXECUTORCH_BUILD_COREAI=ON` with `EXECUTORCH_BUILD_TESTS=ON` registers the diff --git a/backends/apple/coreai/runtime/coreai_pte.cpp b/backends/apple/coreai/runtime/coreai_pte.cpp new file mode 100644 index 00000000000..84fe6a6d939 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_pte.cpp @@ -0,0 +1,222 @@ +/* + * 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 + +#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; + +Result cpu_data(const FreeableBuffer& buffer, size_t size) { + auto data = buffer.data_safe(); + if (!data.ok()) { + return data.error(); + } + if (data.get() == nullptr || buffer.size() < size) { + return Error::InvalidProgram; + } + return data.get(); +} +} // namespace + +Result inspect_coreai_pte(DataLoader& loader) { + auto source_size = loader.size(); + if (!source_size.ok()) { + return source_size.error(); + } + const size_t file_size = source_size.get(); + size_t program_size = file_size; + ExtendedHeader header{}; + bool has_header = false; + if (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()) { + header = parsed.get(); + has_header = true; + 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 > file_size) { + return Error::InvalidProgram; + } + program_size = static_cast(header.program_size); + if (header.segment_base_offset != 0 && + (header.segment_base_offset < program_size || + header.segment_base_offset > file_size)) { + return Error::InvalidProgram; + } + if (header.segment_data_size != 0 && + (header.segment_base_offset == 0 || + header.segment_data_size > file_size - header.segment_base_offset)) { + return Error::InvalidProgram; + } + } else if (parsed.error() != Error::NotFound) { + return Error::InvalidProgram; + } + } + if (program_size == 0 || program_size >= FLATBUFFERS_MAX_BUFFER_SIZE) { + return Error::InvalidProgram; + } + auto loaded = + loader.load(0, program_size, SegmentInfo(SegmentInfo::Type::Program)); + if (!loaded.ok()) { + return loaded.error(); + } + auto data = cpu_data(loaded.get(), program_size); + if (!data.ok()) { + return data.error(); + } + if (reinterpret_cast(data.get()) % alignof(std::max_align_t) != + 0) { + return Error::InvalidProgram; + } + if (program_size >= ExtendedHeader::kNumHeadBytes) { + auto retained_header = + ExtendedHeader::Parse(data.get(), ExtendedHeader::kNumHeadBytes); + if (has_header) { + if (!retained_header.ok() || + retained_header->program_size != header.program_size || + retained_header->segment_base_offset != header.segment_base_offset || + retained_header->segment_data_size != header.segment_data_size) { + return Error::InvalidProgram; + } + } else if ( + retained_header.ok() || retained_header.error() != Error::NotFound) { + return Error::InvalidProgram; + } + } + flatbuffers::Verifier verifier( + static_cast(data.get()), 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; + } + + CoreAIPteData result(std::move(loaded.get())); + std::set method_names; + for (const auto* plan : *program->execution_plan()) { + if (plan->name() == nullptr || plan->non_const_buffer_sizes() == nullptr || + plan->inputs() == nullptr || plan->outputs() == nullptr || + !method_names.emplace(plan->name()->c_str(), plan->name()->size()) + .second) { + return Error::InvalidProgram; + } + if (plan->delegates() == nullptr) { + continue; + } + for (const auto* delegate : *plan->delegates()) { + if (delegate == nullptr || delegate->id() == nullptr) { + return Error::InvalidProgram; + } + constexpr std::string_view backend = "CoreAIBackend"; + const std::string_view id( + delegate->id()->c_str(), delegate->id()->size()); + if (id != backend) { + continue; + } + const auto* reference = delegate->processed(); + if (reference == nullptr) { + return Error::InvalidProgram; + } + const uint32_t index = reference->index(); + if (reference->location() == fb::DataLocation::INLINE) { + const auto* buffers = program->backend_delegate_data(); + if (buffers == nullptr || index >= buffers->size() || + buffers->Get(index) == nullptr) { + return Error::InvalidProgram; + } + const auto* bytes = buffers->Get(index)->data(); + if (bytes == nullptr || bytes->size() == 0) { + return Error::InvalidProgram; + } + result.processed.emplace_back( + bytes->data(), bytes->size(), nullptr, nullptr); + } else if (reference->location() == fb::DataLocation::SEGMENT) { + const auto* segments = program->segments(); + if (!has_header || header.segment_base_offset == 0 || + segments == nullptr || index >= segments->size() || + segments->Get(index) == nullptr) { + return Error::InvalidProgram; + } + const auto* segment = segments->Get(index); + if (segment->offset() > std::numeric_limits::max() || + segment->size() > std::numeric_limits::max() || + segment->size() == 0) { + 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 (offset > file_size - base || size > 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, delegate->id()->c_str())); + if (!buffer.ok()) { + return buffer.error(); + } + auto bytes = cpu_data(buffer.get(), size); + if (!bytes.ok()) { + return bytes.error(); + } + // Preserve the original owner without exposing bytes beyond the + // requested extent. + if (buffer->size() != size) { + return Error::InvalidProgram; + } + result.processed.emplace_back(std::move(buffer.get())); + } else { + return Error::InvalidProgram; + } + } + } + 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..af15f703ffb --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_pte.h @@ -0,0 +1,37 @@ +/* + * 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 { + // Borrowed inline data and segment descriptors must die before the program. + 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; +}; + +// The loader and its data must remain stable until the returned owner is +// released. +runtime::Result inspect_coreai_pte(runtime::DataLoader& loader); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/test/coreai_program_data_test.mm b/backends/apple/coreai/runtime/test/coreai_program_data_test.mm new file mode 100644 index 00000000000..d5aa6a492ad --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_program_data_test.mm @@ -0,0 +1,266 @@ +/* + * 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 + +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 CoreAIProgramDataTest : 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 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(CoreAIProgramDataTest, 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(CoreAIProgramDataTest, 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(CoreAIProgramDataTest, MatchesCompleteBackendIdentifier) { + const char embedded[] = "CoreAIBackend\0suffix"; + CacheDelegateSpec 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}, + 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(CoreAIProgramDataTest, 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(CoreAIProgramDataTest, 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(CoreAIProgramDataTest, 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(CoreAIProgramDataTest, 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(CoreAIProgramDataTest, ReleasesEarlierBuffersOnLateFailure) { + FBSyntheticPTE pte({{"forward", + {{"CoreAIBackend", payload()}, + {"CoreAIBackend", payload("one"), true}, + {"CoreAIBackend", payload("two"), true}}}}); + pte.fault_call = 3; + pte.load_fault = LoadFault::Error; + ASSERT_NO_FATAL_FAILURE(expect_invalid(pte, Error::NotSupported)); + EXPECT_GE(pte.backend_loads, 1); + 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_pte_fixture.h b/backends/apple/coreai/runtime/test/coreai_pte_fixture.h new file mode 100644 index 00000000000..8632f808030 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_pte_fixture.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 + +#import + +#include +#include +#include + +#include +#include +#include +#include + +namespace executorch::backends::coreai::testing { + +struct CacheDelegateSpec { + const char* backend; + NSData* data; + bool segmented = false; + int index_override = -1; + bool missing_processed = false; + std::optional backend_length = std::nullopt; +}; + +struct CacheMethodSpec { + const char* name; + std::vector delegates; +}; + +class FBSyntheticPTE final : public executorch::runtime::DataLoader { + public: + enum class LoadFault { None, Error }; + 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; + int fail_backend_index = -1; + + 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..96ba23f7715 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_pte_fixture.mm @@ -0,0 +1,216 @@ +/* + * 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); + if (static_cast(info.segment_index) == fail_backend_index) { + return Error::AccessFailed; + } + } else if (info.descriptor != nullptr) { + ADD_FAILURE() << "Program segment has a backend descriptor"; + return Error::InvalidArgument; + } + auto* loads = backend ? &backend_loads : &program_loads; + void* allocation = std::malloc(std::max(size, size_t(1))); + if (allocation == nullptr) { + ADD_FAILURE() << "Could not allocate " << size << " bytes"; + return Error::MemoryAllocationFailed; + } + auto* data = static_cast(allocation); + if (size != 0) { + std::memcpy(data, bytes.data() + offset, size); + } + struct Allocation { + const FBSyntheticPTE* loader; + void* base; + bool backend; + }; + auto* context = new Allocation{this, allocation, backend}; + requests.back().data = data; + ++*loads; + return FreeableBuffer( + data, 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