Isolate cached delegates by Module lifetime (pytorch#24462)

Summary:
Pull Request resolved: pytorch#24462

## Summary

The QNN backend caches a `QnnManager` so methods in one ExecuTorch `Module` can share an expensive QNN context. The previous cache key used only the exported context signature, so independent modules loading equivalent programs could share a manager that was placement-allocated from one module's runtime allocator.

Key the cache by both context signature and runtime allocator. Methods in one module continue to share a context, while independent modules receive managers with independent lifetimes.

The change also:

- serializes cache lookup, insertion, membership checks, and removal;
- avoids allocating a manager on cache hits;
- removes cache state before running the non-trivial `QnnManager` destructor.

## Scope

This fixes sequential independent-module lifetime isolation. It does not change exported program contents, backend selection, or model behavior. Broader same-module concurrent execution and initialization races remain outside this change.

Differential Revision: D122851825

Author: alirezamonfared

backends/qualcomm/runtime/QnnExecuTorchBackend.cpp

@@ -65,30 +65,30 @@ 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
   //       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 (DelegateHandle* 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;
+  }
+
+  // Create QnnManager
+  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);
+
   ET_CHECK_OR_RETURN_ERROR(
       qnn_manager->InitBackend() == Error::Ok,
       Internal,
@@ -111,7 +111,7 @@ Result<DelegateHandle*> QnnExecuTorchBackend::init(
           "Fail to allocate tensor");
     }
   }
-  add_cached_delegate(signature, qnn_manager);
+  add_cached_delegate(key, qnn_manager);
 
 #ifndef __hexagon__
   // This backend does not need its processed data after Init.
@@ -126,9 +126,7 @@ Error QnnExecuTorchBackend::execute(
     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,10 +223,8 @@ 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);
+  if (handle != nullptr && erase_cached_delegate(handle)) {
+    static_cast<QnnManager*>(handle)->~QnnManager();
   }
 }
 
@@ -365,23 +361,37 @@ bool QnnExecuTorchBackend::is_available() const {
   return true;
 }
 
+executorch::runtime::DelegateHandle* QnnExecuTorchBackend::get_cached_delegate(
+    const DelegateKey& key) const {
+  std::lock_guard<std::mutex> guard(mutex_);
+  auto iter = delegate_map_.find(key);
+  return iter == delegate_map_.end() ? nullptr : iter->second;
+}
+
 void QnnExecuTorchBackend::add_cached_delegate(
-    const std::int64_t& signature,
+    const DelegateKey& key,
     executorch::runtime::DelegateHandle* handle) const {
   std::lock_guard<std::mutex> guard(mutex_);
-  delegate_map_[signature] = handle;
-  delegate_map_rev_[handle] = signature;
+  delegate_map_[key] = handle;
+  delegate_map_rev_[handle] = key;
 }
 
-void QnnExecuTorchBackend::erase_cached_delegate(
+bool QnnExecuTorchBackend::is_cached_delegate(
+    executorch::runtime::DelegateHandle* handle) const {
+  std::lock_guard<std::mutex> guard(mutex_);
+  return delegate_map_rev_.count(handle) != 0;
+}
+
+bool 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 false;
   }
   delegate_map_.erase(iter->second);
   delegate_map_rev_.erase(handle);
+  return true;
 }
 
 namespace {

backends/qualcomm/runtime/QnnExecuTorchBackend.h

@@ -12,8 +12,10 @@
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/evalue.h>
 
+#include <map>
 #include <mutex>
 #include <unordered_map>
+#include <utility>
 
 namespace executorch {
 namespace backends {
@@ -51,16 +53,25 @@ class QnnExecuTorchBackend final
   bool is_available() const override;
 
  private:
+  // 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*>;
+
+  executorch::runtime::DelegateHandle* get_cached_delegate(
+      const DelegateKey& key) const;
   void add_cached_delegate(
-      const std::int64_t& signature,
+      const DelegateKey& key,
       executorch::runtime::DelegateHandle* handle) const;
-  void erase_cached_delegate(executorch::runtime::DelegateHandle* handle) const;
+  bool is_cached_delegate(executorch::runtime::DelegateHandle* handle) const;
+  bool 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*>
+  mutable std::map<DelegateKey, executorch::runtime::DelegateHandle*>
       delegate_map_;
-  mutable std::unordered_map<executorch::runtime::DelegateHandle*, std::int64_t>
+  mutable std::unordered_map<executorch::runtime::DelegateHandle*, DelegateKey>
       delegate_map_rev_;
 
   RuntimeOption qnn_runtime_log_level_{false, 0};