Isolate cached delegates by Module lifetime (#24462)

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