From 87da567da199cbda5a071351d556ac4a338999ef Mon Sep 17 00:00:00 2001 From: Ali Monfared Date: Mon, 2 Nov 2026 14:26:16 -0800 Subject: [PATCH] Validate and lock SharedBuffer allocations (#24463) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/24463 ## TL;DR Make the QNN `SharedBuffer` manager safe for concurrent callers. Reject invalid `AllocMem` requests before they reach the RPC allocator. ## Problem `SharedBuffer` is a process-wide singleton. Its four maps (aligned-to-original pointer, allocated size, and the tensor-to-custom-memory map and its reverse) had no synchronization. Applications call `QnnExecuTorchAllocCustomMem`, `QnnExecuTorchAddCustomMemTensorAddr` and `QnnExecuTorchFreeCustomMem` while `QnnManager` reads the same maps to register tensor memory. Allocate, register, lookup and free on different threads could therefore race. `AllocMem` also had three validation and bookkeeping bugs: - It accepted a zero size, a zero or non-power-of-two alignment, and sizes whose alignment padding overflows the `int32_t` rpcmem size. - It recorded an allocated size before confirming the pointer insertion. - When that insertion failed, it freed the allocation but still returned the aligned pointer. ## Solution A per-instance `memory_mutex_` now guards every map access. Before calling the RPC allocator, `AllocMem` rejects a zero size, a zero or non-power-of-two alignment, and any request where `bytes + alignment` exceeds `INT32_MAX`. A new allocation stays in a scoped owner until its bookkeeping is installed. If installing the bookkeeping fails after the RPC allocation succeeds, the scoped owner frees the original rpcmem pointer, removes any partial entries, and `AllocMem` returns null. A new `shared_buffer_test` drives the manager through fake RPC callbacks, so it needs no device. Authored with Claude Code and Codex. Differential Revision: D128567553 --- backends/qualcomm/runtime/SharedBuffer.cpp | 46 ++++-- backends/qualcomm/runtime/SharedBuffer.h | 2 + backends/qualcomm/runtime/tests/BUCK | 7 + .../runtime/tests/SharedBufferTest.cpp | 153 ++++++++++++++++++ backends/qualcomm/runtime/tests/targets.bzl | 15 ++ 5 files changed, 214 insertions(+), 9 deletions(-) create mode 100644 backends/qualcomm/runtime/tests/BUCK create mode 100644 backends/qualcomm/runtime/tests/SharedBufferTest.cpp create mode 100644 backends/qualcomm/runtime/tests/targets.bzl diff --git a/backends/qualcomm/runtime/SharedBuffer.cpp b/backends/qualcomm/runtime/SharedBuffer.cpp index 186cae26e2b..ed29f1ff43c 100644 --- a/backends/qualcomm/runtime/SharedBuffer.cpp +++ b/backends/qualcomm/runtime/SharedBuffer.cpp @@ -10,6 +10,8 @@ #include +#include + // Refer to the QNN HTP Shared Buffer Tutorial // in Qualcomm® AI Engine Direct document constexpr uint8_t RPCMEM_HEAP_ID_SYSTEM = 25; @@ -63,6 +65,7 @@ intptr_t alignTo(size_t alignment, intptr_t offset) { std::mutex SharedBuffer::init_mutex_; void* SharedBuffer::GetCustomMemBase(void* buf) { + std::lock_guard lock(memory_mutex_); auto it = tensor_addr_to_custom_mem_.find(buf); if (it == tensor_addr_to_custom_mem_.end()) { return nullptr; @@ -71,6 +74,7 @@ void* SharedBuffer::GetCustomMemBase(void* buf) { } size_t SharedBuffer::GetAllocatedSize(void* buf) { + std::lock_guard lock(memory_mutex_); auto it = allocated_size_map_.find(buf); if (it == allocated_size_map_.end()) { return 0; @@ -107,26 +111,47 @@ void* SharedBuffer::AllocMem(size_t bytes, size_t alignment) { QNN_EXECUTORCH_LOG_ERROR("Shared memory not initialized."); return nullptr; } + const size_t max_bytes = std::numeric_limits::max(); + if (bytes == 0 || alignment == 0 || (alignment & (alignment - 1)) != 0 || + alignment > max_bytes || bytes > max_bytes - alignment) { + QNN_EXECUTORCH_LOG_ERROR("Invalid shared memory size or alignment."); + return nullptr; + } + std::lock_guard lock(memory_mutex_); // do alignment: auto allocate_bytes = static_cast(bytes + alignment); - void* buf = rpc_mem_alloc_( - RPCMEM_HEAP_ID_SYSTEM, RPCMEM_DEFAULT_FLAGS, allocate_bytes); - if (buf == nullptr) { + void* aligned_buf = nullptr; + bool restored = false; + auto release_allocation = [&](void* original) { + if (restored) { + restore_map_.erase(aligned_buf); + allocated_size_map_.erase(aligned_buf); + } + rpc_mem_free_(original); + }; + std::unique_ptr allocation( + rpc_mem_alloc_( + RPCMEM_HEAP_ID_SYSTEM, RPCMEM_DEFAULT_FLAGS, allocate_bytes), + release_allocation); + if (!allocation) { QNN_EXECUTORCH_LOG_WARN("Failed to allocate the tensor by RPC memory."); return nullptr; } - auto aligned_buf = reinterpret_cast( - alignTo(alignment, reinterpret_cast(buf))); - bool status = restore_map_.insert({aligned_buf, buf}).second; - allocated_size_map_.insert({aligned_buf, allocate_bytes}); - if (!status) { + // NOLINTNEXTLINE(performance-no-int-to-ptr) + aligned_buf = reinterpret_cast( + alignTo(alignment, reinterpret_cast(allocation.get()))); + restored = restore_map_.insert({aligned_buf, allocation.get()}).second; + if (!restored) { QNN_EXECUTORCH_LOG_ERROR("Failed to allocate the tensor by RPC memory."); - rpc_mem_free_(buf); + return nullptr; } + allocated_size_map_.insert({aligned_buf, allocate_bytes}); + restore_map_[aligned_buf] = allocation.release(); return aligned_buf; } int32_t SharedBuffer::MemToFd(void* buf) { + std::lock_guard lock(memory_mutex_); int32_t memFd = -1; if (!initialize_) { QNN_EXECUTORCH_LOG_ERROR("Shared memory not initialized."); @@ -137,6 +162,7 @@ int32_t SharedBuffer::MemToFd(void* buf) { } void SharedBuffer::FreeMem(void* buf) { + std::lock_guard lock(memory_mutex_); if (!initialize_) { QNN_EXECUTORCH_LOG_ERROR("Shared memory not initialized."); } else if (restore_map_.count(buf) == 0) { @@ -157,6 +183,7 @@ void SharedBuffer::FreeMem(void* buf) { } bool SharedBuffer::IsAllocated(void* buf) { + std::lock_guard lock(memory_mutex_); return restore_map_.count(buf) != 0U; } @@ -190,6 +217,7 @@ Error SharedBuffer::Load() { } void SharedBuffer::AddCusomMemTensorAddr(void* tensor_addr, void* custom_mem) { + std::lock_guard lock(memory_mutex_); bool status = tensor_addr_to_custom_mem_.insert({tensor_addr, custom_mem}).second; if (!status) { diff --git a/backends/qualcomm/runtime/SharedBuffer.h b/backends/qualcomm/runtime/SharedBuffer.h index 6bf06a6350b..dad855aa188 100644 --- a/backends/qualcomm/runtime/SharedBuffer.h +++ b/backends/qualcomm/runtime/SharedBuffer.h @@ -64,6 +64,7 @@ class SharedBuffer final { void* GetCustomMemBase(void* buf); private: + friend class SharedBufferTest; SharedBuffer() = default; // dlopen RPCMem library and dlysm required functions @@ -88,6 +89,7 @@ class SharedBuffer final { // remain linked to this custom memory. std::unordered_map> custom_mem_to_tensor_addr_; + std::mutex memory_mutex_; std::atomic_bool initialize_{false}; static std::mutex init_mutex_; }; diff --git a/backends/qualcomm/runtime/tests/BUCK b/backends/qualcomm/runtime/tests/BUCK new file mode 100644 index 00000000000..ae833e41d49 --- /dev/null +++ b/backends/qualcomm/runtime/tests/BUCK @@ -0,0 +1,7 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +non_fbcode_target(_kind = define_common_targets) +fbcode_target(_kind = define_common_targets) diff --git a/backends/qualcomm/runtime/tests/SharedBufferTest.cpp b/backends/qualcomm/runtime/tests/SharedBufferTest.cpp new file mode 100644 index 00000000000..e1b2ad8cc08 --- /dev/null +++ b/backends/qualcomm/runtime/tests/SharedBufferTest.cpp @@ -0,0 +1,153 @@ +/* + * 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 SharedBufferTest : public ::testing::Test { + protected: + void SetUp() override { + allocationCalls_ = 0; + buffer_.reset(new SharedBuffer()); + buffer_->rpc_mem_alloc_ = [](int, uint32_t, int bytes) -> void* { + ++allocationCalls_; + return bytes > 0 && bytes <= 1024 * 1024 + ? std::malloc(static_cast(bytes)) + : nullptr; + }; + buffer_->rpc_mem_free_ = [](void* pointer) { std::free(pointer); }; + buffer_->rpc_mem_to_fd_ = [](void*) { return 17; }; + buffer_->SetInitialize(true); + } + + void TearDown() override { + for (const auto& allocation : buffer_->restore_map_) { + std::free(allocation.second); + } + buffer_->SetInitialize(false); + } + + void useOffsetRpcAllocator() { + buffer_->rpc_mem_alloc_ = [](int, uint32_t, int bytes) -> void* { + rpcAllocation_ = + static_cast(std::malloc(static_cast(bytes) + 1)) + + 1; + return rpcAllocation_; + }; + buffer_->rpc_mem_free_ = [](void* pointer) { + rpcFreed_ = pointer; + std::free(static_cast(pointer) - 1); + }; + } + + bool empty() const { + return buffer_->restore_map_.empty() && + buffer_->allocated_size_map_.empty() && + buffer_->tensor_addr_to_custom_mem_.empty() && + buffer_->custom_mem_to_tensor_addr_.empty(); + } + + std::unique_ptr buffer_; + inline static std::atomic allocationCalls_{0}; + inline static void* rpcAllocation_{nullptr}; + inline static void* rpcFreed_{nullptr}; +}; + +TEST_F(SharedBufferTest, AlignsOffsetRpcAllocationAndFreesTheOriginal) { + constexpr size_t kAlignment = 64; + constexpr size_t kBytes = 128; + useOffsetRpcAllocator(); + + auto* aligned = static_cast(buffer_->AllocMem(kBytes, kAlignment)); + ASSERT_NE(aligned, nullptr); + auto* original = static_cast(rpcAllocation_); + EXPECT_EQ(reinterpret_cast(aligned) % kAlignment, 0); + EXPECT_GT(aligned, original); + EXPECT_LE(aligned + kBytes, original + buffer_->GetAllocatedSize(aligned)); + buffer_->FreeMem(aligned); + + EXPECT_EQ(rpcFreed_, original); + EXPECT_TRUE(empty()); +} + +TEST_F(SharedBufferTest, InvalidSizeAndAlignmentDoNotCallRpcAllocator) { + EXPECT_EQ(buffer_->AllocMem(0, alignof(void*)), nullptr); + EXPECT_EQ(buffer_->AllocMem(64, 0), nullptr); + EXPECT_EQ(buffer_->AllocMem(64, 3), nullptr); + EXPECT_EQ( + buffer_->AllocMem(std::numeric_limits::max(), alignof(void*)), + nullptr); + EXPECT_EQ( + buffer_->AllocMem(std::numeric_limits::max(), alignof(void*)), + nullptr); + EXPECT_EQ(allocationCalls_.load(), 0); + EXPECT_TRUE(empty()); +} + +TEST_F(SharedBufferTest, ConcurrentAllocateRegisterLookupAndFree) { + constexpr size_t kThreads = 4; + constexpr size_t kAllocations = 256; + constexpr size_t kBytes = 64; + std::vector buffers(kThreads * kAllocations, nullptr); + std::vector workers; + workers.reserve(kThreads); + std::atomic start{false}; + for (size_t worker = 0; worker < kThreads; ++worker) { + workers.emplace_back([&, worker] { + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + for (size_t index = worker; index < buffers.size(); index += kThreads) { + void* pointer = buffer_->AllocMem(kBytes, alignof(void*)); + ASSERT_NE(pointer, nullptr); + buffers[index] = pointer; + void* alias = static_cast(pointer) + sizeof(uint64_t); + buffer_->AddCusomMemTensorAddr(alias, pointer); + EXPECT_TRUE(buffer_->IsAllocated(pointer)); + EXPECT_EQ(buffer_->GetCustomMemBase(alias), pointer); + EXPECT_GE(buffer_->GetAllocatedSize(pointer), kBytes); + EXPECT_EQ(buffer_->MemToFd(pointer), 17); + } + }); + } + start.store(true, std::memory_order_release); + for (auto& worker : workers) { + worker.join(); + } + EXPECT_EQ(allocationCalls_.load(), buffers.size()); + + workers.clear(); + for (size_t worker = 0; worker < kThreads; ++worker) { + workers.emplace_back([&, worker] { + for (size_t index = worker; index < buffers.size(); index += kThreads) { + void* pointer = buffers[index]; + ASSERT_NE(pointer, nullptr); + void* alias = static_cast(pointer) + sizeof(uint64_t); + buffer_->FreeMem(pointer); + EXPECT_FALSE(buffer_->IsAllocated(pointer)); + EXPECT_EQ(buffer_->GetAllocatedSize(pointer), 0); + EXPECT_EQ(buffer_->GetCustomMemBase(alias), nullptr); + } + }); + } + for (auto& worker : workers) { + worker.join(); + } + EXPECT_TRUE(empty()); +} + +} // namespace executorch::backends::qnn diff --git a/backends/qualcomm/runtime/tests/targets.bzl b/backends/qualcomm/runtime/tests/targets.bzl new file mode 100644 index 00000000000..cc4db4f0787 --- /dev/null +++ b/backends/qualcomm/runtime/tests/targets.bzl @@ -0,0 +1,15 @@ +load("@fbsource//tools/build_defs:default_platform_defs.bzl", "CXX") +load("@fbsource//xplat/executorch/backends/qualcomm/third-party:third_party_libs.bzl", "qnn_third_party_dep") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + runtime.cxx_test( + name = "shared_buffer_test", + srcs = ["SharedBufferTest.cpp"], + compatible_with = [] if runtime.is_oss else ["ovr_config//os:linux"], + platforms = [CXX], + deps = [ + "//executorch/backends/qualcomm/runtime:runtime", + qnn_third_party_dep("api"), + ], + )