Isolate cached delegates by Module lifetime (#24462)

Summary:

The QNN backend caches a `QnnManager` so methods in one ExecuTorch `Module` can share a QNN context. A context signature alone also matched independent modules loading identical program bytes, even though the manager was placement-allocated from only one module's runtime allocator. Key the cache by both signature and allocator to give independent modules independent manager lifetimes.

If a reverse-cache entry has no matching forward entry during destruction, report it and remove the stale reverse entry before returning. The regression deliberately creates this inconsistent state and verifies safe cleanup and a repeated release.

Count each loaded method's reference to its cached manager and destroy the manager after its last method releases it. Protect cache operations with the mutex and use `try_emplace` to preserve the first manager published for a key. A competing initializer returns that manager and destroys its own candidate. A destructor-only ownership guard cleans up every initialization error path, including backend, context, DLC compilation, and tensor allocation failures.

Scope HTP spill-fill groups by runtime allocator and configured buffer size. Each context receives its group handle explicitly, replacing the process-global static handle. Keep the first group manager alive until all group members and pending initializations are gone, and free each member before the first context. Serialize initialization within a group so concurrent loads select one first context.

Separate live modules now allocate separate QNN contexts, which can increase HTP memory use when they load identical blobs. Sharing between methods within one module remains intact. Full `QnnManager` destruction also releases its tensor wrappers and other non-trivial members. FCB export contents and selection are unchanged; FCB weight sharing within an artifact does not provide shared runtime ownership across independent modules.

Export the lifecycle tests so they are visible in the PR. They cover both module destruction orders, execution of multiple methods in one module, cleanup/reload, and method unload/reload on the HTP simulator. C++ tests verify shared cache handle identity, allocator isolation, reference accounting, concurrent cache publication, spill-fill group isolation, cleanup/reload, and owner destruction order. The simulator lifecycle suite also runs with multi-context enabled and a nonzero spill-fill size; target SDK registration still requires device verification. Concurrent execution through one module and execution racing destruction still require client coordination.

Authored with AI assistance from Codex.

Differential Revision: D122851825




Signed-off-by: Ali Monfared <alimonfared@meta.com>

Author: alirezamonfared

backends/qualcomm/runtime/QnnExecuTorchBackend.cpp

@@ -14,6 +14,9 @@
 #include <executorch/backends/qualcomm/runtime/backends/QnnCustomProtocol.h>
 #include <executorch/runtime/backend/interface.h>
 #include <executorch/runtime/backend/options.h>
+
+#include <memory>
+
 namespace executorch {
 namespace backends {
 namespace qnn {
@@ -65,36 +68,58 @@ Result<DelegateHandle*> QnnExecuTorchBackend::init(
       QNN_EXECUTORCH_LOG_WARN("unknown argument: %s", compile_spec.key);
   }
 
-  // Create QnnManager
   MemoryAllocator* runtime_allocator = context.get_runtime_allocator();
-  QnnManager* qnn_manager = runtime_allocator->allocateInstance<QnnManager>();
-  if (qnn_manager == nullptr) {
-    return Error::MemoryAllocationFailed;
-  }
-
-  // NOTE: Since we use placement new and since this type is not trivially
-  // destructible, we must call the destructor manually in destroy().
-  new (qnn_manager) QnnManager(qnn_executorch_options, qnn_context_blob);
-  // TODO: this is a temporal solution for multi-graph support, will be
+  // TODO: this is a temporary solution for multi-graph support, will be
   //       removed once framework starts to accept runtime configuration
   // ---
   // check if current context binary has already been initialized
   // return cached one for reducing memory footprint
-
-  auto iter = delegate_map_.find(signature);
-  if (iter != delegate_map_.end()) {
+  const DelegateKey key{signature, runtime_allocator};
+  if (auto cached_handle = get_cached_delegate(key)) {
     QNN_EXECUTORCH_LOG_INFO(
         "Use cached delegate handle for current method: %s",
         context.get_method_name());
-    return iter->second;
+    return *cached_handle;
+  }
+
+  std::shared_ptr<SpillFillGroup> group;
+  std::unique_lock<std::mutex> group_lock;
+  auto* backend_options = qnn_executorch_options->backend_options();
+  if (backend_options->backend_type() ==
+      QnnExecuTorchBackendType::kHtpBackend) {
+    auto* htp_options = backend_options->htp_options();
+    if (htp_options->use_multi_contexts() &&
+        htp_options->max_sf_buf_size() != 0) {
+      group = get_spill_fill_group(
+          {runtime_allocator, htp_options->max_sf_buf_size()});
+      group_lock = std::unique_lock<std::mutex>(group->mutex);
+      if (auto cached_handle = get_cached_delegate(key)) {
+        return *cached_handle;
+      }
+    }
+  }
+
+  // Create QnnManager
+  QnnManager* qnn_manager = runtime_allocator->allocateInstance<QnnManager>();
+  if (qnn_manager == nullptr) {
+    return Error::MemoryAllocationFailed;
   }
 
+  // The runtime allocator owns the storage, so only invoke the destructor.
+  new (qnn_manager) QnnManager(qnn_executorch_options, qnn_context_blob);
+  auto destroy_manager = [](QnnManager* manager) { manager->~QnnManager(); };
+  std::shared_ptr<QnnManager> manager_guard(qnn_manager, destroy_manager);
+
   ET_CHECK_OR_RETURN_ERROR(
       qnn_manager->InitBackend() == Error::Ok,
       Internal,
       "Fail to initialize Qnn Manager");
   ET_CHECK_OR_RETURN_ERROR(
-      qnn_manager->InitContext() == Error::Ok,
+      qnn_manager->InitContext(
+          std::nullopt,
+          group && group->owner
+              ? static_cast<QnnManager*>(group->owner.get())->GetContextHandle()
+              : nullptr) == Error::Ok,
       Internal,
       "Fail to initialize Qnn Manager");
 
@@ -111,24 +136,26 @@ Result<DelegateHandle*> QnnExecuTorchBackend::init(
           "Fail to allocate tensor");
     }
   }
-  add_cached_delegate(signature, qnn_manager);
+  DelegateHandle* cached_handle =
+      add_cached_delegate(key, qnn_manager, manager_guard, group);
+  if (group && !group->owner && cached_handle == qnn_manager) {
+    group->owner = manager_guard;
+  }
 
 #ifndef __hexagon__
   // This backend does not need its processed data after Init.
   processed->Free();
 #endif
 
-  return qnn_manager;
+  return cached_handle;
 }
 
 Error QnnExecuTorchBackend::execute(
     BackendExecutionContext& context,
     DelegateHandle* handle,
     Span<EValue*> args) const {
   ET_CHECK_OR_RETURN_ERROR(
-      delegate_map_rev_.count(handle) != 0,
-      Internal,
-      "DelegateHandle has been deleted");
+      is_cached_delegate(handle), Internal, "DelegateHandle has been deleted");
   QnnManager* qnn_manager = static_cast<QnnManager*>(handle);
 
   std::string method_name = context.get_method_name();
@@ -225,11 +252,9 @@ Error QnnExecuTorchBackend::execute(
 }
 
 void QnnExecuTorchBackend::destroy(DelegateHandle* handle) const {
-  if (handle != nullptr && delegate_map_rev_.count(handle)) {
-    QnnManager* qnn_manager = static_cast<QnnManager*>(handle);
-    qnn_manager->Destroy();
-    erase_cached_delegate(handle);
-  }
+  // Releasing the final shared owner invokes ~QnnManager(), which calls
+  // Destroy(). The returned entry releases owners after mutex_ is unlocked.
+  (void)erase_cached_delegate(handle);
 }
 
 executorch::runtime::Error QnnExecuTorchBackend::set_option(
@@ -365,41 +390,98 @@ bool QnnExecuTorchBackend::is_available() const {
   return true;
 }
 
-void QnnExecuTorchBackend::add_cached_delegate(
-    const std::int64_t& signature,
+std::optional<executorch::runtime::DelegateHandle*>
+QnnExecuTorchBackend::get_cached_delegate(const DelegateKey& key) const {
+  std::lock_guard<std::mutex> guard(mutex_);
+  auto iter = delegate_map_.find(key);
+  if (iter == delegate_map_.end()) {
+    return std::nullopt;
+  }
+  ++iter->second.users;
+  return iter->second.handle;
+}
+
+std::shared_ptr<QnnExecuTorchBackend::SpillFillGroup>
+QnnExecuTorchBackend::get_spill_fill_group(const SpillFillKey& key) const {
+  std::lock_guard<std::mutex> guard(mutex_);
+  for (auto iter = spill_fill_groups_.begin();
+       iter != spill_fill_groups_.end();) {
+    if (iter->second.expired()) {
+      iter = spill_fill_groups_.erase(iter);
+    } else {
+      ++iter;
+    }
+  }
+  auto& weak_group = spill_fill_groups_[key];
+  auto group = weak_group.lock();
+  if (!group) {
+    group = std::make_shared<SpillFillGroup>();
+    weak_group = group;
+  }
+  return group;
+}
+
+executorch::runtime::DelegateHandle* QnnExecuTorchBackend::add_cached_delegate(
+    const DelegateKey& key,
+    executorch::runtime::DelegateHandle* handle,
+    std::shared_ptr<executorch::runtime::DelegateHandle> manager,
+    std::shared_ptr<SpillFillGroup> group) const {
+  std::lock_guard<std::mutex> guard(mutex_);
+  auto [iter, inserted] = delegate_map_.try_emplace(
+      key, DelegateEntry{handle, 1, std::move(group), std::move(manager)});
+  if (inserted) {
+    delegate_map_rev_.emplace(handle, key);
+  } else {
+    ++iter->second.users;
+  }
+  return iter->second.handle;
+}
+
+bool QnnExecuTorchBackend::is_cached_delegate(
     executorch::runtime::DelegateHandle* handle) const {
   std::lock_guard<std::mutex> guard(mutex_);
-  delegate_map_[signature] = handle;
-  delegate_map_rev_[handle] = signature;
+  return delegate_map_rev_.count(handle) != 0;
 }
 
-void QnnExecuTorchBackend::erase_cached_delegate(
+std::optional<QnnExecuTorchBackend::DelegateEntry>
+QnnExecuTorchBackend::erase_cached_delegate(
     executorch::runtime::DelegateHandle* handle) const {
   std::lock_guard<std::mutex> guard(mutex_);
   auto iter = delegate_map_rev_.find(handle);
   if (iter == delegate_map_rev_.end()) {
-    return;
+    return std::nullopt;
   }
-  delegate_map_.erase(iter->second);
-  delegate_map_rev_.erase(handle);
+  auto delegate_iter = delegate_map_.find(iter->second);
+  if (delegate_iter == delegate_map_.end()) {
+    QNN_EXECUTORCH_LOG_ERROR("Cached delegate has no matching entry.");
+    delegate_map_rev_.erase(iter);
+    return std::nullopt;
+  }
+  if (--delegate_iter->second.users != 0) {
+    return std::nullopt;
+  }
+  DelegateEntry entry = std::move(delegate_iter->second);
+  delegate_map_.erase(delegate_iter);
+  delegate_map_rev_.erase(iter);
+  return entry;
 }
 
 namespace {
+// Loaded modules may outlive static teardown. Their managers must be released
+// by destroy(), while their allocators and QNN libraries are still alive.
+const executorch::runtime::Backend& qnn_backend() {
+  static const executorch::runtime::Backend backend{
+      QNN_BACKEND, new QnnExecuTorchBackend()};
+  return backend;
+}
 #if !defined(_MSC_VER)
-auto cls = QnnExecuTorchBackend();
-executorch::runtime::Backend backend{QNN_BACKEND, &cls};
 // On non-Windows platforms a single copy of executorch_core is shared between
 // the DLL and the exe, so the static initializer writes into the correct
 // registered_backends[] array.
-static auto success_with_compiler = register_backend(backend);
+const auto success_with_compiler = register_backend(qnn_backend());
 #endif
 } // namespace
 
-#if defined(_MSC_VER)
-static QnnExecuTorchBackend cls_msvc;
-static executorch::runtime::Backend backend_msvc{QNN_BACKEND, &cls_msvc};
-#endif
-
 } // namespace qnn
 } // namespace backends
 } // namespace executorch
@@ -417,7 +499,7 @@ void QnnExecuTorchBackendRegister(void* register_fn) {
     return;
   }
   executorch::runtime::Error err = reinterpret_cast<RegisterFn>(register_fn)(
-      executorch::backends::qnn::backend_msvc);
+      executorch::backends::qnn::qnn_backend());
   if (err != executorch::runtime::Error::Ok) {
     ET_LOG(
         Error,

backends/qualcomm/runtime/QnnExecuTorchBackend.h

@@ -12,8 +12,12 @@
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/evalue.h>
 
+#include <map>
+#include <memory>
 #include <mutex>
+#include <optional>
 #include <unordered_map>
+#include <utility>
 
 namespace executorch {
 namespace backends {
@@ -22,7 +26,7 @@ namespace qnn {
 class QnnExecuTorchBackend final
     : public ::executorch::runtime::BackendInterface {
  public:
-  ~QnnExecuTorchBackend(){};
+  ~QnnExecuTorchBackend() {};
 
   executorch::runtime::Result<executorch::runtime::DelegateHandle*> init(
       executorch::runtime::BackendInitContext& context,
@@ -51,16 +55,49 @@ class QnnExecuTorchBackend final
   bool is_available() const override;
 
  private:
-  void add_cached_delegate(
-      const std::int64_t& signature,
+  friend class QnnExecuTorchBackendTest;
+
+  // A context binary signature plus the runtime allocator that the methods of
+  // one Module share. Other Modules can be destroyed independently, so they
+  // must never receive this Module's QnnManager.
+  using DelegateKey =
+      std::pair<std::int64_t, const executorch::runtime::MemoryAllocator*>;
+
+  struct SpillFillGroup {
+    std::mutex mutex;
+    // HTP requires the first context to outlive every other group member.
+    std::shared_ptr<executorch::runtime::DelegateHandle> owner;
+  };
+  using SpillFillKey =
+      std::pair<const executorch::runtime::MemoryAllocator*, uint64_t>;
+
+  struct DelegateEntry {
+    executorch::runtime::DelegateHandle* handle;
+    size_t users;
+    // Destroy each member context before releasing the group's first context.
+    std::shared_ptr<SpillFillGroup> group;
+    std::shared_ptr<executorch::runtime::DelegateHandle> manager;
+  };
+
+  std::shared_ptr<SpillFillGroup> get_spill_fill_group(
+      const SpillFillKey& key) const;
+  std::optional<executorch::runtime::DelegateHandle*> get_cached_delegate(
+      const DelegateKey& key) const;
+  executorch::runtime::DelegateHandle* add_cached_delegate(
+      const DelegateKey& key,
+      executorch::runtime::DelegateHandle* handle,
+      std::shared_ptr<executorch::runtime::DelegateHandle> manager = nullptr,
+      std::shared_ptr<SpillFillGroup> group = nullptr) const;
+  bool is_cached_delegate(executorch::runtime::DelegateHandle* handle) const;
+  std::optional<DelegateEntry> erase_cached_delegate(
       executorch::runtime::DelegateHandle* handle) const;
-  void erase_cached_delegate(executorch::runtime::DelegateHandle* handle) const;
 
   mutable std::mutex mutex_;
   mutable std::mutex runtime_option_mutex_;
-  mutable std::unordered_map<int64_t, executorch::runtime::DelegateHandle*>
-      delegate_map_;
-  mutable std::unordered_map<executorch::runtime::DelegateHandle*, std::int64_t>
+  mutable std::map<DelegateKey, DelegateEntry> delegate_map_;
+  mutable std::map<SpillFillKey, std::weak_ptr<SpillFillGroup>>
+      spill_fill_groups_;
+  mutable std::unordered_map<executorch::runtime::DelegateHandle*, DelegateKey>
       delegate_map_rev_;
 
   RuntimeOption qnn_runtime_log_level_{false, 0};

backends/qualcomm/runtime/QnnManager.cpp

@@ -342,7 +342,8 @@ Error QnnManager::InitBackend() {
 }
 
 Error QnnManager::InitContext(
-    std::optional<std::vector<std::string>> graph_names) {
+    const std::optional<std::vector<std::string>>& graph_names,
+    Qnn_ContextHandle_t spill_fill_handle) {
   if (backend_params_ptr_->backend_init_state_ ==
       BackendInitializeState::UNINITIALIZED) {
     QNN_EXECUTORCH_LOG_INFO(
@@ -356,7 +357,8 @@ Error QnnManager::InitContext(
         backend_bundle_ptr_->qnn_device_ptr.get(),
         qnn_context_blob_,
         options_,
-        qnn_dlc_manager_.get());
+        qnn_dlc_manager_.get(),
+        spill_fill_handle);
     ET_CHECK_OR_RETURN_ERROR(
         backend_params_ptr_ != nullptr,
         Internal,

backends/qualcomm/runtime/QnnManager.h

@@ -38,7 +38,12 @@ class QnnManager {
   // provided graph_names. Note: For online_prepare or deserialization, the
   // graph name will be obtained from the binary.
   executorch::runtime::Error InitContext(
-      std::optional<std::vector<std::string>> graph_names = std::nullopt);
+      const std::optional<std::vector<std::string>>& graph_names = std::nullopt,
+      Qnn_ContextHandle_t spill_fill_handle = nullptr);
+
+  Qnn_ContextHandle_t GetContextHandle() const {
+    return backend_params_ptr_->qnn_context_ptr_->GetHandle();
+  }
   // This function only initialize the context cache to get spill fill buffer
   // size
   executorch::runtime::Error InitContextCache();

backends/qualcomm/runtime/backends/QnnBackendFactory.cpp

@@ -20,7 +20,8 @@ std::unique_ptr<BackendConfigParameters> QnnBackendFactory::Create(
     QnnDevice* qnn_device_ptr,
     const QnnExecuTorchContextBinary& qnn_context_blob,
     const QnnExecuTorchOptions* options,
-    QnnDlcManager* qnn_dlc_manager) {
+    QnnDlcManager* qnn_dlc_manager,
+    Qnn_ContextHandle_t spill_fill_handle) {
   auto backend_params = std::make_unique<BackendConfigParameters>();
 
   switch (options->backend_options()->backend_type()) {
@@ -79,7 +80,8 @@ std::unique_ptr<BackendConfigParameters> QnnBackendFactory::Create(
           htp_options,
           options->target_options(),
           qnn_dlc_manager,
-          get_option(options->profile_level(), QNN_RUNTIME_PROFILE_LEVEL));
+          get_option(options->profile_level(), QNN_RUNTIME_PROFILE_LEVEL),
+          spill_fill_handle);
 
       backend_params->qnn_graph_ptr_ = std::make_unique<HtpGraph>(
           implementation_ptr,

backends/qualcomm/runtime/backends/QnnBackendFactory.h

@@ -68,7 +68,8 @@ class QnnBackendFactory {
       QnnDevice* qnn_device_ptr,
       const QnnExecuTorchContextBinary& qnn_context_blob,
       const QnnExecuTorchOptions* options,
-      QnnDlcManager* qnn_dlc_manager);
+      QnnDlcManager* qnn_dlc_manager,
+      Qnn_ContextHandle_t spill_fill_handle = nullptr);
 };
 } // namespace qnn
 } // namespace backends