From bb9fa672366fb6ccd05a1054aaf9c2d3189bf421 Mon Sep 17 00:00:00 2001 From: lucylq Date: Mon, 28 Apr 2025 15:06:42 -0700 Subject: [PATCH] Backend data separation test Add backend data separation test with demo ExecutorBackend. Note: ExecutorBackend is a wrapper around a portable PTE file. It doesn't support delegated program-data separation in the sense of tagging tensors as external and placing them into the named data store as such. Rather, this test: 1. Creates a linear PTE file with data removed, using the portable flow. 2. Packages (1) into preprocessed blob as a delegate, using ExecutorBackend. Note: this discards the PTD portion. 3. Re-create the PTD portion via export_program (portable flow again). 4. Runs the delegated ExecutorBackend linear with portable linear.ptd file. Caveat: this means that LinearModule definition in export_program and export_delegated_program must stay in sync, as we get the PTE artifact from export_delegated_program (wrapped in ExecutorBackend), and the PTD artifact from export_program. Differential Revision: [D73679733](https://our.internmc.facebook.com/intern/diff/D73679733/) [ghstack-poisoned] --- .../test/demos/rpc/ExecutorBackend.cpp | 10 +- exir/backend/test/demos/rpc/TARGETS | 1 + .../demos/rpc/executor_backend_preprocess.py | 8 +- exir/backend/test/demos/rpc/targets.bzl | 1 + runtime/executor/method.cpp | 2 + .../test/backend_data_separation_test.cpp | 95 +++++++++++++++++++ runtime/executor/test/targets.bzl | 22 +++++ test/models/export_delegated_program.py | 16 +++- test/models/export_program.py | 13 +++ test/models/targets.bzl | 25 ++++- 10 files changed, 187 insertions(+), 6 deletions(-) create mode 100644 runtime/executor/test/backend_data_separation_test.cpp diff --git a/exir/backend/test/demos/rpc/ExecutorBackend.cpp b/exir/backend/test/demos/rpc/ExecutorBackend.cpp index 7dc0d2b2373..7632e4ad33c 100644 --- a/exir/backend/test/demos/rpc/ExecutorBackend.cpp +++ b/exir/backend/test/demos/rpc/ExecutorBackend.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #include @@ -37,6 +38,7 @@ using ::executorch::runtime::MemoryAllocator; using ::executorch::runtime::MemoryManager; using ::executorch::runtime::Method; using ::executorch::runtime::MethodMeta; +using ::executorch::runtime::NamedDataMap; using ::executorch::runtime::Program; using ::executorch::runtime::Result; using ::executorch::runtime::Span; @@ -156,9 +158,13 @@ class ExecutorBackend final : public ::executorch::runtime::BackendInterface { new (client_memory_manager) MemoryManager(client_method_allocator, client_planned_memory); + const NamedDataMap* named_data_map = context.get_named_data_map(); // Construct the client Method - Result method_res = - client_program->load_method("forward", client_memory_manager); + Result method_res = client_program->load_method( + "forward", + client_memory_manager, + /*event_tracer=*/nullptr, + named_data_map); if (!method_res.ok()) { ET_LOG( Error, diff --git a/exir/backend/test/demos/rpc/TARGETS b/exir/backend/test/demos/rpc/TARGETS index 3fdb1d4360a..d8fb426ba6a 100644 --- a/exir/backend/test/demos/rpc/TARGETS +++ b/exir/backend/test/demos/rpc/TARGETS @@ -11,6 +11,7 @@ runtime.python_library( ], visibility = [ "//executorch/exir/backend/test/...", + "//executorch/test/...", ], deps = [ "//caffe2:torch", diff --git a/exir/backend/test/demos/rpc/executor_backend_preprocess.py b/exir/backend/test/demos/rpc/executor_backend_preprocess.py index 0e5b8a8d3d5..19f7a25bd73 100644 --- a/exir/backend/test/demos/rpc/executor_backend_preprocess.py +++ b/exir/backend/test/demos/rpc/executor_backend_preprocess.py @@ -8,6 +8,8 @@ from typing import final, List +from executorch.exir import ExecutorchBackendConfig + from executorch.exir.backend.backend_details import ( BackendDetails, ExportedProgram, @@ -24,10 +26,14 @@ def preprocess( edge_program: ExportedProgram, compile_specs: List[CompileSpec], ) -> PreprocessResult: + config = ExecutorchBackendConfig() + for spec in compile_specs: + if spec.key == "external_constants": + config.external_constants = True return PreprocessResult( processed_bytes=EdgeProgramManager( edge_programs=edge_program, ) - .to_executorch() + .to_executorch(config) .buffer, ) diff --git a/exir/backend/test/demos/rpc/targets.bzl b/exir/backend/test/demos/rpc/targets.bzl index c5cfb343a6c..486444e400d 100644 --- a/exir/backend/test/demos/rpc/targets.bzl +++ b/exir/backend/test/demos/rpc/targets.bzl @@ -40,6 +40,7 @@ def define_common_targets(): ], visibility = [ "//executorch/exir/backend/test/...", + "//executorch/runtime/executor/test/...", ], deps = [ ":executor_backend", diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index 0709e0fac1f..8191f26c118 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -329,6 +329,7 @@ Result Method::get_num_external_constants() { } Error Method::parse_external_constants(const NamedDataMap* named_data_map) { + ET_CHECK_OR_RETURN_ERROR(named_data_map != nullptr, InvalidState, "named_data_map is null"); auto flatbuffer_values = serialization_plan_->values(); size_t n_value = flatbuffer_values->size(); @@ -372,6 +373,7 @@ Error Method::parse_external_constants(const NamedDataMap* named_data_map) { Result tensor_layout = named_data_map->get_metadata(key); if (!tensor_layout.ok()) { + ET_LOG(Info, "Failed to get metadata for key %s", key); return tensor_layout.error(); } // Check external tensor compatibility. diff --git a/runtime/executor/test/backend_data_separation_test.cpp b/runtime/executor/test/backend_data_separation_test.cpp new file mode 100644 index 00000000000..3b77acca711 --- /dev/null +++ b/runtime/executor/test/backend_data_separation_test.cpp @@ -0,0 +1,95 @@ + +/* + * 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 + +using namespace ::testing; +using executorch::extension::FlatTensorDataMap; +using executorch::runtime::DataLoader; +using executorch::runtime::Error; +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 BackendDataSeparationTest : 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(); + + // Make sure that the backend has been registered. Safe to call multiple + // times. Doing this at runtime ensures that it's only registered if these + // tests are run. + ASSERT_EQ(example::register_executor_backend(), Error::Ok); + + // Create data loaders. + Result linear_program_loader = FileDataLoader::from( + std::getenv("ET_MODULE_LINEAR_DELEGATE_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_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())); + + ET_LOG(Info, "setup done, named_data_map_ = %lu", linear_data_map_->get_num_keys().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(BackendDataSeparationTest, TestSeparation) { + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = linear_program_->load_method( + "forward", &mmm.get(), /*event_tracer=*/nullptr, /*named_data_map=*/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/runtime/executor/test/targets.bzl b/runtime/executor/test/targets.bzl index 0ecff00eb8a..8ecdccccdcb 100644 --- a/runtime/executor/test/targets.bzl +++ b/runtime/executor/test/targets.bzl @@ -240,6 +240,28 @@ def define_common_targets(is_fbcode = False): }, ) + runtime.cxx_test( + name = "backend_data_separation_test", + srcs = [ + "backend_data_separation_test.cpp", + ], + deps = [ + ":managed_memory_manager", + "//executorch/runtime/executor:program", + "//executorch/extension/data_loader:file_data_loader", + "//executorch/exir/backend/test/demos/rpc:executor_backend", + "//executorch/exir/backend/test/demos/rpc:executor_backend_register", + "//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_DELEGATE_PROGRAM_PATH": "$(location fbcode//executorch/test/models:exported_executor_backend_program_and_data[ModuleLinear-e.pte])", + "ET_MODULE_LINEAR_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleLinear.ptd])", + }, + ) runtime.cxx_test( name = "memory_manager_test", srcs = [ diff --git a/test/models/export_delegated_program.py b/test/models/export_delegated_program.py index f7b2f354373..5d66e06f83d 100644 --- a/test/models/export_delegated_program.py +++ b/test/models/export_delegated_program.py @@ -20,11 +20,15 @@ from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower from executorch.exir.backend.backend_api import to_backend from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult +from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.backend.test.backend_with_compiler_demo import ( BackendWithCompilerDemo, ) from executorch.exir.passes.external_constants_pass import ( delegate_external_constants_pass, +) +from executorch.exir.backend.test.demos.rpc.executor_backend_preprocess import ( + ExecutorBackend, ) from executorch.exir.program import ExecutorchProgramManager from torch import nn @@ -150,13 +154,18 @@ def __init__(self, fn, method_name=method_name): def forward(self, *args, **kwargs): return getattr(self.fn, self.method_name)(*args, **kwargs) - exported_program = export(WrapperModule(eager_module), args=inputs, strict=True) + if method_name != "forward": + # Only require wrapper module if we're exporting a specific method other than forward. + exported_program = export(WrapperModule(eager_module), args=inputs, strict=True) + else: + exported_program = export(eager_module, args=inputs, strict=True) edge_config = EdgeCompileConfig(_check_ir_validity=False) et_config = exir.ExecutorchBackendConfig( extract_delegate_segments=extract_delegate_segments, constant_tensor_alignment=constant_tensor_alignment, delegate_alignment=delegate_alignment, + external_constants=external_constants, ) if backend_id == "XnnpackBackend": @@ -181,7 +190,10 @@ def forward(self, *args, **kwargs): else: edge: exir.EdgeProgramManager = to_edge(exported_program) lowered_module = to_backend( # type: ignore[call-arg] - backend_id, edge.exported_program(), compile_specs=[] + backend_id, + edge.exported_program(), + # Just for the demo executor_backend. + compile_specs=[CompileSpec(key="external_constants", value=b"")], ) class CompositeModule(nn.Module): diff --git a/test/models/export_program.py b/test/models/export_program.py index 5387df24aad..e13b63eaf74 100644 --- a/test/models/export_program.py +++ b/test/models/export_program.py @@ -146,6 +146,19 @@ def get_random_inputs(self): return (torch.ones(2, 2, dtype=torch.float),) +# Used for program-data-separation. +class ModuleLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(3, 3) + + def forward(self, x: torch.Tensor): + return self.linear(x) + + def get_random_inputs(self): + return (torch.randn(3),) + + class ModuleMultipleEntry(torch.nn.Module): def __init__(self): super().__init__() diff --git a/test/models/targets.bzl b/test/models/targets.bzl index db0f410d727..9e26a6c123b 100644 --- a/test/models/targets.bzl +++ b/test/models/targets.bzl @@ -95,6 +95,7 @@ def define_common_targets(): # Class names of nn.Modules for :exported_programs to export. MODULES_AND_DATA_TO_EXPORT = [ "ModuleAddMul", + "ModuleLinear", "ModuleSimpleTrain", ] @@ -104,6 +105,8 @@ def define_common_targets(): outs = { "ModuleAddMul.pte": ["ModuleAddMulProgram.pte"], "ModuleAddMul.ptd": ["ModuleAddMulProgram.ptd"], + "ModuleLinear.pte": ["ModuleLinearProgram.pte"], + "ModuleLinear.ptd": ["ModuleLinearProgram.ptd"], "ModuleSimpleTrainProgram.pte": ["ModuleSimpleTrainProgram.pte"], "ModuleSimpleTrain.ptd": ["ModuleSimpleTrainProgram.ptd"], }, @@ -146,7 +149,7 @@ def define_common_targets(): deps = [ ":export_delegated_program_lib", "//executorch/backends/xnnpack/partition:xnnpack_partitioner", - + "//executorch/exir/backend/test/demos/rpc:executor_backend_preprocess", ], visibility = [], # Private ) @@ -225,3 +228,23 @@ def define_common_targets(): "//executorch/test/...", ], ) + + # Export with demo ExecutorBackend for program-data separation test. + runtime.genrule( + name = "exported_executor_backend_program_and_data", + cmd = "$(exe :export_delegated_program)" + + " --modules ModuleLinear" + + " --backend_id ExecutorBackend" + + " --external_constants" + + " --outdir $OUT", + + outs = { + "ModuleLinear-e.pte": ["ModuleLinear-e.pte"], + }, + default_outs = ["."], + visibility = [ + "//executorch/runtime/executor/test/...", + "//executorch/extension/flat_tensor/test/...", + "//executorch/test/...", + ], + )