Isolate cached delegates by Module lifetime (#24462)

Summary:
Pull Request resolved: #24462

## TL;DR

Give independently loaded ExecuTorch `Module`s independent QNN manager lifetimes while preserving context sharing between methods within one module. Key the delegate cache by `(context signature, runtime allocator)`, count method references, and release complete manager state on final destruction or initialization failure.

## Context

`QnnExecuTorchBackend` is process-global and caches expensive `QnnManager` instances across methods. A manager is placement-allocated from a module's runtime allocator; methods in that module share its allocator, but independent modules can load identical `.pte` bytes and be destroyed separately. D122851825 hardens this backend contract without changing FCB export contents or artifact selection.

## Problem

A context signature alone also matched independent modules, allowing one module to destroy manager state still used by another. Calling `Destroy()` alone also omitted non-trivial manager member destruction, and competing initializers could publish or abandon duplicate candidates. A process-global HTP spill-fill handle did not represent the lifetime of its owning context.

```text
Before: process cache[signature] -> Manager M (allocator A)

Module A methods ----> M <---- Module B, same program bytes
release Module A ----> M destroyed
Module B executes ---> stale context
```

## Solution

Use a cache key containing the signature and allocator, with mutex-protected lookup, membership, publication and release. Each loaded method holds a reference; `try_emplace` preserves the first published manager and competing initializers destroy their unused candidate. A destructor-only ownership guard covers backend, context, DLC compilation and tensor allocation failures. Full `QnnManager` destruction also releases its input/output tensor wrappers and other non-trivial members; their storage was not released by calling `Destroy()` alone. On final release, erase cache state under the mutex and destroy the manager after unlocking. A missing forward entry is reported and its stale reverse entry removed safely.

```text
After: cache[(signature, allocator A)] -> Manager A -> A's methods
       cache[(signature, allocator B)] -> Manager B -> B's methods

release A's last method -> unlock cache -> ~Manager A -> Destroy()
B's methods            -> Manager B remains alive
```

Scope HTP spill-fill groups by allocator and configured buffer size. Thread the group handle explicitly through context initialization instead of using a process-global handle. Serialize initialization within each group, retain its first manager until all members and pending initializations are gone, and destroy member contexts before the owner. The registry holds weak group references, preventing an ownership cycle and allowing cleanup followed by a fresh group.

```text
weak registry[(allocator, spill-fill bytes)] -> Group
method entry -> member manager + Group
Group        -> first manager (HTP owner)

last member/pending initializer releases Group
  -> member context destroyed
  -> first context destroyed
  -> registry's weak reference expires
```

Keep the registered backend alive for the process lifetime so loaded modules cannot destroy QNN contexts through global cache teardown after allocators or the QNN library have finalized. Normal method destruction still releases manager state; the explicit defaulted backend destructor avoids conflicting empty-body formatting between local and OSS ClangFormat versions. Python lifecycle imports follow the ordering reported by OSS UFMT.

## Compatibility and evidence limits

Separate live modules allocate separate QNN contexts, which can increase HTP memory use for duplicate blobs. Methods within one module still share a manager. FCB weight sharing inside an artifact does not supply shared runtime ownership across independent modules. Concurrent execution through one module and execution racing destruction require client coordination.

### QAI Hub shard spill-fill cost

`gen_pte_from_ctx_bin()` calls `update_spill_fill_size()` over the shard programs before writing separate `.pte` files. The helper sets each participating context's `max_sf_buf_size` to the largest requirement in the export group. The QAI Hub Llama runner constructs a separate `Module` for each shard and loads every shard method; default `Module` construction gives each one a separate runtime allocator.

On device, `HtpContextCustomConfig::CreateContextCustomConfig()` registers an HTP group using `firstGroupHandle` and `maxSpillFillBuffer`. The former process-global handle allowed shard contexts to join one group. With this diff, the `(allocator, max_sf_buf_size)` key isolates those independently allocated Modules, so each shard starts its own spill-fill group even when the exported size matches.

```text
N resident shard Modules, common configured S bytes (S > 0)

Before: shard 1 ... shard N -> one HTP spill-fill group -> S bytes
After:  shard i / allocator i -> group i              -> S bytes each

Expected configured reservation: S -> N * S bytes
Expected additional reservation: (N - 1) * S bytes
```

This is a projection of the spill-fill reservation for the described shard layout, not measured RSS, PSS, DSP heap usage or peak device memory. SDK allocation behavior, alignment and other context/model memory affect the observed totals. The existing host simulator tests do not issue target spill-fill group registration and cannot measure this device cost. No connected physical device was available for this follow-up, so the magnitude on a representative QAI Hub shard model remains unmeasured.

If device evidence shows a material cost, an opt-in cross-Module group ID is a possible follow-up. Such a group must keep its first context and allocator-owned manager storage alive until all shard Modules and pending initializations release it; changing the key to an ID alone would not satisfy that lifetime requirement. This diff does not introduce that opt-in.

Exported lifecycle tests cover both module destruction orders, multiple methods, cleanup/reload and sibling unload/reload on the HTP simulator, including multi-context mode with a nonzero spill-fill size. C++ tests check handle identity, allocator isolation, reference accounting, concurrent publication, stale-map cleanup, spill-fill ownership and process exit. Target SDK spill-fill registration still requires device verification; the recorded FCB AOT failure occurred before runtime module loading and does not establish FCB runtime compatibility.

Authored with AI assistance from Codex.

Differential Revision: D122851825

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() override = default;
 
   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

@@ -338,7 +338,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(
@@ -352,7 +353,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