Validate and lock SharedBuffer allocations (pytorch#24463)

Summary:

## 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.

When the standard library supports throwing allocation failures, `AllocMem` now catches `std::bad_alloc` from either bookkeeping map insertion and returns `nullptr` after the scoped owner rolls back the allocation. The recovery test forces each insertion to fail and verifies that the RPC buffer is freed, both maps are empty, and allocation succeeds afterward. Builds without exception support retain their existing allocation-failure behavior.

Authored with Claude Code and Codex.

Differential Revision: D128567553

Author: alirezamonfared

backends/qualcomm/runtime/SharedBuffer.cpp

@@ -7,9 +7,13 @@
  */
 #include <executorch/backends/qualcomm/runtime/Logging.h>
 #include <executorch/backends/qualcomm/runtime/SharedBuffer.h>
+#include <executorch/runtime/platform/compiler.h>
 
 #include <pal/DynamicLoading.h>
 
+#include <limits>
+#include <new>
+
 // Refer to the QNN HTP Shared Buffer Tutorial
 // in Qualcomm® AI Engine Direct document
 constexpr uint8_t RPCMEM_HEAP_ID_SYSTEM = 25;
@@ -63,6 +67,7 @@ intptr_t alignTo(size_t alignment, intptr_t offset) {
 std::mutex SharedBuffer::init_mutex_;
 
 void* SharedBuffer::GetCustomMemBase(void* buf) {
+  std::lock_guard<std::mutex> lock(memory_mutex_);
   auto it = tensor_addr_to_custom_mem_.find(buf);
   if (it == tensor_addr_to_custom_mem_.end()) {
     return nullptr;
@@ -71,6 +76,7 @@ void* SharedBuffer::GetCustomMemBase(void* buf) {
 }
 
 size_t SharedBuffer::GetAllocatedSize(void* buf) {
+  std::lock_guard<std::mutex> lock(memory_mutex_);
   auto it = allocated_size_map_.find(buf);
   if (it == allocated_size_map_.end()) {
     return 0;
@@ -107,26 +113,56 @@ 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<int32_t>::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<std::mutex> lock(memory_mutex_);
   // do alignment:
   auto allocate_bytes = static_cast<int32_t>(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<void, decltype(release_allocation)> 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<void*>(
-      alignTo(alignment, reinterpret_cast<intptr_t>(buf)));
-  bool status = restore_map_.insert({aligned_buf, buf}).second;
-  allocated_size_map_.insert({aligned_buf, allocate_bytes});
-  if (!status) {
-    QNN_EXECUTORCH_LOG_ERROR("Failed to allocate the tensor by RPC memory.");
-    rpc_mem_free_(buf);
+  // NOLINTNEXTLINE(performance-no-int-to-ptr)
+  aligned_buf = reinterpret_cast<void*>(
+      alignTo(alignment, reinterpret_cast<intptr_t>(allocation.get())));
+#if ET_HAS_EXCEPTIONS
+  try {
+#endif
+    restored = restore_map_.insert({aligned_buf, allocation.get()}).second;
+    if (!restored) {
+      QNN_EXECUTORCH_LOG_ERROR("Failed to allocate the tensor by RPC memory.");
+      return nullptr;
+    }
+    allocated_size_map_.insert({aligned_buf, allocate_bytes});
+#if ET_HAS_EXCEPTIONS
+  } catch (const std::bad_alloc&) {
+    QNN_EXECUTORCH_LOG_ERROR("Failed to allocate shared memory bookkeeping.");
+    return nullptr;
   }
+#endif
+  restore_map_[aligned_buf] = allocation.release();
   return aligned_buf;
 }
 
 int32_t SharedBuffer::MemToFd(void* buf) {
+  std::lock_guard<std::mutex> lock(memory_mutex_);
   int32_t memFd = -1;
   if (!initialize_) {
     QNN_EXECUTORCH_LOG_ERROR("Shared memory not initialized.");
@@ -137,6 +173,7 @@ int32_t SharedBuffer::MemToFd(void* buf) {
 }
 
 void SharedBuffer::FreeMem(void* buf) {
+  std::lock_guard<std::mutex> lock(memory_mutex_);
   if (!initialize_) {
     QNN_EXECUTORCH_LOG_ERROR("Shared memory not initialized.");
   } else if (restore_map_.count(buf) == 0) {
@@ -157,6 +194,7 @@ void SharedBuffer::FreeMem(void* buf) {
 }
 
 bool SharedBuffer::IsAllocated(void* buf) {
+  std::lock_guard<std::mutex> lock(memory_mutex_);
   return restore_map_.count(buf) != 0U;
 }
 
@@ -190,6 +228,7 @@ Error SharedBuffer::Load() {
 }
 
 void SharedBuffer::AddCusomMemTensorAddr(void* tensor_addr, void* custom_mem) {
+  std::lock_guard<std::mutex> lock(memory_mutex_);
   bool status =
       tensor_addr_to_custom_mem_.insert({tensor_addr, custom_mem}).second;
   if (!status) {

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<void*, std::unordered_set<void*>>
       custom_mem_to_tensor_addr_;
+  std::mutex memory_mutex_;
   std::atomic_bool initialize_{false};
   static std::mutex init_mutex_;
 };

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)

backends/qualcomm/runtime/tests/SharedBufferTest.cpp

@@ -0,0 +1,227 @@
+/*
+ * 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 <executorch/backends/qualcomm/runtime/SharedBuffer.h>
+#include <executorch/runtime/platform/compiler.h>
+#include <gtest/gtest.h>
+#include <atomic>
+#include <cstdint>
+#include <cstdlib>
+#include <limits>
+#include <memory>
+#include <new>
+#include <thread>
+#include <vector>
+
+#if ET_HAS_EXCEPTIONS && !defined(_GLIBCXX_RESTRICT_EXCEPTIONS)
+namespace {
+thread_local int allocationsUntilFailure = -1;
+} // namespace
+
+void* operator new(std::size_t size) {
+  if (allocationsUntilFailure >= 0 && allocationsUntilFailure-- == 0) {
+    throw std::bad_alloc();
+  }
+  if (void* pointer = std::malloc(size == 0 ? 1 : size)) {
+    return pointer;
+  }
+  throw std::bad_alloc();
+}
+
+void operator delete(void* pointer) noexcept {
+  std::free(pointer);
+}
+
+void operator delete(void* pointer, std::size_t) noexcept {
+  std::free(pointer);
+}
+
+void* operator new(std::size_t size, const std::nothrow_t&) noexcept {
+  try {
+    return ::operator new(size);
+  } catch (const std::bad_alloc&) {
+    return nullptr;
+  }
+}
+
+void operator delete(void* pointer, const std::nothrow_t&) noexcept {
+  std::free(pointer);
+}
+#endif
+
+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<size_t>(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<uint8_t*>(std::malloc(static_cast<size_t>(bytes) + 1)) +
+          1;
+      return rpcAllocation_;
+    };
+    buffer_->rpc_mem_free_ = [](void* pointer) {
+      rpcFreed_ = pointer;
+      std::free(static_cast<uint8_t*>(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();
+  }
+
+#if ET_HAS_EXCEPTIONS && !defined(_GLIBCXX_RESTRICT_EXCEPTIONS)
+  void failBookkeepingAllocation(int successfulAllocations) {
+    buffer_->restore_map_.reserve(1);
+    buffer_->allocated_size_map_.reserve(1);
+    buffer_->rpc_mem_alloc_ = [](int, uint32_t, int bytes) -> void* {
+      rpcAllocation_ = std::malloc(static_cast<size_t>(bytes));
+      return rpcAllocation_;
+    };
+    buffer_->rpc_mem_free_ = [](void* pointer) {
+      rpcFreed_ = pointer;
+      std::free(pointer);
+    };
+    rpcFreed_ = nullptr;
+    allocationsUntilFailure = successfulAllocations;
+  }
+#endif
+
+  std::unique_ptr<SharedBuffer> buffer_;
+  inline static std::atomic<size_t> 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<uint8_t*>(buffer_->AllocMem(kBytes, kAlignment));
+  ASSERT_NE(aligned, nullptr);
+  auto* original = static_cast<uint8_t*>(rpcAllocation_);
+  EXPECT_EQ(reinterpret_cast<uintptr_t>(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<int32_t>::max(), alignof(void*)),
+      nullptr);
+  EXPECT_EQ(
+      buffer_->AllocMem(std::numeric_limits<size_t>::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<void*> buffers(kThreads * kAllocations, nullptr);
+  std::vector<std::thread> workers;
+  workers.reserve(kThreads);
+  std::atomic<bool> 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<uint8_t*>(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<uint8_t*>(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());
+}
+
+#if ET_HAS_EXCEPTIONS && !defined(_GLIBCXX_RESTRICT_EXCEPTIONS)
+TEST_F(SharedBufferTest, BookkeepingAllocationFailureRollsBackAndReturnsNull) {
+  // Reserved buckets leave one node allocation per map insertion.
+  for (int successfulAllocations : {0, 1}) {
+    failBookkeepingAllocation(successfulAllocations);
+    void* result = buffer_->AllocMem(64, alignof(void*));
+    allocationsUntilFailure = -1;
+    EXPECT_EQ(result, nullptr);
+    EXPECT_EQ(rpcFreed_, rpcAllocation_);
+    EXPECT_TRUE(empty());
+
+    void* retry = buffer_->AllocMem(64, alignof(void*));
+    ASSERT_NE(retry, nullptr);
+    buffer_->FreeMem(retry);
+    EXPECT_TRUE(empty());
+  }
+}
+#endif
+
+} // namespace executorch::backends::qnn

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"),
+        ],
+    )