diff --git a/backends/xnnpack/test/runtime/test_xnn_data_separation.cpp b/backends/xnnpack/test/runtime/test_xnn_data_separation.cpp new file mode 100644 index 00000000000..342e3478e0f --- /dev/null +++ b/backends/xnnpack/test/runtime/test_xnn_data_separation.cpp @@ -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. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +using namespace ::testing; +using executorch::extension::FlatTensorDataMap; +using executorch::runtime::DataLoader; +using executorch::runtime::Error; +using executorch::runtime::FreeableBuffer; +using executorch::runtime::Method; +using executorch::runtime::Program; +using executorch::runtime::Result; +using executorch::runtime::testing::ManagedMemoryManager; +using torch::executor::util::FileDataLoader; + +constexpr size_t kDefaultNonConstMemBytes = 32 * 1024U; +constexpr size_t kDefaultRuntimeMemBytes = 32 * 1024U; + +class DataSeparationTest : public ::testing::Test { + protected: + void SetUp() override { + // Since these tests cause ET_LOG to be called, the PAL must be initialized + // first. + executorch::runtime::runtime_init(); + + // Create data loaders. + Result linear_program_loader = + FileDataLoader::from(std::getenv("ET_MODULE_LINEAR_XNN_PROGRAM_PATH")); + ASSERT_EQ(linear_program_loader.error(), Error::Ok); + linear_program_loader_ = std::make_unique( + std::move(linear_program_loader.get())); + + Result linear_data_loader = + FileDataLoader::from(std::getenv("ET_MODULE_LINEAR_XNN_DATA_PATH")); + ASSERT_EQ(linear_data_loader.error(), Error::Ok); + linear_data_loader_ = + std::make_unique(std::move(linear_data_loader.get())); + + // Create programs. + Result linear_program = Program::load( + linear_program_loader_.get(), + Program::Verification::InternalConsistency); + ASSERT_EQ(linear_program.error(), Error::Ok); + linear_program_ = + std::make_unique(std::move(linear_program.get())); + + Result linear_data_map = + FlatTensorDataMap::load(linear_data_loader_.get()); + EXPECT_EQ(linear_data_map.error(), Error::Ok); + linear_data_map_ = + std::make_unique(std::move(linear_data_map.get())); + } + + private: + std::unique_ptr linear_program_loader_; + std::unique_ptr linear_data_loader_; + + protected: + std::unique_ptr linear_program_; + std::unique_ptr linear_data_map_; +}; + +TEST_F(DataSeparationTest, TestExternalData) { + FlatTensorDataMap* data_map = linear_data_map_.get(); + EXPECT_EQ(data_map->get_num_keys().get(), 2); + + Result key0 = data_map->get_key(0); + EXPECT_EQ(key0.error(), Error::Ok); + Result key1 = data_map->get_key(1); + EXPECT_EQ(key1.error(), Error::Ok); + + // Check that accessing keys out of bounds fails. + EXPECT_EQ(data_map->get_key(2).error(), Error::InvalidArgument); + + // Linear.weight + Result data0 = data_map->get_data(key0.get()); + EXPECT_EQ(data0.error(), Error::Ok); + EXPECT_EQ(data0.get().size(), 36); // 3*3*4 (3*3 matrix, 4 bytes per float) + + // Linear.bias + Result data1 = data_map->get_data(key1.get()); + EXPECT_EQ(data1.error(), Error::Ok); + EXPECT_EQ(data1.get().size(), 12); // 3*4 (3 vector, 4 bytes per float) + + // Check that accessing non-existent data fails. + Result data2 = data_map->get_data("nonexistent"); + EXPECT_EQ(data2.error(), Error::NotFound); +} + +TEST_F(DataSeparationTest, TestE2E) { + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = linear_program_->load_method( + "forward", &mmm.get(), nullptr, linear_data_map_.get()); + ASSERT_EQ(method.error(), Error::Ok); + + // Can execute the method. + Error err = method->execute(); + ASSERT_EQ(err, Error::Ok); +} diff --git a/backends/xnnpack/test/targets.bzl b/backends/xnnpack/test/targets.bzl index 58589b70607..f175e9655ea 100644 --- a/backends/xnnpack/test/targets.bzl +++ b/backends/xnnpack/test/targets.bzl @@ -43,3 +43,23 @@ def define_common_targets(): "//executorch/schema:program", ], ) + + runtime.cxx_test( + name = "test_xnn_data_separation", + srcs = ["runtime/test_xnn_data_separation.cpp"], + deps = [ + "//executorch/runtime/executor/test:managed_memory_manager", + "//executorch/runtime/executor:program", + "//executorch/extension/data_loader:file_data_loader", + "//executorch/backends/xnnpack:xnnpack_backend", + "//executorch/extension/flat_tensor:flat_tensor_data_map", + ], + env = { + # The tests use these vars to find the program files to load. + # Uses an fbcode target path because the authoring/export tools + # intentionally don't work in xplat (since they're host-only + # tools). + "ET_MODULE_LINEAR_XNN_PROGRAM_PATH": "$(location fbcode//executorch/test/models:exported_xnnpack_program_and_data[ModuleLinear-e.pte])", + "ET_MODULE_LINEAR_XNN_DATA_PATH": "$(location fbcode//executorch/test/models:exported_xnnpack_program_and_data[ModuleLinear.ptd])", + }, + ) diff --git a/test/models/targets.bzl b/test/models/targets.bzl index 9e26a6c123b..391ce230ab8 100644 --- a/test/models/targets.bzl +++ b/test/models/targets.bzl @@ -98,7 +98,7 @@ def define_common_targets(): "ModuleLinear", "ModuleSimpleTrain", ] - + runtime.genrule( name = "exported_program_and_data", cmd = "$(exe :export_program) --modules " + ",".join(MODULES_AND_DATA_TO_EXPORT) + " --external-constants --outdir $OUT", @@ -213,18 +213,18 @@ def define_common_targets(): runtime.genrule( name = "exported_xnnpack_program_and_data", cmd = "$(exe :export_delegated_program)" + - " --modules ModuleLinear" + + " --modules ModuleLinear" + " --backend_id XnnpackBackend" + " --external_constants" + " --outdir $OUT", - + outs = { "ModuleLinear-e.pte": ["ModuleLinear-e.pte"], "ModuleLinear.ptd": ["ModuleLinear.ptd"], }, default_outs = ["."], visibility = [ - "//executorch/runtime/executor/test/...", + "//executorch/backends/xnnpack/test/...", "//executorch/test/...", ], ) @@ -233,11 +233,11 @@ def define_common_targets(): runtime.genrule( name = "exported_executor_backend_program_and_data", cmd = "$(exe :export_delegated_program)" + - " --modules ModuleLinear" + + " --modules ModuleLinear" + " --backend_id ExecutorBackend" + " --external_constants" + " --outdir $OUT", - + outs = { "ModuleLinear-e.pte": ["ModuleLinear-e.pte"], },