4/x: Add LoadBackendOptionsMap support to Module::load()

metascroy · web-flow · commit 4528ae26e297 · 2026-02-27T03:06:10.000Z
Differential Revision: D92358607 Pull Request resolved: #17687
diff --git a/extension/module/module.cpp b/extension/module/module.cpp
@@ -176,6 +176,17 @@ Module::Module(
 }
 
 runtime::Error Module::load(const Program::Verification verification) {
+  return load_internal(verification);
+}
+
+runtime::Error Module::load(
+    const LoadBackendOptionsMap& backend_options,
+    const Program::Verification verification) {
+  backend_options_ = &backend_options;
+  return load_internal(verification);
+}
+
+runtime::Error Module::load_internal(const Program::Verification verification) {
   if (!is_loaded()) {
     if (!data_loader_) {
       auto data_loader_result = make_data_loader(file_path_, load_mode_);
@@ -256,10 +267,15 @@ runtime::Result<std::unordered_set<std::string>> Module::method_names() {
 runtime::Error Module::load_method(
     const std::string& method_name,
     runtime::HierarchicalAllocator* planned_memory,
-    torch::executor::EventTracer* event_tracer) {
+    torch::executor::EventTracer* event_tracer,
+    const LoadBackendOptionsMap* backend_options) {
   if (!is_method_loaded(method_name)) {
     ET_CHECK_OK_OR_RETURN_ERROR(load());
 
+    // Use passed backend_options, or fall back to stored one from load()
+    const LoadBackendOptionsMap* effective_backend_options =
+        backend_options ? backend_options : backend_options_;
+
     MethodHolder method_holder;
 
     if (!planned_memory) {
@@ -292,7 +308,8 @@ runtime::Error Module::load_method(
         method_name.c_str(),
         method_holder.memory_manager.get(),
         event_tracer ? event_tracer : this->event_tracer(),
-        merged_data_map_.get());
+        merged_data_map_.get(),
+        effective_backend_options);
     if (!res_method.ok()) {
       return res_method.error();
     }
diff --git a/extension/module/module.h b/extension/module/module.h
@@ -29,6 +29,7 @@ using ET_RUNTIME_NAMESPACE::Method;
 using ET_RUNTIME_NAMESPACE::MethodMeta;
 using ET_RUNTIME_NAMESPACE::NamedDataMap;
 using ET_RUNTIME_NAMESPACE::Program;
+using runtime::LoadBackendOptionsMap;
 
 class ExecuTorchJni;
 
@@ -153,6 +154,22 @@ class Module {
       const Program::Verification verification =
           Program::Verification::Minimal);
 
+  /**
+   * Loads the program with per-delegate runtime options.
+   *
+   * @param[in] backend_options A LoadBackendOptionsMap containing per-delegate
+   * load-time configuration options. The caller must ensure this object
+   * outlives any methods loaded with these options.
+   * @param[in] verification The type of verification to do before returning
+   * success.
+   *
+   * @returns An Error to indicate success or failure of the loading process.
+   */
+  ET_NODISCARD virtual runtime::Error load(
+      const LoadBackendOptionsMap& backend_options,
+      const Program::Verification verification =
+          Program::Verification::Minimal);
+
   /**
    * Checks if the program is loaded.
    *
@@ -207,12 +224,13 @@ class Module {
   runtime::Error load_method(
       const std::string& method_name,
       runtime::HierarchicalAllocator* planned_memory = nullptr,
-      torch::executor::EventTracer* event_tracer = nullptr);
+      torch::executor::EventTracer* event_tracer = nullptr,
+      const LoadBackendOptionsMap* backend_options = nullptr);
 
   ET_DEPRECATED ET_NODISCARD runtime::Error inline load_method(
       const std::string& method_name,
       torch::executor::EventTracer* event_tracer) {
-    return load_method(method_name, nullptr, event_tracer);
+    return load_method(method_name, nullptr, event_tracer, nullptr);
   }
 
   /**
@@ -254,13 +272,15 @@ class Module {
    */
   ET_NODISCARD inline runtime::Error load_forward(
       runtime::HierarchicalAllocator* planned_memory = nullptr,
-      torch::executor::EventTracer* event_tracer = nullptr) {
-    return load_method("forward", planned_memory, event_tracer);
+      torch::executor::EventTracer* event_tracer = nullptr,
+      const LoadBackendOptionsMap* backend_options = nullptr) {
+    return load_method(
+        "forward", planned_memory, event_tracer, backend_options);
   }
 
   ET_DEPRECATED ET_NODISCARD inline runtime::Error load_forward(
       torch::executor::EventTracer* event_tracer) {
-    return load_forward(nullptr, event_tracer);
+    return load_forward(nullptr, event_tracer, nullptr);
   }
 
   /**
@@ -650,6 +670,10 @@ class Module {
   std::vector<std::unique_ptr<NamedDataMap>> named_data_maps_;
   std::unique_ptr<NamedDataMap> merged_data_map_;
   ET_DEPRECATED std::vector<uint8_t> debug_buffer_;
+  const LoadBackendOptionsMap* backend_options_ = nullptr;
+
+  ET_NODISCARD runtime::Error load_internal(
+      const Program::Verification verification);
 
  protected:
   std::unordered_map<std::string, MethodHolder> methods_;
diff --git a/extension/module/test/module_test.cpp b/extension/module/test/module_test.cpp
@@ -15,6 +15,8 @@
 
 #include <executorch/extension/data_loader/file_data_loader.h>
 #include <executorch/extension/tensor/tensor.h>
+#include <executorch/runtime/backend/backend_options_map.h>
+#include <executorch/runtime/backend/options.h>
 #include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
 
 using namespace ::executorch::extension;
@@ -554,3 +556,165 @@ TEST_F(ModuleTest, TestPTD_Multiple) {
   auto tensor2 = make_tensor_ptr({3}, {2.f, 3.f, 4.f});
   ASSERT_EQ(module_linear.forward(tensor2).error(), Error::Ok);
 }
+
+// =============================================================================
+// LoadBackendOptionsMap / RuntimeSpec Tests
+// =============================================================================
+
+TEST_F(ModuleTest, TestLoadWithLoadBackendOptionsMap) {
+  Module module(model_path_);
+
+  // Create a LoadBackendOptionsMap with some options for a hypothetical backend
+  LoadBackendOptionsMap backend_options;
+  BackendOptions<4> options;
+  options.set_option("compute_unit", "cpu_only");
+  options.set_option("debug_mode", true);
+  ASSERT_EQ(
+      backend_options.set_options("TestBackend", options.view()), Error::Ok);
+
+  // Load with backend options - should succeed even though the model
+  // doesn't use this backend (options are simply passed through)
+  EXPECT_FALSE(module.is_loaded());
+  const auto error = module.load(backend_options);
+  EXPECT_EQ(error, Error::Ok);
+  EXPECT_TRUE(module.is_loaded());
+}
+
+TEST_F(ModuleTest, TestLoadWithLoadBackendOptionsMapThenExecute) {
+  Module module(model_path_);
+
+  LoadBackendOptionsMap backend_options;
+  BackendOptions<2> options;
+  options.set_option("key1", "value1");
+  ASSERT_EQ(
+      backend_options.set_options("SomeBackend", options.view()), Error::Ok);
+
+  const auto load_error = module.load(backend_options);
+  EXPECT_EQ(load_error, Error::Ok);
+
+  // Execute should work normally
+  auto tensor = make_tensor_ptr({2, 2}, {1.f, 2.f, 3.f, 4.f});
+  const auto result = module.execute("forward", {tensor, tensor, 1.0});
+  EXPECT_EQ(result.error(), Error::Ok);
+
+  const auto expected = make_tensor_ptr({2, 2}, {2.f, 4.f, 6.f, 8.f});
+  EXPECT_TENSOR_CLOSE(result->at(0).toTensor(), *expected.get());
+}
+
+TEST_F(ModuleTest, TestLoadMethodWithLoadBackendOptionsMap) {
+  Module module(model_path_);
+
+  LoadBackendOptionsMap backend_options;
+  BackendOptions<2> options;
+  options.set_option("option1", 42);
+  ASSERT_EQ(
+      backend_options.set_options("AnotherBackend", options.view()), Error::Ok);
+
+  EXPECT_FALSE(module.is_method_loaded("forward"));
+  const auto error =
+      module.load_method("forward", nullptr, nullptr, &backend_options);
+  EXPECT_EQ(error, Error::Ok);
+  EXPECT_TRUE(module.is_method_loaded("forward"));
+  EXPECT_TRUE(module.is_loaded());
+}
+
+TEST_F(ModuleTest, TestLoadForwardWithLoadBackendOptionsMap) {
+  Module module(model_path_);
+
+  LoadBackendOptionsMap backend_options;
+  BackendOptions<2> options;
+  options.set_option("setting", "enabled");
+  ASSERT_EQ(
+      backend_options.set_options("ForwardBackend", options.view()), Error::Ok);
+
+  EXPECT_FALSE(module.is_method_loaded("forward"));
+  const auto error = module.load_forward(nullptr, nullptr, &backend_options);
+  EXPECT_EQ(error, Error::Ok);
+  EXPECT_TRUE(module.is_method_loaded("forward"));
+}
+
+TEST_F(ModuleTest, TestLoadWithEmptyLoadBackendOptionsMap) {
+  Module module(model_path_);
+
+  // Empty LoadBackendOptionsMap should work fine
+  LoadBackendOptionsMap backend_options;
+
+  const auto error = module.load(backend_options);
+  EXPECT_EQ(error, Error::Ok);
+  EXPECT_TRUE(module.is_loaded());
+}
+
+TEST_F(ModuleTest, TestLoadBackendOptionsMapPersistedAcrossLoadMethod) {
+  Module module(model_path_);
+
+  // Set backend options via load()
+  LoadBackendOptionsMap backend_options;
+  BackendOptions<2> options;
+  options.set_option("persist_test", true);
+  ASSERT_EQ(
+      backend_options.set_options("PersistBackend", options.view()), Error::Ok);
+
+  const auto load_error = module.load(backend_options);
+  EXPECT_EQ(load_error, Error::Ok);
+
+  // load_method without explicit backend_options should use the stored ones
+  const auto method_error = module.load_method("forward");
+  EXPECT_EQ(method_error, Error::Ok);
+  EXPECT_TRUE(module.is_method_loaded("forward"));
+}
+
+TEST_F(ModuleTest, TestLoadMethodOverridesStoredBackendOptions) {
+  Module module(model_path_);
+
+  // Set initial backend options via load()
+  LoadBackendOptionsMap initial_options;
+  BackendOptions<2> opts1;
+  opts1.set_option("source", "load");
+  ASSERT_EQ(
+      initial_options.set_options("TestBackend", opts1.view()), Error::Ok);
+
+  const auto load_error = module.load(initial_options);
+  EXPECT_EQ(load_error, Error::Ok);
+
+  // Unload and reload with different options passed to load_method
+  module.unload_method("forward");
+
+  LoadBackendOptionsMap override_options;
+  BackendOptions<2> opts2;
+  opts2.set_option("source", "load_method");
+  ASSERT_EQ(
+      override_options.set_options("TestBackend", opts2.view()), Error::Ok);
+
+  // The override_options should take precedence
+  const auto method_error =
+      module.load_method("forward", nullptr, nullptr, &override_options);
+  EXPECT_EQ(method_error, Error::Ok);
+  EXPECT_TRUE(module.is_method_loaded("forward"));
+}
+
+TEST_F(ModuleTest, TestMultipleBackendsInOptionsMap) {
+  Module module(model_path_);
+
+  LoadBackendOptionsMap backend_options;
+
+  // Add options for multiple backends
+  BackendOptions<2> backend1_opts;
+  backend1_opts.set_option("compute_unit", "gpu");
+  ASSERT_EQ(
+      backend_options.set_options("Backend1", backend1_opts.view()), Error::Ok);
+
+  BackendOptions<3> backend2_opts;
+  backend2_opts.set_option("optimization_level", 3);
+  backend2_opts.set_option("debug", false);
+  ASSERT_EQ(
+      backend_options.set_options("Backend2", backend2_opts.view()), Error::Ok);
+
+  const auto error = module.load(backend_options);
+  EXPECT_EQ(error, Error::Ok);
+  EXPECT_TRUE(module.is_loaded());
+
+  // Should still execute normally
+  auto tensor = make_tensor_ptr({2, 2}, {1.f, 2.f, 3.f, 4.f});
+  const auto result = module.forward({tensor, tensor, 1.0});
+  EXPECT_EQ(result.error(), Error::Ok);
+}
diff --git a/extension/module/test/targets.bzl b/extension/module/test/targets.bzl
@@ -37,6 +37,8 @@ def define_common_targets(is_fbcode=False):
                     "//executorch/extension/data_loader:file_data_loader",
                     "//executorch/extension/module:module" + aten_suffix,
                     "//executorch/extension/tensor:tensor" + aten_suffix,
+                    "//executorch/runtime/backend:backend_options",
+                    "//executorch/runtime/backend:backend_options_map",
                     "//executorch/runtime/core/exec_aten/testing_util:tensor_util" + aten_suffix,
                 ],
                 env = modules_env,
