From 2f58db187af391d31935934978e8905e096c369b Mon Sep 17 00:00:00 2001 From: Ali Monfared Date: Tue, 3 Nov 2026 21:40:36 -0800 Subject: [PATCH] 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. 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, same-module sharing, cleanup/reload, and method unload/reload on the HTP simulator. C++ tests cover 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 --- .../qualcomm/runtime/QnnExecuTorchBackend.cpp | 143 ++++++++++--- .../qualcomm/runtime/QnnExecuTorchBackend.h | 50 ++++- 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 | 191 ++++++++++++++++++ .../tests/test_qnn_shared_delegate.py | 160 +++++++++++++++ 11 files changed, 554 insertions(+), 60 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..a034a836fb5 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); - } + // The returned entry releases its context and spill-fill group outside + // mutex_. + (void)erase_cached_delegate(handle); } executorch::runtime::Error QnnExecuTorchBackend::set_option( @@ -365,23 +390,75 @@ 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; + } + auto delegate_iter = delegate_map_.find(iter->second); + if (--delegate_iter->second.users != 0) { + return std::nullopt; } - delegate_map_.erase(iter->second); - delegate_map_rev_.erase(handle); + DelegateEntry entry = std::move(delegate_iter->second); + delegate_map_.erase(delegate_iter); + delegate_map_rev_.erase(iter); + return entry; } namespace { diff --git a/backends/qualcomm/runtime/QnnExecuTorchBackend.h b/backends/qualcomm/runtime/QnnExecuTorchBackend.h index e3548c8752b..62c8d1b5250 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() {}; executorch::runtime::Result init( executorch::runtime::BackendInitContext& context, @@ -51,16 +55,48 @@ 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; + 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 ff9fee63178..76a920cf6ae 100644 --- a/backends/qualcomm/runtime/QnnManager.cpp +++ b/backends/qualcomm/runtime/QnnManager.cpp @@ -342,7 +342,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( @@ -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, 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..77b2adcc545 --- /dev/null +++ b/backends/qualcomm/tests/QnnExecuTorchBackendTest.cpp @@ -0,0 +1,191 @@ +/* + * 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 + +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(); + } + + QnnExecuTorchBackend backend_; +}; + +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, 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..98a924c884b --- /dev/null +++ b/backends/qualcomm/tests/test_qnn_shared_delegate.py @@ -0,0 +1,160 @@ +# 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_share_delegate(self) -> None: + module = _load_for_executorch(self.pte_path) + self.assert_forward_matches(module) + + with tempfile.TemporaryFile() as captured: + saved_stderr = os.dup(2) + try: + os.dup2(captured.fileno(), 2) + (output,) = module.run_method("shift", [self.sample_input]) + finally: + os.dup2(saved_stderr, 2) + os.close(saved_stderr) + captured.seek(0) + log = captured.read().decode("utf-8", "replace") + + torch.testing.assert_close(output, self.expected, atol=1e-1, rtol=1e-1) + self.assertIn( + "Use cached delegate handle for current method: shift", + log, + ) + + 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()