@@ -986,6 +986,9 @@ jobs:
         # Build QNN backend with Buck
         buck2 build //backends/qualcomm/...
 
+        # Run the host allocator regressions without a device or RPC library.
+        buck2 test //backends/qualcomm/runtime/tests:shared_buffer_test
+
   test-arm-backend-docgen:
     needs: [docker-image, changed-files]
     if: |
 
@@ -89,7 +89,8 @@ QNN_EXECUTORCH_EXPORT void QnnExecuTorchAddCustomMemTensorAddr(
     void* tensor_addr,
     void* custom_mem);
 
-/// Free the allocated shared memory.
+/// Free the allocated shared memory. The caller must wait for all registration
+/// and execution using this allocation or its tensor aliases to finish.
 QNN_EXECUTORCH_EXPORT void QnnExecuTorchFreeCustomMem(void* buffer_ptr);
 
 /// Register the QNN backend with the ExecuTorch runtime living in the caller's
 
@@ -214,17 +214,13 @@ Error QnnManager::RegisterCustomMem(
         tensor_wrapper, data_ptr, pre_registered_handle);
   }
 
-  SharedBuffer& shared_buffer_manager = SharedBuffer::GetSharedBufferManager();
-
   size_t tensor_offset = info.pos;
-  size_t total_custom_mem_size =
-      shared_buffer_manager.GetAllocatedSize(custom_mem_base);
-
-  int32_t mem_fd = shared_buffer_manager.MemToFd(custom_mem_base);
+  SharedBuffer::AllocationInfo allocation =
+      SharedBuffer::GetSharedBufferManager().GetAllocationInfo(custom_mem_base);
   // Note: If obtaining the file descriptor fails, it may be due to memory not
   // being released with QnnExecuTorchFreeCustomMem. In this situation, we could
   // consider adding a map to monitor it.
-  if (mem_fd == -1) {
+  if (allocation.fd == -1) {
     QNN_EXECUTORCH_LOG_WARN(
         "Tensor name %s failed to get file descriptor.",
         tensor_wrapper->GetName().c_str());
@@ -234,9 +230,9 @@ Error QnnManager::RegisterCustomMem(
   ET_CHECK_OR_RETURN_ERROR(
       backend_params_ptr_->qnn_mem_manager_ptr_->RegisterCustomMem(
           tensor_wrapper,
-          mem_fd,
+          allocation.fd,
           data_ptr,
-          total_custom_mem_size,
+          allocation.allocated_size,
           tensor_offset,
           info) == Error::Ok,
       Internal,
 
@@ -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,19 +67,21 @@ 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;
   }
   return it->second;
 }
 
-size_t SharedBuffer::GetAllocatedSize(void* buf) {
+SharedBuffer::AllocationInfo SharedBuffer::GetAllocationInfo(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;
+    return {};
   }
-  return it->second;
+  return {it->second, rpc_mem_to_fd_(buf)};
 }
 
 SharedBuffer& SharedBuffer::GetSharedBufferManager() {
@@ -109,6 +115,13 @@ 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 (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_(
@@ -117,18 +130,34 @@ void* SharedBuffer::AllocMem(size_t bytes, size_t alignment) {
     QNN_EXECUTORCH_LOG_WARN("Failed to allocate the tensor by RPC memory.");
     return nullptr;
   }
+  // NOLINTNEXTLINE(performance-no-int-to-ptr)
   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.");
+  if (restore_map_.count(aligned_buf) != 0) {
+    // The RPC allocator returned memory that is still live, so freeing buf
+    // would release the existing allocation.
+    QNN_EXECUTORCH_LOG_ERROR(
+        "RPC memory %p is already allocated.", aligned_buf);
+    return nullptr;
+  }
+#if ET_HAS_EXCEPTIONS
+  try {
+#endif
+    restore_map_.insert({aligned_buf, buf});
+    allocated_size_map_.insert({aligned_buf, allocate_bytes});
+#if ET_HAS_EXCEPTIONS
+  } catch (const std::bad_alloc&) {
+    restore_map_.erase(aligned_buf);
     rpc_mem_free_(buf);
+    QNN_EXECUTORCH_LOG_ERROR("Failed to allocate shared memory bookkeeping.");
+    return nullptr;
   }
+#endif
   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.");
@@ -139,6 +168,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) {
@@ -159,6 +189,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;
 }
 
@@ -197,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) {
 
@@ -59,11 +59,19 @@ class SharedBuffer final {
   // memory handle is registered during execution
   void AddCusomMemTensorAddr(void* tensor_addr, void* custom_mem);
 
-  size_t GetAllocatedSize(void* buf);
+  struct AllocationInfo {
+    size_t allocated_size{0};
+    int32_t fd{-1};
+  };
+  // Reads the size and file descriptor of an allocation under one lock, so a
+  // concurrent FreeMem cannot leave them describing different states. Returns
+  // {0, -1} if buf is not allocated.
+  AllocationInfo GetAllocationInfo(void* buf);
 
   void* GetCustomMemBase(void* buf);
 
  private:
+  friend class SharedBufferTest;
   SharedBuffer() = default;
 
   // dlopen RPCMem library and dlysm required functions
@@ -88,6 +96,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_;
 };
 
@@ -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)