From 5feb8d5ebbd4c7513890d976e7bda3f45de59d44 Mon Sep 17 00:00:00 2001 From: zhenyanzhang Date: Thu, 24 Apr 2025 14:29:31 -0700 Subject: [PATCH 1/3] [ExecuTorch][#10375] Add `extension.BundledModule` to Wrap `extension.Module` with Bundled Program Logic https://github.com/pytorch/executorch/issues/10375 # Context This issue is a step of https://github.com/pytorch/executorch/discussions/9638. In https://github.com/pytorch/executorch/discussions/9638, we want to have `extension.Module` as the single source of implementation in `pybindings`, which means that `pybindings.PyModule` should use `extension.Module` rather than its own `pybindings.Module`. The issue is that `pybindings.PyModule` is dependent on the `method` getter from `pybindings.Module`, which `extension.Module` do not have. Since we don't want to expose `method` getter in `extension.Module`, we have to protect the getter, wrap the functions that is dependent on it and use the protected getter there, ultimately decouple `pybindings` from a `method` getter. # Proposal Now that we have a protected `method` getter, we can introduce a `extension.BundledModule`, a child class inheriting `extension.Module` which wraps up bundled program logic that is dependent on the `method` getter. Differential Revision: [D73564125](https://our.internmc.facebook.com/intern/diff/D73564125/) [ghstack-poisoned] --- devtools/bundled_program/schema/targets.bzl | 1 + extension/module/module.cpp | 62 ++++++++++++++-- extension/module/module.h | 82 +++++++++++++++++++-- extension/module/targets.bzl | 3 + 4 files changed, 132 insertions(+), 16 deletions(-) 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/module.cpp b/extension/module/module.cpp index ec01323edc7..d2917a54de6 100644 --- a/extension/module/module.cpp +++ b/extension/module/module.cpp @@ -8,6 +8,9 @@ #include +#include +#include +#include #include #include #include @@ -302,14 +305,57 @@ 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 { +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::Error BundledModule::load_bundled_input( + 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; + auto& inputs = methods_.at(method_name).inputs; + + auto status = executorch::BUNDLED_PROGRAM_NAMESPACE::load_bundled_input( + *method, bundled_program_ptr_, testset_idx); + ET_CHECK_OK_OR_RETURN_ERROR( + status, + "Bundled Program's load_bundled_input failed with status 0x%" PRIx32, + static_cast(status)); + + return method->get_inputs(inputs.data(), inputs.size()); +} + +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 diff --git a/extension/module/module.h b/extension/module/module.h index 8cedb79c06e..dd0b3d8d0f4 100644 --- a/extension/module/module.h +++ b/extension/module/module.h @@ -493,19 +493,85 @@ class Module { std::unique_ptr data_map_; protected: + std::unordered_map methods_; + + friend class ExecuTorchJni; +}; + +/** + * A facade class for loading bundled programs and executing methods within + * them. + */ +class BundledModule : public Module { + public: /** - * Get a method by method name. + * Constructs an instance with the bundled program buffer pointer. * - * @param[in] method_name The name of the method to get. + * 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. * - * @returns A Result object containing either a pointer to the requested - * method or an error to indicate failure. + * @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. */ - ET_NODISCARD inline runtime::Result get_method( - const std::string& method_name); - std::unordered_map methods_; + 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); - friend class ExecuTorchJni; + BundledModule(const BundledModule&) = delete; + BundledModule& operator=(const BundledModule&) = delete; + BundledModule(Module&&) = delete; + BundledModule& operator=(BundledModule&&) = delete; + + /** + * Execute a specific method with the input value at the given testset_idx + * from the program bundle. + * + * Before execution, this function loads the program and method with + * 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 A Result object containing either a vector of output values + * from the method or an error to indicate failure. + */ + ET_NODISCARD + runtime::Error load_bundled_input( + 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_; }; } // namespace extension diff --git a/extension/module/targets.bzl b/extension/module/targets.bzl index d8019ce9c4e..a4ed383cb6b 100644 --- a/extension/module/targets.bzl +++ b/extension/module/targets.bzl @@ -22,6 +22,9 @@ def define_common_targets(): "@EXECUTORCH_CLIENTS", ], deps = [ + "//executorch/extension/data_loader:buffer_data_loader", + "//executorch/devtools/bundled_program:runtime", + "//executorch/devtools/bundled_program/schema:bundled_program_schema_fbs", "//executorch/extension/memory_allocator:malloc_memory_allocator", "//executorch/extension/data_loader:file_data_loader", "//executorch/extension/data_loader:mmap_data_loader", From f6ba0bd64f2ccc8107dbaec8379d0001ee2d5248 Mon Sep 17 00:00:00 2001 From: zhenyanzhang Date: Thu, 24 Apr 2025 14:52:30 -0700 Subject: [PATCH 2/3] Update on "[ExecuTorch][#10375] Add `extension.BundledModule` to Wrap `extension.Module` with Bundled Program Logic" https://github.com/pytorch/executorch/issues/10375 # Context This issue is a step of https://github.com/pytorch/executorch/discussions/9638. In https://github.com/pytorch/executorch/discussions/9638, we want to have `extension.Module` as the single source of implementation in `pybindings`, which means that `pybindings.PyModule` should use `extension.Module` rather than its own `pybindings.Module`. The issue is that `pybindings.PyModule` is dependent on the `method` getter from `pybindings.Module`, which `extension.Module` do not have. Since we don't want to expose `method` getter in `extension.Module`, we have to protect the getter, wrap the functions that is dependent on it and use the protected getter there, ultimately decouple `pybindings` from a `method` getter. # Proposal Now that we have a protected `method` getter, we can introduce a `extension.BundledModule`, a child class inheriting `extension.Module` which wraps up bundled program logic that is dependent on the `method` getter. Differential Revision: [D73564125](https://our.internmc.facebook.com/intern/diff/D73564125/) [ghstack-poisoned] --- extension/module/module.h | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/extension/module/module.h b/extension/module/module.h index dd0b3d8d0f4..f4f8e29b267 100644 --- a/extension/module/module.h +++ b/extension/module/module.h @@ -531,18 +531,17 @@ class BundledModule : public Module { BundledModule& operator=(BundledModule&&) = delete; /** - * Execute a specific method with the input value at the given testset_idx - * from the program bundle. + * Load a specific method with the input value at the given `testset_idx` from + * the bundle to the method. * - * Before execution, this function loads the program and method with - * load_bundled_input in bundled_program. + * 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. + * @param[in] testset_idx The index of the input value to be passed to the + * method. * - * @returns A Result object containing either a vector of output values - * from the method or an error to indicate failure. + * @returns Return Error::Ok on a successful load, or the error happens during + * execution. */ ET_NODISCARD runtime::Error load_bundled_input( @@ -551,9 +550,9 @@ class BundledModule : public Module { /** * Verify the output of a specific method with the expected output from the - * program bundle at the given testset_idx. + * program bundle at the given `testset_idx`. * - * This function is a wrapper of verify_method_outputs in bundled_program. + * 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. From 7fdd05d82a6d12dc8bb99acf2c06511dd8fd8ba4 Mon Sep 17 00:00:00 2001 From: zhenyanzhang Date: Thu, 24 Apr 2025 16:24:55 -0700 Subject: [PATCH 3/3] Update on "[ExecuTorch][#10375] Add `extension.BundledModule` to Wrap `extension.Module` with Bundled Program Logic" https://github.com/pytorch/executorch/issues/10375 # Context This issue is a step of https://github.com/pytorch/executorch/discussions/9638. In https://github.com/pytorch/executorch/discussions/9638, we want to have `extension.Module` as the single source of implementation in `pybindings`, which means that `pybindings.PyModule` should use `extension.Module` rather than its own `pybindings.Module`. The issue is that `pybindings.PyModule` is dependent on the `method` getter from `pybindings.Module`, which `extension.Module` do not have. Since we don't want to expose `method` getter in `extension.Module`, we have to protect the getter, wrap the functions that is dependent on it and use the protected getter there, ultimately decouple `pybindings` from a `method` getter. # Proposal Now that we have a protected `method` getter, we can introduce a `extension.BundledModule`, a child class inheriting `extension.Module` which wraps up bundled program logic that is dependent on the `method` getter. Differential Revision: [D73564125](https://our.internmc.facebook.com/intern/diff/D73564125/) [ghstack-poisoned] --- extension/module/targets.bzl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/extension/module/targets.bzl b/extension/module/targets.bzl index a4ed383cb6b..73d59e7547e 100644 --- a/extension/module/targets.bzl +++ b/extension/module/targets.bzl @@ -23,7 +23,7 @@ def define_common_targets(): ], deps = [ "//executorch/extension/data_loader:buffer_data_loader", - "//executorch/devtools/bundled_program:runtime", + "//executorch/devtools/bundled_program:runtime" + aten_suffix, "//executorch/devtools/bundled_program/schema:bundled_program_schema_fbs", "//executorch/extension/memory_allocator:malloc_memory_allocator", "//executorch/extension/data_loader:file_data_loader",