Module deep-copies LoadBackendOptionsMap on load (pytorch#19673)

Differential Revision: D105100372

Pull Request resolved: pytorch#19673

Author: metascroy

extension/module/module.cpp

@@ -203,8 +203,49 @@ runtime::Error Module::load(const Program::Verification verification) {
 runtime::Error Module::load(
     const LoadBackendOptionsMap& backend_options,
     const Program::Verification verification) {
-  backend_options_ = &backend_options;
-  return load_internal(verification);
+  // load_internal does not read backend options, so run it first; on
+  // failure we skip the deep-copy work entirely and leave the prior
+  // installed options (if any) in place.
+  ET_CHECK_OK_OR_RETURN_ERROR(load_internal(verification));
+
+  // Deep-copy the input into local storage so the Module owns the
+  // BackendOption arrays for the lifetime of any methods loaded with
+  // these options. Build BOTH the storage and the map in locals so any
+  // mid-loop failure (or exception from emplace) leaves the prior
+  // installed state untouched -- the two members are only committed
+  // together at the end on full success.
+  //
+  // local_storage is reserve()'d up front, so emplace_back() never
+  // reallocates the outer buffer and the inner vectors keep stable
+  // addresses while we build local_map. The final move of
+  // local_storage into backend_options_storage_ uses std::vector's
+  // O(1) buffer transfer, so the heap buffers that local_map's spans
+  // point into remain valid after the move; the static_assert documents
+  // the inner-vector property we rely on for that span stability.
+  static_assert(
+      std::is_nothrow_move_constructible_v<std::vector<runtime::BackendOption>>,
+      "Moving local_storage must not move-construct the inner vectors; "
+      "local_map's spans reference their heap buffers.");
+
+  std::vector<std::vector<runtime::BackendOption>> local_storage;
+  local_storage.reserve(backend_options.size());
+  LoadBackendOptionsMap local_map;
+  for (size_t i = 0; i < backend_options.size(); ++i) {
+    const auto entry = backend_options.entry_at(i);
+    local_storage.emplace_back(entry.options.begin(), entry.options.end());
+    auto& owned = local_storage.back();
+    // The input map was already valid, so set_options should not fail
+    // here; assert it loudly rather than leaving partial state behind.
+    ET_CHECK_OK_OR_RETURN_ERROR(local_map.set_options(
+        entry.backend_id,
+        runtime::Span<runtime::BackendOption>(owned.data(), owned.size())));
+  }
+
+  // Single commit point: both members updated together.
+  backend_options_storage_ = std::move(local_storage);
+  backend_options_map_ = std::move(local_map);
+
+  return runtime::Error::Ok;
 }
 
 runtime::Error Module::load_internal(const Program::Verification verification) {
@@ -369,9 +410,12 @@ runtime::Error Module::load_method(
   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_;
+    // Use passed backend_options, or fall back to stored ones from load().
+    // An empty stored map behaves identically to nullptr downstream, so we
+    // only forward the stored map when it actually has entries.
+    const LoadBackendOptionsMap* effective_backend_options = backend_options
+        ? backend_options
+        : (backend_options_map_.size() > 0 ? &backend_options_map_ : nullptr);
 
     MethodHolder method_holder;
 

extension/module/module.h

@@ -14,6 +14,8 @@
 #include <unordered_set>
 #include <vector>
 
+#include <executorch/runtime/backend/backend_options_map.h>
+#include <executorch/runtime/backend/options.h>
 #include <executorch/runtime/executor/program.h>
 
 #ifdef USE_ATEN_LIB
@@ -187,9 +189,18 @@ class Module {
   /**
    * 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.
+   * The Module deep-copies `backend_options` into internal storage, so the
+   * caller may release the input (and any backing BackendOption arrays its
+   * Spans referenced) immediately after this call returns. Future lazy
+   * `load_method` calls (e.g. triggered by `forward`) consume the
+   * Module-owned copy.
+   *
+   * Transactional: on failure, the previously-installed backend options
+   * (if any) are left in place; the input is not committed.
+   *
+   * @param[in] backend_options A LoadBackendOptionsMap containing
+   * per-delegate load-time configuration options. Deep-copied into the
+   * Module on success; not retained on failure.
    * @param[in] verification The type of verification to do before returning
    * success.
    *
@@ -200,6 +211,21 @@ class Module {
       const Program::Verification verification =
           Program::Verification::Minimal);
 
+  /**
+   * Returns the deep-copied LoadBackendOptionsMap most recently installed
+   * via `load(LoadBackendOptionsMap, ...)`. The returned reference is owned
+   * by the Module and remains valid until the next call to
+   * `load(LoadBackendOptionsMap, ...)` or until the Module is destroyed.
+   *
+   * If `load(LoadBackendOptionsMap, ...)` has never been called, returns a
+   * default-constructed (empty, `size() == 0`) map.
+   *
+   * @returns Const reference to the Module-owned LoadBackendOptionsMap.
+   */
+  inline const LoadBackendOptionsMap& backend_options() const {
+    return backend_options_map_;
+  }
+
   /**
    * Checks if the program is loaded.
    *
@@ -713,7 +739,14 @@ class Module {
   std::unique_ptr<NamedDataMap> merged_data_map_;
   std::vector<std::vector<uint8_t>> shared_arenas_;
   ET_DEPRECATED std::vector<uint8_t> debug_buffer_;
-  const LoadBackendOptionsMap* backend_options_ = nullptr;
+  // Module-owned deep-copy of the backend options most recently installed
+  // via load(LoadBackendOptionsMap, ...). `backend_options_storage_` owns
+  // the per-backend BackendOption arrays; `backend_options_map_` is a
+  // LoadBackendOptionsMap whose Spans reference those owned arrays. An
+  // empty map (`size() == 0`) is observationally indistinguishable from
+  // "never set" by downstream consumers, so we don't track that bit.
+  std::vector<std::vector<runtime::BackendOption>> backend_options_storage_;
+  LoadBackendOptionsMap backend_options_map_;
   bool share_memory_arenas_;
 
   ET_NODISCARD runtime::Error load_internal(

extension/module/targets.bzl

@@ -27,6 +27,8 @@ def define_common_targets():
                 "//executorch/extension/named_data_map:merged_data_map" + aten_suffix,
             ],
             exported_deps = [
+                "//executorch/runtime/backend:backend_options",
+                "//executorch/runtime/backend:backend_options_map",
                 "//executorch/runtime/executor:program_no_prim_ops" + aten_suffix,
             ],
         )

extension/module/test/module_test.cpp

@@ -9,7 +9,9 @@
 #include <executorch/extension/module/module.h>
 
 #include <array>
+#include <cstring>
 #include <thread>
+#include <variant>
 
 #include <gtest/gtest.h>
 
@@ -663,6 +665,115 @@ TEST_F(ModuleTest, TestLoadWithEmptyLoadBackendOptionsMap) {
   EXPECT_TRUE(module.is_loaded());
 }
 
+TEST_F(ModuleTest, TestLoadWithBackendOptionsRollbackOnFailure) {
+  // Module pointed at a non-existent file so `load_internal` will fail.
+  Module module("/this/path/should/not/exist.pte");
+
+  {
+    // `bo1` lives only in this scope. The Module deep-copies the input,
+    // so dropping `bo1` is always safe regardless of whether the load
+    // succeeded, but on the failure path we additionally verify the
+    // Module did NOT install the input options (transactional contract).
+    LoadBackendOptionsMap bo1;
+    BackendOptions<2> opts;
+    opts.set_option("rollback_test", true);
+    ASSERT_EQ(bo1.set_options("RollbackBackend", opts.view()), Error::Ok);
+
+    const auto load_error = module.load(bo1);
+    EXPECT_NE(load_error, Error::Ok);
+    EXPECT_FALSE(module.is_loaded());
+  }
+  // `bo1` is destroyed. Module must remain in a usable state and a
+  // subsequent `load_method` should fail with the same load-time error
+  // (file not found) rather than crashing.
+  EXPECT_FALSE(module.is_loaded());
+  const auto method_error = module.load_method("forward");
+  EXPECT_NE(method_error, Error::Ok);
+  EXPECT_FALSE(module.is_method_loaded("forward"));
+}
+
+TEST_F(ModuleTest, TestLoadDeepCopiesBackendOptionsInputCanBeReleased) {
+  // Pin the deep-copy contract: the caller may release the input
+  // LoadBackendOptionsMap (and the BackendOption arrays its Spans
+  // referenced) immediately after `load()` returns. A subsequent
+  // `load_method` must use the Module-owned copy via the fallback path,
+  // not dereference the released input.
+  Module module(model_path_);
+
+  {
+    LoadBackendOptionsMap bo;
+    BackendOptions<2> opts;
+    opts.set_option("persist_test", true);
+    ASSERT_EQ(bo.set_options("PersistBackend", opts.view()), Error::Ok);
+
+    ASSERT_EQ(module.load(bo), Error::Ok);
+    // `bo` and `opts` go out of scope here; their storage is freed.
+  }
+
+  // load_method without explicit backend_options falls back to the
+  // Module's stored copy. With the old borrowed-pointer design this
+  // would have been a use-after-free; with deep-copy it is safe.
+  EXPECT_EQ(module.load_method("forward"), Error::Ok);
+  EXPECT_TRUE(module.is_method_loaded("forward"));
+
+  // Forward should still execute correctly using the Module-owned
+  // backend options.
+  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);
+}
+
+TEST_F(ModuleTest, TestLoadStoresBackendOptionsForReadback) {
+  // Verify that Module deep-copies the input LoadBackendOptionsMap into
+  // its own storage so callers can both (a) release the input
+  // immediately and (b) read back exactly what was stored via the
+  // public `backend_options()` accessor.
+  Module module(model_path_);
+
+  // Default-constructed: no options stored yet.
+  EXPECT_EQ(module.backend_options().size(), 0u);
+
+  {
+    LoadBackendOptionsMap bo;
+    BackendOptions<2> opts;
+    opts.set_option("num_threads", 8);
+    opts.set_option("enable_profiling", true);
+    ASSERT_EQ(bo.set_options("MyBackend", opts.view()), Error::Ok);
+
+    ASSERT_EQ(module.load(bo), Error::Ok);
+    // `bo` and `opts` go out of scope here; their backing storage is
+    // freed. Anything we read back from `module.backend_options()` must
+    // therefore live in Module-owned storage.
+  }
+
+  const auto& stored = module.backend_options();
+  ASSERT_EQ(stored.size(), 1u);
+
+  const auto entry = stored.entry_at(0);
+  EXPECT_STREQ(entry.backend_id, "MyBackend");
+  ASSERT_EQ(entry.options.size(), 2u);
+
+  // Look up each option by key so the value assertions are direct and
+  // independent of insertion order.
+  const BackendOption* num_threads_opt = nullptr;
+  const BackendOption* enable_profiling_opt = nullptr;
+  for (const auto& opt : entry.options) {
+    if (std::strcmp(opt.key, "num_threads") == 0) {
+      num_threads_opt = &opt;
+    } else if (std::strcmp(opt.key, "enable_profiling") == 0) {
+      enable_profiling_opt = &opt;
+    }
+  }
+
+  ASSERT_NE(num_threads_opt, nullptr);
+  ASSERT_TRUE(std::holds_alternative<int>(num_threads_opt->value));
+  EXPECT_EQ(std::get<int>(num_threads_opt->value), 8);
+
+  ASSERT_NE(enable_profiling_opt, nullptr);
+  ASSERT_TRUE(std::holds_alternative<bool>(enable_profiling_opt->value));
+  EXPECT_TRUE(std::get<bool>(enable_profiling_opt->value));
+}
+
 TEST_F(ModuleTest, TestLoadBackendOptionsMapPersistedAcrossLoadMethod) {
   Module module(model_path_);