diff --git a/devtools/bundled_program/bundled_program.cpp b/devtools/bundled_program/bundled_program.cpp index df4124e0038..913c349a53a 100644 --- a/devtools/bundled_program/bundled_program.cpp +++ b/devtools/bundled_program/bundled_program.cpp @@ -260,9 +260,16 @@ ET_NODISCARD Error load_bundled_input( if (!method_test.ok()) { return method_test.error(); } - + auto test_cases = method_test.get()->test_cases(); + ET_CHECK_OR_RETURN_ERROR( + testset_idx < test_cases->size(), + InvalidArgument, + "testset_idx %zu is out of range [0, %u]", + testset_idx, + test_cases->size()); auto bundled_inputs = - method_test.get()->test_cases()->Get(testset_idx)->inputs(); + test_cases->Get(static_cast(testset_idx)) + ->inputs(); for (size_t input_idx = 0; input_idx < method.inputs_size(); input_idx++) { auto bundled_input = bundled_inputs->GetMutableObject(input_idx); @@ -359,8 +366,16 @@ ET_NODISCARD Error verify_method_outputs( return method_test.error(); } + auto test_cases = method_test.get()->test_cases(); + ET_CHECK_OR_RETURN_ERROR( + testset_idx < test_cases->size(), + InvalidArgument, + "testset_idx %zu is out of range [0, %u]", + testset_idx, + test_cases->size()); auto bundled_expected_outputs = - method_test.get()->test_cases()->Get(testset_idx)->expected_outputs(); + test_cases->Get(static_cast(testset_idx)) + ->expected_outputs(); if (bundled_expected_outputs->size() == 0) { // No bundled expected outputs, so we can't verify the method outputs. diff --git a/devtools/bundled_program/schema/targets.bzl b/devtools/bundled_program/schema/targets.bzl index 532a01e039e..1201458b42f 100644 --- a/devtools/bundled_program/schema/targets.bzl +++ b/devtools/bundled_program/schema/targets.bzl @@ -74,6 +74,7 @@ def define_common_targets(): visibility = [ "//executorch/devtools/bundled_program/...", "//executorch/extension/pybindings/...", + "//executorch/extension/module/...", ], exported_headers = { OUTPUT_BUNDLED_HEADER: ":{}[{}]".format(BUNDLED_GEN_RULE_NAME, OUTPUT_BUNDLED_HEADER), diff --git a/extension/module/bundled_module.cpp b/extension/module/bundled_module.cpp new file mode 100644 index 00000000000..083aef141a0 --- /dev/null +++ b/extension/module/bundled_module.cpp @@ -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 + +#include +#include +#include +#include + +namespace executorch { +namespace extension { + +namespace { +std::unique_ptr program_data_loader( + const void* bundled_program_ptr) { + auto bundled_program = + bundled_program_flatbuffer::GetBundledProgram(bundled_program_ptr); + // the program inside the bundled program + auto program = bundled_program->program(); + return std::make_unique(program->data(), program->size()); +} +} // namespace + +BundledModule::BundledModule( + const void* bundled_program_ptr, + std::unique_ptr memory_allocator, + std::unique_ptr temp_allocator, + std::unique_ptr event_tracer, + std::unique_ptr data_map_loader) + : Module( + program_data_loader(bundled_program_ptr), + std::move(memory_allocator), + std::move(temp_allocator), + std::move(event_tracer), + std::move(data_map_loader)), + bundled_program_ptr_(bundled_program_ptr) {} + +runtime::Result> BundledModule::from_file( + const std::string& file_path, + std::unique_ptr memory_allocator, + std::unique_ptr temp_allocator, + std::unique_ptr event_tracer, + std::unique_ptr data_map_loader) { + auto data_loader_result = FileDataLoader::from(file_path.c_str()); + if (!data_loader_result.ok()) { + return data_loader_result.error(); + } + + auto file_size_result = data_loader_result->size(); + if (!file_size_result.ok()) { + return file_size_result.error(); + } + + size_t file_size = file_size_result.get(); + auto file_data = std::make_unique(file_size); + auto buffer_result = + data_loader_result->load_into(0, file_size, {}, file_data.get()); + if (buffer_result != runtime::Error::Ok) { + return buffer_result; + } + + // Pass ownership of the data to BundledModule + auto bm = std::make_unique( + file_data.release(), + std::move(memory_allocator), + std::move(temp_allocator), + std::move(event_tracer), + std::move(data_map_loader)); + + bm->is_loaded_from_file_ = true; + + return bm; +} + +runtime::Result> BundledModule::execute( + const std::string& method_name, + const size_t testset_idx) { + ET_CHECK_OK_OR_RETURN_ERROR(load_method(method_name)); + auto& method = methods_.at(method_name).method; + + ET_CHECK_OK_OR_RETURN_ERROR( + executorch::BUNDLED_PROGRAM_NAMESPACE::load_bundled_input( + *method, bundled_program_ptr_, testset_idx)); + ET_CHECK_OK_OR_RETURN_ERROR(method->execute()); + + const auto outputs_size = method->outputs_size(); + std::vector outputs(outputs_size); + ET_CHECK_OK_OR_RETURN_ERROR( + method->get_outputs(outputs.data(), outputs_size)); + + return outputs; +} + +runtime::Error BundledModule::verify_method_outputs( + const std::string& method_name, + const size_t testset_idx, + double rtol, + double atol) { + ET_CHECK_OK_OR_RETURN_ERROR(load_method(method_name)); + auto& method = methods_.at(method_name).method; + return executorch::BUNDLED_PROGRAM_NAMESPACE::verify_method_outputs( + *method, bundled_program_ptr_, testset_idx, rtol, atol); +} + +} // namespace extension +} // namespace executorch diff --git a/extension/module/bundled_module.h b/extension/module/bundled_module.h new file mode 100644 index 00000000000..d254a2cdcb5 --- /dev/null +++ b/extension/module/bundled_module.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 + +namespace executorch { +namespace extension { + +/** + * A facade class for loading bundled programs and executing methods within + * them. + */ +class BundledModule : public Module { + public: + /** + * Constructs an instance with the bundled program buffer pointer. + * + * This constructor reads the program from bundled program buffer to load the + * module with data loader. The bundled program pointer is preserved so that + * the portion outside of program is accessible. + * + * @param[in] bundled_program_ptr A DataLoader used for loading program data. + * @param[in] memory_allocator A MemoryAllocator used for memory management. + * @param[in] temp_allocator A MemoryAllocator to use when allocating + * temporary data during kernel or delegate execution. + * @param[in] event_tracer A EventTracer used for tracking and logging events. + * @param[in] data_map_loader A DataLoader used for loading external weights. + */ + explicit BundledModule( + const void* bundled_program_ptr, + std::unique_ptr memory_allocator = nullptr, + std::unique_ptr temp_allocator = nullptr, + std::unique_ptr event_tracer = nullptr, + std::unique_ptr data_map_loader = nullptr); + + // Disallow copying + BundledModule(const BundledModule&) = delete; + BundledModule& operator=(const BundledModule&) = delete; + // Disallow copying + BundledModule(BundledModule&&) = delete; + BundledModule& operator=(BundledModule&&) = delete; + // Default destructor + ~BundledModule() { + if (is_loaded_from_file_) { + delete[] static_cast(bundled_program_ptr_); + } + } + + /** + * Constructs an instance by loading a bundled program from a file with + * specified memory locking behavior. + * + * @param[in] file_path The path to the ExecuTorch bundled program file to + * load. + * @param[in] memory_allocator A MemoryAllocator used for memory management. + * @param[in] temp_allocator A MemoryAllocator to use when allocating + * temporary data during kernel or delegate execution. + * @param[in] event_tracer A EventTracer used for tracking and logging events. + * @param[in] data_map_loader A DataLoader used for loading external weights. + */ + ET_NODISCARD static runtime::Result> from_file( + const std::string& file_path, + std::unique_ptr memory_allocator = nullptr, + std::unique_ptr temp_allocator = nullptr, + std::unique_ptr event_tracer = nullptr, + std::unique_ptr data_map_loader = nullptr); + + using Module::execute; + + /** + * Execute a specific method with the input value at the given `testset_idx` + * from the bundle to the method. Loads the program and method before + * executing if needed. + * + * This function is a wrapper of `load_bundled_input` in `bundled_program`. + * + * @param[in] method_name The name of the method to execute. + * @param[in] testset_idx The index of the input value to be passed to the + * method. + * + * @returns Return Error::Ok on a successful load, or the error happens during + * execution. + */ + ET_NODISCARD + runtime::Result> execute( + const std::string& method_name, + const size_t testset_idx); + + /** + * Verify the output of a specific method with the expected output from the + * program bundle at the given `testset_idx`. + * + * This function is a wrapper of `verify_method_outputs` in `bundled_program`. + * + * @param[in] method_name The name of the method to extract outputs from. + * @param[in] testset_idx The index of expected output needs to be compared. + * @param[in] rtol Relative tolerance used for data comparsion. + * @param[in] atol Absolute tolerance used for data comparsion. + * + * @returns Return Error::Ok if two outputs match, or the error happens during + * execution. + */ + ET_NODISCARD + runtime::Error verify_method_outputs( + const std::string& method_name, + const size_t testset_idx, + double rtol = 1e-5, + double atol = 1e-8); + + private: + const void* bundled_program_ptr_; + bool is_loaded_from_file_ = false; +}; + +} // namespace extension +} // namespace executorch diff --git a/extension/module/module.cpp b/extension/module/module.cpp index ec01323edc7..6c534b8d560 100644 --- a/extension/module/module.cpp +++ b/extension/module/module.cpp @@ -302,15 +302,5 @@ runtime::Error Module::set_output( output_tensor.mutable_data_ptr(), output_tensor.nbytes(), output_index); } -ET_NODISCARD inline runtime::Result Module::get_method( - const std::string& method_name) { - ET_CHECK_OR_RETURN_ERROR( - methods_.count(method_name) > 0, - InvalidArgument, - "no such method in program: %s", - method_name.c_str()); - return methods_[method_name].method.get(); -} - } // namespace extension } // namespace executorch diff --git a/extension/module/module.h b/extension/module/module.h index 201887b9ccc..0c4d4779bea 100644 --- a/extension/module/module.h +++ b/extension/module/module.h @@ -491,16 +491,6 @@ class Module { std::unique_ptr data_map_; protected: - /** - * Get a method by method name. - * - * @param[in] method_name The name of the method to get. - * - * @returns A Result object containing either a pointer to the requested - * method or an error to indicate failure. - */ - ET_NODISCARD inline runtime::Result get_method( - const std::string& method_name); std::unordered_map methods_; friend class ExecuTorchJni; diff --git a/extension/module/targets.bzl b/extension/module/targets.bzl index 78077df6387..3e449da5e14 100644 --- a/extension/module/targets.bzl +++ b/extension/module/targets.bzl @@ -31,3 +31,25 @@ def define_common_targets(): "//executorch/runtime/executor:program_no_prim_ops" + aten_suffix, ], ) + + runtime.cxx_library( + name = "bundled_module" + aten_suffix, + srcs = [ + "bundled_module.cpp", + ], + exported_headers = [ + "bundled_module.h", + ], + visibility = [ + "@EXECUTORCH_CLIENTS", + ], + deps = [ + "//executorch/extension/data_loader:buffer_data_loader", + "//executorch/extension/data_loader:file_data_loader", + "//executorch/devtools/bundled_program:runtime" + aten_suffix, + "//executorch/devtools/bundled_program/schema:bundled_program_schema_fbs", + ], + exported_deps = [ + "//executorch/extension/module:module" + aten_suffix, + ], + ) diff --git a/extension/module/test/bundled_module_test.cpp b/extension/module/test/bundled_module_test.cpp new file mode 100644 index 00000000000..a07c5dd5486 --- /dev/null +++ b/extension/module/test/bundled_module_test.cpp @@ -0,0 +1,110 @@ +/* + * 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 + +using namespace ::executorch::extension; +using namespace ::executorch::runtime; + +class BundledModuleTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + std::string resources_path; + if (const char* env = std::getenv("RESOURCES_PATH")) { + resources_path = env; + } + pte_path_ = std::getenv("ET_MODULE_PTE_PATH"); + bpte_path_ = resources_path + "/bundled_program.bpte"; + } + + static inline std::string bpte_path_; + static inline std::string pte_path_; +}; + +TEST_F(BundledModuleTest, TestExecute) { + auto bundled_module_output = BundledModule::from_file(bpte_path_.c_str()); + EXPECT_EQ(bundled_module_output.error(), Error::Ok); + auto& bundled_module = bundled_module_output.get(); + + auto outputs = bundled_module->execute("forward", /*testset_idx=*/0); + EXPECT_EQ(bundled_module->Module::is_loaded(), true); + EXPECT_EQ(outputs.error(), Error::Ok); + + auto status = + bundled_module->verify_method_outputs("forward", /*testset_idx=*/0); + EXPECT_EQ(status, Error::Ok); +} + +TEST_F(BundledModuleTest, TestNonExistBPFile) { + auto bundled_module_output = + BundledModule::from_file("/path/to/nonexistent/file.bpte"); + EXPECT_EQ(bundled_module_output.error(), Error::AccessFailed); +} + +TEST_F(BundledModuleTest, TestNonBPFile) { + auto bundled_module_output = BundledModule::from_file(pte_path_.c_str()); + EXPECT_EQ(bundled_module_output.error(), Error::Ok); + + auto& bundled_module = bundled_module_output.get(); + + auto outputs = bundled_module->execute("forward", /*testset_idx=*/0); + EXPECT_EQ(bundled_module->Module::is_loaded(), false); + EXPECT_EQ(outputs.error(), Error::InvalidArgument); + + auto status = + bundled_module->verify_method_outputs("forward", /*testset_idx=*/0); + EXPECT_EQ(status, Error::InvalidArgument); +} + +TEST_F(BundledModuleTest, TestExecuteInvalidMethod) { + auto bundled_module_output = BundledModule::from_file(bpte_path_.c_str()); + EXPECT_EQ(bundled_module_output.error(), Error::Ok); + auto& bundled_module = bundled_module_output.get(); + + auto outputs = + bundled_module->execute("non_existent_method", /*testset_idx=*/0); + EXPECT_EQ(outputs.error(), Error::InvalidArgument); +} + +TEST_F(BundledModuleTest, TestExecuteInvalidIdx) { + auto bundled_module_output = BundledModule::from_file(bpte_path_.c_str()); + EXPECT_EQ(bundled_module_output.error(), Error::Ok); + auto& bundled_module = bundled_module_output.get(); + + auto outputs = bundled_module->execute("forward", /*testset_idx=*/10000); + EXPECT_EQ(outputs.error(), Error::InvalidArgument); +} + +TEST_F(BundledModuleTest, TestVerifyInvalidMethod) { + auto bundled_module_output = BundledModule::from_file(bpte_path_.c_str()); + EXPECT_EQ(bundled_module_output.error(), Error::Ok); + auto& bundled_module = bundled_module_output.get(); + + auto outputs = bundled_module->execute("forward", /*testset_idx=*/0); + EXPECT_EQ(bundled_module->Module::is_loaded(), true); + EXPECT_EQ(outputs.error(), Error::Ok); + + auto status = bundled_module->verify_method_outputs( + "non_existent_method", /*testset_idx=*/0); + EXPECT_EQ(status, Error::InvalidArgument); +} + +TEST_F(BundledModuleTest, TestVerifyInvalidIdx) { + auto bundled_module_output = BundledModule::from_file(bpte_path_.c_str()); + EXPECT_EQ(bundled_module_output.error(), Error::Ok); + auto& bundled_module = bundled_module_output.get(); + + auto outputs = bundled_module->execute("forward", /*testset_idx=*/0); + EXPECT_EQ(bundled_module->Module::is_loaded(), true); + EXPECT_EQ(outputs.error(), Error::Ok); + + auto status = + bundled_module->verify_method_outputs("forward", /*testset_idx=*/10000); + EXPECT_EQ(status, Error::InvalidArgument); +} diff --git a/extension/module/test/resources/README.md b/extension/module/test/resources/README.md new file mode 100644 index 00000000000..026042ab121 --- /dev/null +++ b/extension/module/test/resources/README.md @@ -0,0 +1,7 @@ +## Resources + +### bundled_program.bpte + + ``` + python3 extension/module/test/resources/gen_bundled_program.py + ``` diff --git a/extension/module/test/resources/bundled_program.bpte b/extension/module/test/resources/bundled_program.bpte new file mode 100644 index 00000000000..ea5b42c03d0 Binary files /dev/null and b/extension/module/test/resources/bundled_program.bpte differ diff --git a/extension/module/test/resources/gen_bundled_program.py b/extension/module/test/resources/gen_bundled_program.py new file mode 100644 index 00000000000..f1fa0a4a7e3 --- /dev/null +++ b/extension/module/test/resources/gen_bundled_program.py @@ -0,0 +1,99 @@ +import torch + +from executorch.devtools import BundledProgram + +from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite +from executorch.devtools.bundled_program.serialize import ( + serialize_from_bundled_program_to_flatbuffer, +) + +from executorch.exir import to_edge_transform_and_lower +from torch.export import export, export_for_training + + +# Step 1: ExecuTorch Program Export +class SampleModel(torch.nn.Module): + """An example model with multi-methods. Each method has multiple input and single output""" + + def __init__(self) -> None: + super().__init__() + self.register_buffer("a", 3 * torch.ones(2, 2, dtype=torch.int32)) + self.register_buffer("b", 2 * torch.ones(2, 2, dtype=torch.int32)) + + def forward(self, x: torch.Tensor, q: torch.Tensor) -> torch.Tensor: + z = x.clone() + torch.mul(self.a, x, out=z) + y = x.clone() + torch.add(z, self.b, out=y) + torch.add(y, q, out=y) + return y + + +def main() -> None: + """Sample code to generate bundled program and save it to file. It is the same as in https://pytorch.org/executorch/0.6/bundled-io.html#emit-example""" + # Inference method name of SampleModel we want to bundle testcases to. + # Notices that we do not need to bundle testcases for every inference methods. + method_name = "forward" + model = SampleModel() + + # Inputs for graph capture. + capture_input = ( + (torch.rand(2, 2) - 0.5).to(dtype=torch.int32), + (torch.rand(2, 2) - 0.5).to(dtype=torch.int32), + ) + + # Export method's FX Graph. + method_graph = export( + export_for_training(model, capture_input).module(), + capture_input, + ) + + # Emit the traced method into ET Program. + et_program = to_edge_transform_and_lower(method_graph).to_executorch() + + # Step 2: Construct MethodTestSuite for Each Method + + # Prepare the Test Inputs. + + # Number of input sets to be verified + n_input = 10 + + # Input sets to be verified. + inputs = [ + # Each list below is a individual input set. + # The number of inputs, dtype and size of each input follow Program's spec. + [ + (torch.rand(2, 2) - 0.5).to(dtype=torch.int32), + (torch.rand(2, 2) - 0.5).to(dtype=torch.int32), + ] + for _ in range(n_input) + ] + + # Generate Test Suites + method_test_suites = [ + MethodTestSuite( + method_name=method_name, + test_cases=[ + MethodTestCase( + inputs=input, + expected_outputs=(getattr(model, method_name)(*input),), + ) + for input in inputs + ], + ), + ] + + # Step 3: Generate BundledProgram + bundled_program = BundledProgram(et_program, method_test_suites) + + # Step 4: Serialize BundledProgram to flatbuffer. + serialized_bundled_program = serialize_from_bundled_program_to_flatbuffer( + bundled_program + ) + save_path = "bundled_program.bpte" + with open(save_path, "wb") as f: + f.write(serialized_bundled_program) + + +if __name__ == "__main__": + main() diff --git a/extension/module/test/targets.bzl b/extension/module/test/targets.bzl index e308ca89c30..e09b43e356d 100644 --- a/extension/module/test/targets.bzl +++ b/extension/module/test/targets.bzl @@ -42,3 +42,30 @@ def define_common_targets(is_fbcode=False): "-Wno-error=deprecated-declarations", ], ) + + runtime.cxx_test( + name = "bundled_test" + aten_suffix, + srcs = [ + "bundled_module_test.cpp", + ], + deps = [ + "//executorch/kernels/portable:generated_lib" + aten_suffix, + "//executorch/extension/module:bundled_module" + aten_suffix, + "//executorch/extension/tensor:tensor" + aten_suffix, + ], + env = { + "RESOURCES_PATH": "$(location :resources)/resources", + "ET_MODULE_PTE_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleAdd.pte])", + }, + platforms = [CXX, ANDROID], # Cannot bundle resources on Apple platform. + compiler_flags = [ + "-Wno-error=deprecated-declarations", + ], + ) + + runtime.filegroup( + name = "resources", + srcs = native.glob([ + "resources/**", + ]), + )