From c44220f672f1356c9c346a5cefe8177cd46d4e82 Mon Sep 17 00:00:00 2001 From: Ali Monfared Date: Mon, 9 Nov 2026 19:29:56 -0800 Subject: [PATCH] Isolate cached delegates by Module lifetime (#24462) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/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 --- .../qualcomm/runtime/QnnExecuTorchBackend.cpp | 166 +++++++++--- .../qualcomm/runtime/QnnExecuTorchBackend.h | 51 +++- backends/qualcomm/runtime/QnnManager.cpp | 6 +- backends/qualcomm/runtime/QnnManager.h | 7 +- .../runtime/backends/QnnBackendFactory.cpp | 6 +- .../runtime/backends/QnnBackendFactory.h | 3 +- .../runtime/backends/htp/HtpContext.cpp | 9 - .../runtime/backends/htp/HtpContext.h | 10 +- backends/qualcomm/tests/BUCK | 29 ++ .../tests/QnnExecuTorchBackendTest.cpp | 252 ++++++++++++++++++ .../tests/test_qnn_shared_delegate.py | 142 ++++++++++ 11 files changed, 612 insertions(+), 69 deletions(-) create mode 100644 backends/qualcomm/tests/QnnExecuTorchBackendTest.cpp create mode 100644 backends/qualcomm/tests/test_qnn_shared_delegate.py diff --git a/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp b/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp index 76b67f911cd..ea372cbda40 100644 --- a/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp +++ b/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp @@ -14,6 +14,9 @@ #include #include #include + +#include + namespace executorch { namespace backends { namespace qnn { @@ -65,36 +68,58 @@ Result 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(); - 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 group; + std::unique_lock 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(group->mutex); + if (auto cached_handle = get_cached_delegate(key)) { + return *cached_handle; + } + } + } + + // Create QnnManager + QnnManager* qnn_manager = runtime_allocator->allocateInstance(); + 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 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(group->owner.get())->GetContextHandle() + : nullptr) == Error::Ok, Internal, "Fail to initialize Qnn Manager"); @@ -111,14 +136,18 @@ Result 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( @@ -126,9 +155,7 @@ Error QnnExecuTorchBackend::execute( DelegateHandle* handle, Span 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(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(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 +QnnExecuTorchBackend::get_cached_delegate(const DelegateKey& key) const { + std::lock_guard 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::get_spill_fill_group(const SpillFillKey& key) const { + std::lock_guard 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(); + weak_group = group; + } + return group; +} + +executorch::runtime::DelegateHandle* QnnExecuTorchBackend::add_cached_delegate( + const DelegateKey& key, + executorch::runtime::DelegateHandle* handle, + std::shared_ptr manager, + std::shared_ptr group) const { + std::lock_guard 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 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::erase_cached_delegate( executorch::runtime::DelegateHandle* handle) const { std::lock_guard 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(register_fn)( - executorch::backends::qnn::backend_msvc); + executorch::backends::qnn::qnn_backend()); if (err != executorch::runtime::Error::Ok) { ET_LOG( Error, diff --git a/backends/qualcomm/runtime/QnnExecuTorchBackend.h b/backends/qualcomm/runtime/QnnExecuTorchBackend.h index e3548c8752b..2f36bccb3da 100644 --- a/backends/qualcomm/runtime/QnnExecuTorchBackend.h +++ b/backends/qualcomm/runtime/QnnExecuTorchBackend.h @@ -12,8 +12,12 @@ #include #include +#include +#include #include +#include #include +#include 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 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; + + struct SpillFillGroup { + std::mutex mutex; + // HTP requires the first context to outlive every other group member. + std::shared_ptr owner; + }; + using SpillFillKey = + std::pair; + + struct DelegateEntry { + executorch::runtime::DelegateHandle* handle; + size_t users; + // Destroy each member context before releasing the group's first context. + std::shared_ptr group; + std::shared_ptr manager; + }; + + std::shared_ptr get_spill_fill_group( + const SpillFillKey& key) const; + std::optional get_cached_delegate( + const DelegateKey& key) const; + executorch::runtime::DelegateHandle* add_cached_delegate( + const DelegateKey& key, + executorch::runtime::DelegateHandle* handle, + std::shared_ptr manager = nullptr, + std::shared_ptr group = nullptr) const; + bool is_cached_delegate(executorch::runtime::DelegateHandle* handle) const; + std::optional 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 - delegate_map_; - mutable std::unordered_map + mutable std::map delegate_map_; + mutable std::map> + spill_fill_groups_; + mutable std::unordered_map delegate_map_rev_; RuntimeOption qnn_runtime_log_level_{false, 0}; diff --git a/backends/qualcomm/runtime/QnnManager.cpp b/backends/qualcomm/runtime/QnnManager.cpp index 95aa037cb31..8a3e5f0d462 100644 --- a/backends/qualcomm/runtime/QnnManager.cpp +++ b/backends/qualcomm/runtime/QnnManager.cpp @@ -338,7 +338,8 @@ Error QnnManager::InitBackend() { } Error QnnManager::InitContext( - std::optional> graph_names) { + const std::optional>& 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, diff --git a/backends/qualcomm/runtime/QnnManager.h b/backends/qualcomm/runtime/QnnManager.h index 58a593a733a..8564ea4ca25 100644 --- a/backends/qualcomm/runtime/QnnManager.h +++ b/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> graph_names = std::nullopt); + const std::optional>& 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(); diff --git a/backends/qualcomm/runtime/backends/QnnBackendFactory.cpp b/backends/qualcomm/runtime/backends/QnnBackendFactory.cpp index 5b8f96b84b8..1d71de9d00f 100644 --- a/backends/qualcomm/runtime/backends/QnnBackendFactory.cpp +++ b/backends/qualcomm/runtime/backends/QnnBackendFactory.cpp @@ -20,7 +20,8 @@ std::unique_ptr 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(); switch (options->backend_options()->backend_type()) { @@ -79,7 +80,8 @@ std::unique_ptr 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( implementation_ptr, diff --git a/backends/qualcomm/runtime/backends/QnnBackendFactory.h b/backends/qualcomm/runtime/backends/QnnBackendFactory.h index a5b9af05029..df552aab321 100644 --- a/backends/qualcomm/runtime/backends/QnnBackendFactory.h +++ b/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 diff --git a/backends/qualcomm/runtime/backends/htp/HtpContext.cpp b/backends/qualcomm/runtime/backends/htp/HtpContext.cpp index 0056a2c0917..dda2525f7f1 100644 --- a/backends/qualcomm/runtime/backends/htp/HtpContext.cpp +++ b/backends/qualcomm/runtime/backends/htp/HtpContext.cpp @@ -36,15 +36,6 @@ Error HtpContext::MakeConfig(std::vector& config) { return Error::Ok; } -Error HtpContext::AfterConfigure() { - // update sf_handle with first context handle encounterded as group handle - // TODO: should handle the thread safety if needed - if (sf_handle_ == 0x0) { - sf_handle_ = GetHandle(); - } - return Error::Ok; -} - } // namespace qnn } // namespace backends } // namespace executorch diff --git a/backends/qualcomm/runtime/backends/htp/HtpContext.h b/backends/qualcomm/runtime/backends/htp/HtpContext.h index a41e09c0986..e5bca85d76c 100644 --- a/backends/qualcomm/runtime/backends/htp/HtpContext.h +++ b/backends/qualcomm/runtime/backends/htp/HtpContext.h @@ -28,7 +28,8 @@ class HtpContext : public QnnContext { const QnnExecuTorchHtpBackendOptions* htp_options, const QnnExecuTorchTargetOptions* target_options, QnnDlcManager* qnn_dlc_manager, - const QnnExecuTorchProfileLevel& profile_level) + const QnnExecuTorchProfileLevel& profile_level, + Qnn_ContextHandle_t spill_fill_handle = nullptr) : QnnContext( implementation, system_implementation, @@ -36,7 +37,8 @@ class HtpContext : public QnnContext { device, cache, qnn_dlc_manager, - profile_level) { + profile_level), + sf_handle_(spill_fill_handle) { htp_context_custom_config_ = std::make_unique( this, htp_options, target_options, profile_level); } @@ -49,13 +51,11 @@ class HtpContext : public QnnContext { protected: executorch::runtime::Error MakeConfig( std::vector& config) override; - executorch::runtime::Error AfterConfigure() override; private: std::vector context_config_; std::unique_ptr htp_context_custom_config_; - // this is shared between all partitions inside one pte - static inline Qnn_ContextHandle_t sf_handle_ = 0x0; + Qnn_ContextHandle_t sf_handle_; }; } // namespace qnn diff --git a/backends/qualcomm/tests/BUCK b/backends/qualcomm/tests/BUCK index 7676953c093..bd5cf14ffe6 100644 --- a/backends/qualcomm/tests/BUCK +++ b/backends/qualcomm/tests/BUCK @@ -1,5 +1,6 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +load("@fbsource//tools/build_defs:default_platform_defs.bzl", "CXX") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") load("@fbsource//xplat/executorch/backends/qualcomm/qnn_version.bzl", "get_qnn_library_version") @@ -143,3 +144,31 @@ None if runtime.is_oss else fbcode_target( ":test_qnn_delegate", ], ) + +None if runtime.is_oss else fbcode_target(_kind = runtime.python_test, + name = "test_qnn_shared_delegate", + srcs = ["test_qnn_shared_delegate.py"], + env = { + "LD_LIBRARY_PATH": "$(location fbsource//third-party/qualcomm/qnn/qnn-{0}:qnn_offline_compile_libs)".format(get_qnn_library_version()), + "QNN_SDK_ROOT": "$(location fbsource//third-party/qualcomm/qnn/qnn-{0}:__dir__)".format(get_qnn_library_version()), + }, + deps = [ + "//caffe2:torch", + "//caffe2/functorch:functorch_src", + "//executorch/backends/qualcomm/partition:partition", + "//executorch/backends/qualcomm/serialization:serialization", + "//executorch/backends/qualcomm/utils:utils", + "//executorch/exir:lib", + "//executorch/extension/pybindings:portable_lib", # @manual + ], +) + +fbcode_target(_kind = runtime.cxx_test, + name = "qnn_delegate_cache_test", + srcs = ["QnnExecuTorchBackendTest.cpp"], + platforms = [CXX], + deps = [ + "//executorch/backends/qualcomm/runtime:runtime", + "//executorch/runtime/core:memory_allocator", + ], +) diff --git a/backends/qualcomm/tests/QnnExecuTorchBackendTest.cpp b/backends/qualcomm/tests/QnnExecuTorchBackendTest.cpp new file mode 100644 index 00000000000..a99ed7a70da --- /dev/null +++ b/backends/qualcomm/tests/QnnExecuTorchBackendTest.cpp @@ -0,0 +1,252 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::qnn { + +class QnnExecuTorchBackendTest : public ::testing::Test { + protected: + using DelegateKey = QnnExecuTorchBackend::DelegateKey; + using SpillFillKey = QnnExecuTorchBackend::SpillFillKey; + using SpillFillGroup = QnnExecuTorchBackend::SpillFillGroup; + + std::optional acquire(const DelegateKey& key) { + return backend_.get_cached_delegate(key); + } + + void* publish( + const DelegateKey& key, + void* handle, + std::shared_ptr manager = nullptr, + std::shared_ptr group = nullptr) { + return backend_.add_cached_delegate(key, handle, manager, group); + } + + std::shared_ptr spill_fill_group(const SpillFillKey& key) { + return backend_.get_spill_fill_group(key); + } + + bool contains(void* handle) { + return backend_.is_cached_delegate(handle); + } + + bool release(void* handle) { + return backend_.erase_cached_delegate(handle).has_value(); + } + + void leave_registered_delegate_alive() { + auto* registered = static_cast( + runtime::get_backend_class("QnnBackend")); + std::shared_ptr manager(new int, [](void* handle) { + delete static_cast(handle); + std::_Exit(42); + }); + registered->add_cached_delegate({42, nullptr}, manager.get(), manager); + } + + QnnExecuTorchBackend backend_; + decltype(QnnExecuTorchBackend::delegate_map_rev_)& reverse_map_ = + backend_.delegate_map_rev_; +}; + +TEST_F( + QnnExecuTorchBackendTest, + RegisteredBackendSurvivesProcessExitWithLiveDelegate) { + ASSERT_NE(runtime::get_backend_class("QnnBackend"), nullptr); + EXPECT_EXIT( + { + leave_registered_delegate_alive(); + std::exit(0); + }, + ::testing::ExitedWithCode(0), + ""); +} + +TEST_F( + QnnExecuTorchBackendTest, + RemovesStaleReverseEntryWithoutDereferencingEnd) { + int storage; + DelegateKey key{42, nullptr}; + ASSERT_TRUE(reverse_map_.emplace(&storage, key).second); + EXPECT_TRUE(contains(&storage)); + EXPECT_FALSE(release(&storage)); + EXPECT_FALSE(contains(&storage)); + EXPECT_EQ(acquire(key), std::nullopt); + EXPECT_FALSE(release(&storage)); +} + +TEST_F(QnnExecuTorchBackendTest, SharesUntilLastMethodIsReleased) { + int storage; + DelegateKey key{42, nullptr}; + ASSERT_EQ(acquire(key), std::nullopt); + ASSERT_EQ(publish(key, &storage), &storage); + ASSERT_EQ(acquire(key), &storage); + EXPECT_FALSE(release(&storage)); + EXPECT_TRUE(contains(&storage)); + EXPECT_TRUE(release(&storage)); + EXPECT_FALSE(contains(&storage)); + EXPECT_EQ(acquire(key), std::nullopt); + EXPECT_FALSE(release(&storage)); + EXPECT_EQ(publish(key, &storage), &storage); + EXPECT_TRUE(release(&storage)); +} + +TEST_F(QnnExecuTorchBackendTest, DestroyReleasesLastManagerOutsideCacheMutex) { + int storage; + size_t destroyed = 0; + std::shared_ptr manager(&storage, [&](void*) { + ++destroyed; + EXPECT_FALSE(contains(&storage)); + }); + DelegateKey key{42, nullptr}; + ASSERT_EQ(publish(key, &storage, manager), &storage); + ASSERT_EQ(acquire(key), &storage); + manager.reset(); + + backend_.destroy(&storage); + EXPECT_EQ(destroyed, 0); + EXPECT_TRUE(contains(&storage)); + backend_.destroy(&storage); + EXPECT_EQ(destroyed, 1); + EXPECT_FALSE(contains(&storage)); + backend_.destroy(&storage); + EXPECT_EQ(destroyed, 1); +} + +TEST_F(QnnExecuTorchBackendTest, IsolatesAllocatorsWithIdenticalSignatures) { + std::array storage{}; + runtime::MemoryAllocator first_allocator(0, nullptr); + runtime::MemoryAllocator second_allocator(0, nullptr); + DelegateKey first_key{42, &first_allocator}; + DelegateKey second_key{42, &second_allocator}; + ASSERT_EQ(publish(first_key, &storage[0]), &storage[0]); + ASSERT_EQ(publish(second_key, &storage[1]), &storage[1]); + EXPECT_TRUE(release(&storage[0])); + EXPECT_FALSE(contains(&storage[0])); + EXPECT_TRUE(contains(&storage[1])); + EXPECT_EQ(acquire(second_key), &storage[1]); + EXPECT_FALSE(release(&storage[1])); + EXPECT_TRUE(release(&storage[1])); +} + +TEST_F(QnnExecuTorchBackendTest, ConcurrentPublicationKeepsOneManager) { + constexpr size_t count = 16; + std::array storage{}; + std::array published{}; + std::vector threads; + threads.reserve(count); + DelegateKey key{42, nullptr}; + for (size_t i = 0; i < count; ++i) { + threads.emplace_back([&, i] { published[i] = publish(key, &storage[i]); }); + } + for (auto& thread : threads) { + thread.join(); + } + void* winner = published[0]; + ASSERT_NE(winner, nullptr); + for (size_t i = 0; i < count; ++i) { + EXPECT_EQ(published[i], winner); + EXPECT_EQ(contains(&storage[i]), &storage[i] == winner); + EXPECT_EQ(release(winner), i == count - 1); + } + EXPECT_EQ(acquire(key), std::nullopt); + EXPECT_FALSE(contains(winner)); +} + +TEST_F(QnnExecuTorchBackendTest, SpillFillOwnerOutlivesAllGroupMembers) { + for (bool owner_first : {true, false}) { + std::array storage{}; + std::vector destroyed; + auto group = spill_fill_group({nullptr, 1024}); + std::shared_ptr owner( + &storage[0], [&](void*) { destroyed.push_back(0); }); + std::shared_ptr member( + &storage[1], [&](void*) { destroyed.push_back(1); }); + group->owner = owner; + ASSERT_EQ(publish({42, nullptr}, owner.get(), owner, group), &storage[0]); + ASSERT_EQ(publish({43, nullptr}, member.get(), member, group), &storage[1]); + std::weak_ptr previous_group = group; + owner.reset(); + member.reset(); + group.reset(); + + ASSERT_TRUE(release(&storage[owner_first ? 0 : 1])); + EXPECT_EQ( + destroyed, owner_first ? std::vector{} : std::vector{1}); + auto live_group = spill_fill_group({nullptr, 1024}); + EXPECT_EQ(live_group->owner.get(), &storage[0]); + live_group.reset(); + ASSERT_TRUE(release(&storage[owner_first ? 1 : 0])); + EXPECT_EQ(destroyed, (std::vector{1, 0})); + EXPECT_TRUE(previous_group.expired()); + EXPECT_EQ(spill_fill_group({nullptr, 1024})->owner, nullptr); + } +} + +TEST_F(QnnExecuTorchBackendTest, SpillFillGroupsAreIndependentAcrossModules) { + runtime::MemoryAllocator first_allocator(0, nullptr); + runtime::MemoryAllocator second_allocator(0, nullptr); + std::array storage{}; + std::vector destroyed; + auto first_group = spill_fill_group({&first_allocator, 1024}); + auto second_group = spill_fill_group({&second_allocator, 1024}); + EXPECT_NE(first_group, second_group); + EXPECT_NE(first_group, spill_fill_group({&first_allocator, 2048})); + std::shared_ptr first( + &storage[0], [&](void*) { destroyed.push_back(0); }); + std::shared_ptr second( + &storage[1], [&](void*) { destroyed.push_back(1); }); + first_group->owner = first; + second_group->owner = second; + ASSERT_EQ( + publish({42, &first_allocator}, first.get(), first, first_group), + &storage[0]); + ASSERT_EQ( + publish({42, &second_allocator}, second.get(), second, second_group), + &storage[1]); + first.reset(); + second.reset(); + first_group.reset(); + second_group.reset(); + ASSERT_TRUE(release(&storage[0])); + EXPECT_EQ(destroyed, (std::vector{0})); + EXPECT_TRUE(contains(&storage[1])); + EXPECT_EQ( + spill_fill_group({&second_allocator, 1024})->owner.get(), &storage[1]); + ASSERT_TRUE(release(&storage[1])); + EXPECT_EQ(destroyed, (std::vector{0, 1})); +} + +TEST_F(QnnExecuTorchBackendTest, ConcurrentLoadsAcquireOneSpillFillGroup) { + constexpr size_t count = 16; + std::array, count> groups; + std::vector threads; + threads.reserve(count); + for (size_t i = 0; i < count; ++i) { + threads.emplace_back( + [&, i] { groups[i] = spill_fill_group({nullptr, 1024}); }); + } + for (auto& thread : threads) { + thread.join(); + } + for (const auto& group : groups) { + EXPECT_EQ(group, groups[0]); + } +} + +} // namespace executorch::backends::qnn diff --git a/backends/qualcomm/tests/test_qnn_shared_delegate.py b/backends/qualcomm/tests/test_qnn_shared_delegate.py new file mode 100644 index 00000000000..4b108e347e4 --- /dev/null +++ b/backends/qualcomm/tests/test_qnn_shared_delegate.py @@ -0,0 +1,142 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import gc +import os +import tempfile +import unittest + +import torch + +from executorch.backends.qualcomm.serialization.qc_schema import QcomChipset +from executorch.backends.qualcomm.utils.utils import ( + generate_htp_compiler_spec, + generate_qnn_executorch_compiler_spec, + get_qnn_context_binary_alignment, + to_edge_transform_and_lower_to_qnn, +) +from executorch.exir import ExecutorchBackendConfig +from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch, + _load_program, +) + + +class ScaleShift(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * 2.0 + 1.0 + + +class TestQnnSharedDelegate(unittest.TestCase): + """Modules that load the same program must each keep running, whichever of + them is destroyed first.""" + + use_multi_contexts = False + + @classmethod + def setUpClass(cls) -> None: + cls.sample_input = torch.randn(1, 3, 8, 8) + model = ScaleShift() + cls.expected = model(cls.sample_input) + backend_options = generate_htp_compiler_spec( + use_fp16=True, use_multi_contexts=cls.use_multi_contexts + ) + if cls.use_multi_contexts: + backend_options.htp_options.max_sf_buf_size = 1024 * 1024 + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=backend_options, + ) + program = to_edge_transform_and_lower_to_qnn( + {"forward": model, "shift": model}, + { + "forward": (cls.sample_input,), + "shift": (cls.sample_input,), + }, + {"forward": compiler_specs, "shift": compiler_specs}, + ).to_executorch( + ExecutorchBackendConfig( + segment_alignment=get_qnn_context_binary_alignment() + ) + ) + cls.tmp_dir = tempfile.TemporaryDirectory() + cls.pte_path = os.path.join(cls.tmp_dir.name, "scale_shift.pte") + with open(cls.pte_path, "wb") as f: + f.write(program.buffer) + + @classmethod + def tearDownClass(cls) -> None: + cls.tmp_dir.cleanup() + + def assert_forward_matches(self, module) -> None: + (output,) = module.forward([self.sample_input]) + torch.testing.assert_close(output, self.expected, atol=1e-1, rtol=1e-1) + + def load_two_modules(self): + first = _load_for_executorch(self.pte_path) + self.assert_forward_matches(first) + second = _load_for_executorch(self.pte_path) + self.assert_forward_matches(second) + return first, second + + def test_methods_in_one_module_execute(self) -> None: + module = _load_for_executorch(self.pte_path) + self.assert_forward_matches(module) + (output,) = module.run_method("shift", [self.sample_input]) + torch.testing.assert_close(output, self.expected, atol=1e-1, rtol=1e-1) + + def test_second_module_runs_after_first_is_destroyed(self) -> None: + first, second = self.load_two_modules() + del first + gc.collect() + + for _ in range(3): + self.assert_forward_matches(second) + + def test_first_module_runs_after_second_is_destroyed(self) -> None: + first, second = self.load_two_modules() + del second + gc.collect() + + for _ in range(3): + self.assert_forward_matches(first) + + def test_new_module_runs_after_previous_ones_are_destroyed(self) -> None: + first, second = self.load_two_modules() + del first, second + gc.collect() + + self.assert_forward_matches(_load_for_executorch(self.pte_path)) + + def test_method_survives_sibling_unload_and_reload(self) -> None: + for first_name, second_name in (("forward", "shift"), ("shift", "forward")): + with self.subTest(first=first_name): + program = _load_program(self.pte_path) + first = program.load_method(first_name) + second = program.load_method(second_name) + del first + gc.collect() + (output,) = second([self.sample_input]) + torch.testing.assert_close(output, self.expected, atol=1e-1, rtol=1e-1) + first = program.load_method(first_name) + del second + gc.collect() + (output,) = first([self.sample_input]) + torch.testing.assert_close(output, self.expected, atol=1e-1, rtol=1e-1) + del first + gc.collect() + reloaded = program.load_method(second_name) + (output,) = reloaded([self.sample_input]) + torch.testing.assert_close(output, self.expected, atol=1e-1, rtol=1e-1) + del reloaded, program + + +class TestQnnSpillFillSharedDelegate(TestQnnSharedDelegate): + use_multi_contexts = True + + +if __name__ == "__main__": + unittest.main()