Validate and lock SharedBuffer allocations (#24463)

backends/qualcomm/runtime/SharedBuffer.cpp

@@ -10,6 +10,8 @@
 
 #include <pal/DynamicLoading.h>
 
+#include <limits>
+
 // 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<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 +74,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 +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<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) {
+  // NOLINTNEXTLINE(performance-no-int-to-ptr)
+  aligned_buf = reinterpret_cast<void*>(
+      alignTo(alignment, reinterpret_cast<intptr_t>(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<std::mutex> 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<std::mutex> 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<std::mutex> 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<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,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 <executorch/backends/qualcomm/runtime/SharedBuffer.h>
+#include <gtest/gtest.h>
+#include <atomic>
+#include <cstdint>
+#include <cstdlib>
+#include <limits>
+#include <memory>
+#include <thread>
+#include <vector>
+
+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();
+  }
+
+  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());
+}
+
+} // 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"),
+        ],
+    )