[6/N] Add update in module

Pull Request resolved: pytorch/executorch#11533

- Expose the update API in module. Inside the module. invoke the method. Update API
- Did a bit refactor for the `StubBackend` test code, such that it can shared between `backend_interface_update_test.cpp` and `module_test.cpp`
ghstack-source-id: 290372284
@exported-using-ghexport

Differential Revision: [D76172680](https://our.internmc.facebook.com/intern/diff/D76172680/)

Author: cccclai

extension/module/module.cpp

@@ -309,6 +309,13 @@ runtime::Error Module::set_output(
       output_tensor.mutable_data_ptr(), output_tensor.nbytes(), output_index);
 }
 
-} // namespace ET_MODULE_NAMESPACE
+runtime::Error Module::update(
+    const std::string& method_name,
+    runtime::ArrayRef<runtime::Entry> backend_options) {
+  ET_CHECK_OK_OR_RETURN_ERROR(load_method(method_name));
+  auto& method = methods_.at(method_name).method;
+  return method->update(backend_options);
+}
+
 } // namespace extension
 } // namespace executorch

extension/module/module.h

@@ -466,6 +466,32 @@ class Module {
     return set_output("forward", std::move(output_value), output_index);
   }
 
+  /**
+   * EXPERIMENTAL: Updates backend options for a specific method.
+   * Loads the program and method before updating if needed.
+   *
+   * @param[in] method_name The name of the method to update.
+   * @param[in] backend_options The backend options to update the method with.
+   *
+   * @returns An Error to indicate success or failure.
+   */
+  ET_EXPERIMENTAL ET_NODISCARD runtime::Error update(
+      const std::string& method_name,
+      runtime::ArrayRef<runtime::Entry> backend_options);
+
+  /**
+   * EXPERIMENTAL: Updates backend options for the 'forward' method.
+   * Loads the program and method before updating if needed.
+   *
+   * @param[in] backend_options The backend options to update the method with.
+   *
+   * @returns An Error to indicate success or failure.
+   */
+  ET_EXPERIMENTAL ET_NODISCARD inline runtime::Error update(
+      runtime::ArrayRef<runtime::Entry> backend_options) {
+    return update("forward", backend_options);
+  }
+
   /**
    * Retrieves the EventTracer instance being used by the Module.
    * EventTracer is used for tracking and logging events during the execution

extension/module/test/module_test.cpp

@@ -16,21 +16,32 @@
 #include <executorch/extension/data_loader/file_data_loader.h>
 #include <executorch/extension/tensor/tensor.h>
 #include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
+#include <executorch/runtime/backend/backend_options.h>
+#include <executorch/runtime/backend/backend_options_map.h>
+#include <executorch/runtime/executor/test/stub_backend.h>
 
 using namespace ::executorch::extension;
 using namespace ::executorch::runtime;
+using executorch::runtime::BackendOptions;
+using executorch::runtime::Entry;
+using executorch::runtime::IntKey;
 
 class ModuleTest : public ::testing::Test {
  protected:
   static void SetUpTestSuite() {
     model_path_ = std::getenv("ET_MODULE_ADD_PATH");
     add_mul_path_ = std::getenv("ET_MODULE_ADD_MUL_PROGRAM_PATH");
     add_mul_data_path_ = std::getenv("ET_MODULE_ADD_MUL_DATA_PATH");
+    stub_model_path_ = std::getenv("ET_MODULE_ADD_MUL_DELEGATED_PATH");
+    
+    // Register the StubBackend for testing
+    StubBackend::register_singleton();
   }
 
   static inline std::string model_path_;
   static inline std::string add_mul_path_;
   static inline std::string add_mul_data_path_;
+  static inline std::string stub_model_path_;
 };
 
 TEST_F(ModuleTest, TestLoad) {
@@ -466,3 +477,34 @@ TEST_F(ModuleTest, TestPTD) {
   auto tensor = make_tensor_ptr({2, 2}, {2.f, 3.f, 4.f, 2.f});
   ASSERT_EQ(module.forward(tensor).error(), Error::Ok);
 }
+
+TEST_F(ModuleTest, TestUpdate) {
+  Module module(stub_model_path_);
+
+  BackendOptionsMap<3> map;
+  BackendOptions<1> backend_options;
+  int new_num_threads = 4;
+  backend_options.set_option(IntKey("NumberOfThreads"), new_num_threads);
+  map.add("StubBackend", backend_options.view());
+
+  // Test update method with specific method name
+  const auto update_result = module.update("forward", map.entries());
+  EXPECT_EQ(update_result, Error::Ok);
+
+  ASSERT_EQ(StubBackend::singleton().num_threads(), new_num_threads);
+  
+}
+
+TEST_F(ModuleTest, TestUpdateNonExistentMethod) {
+  Module module(stub_model_path_);
+
+  BackendOptionsMap<3> map;
+  BackendOptions<1> backend_options;
+  int new_num_threads = 4;
+  backend_options.set_option(IntKey("NumberOfThreads"), new_num_threads);
+  map.add("StubBackend", backend_options.view());
+  
+  // Test update method with non-existent method name
+  const auto update_result = module.update("nonexistent", map.entries());
+  EXPECT_NE(update_result, Error::Ok);
+}

extension/module/test/targets.bzl

@@ -19,6 +19,7 @@ def define_common_targets(is_fbcode=False):
             "ET_MODULE_ADD_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleAdd.pte])",
             "ET_MODULE_ADD_MUL_PROGRAM_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleAddMul.pte])",
             "ET_MODULE_ADD_MUL_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleAddMul.ptd])",
+            "ET_MODULE_ADD_MUL_DELEGATED_PATH": "$(location fbcode//executorch/test/models:exported_delegated_add_mul[ModuleAddMul.pte])",
         }
 
         for aten_mode in get_aten_mode_options():
@@ -35,6 +36,7 @@ def define_common_targets(is_fbcode=False):
                     "//executorch/extension/module:module" + aten_suffix,
                     "//executorch/extension/tensor:tensor" + aten_suffix,
                     "//executorch/runtime/core/exec_aten/testing_util:tensor_util" + aten_suffix,
+                    "//executorch/runtime/executor/test:stub_backend",
                 ],
                 env = modules_env,
                 platforms = [CXX, ANDROID],  # Cannot bundle resources on Apple platform.

runtime/executor/method.cpp

@@ -1513,7 +1513,8 @@ Error Method::experimental_step() {
   return step();
 }
 
-Error Method::update(executorch::runtime::ArrayRef<executorch::runtime::Entry> backend_option) {
+Error Method::update(
+    executorch::runtime::ArrayRef<executorch::runtime::Entry> backend_option) {
   for (const auto& entry : backend_option) {
     const char* backend_name = entry.backend_name;
     auto backend_options = entry.options;

runtime/executor/method.h

@@ -241,13 +241,14 @@ class Method final {
   /// DEPRECATED: Use `reset_execution()` instead.
   ET_DEPRECATED ET_NODISCARD Error experimental_reset_execution();
 
-    /**
+  /**
    * EXPERIMENTAL: Update backend options, which will be dispatched to different backends.
    *
    * @retval Error::Ok step succeeded
    * @retval non-Ok Method update fails
    */
-   ET_EXPERIMENTAL ET_NODISCARD Error update(executorch::runtime::ArrayRef<executorch::runtime::Entry> backend_option);
+  ET_EXPERIMENTAL ET_NODISCARD Error update(
+      executorch::runtime::ArrayRef<executorch::runtime::Entry> backend_option);
 
   /**
    * Returns the MethodMeta that corresponds to the calling Method.