Validate and lock SharedBuffer allocations (#24463) by alirezamonfared · Pull Request #24463 · pytorch/executorch

alirezamonfared · 2026-11-02T20:35:35Z

Summary:

TL;DR

Serialize the QNN shared-buffer manager's bookkeeping, preserve zero-byte allocations, and return nullptr for invalid requests or recoverable bookkeeping failures. A duplicate pointer returned by the RPC allocator must leave the existing live allocation intact. Host regressions cover these contracts, and the existing OSS QNN Buck job now executes the suite instead of only building it.

Context

QnnExecuTorchAllocCustomMem allocates RPC-backed memory through the process-wide SharedBuffer singleton. Clients can associate tensor addresses with an allocation using QnnExecuTorchAddCustomMemTensorAddr; QnnManager::RegisterMem then selects direct ION registration or custom-memory registration during execution. QnnExecuTorchFreeCustomMem frees the original RPC pointer and removes the allocation and its tensor aliases.

The manager maintains four maps: aligned pointer to original RPC pointer, aligned pointer to allocated size, tensor address to custom-memory base, and custom-memory base to tensor addresses. These maps describe ownership and registration metadata; they are not owners of the client's execution lifetime. The client must retain an allocation and its aliases until every registration and execution using them has finished.

client allocation                           client execution
QnnExecuTorchAllocCustomMem                  QnnManager::RegisterMem
  -> SharedBuffer::AllocMem                    -> alias lookup
  -> RPC allocation [bytes + alignment]        -> RegisterCustomMem / RegisterIonMem
  -> aligned address + bookkeeping             -> QNN execution

client teardown, after registration and execution have finished
QnnExecuTorchFreeCustomMem -> FreeMem -> RPC free + removal of all aliases

This is allocator hardening for every QNN custom-memory client. It does not enable shared I/O, change exported model contents, or establish safe concurrent execution and freeing of the same allocation.

Problem

Map reads and writes previously had no common mutex. Calls on independent threads could mutate the same containers while another thread looked up registration metadata. Even locking each lookup separately does not make a series of size and file-descriptor reads one consistent observation.

Allocation bookkeeping also had failure cases. The RPC allocator takes a signed 32-bit byte count, but the requested size and alignment are size_t. For example, bytes = INT32_MAX - 63 and alignment = 64 require 2,147,483,648 bytes, which exceeds its limit by one byte. Alignment zero and non-power-of-two alignment are invalid. Zero bytes, however, were historically valid: the alignment padding still makes the RPC request positive, and rejecting that request breaks callers with empty I/O tensors.

The duplicate-address failure is distinct from allocation failure. If an allocator returns the same still-live pointer twice, failure to insert the second mapping must not free that pointer or overwrite the first allocation's size. A successful RPC allocation can also be followed by std::bad_alloc from either map insertion; that must release the new allocation and return nullptr when the standard library supports throwing allocation failures.

duplicate RPC address P

before: existing P -> first size
        second allocation returns P
        insertion fails -> free(P) -> first allocation now refers to freed memory

after:  existing P -> first size
        second allocation returns P
        reject with nullptr -> no free, no map changes
        first allocation's eventual FreeMem(P) frees it once

Solution

One memory_mutex_ protects allocation, freeing, alias bookkeeping and lookup operations. GetAllocationInfo obtains the allocated size and RPC file descriptor under one lock and returns {0, -1} for an allocation that is no longer recorded. RegisterCustomMem consumes that snapshot instead of taking independent size and descriptor snapshots. This makes the snapshot consistent; it does not retain memory after the lock is released. The public free API documents the client's lifetime obligation explicitly. Remove the obsolete GetAllocatedSize accessor now that runtime registration uses GetAllocationInfo. Tests read allocated_size from that same aggregate API; the live-allocation test checks it directly against requested bytes plus alignment, preserving an independent expected value.

Validate alignment and require bytes <= INT32_MAX - alignment before calling RPC. This avoids overflowing the sum and accepts the exact limit bytes + alignment == INT32_MAX. AllocMem(0, alignment) allocates only the padding. Reject a duplicate aligned address before taking ownership of it, because freeing it would destroy the existing live allocation.

For a new address, insert the original-pointer and size mappings while holding the mutex. If an insertion throws std::bad_alloc, erase the partial pointer mapping, free the new RPC allocation, log the failure and return nullptr. There is no redundant assignment after insertion. Toolchains without throwing allocation support retain their allocation-failure behavior; the regression that injects throwing operator new is enabled only where those exceptions are supported.

AllocMem(bytes, alignment), size units: bytes

validate -> lock -> RPC allocation -> align address
                                    |
                      address already live? -> nullptr; preserve live owner
                                    |
                      insert original-pointer mapping
                                    |
                      insert allocated-size mapping
                                    |
                           return aligned pointer

either insertion throws bad_alloc
  -> erase partial bookkeeping -> RPC free -> nullptr

The host tests exercise offset alignment and freeing the original pointer, zero-byte compatibility, invalid requests without an RPC call, the exact signed-32-bit limit, duplicate live addresses, allocation snapshots before and after free, and concurrent bookkeeping on distinct allocations. Where supported, an injected allocation failure exercises each insertion separately and checks cleanup and successful retry. Test teardown uses the configured RPC free callback, so an offset allocator is cleaned up correctly even after a failed assertion.

The OSS test-qnn-buck-build-linux job already installs the QNN SDK and builds the backend. It now runs buck2 test //backends/qualcomm/runtime/tests:shared_buffer_test. The suite uses fake RPC callbacks and needs no device or loaded RPC library; an actual green OSS run remains a separate verification gate.

Compatibility and evidence limits

No exported C function signature changes. Empty tensors keep their previous allocation behavior. Invalid alignment and requests exceeding the signed RPC allocation limit return nullptr; duplicate allocator addresses preserve the original allocation. The aggregate lookup is a C++ backend helper, not a change to model serialization.

The mutex protects bookkeeping and metadata snapshots, not a registration-to-execution lease. Clients must coordinate FreeMem with all users of that allocation. Host tests verify allocator contracts; they do not establish device registration behavior or performance. Concrete results and earlier device evidence are recorded separately in the Test Plan.

Authored with AI assistance from Codex and Claude Code.

Differential Revision: D128567553

pytorch-bot · 2026-11-02T20:35:40Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24463

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

linux-foundation-easycla · 2026-11-02T20:35:44Z

The committers listed above are authorized under a signed CLA.

✅ login: alirezamonfared / name: Ali Monfared (c0c300f)

meta-codesync · 2026-11-02T20:35:46Z

@alirezamonfared has exported this pull request. If you are a Meta employee, you can view the originating Diff in D128567553.

github-actions · 2026-11-02T20:36:37Z

This PR needs a `release notes:` label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Summary: Pull Request resolved: pytorch#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

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 Signed-off-by: Ali Monfared <alimonfared@meta.com>

psiddh · 2026-11-04T18:17:06Z

cc @qti-horodnic can you review this one

@alirezamonfared easycla is needed

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 Signed-off-by: Ali Monfared <alimonfared@meta.com>

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

qti-horodnic

Thanks for the fixes @alirezamonfared, please see the comments below. Mostly minor and testing-related

qti-horodnic · 2026-11-06T23:02:44Z

+    if (!restored) {
+      QNN_EXECUTORCH_LOG_ERROR("Failed to allocate the tensor by RPC memory.");
+      return nullptr;


Can we add a test where the fake rpc_mem_alloc_ returns the same pointer two times. Expect nullptr on the second call, the first entry still in both maps, and no free of the live buffer?

qti-horodnic · 2026-11-06T23:03:29Z

    return nullptr;
  }
+  const size_t max_bytes = std::numeric_limits<int32_t>::max();
+  if (bytes == 0 || alignment == 0 || (alignment & (alignment - 1)) != 0 ||


Before this PR, AllocMem(0, align) was successful. Now it returns nullptr, and qnn_executor_runner.cpp:442-447 aborts at ET_CHECK_MSG for a zero-byte I/O tensor with --shared_buffer. Can we keep bytes == 0 valid and allocate only the alignment padding?

qti-horodnic · 2026-11-06T23:04:19Z

  // remain linked to this custom memory.
  std::unordered_map<void*, std::unordered_set<void*>>
      custom_mem_to_tensor_addr_;
+  std::mutex memory_mutex_;


This lock makes each call atomic, but not a sequence of calls. QnnManager::RegisterCustomMem calls GetCustomMemBase, GetAllocatedSize and MemToFd with three separate locks, so a concurrent FreeMem can still happen between them

qti-horodnic · 2026-11-06T23:05:01Z

+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+def define_common_targets():
+    runtime.cxx_test(


@psiddh Can we add this test to the OSS CI?

qti-horodnic · 2026-11-06T23:05:44Z

+    return nullptr;
  }
+#endif
+  restore_map_[aligned_buf] = allocation.release();


Minor: The map already has this value from line 148.

qti-horodnic · 2026-11-06T23:06:29Z

-      RPCMEM_HEAP_ID_SYSTEM, RPCMEM_DEFAULT_FLAGS, allocate_bytes);
-  if (buf == nullptr) {
+  void* aligned_buf = nullptr;
+  bool restored = false;


Minor: can we rename this inserted? Would be more representative of its purpose

qti-horodnic · 2026-11-06T23:07:49Z

+  EXPECT_TRUE(empty());
+}
+
+TEST_F(SharedBufferTest, InvalidSizeAndAlignmentDoNotCallRpcAllocator) {


Can we add tests for alignment > INT32_MAX, and for the accepted limit bytes + alignment == INT32_MAX? Maybe add AllocMem(64, size_t{1} << 31) == nullptrand a limit case with a fake allocator that has no 1 MiB cap and expectallocationCalls_ == 1`.

qti-horodnic · 2026-11-06T23:08:24Z

+  }
+
+  void TearDown() override {
+    for (const auto& allocation : buffer_->restore_map_) {


After useOffsetRpcAllocator, a leaked entry has a malloc + 1 pointer. Thus std::free on this pointer is not valid. This occurs only after a test has already failed

qti-horodnic · 2026-11-06T23:08:53Z

+#include <thread>
+#include <vector>
+
+#if ET_HAS_EXCEPTIONS && !defined(_GLIBCXX_RESTRICT_EXCEPTIONS)


Minor: The local libstdc++ headers do not contain _GLIBCXX_RESTRICT_EXCEPTIONS

Summary: ## TL;DR Serialize the QNN shared-buffer manager's bookkeeping, preserve zero-byte allocations, and return `nullptr` for invalid requests or recoverable bookkeeping failures. A duplicate pointer returned by the RPC allocator must leave the existing live allocation intact. Host regressions cover these contracts, and the existing OSS QNN Buck job now executes the suite instead of only building it. ## Context `QnnExecuTorchAllocCustomMem` allocates RPC-backed memory through the process-wide `SharedBuffer` singleton. Clients can associate tensor addresses with an allocation using `QnnExecuTorchAddCustomMemTensorAddr`; `QnnManager::RegisterMem` then selects direct ION registration or custom-memory registration during execution. `QnnExecuTorchFreeCustomMem` frees the original RPC pointer and removes the allocation and its tensor aliases. The manager maintains four maps: aligned pointer to original RPC pointer, aligned pointer to allocated size, tensor address to custom-memory base, and custom-memory base to tensor addresses. These maps describe ownership and registration metadata; they are not owners of the client's execution lifetime. The client must retain an allocation and its aliases until every registration and execution using them has finished. ```text client allocation client execution QnnExecuTorchAllocCustomMem QnnManager::RegisterMem -> SharedBuffer::AllocMem -> alias lookup -> RPC allocation [bytes + alignment] -> RegisterCustomMem / RegisterIonMem -> aligned address + bookkeeping -> QNN execution client teardown, after registration and execution have finished QnnExecuTorchFreeCustomMem -> FreeMem -> RPC free + removal of all aliases ``` This is allocator hardening for every QNN custom-memory client. It does not enable shared I/O, change exported model contents, or establish safe concurrent execution and freeing of the same allocation. ## Problem Map reads and writes previously had no common mutex. Calls on independent threads could mutate the same containers while another thread looked up registration metadata. Even locking each lookup separately does not make a series of size and file-descriptor reads one consistent observation. Allocation bookkeeping also had failure cases. The RPC allocator takes a signed 32-bit byte count, but the requested size and alignment are `size_t`. For example, `bytes = INT32_MAX - 63` and `alignment = 64` require `2,147,483,648` bytes, which exceeds its limit by one byte. Alignment zero and non-power-of-two alignment are invalid. Zero **bytes**, however, were historically valid: the alignment padding still makes the RPC request positive, and rejecting that request breaks callers with empty I/O tensors. The duplicate-address failure is distinct from allocation failure. If an allocator returns the same still-live pointer twice, failure to insert the second mapping must not free that pointer or overwrite the first allocation's size. A successful RPC allocation can also be followed by `std::bad_alloc` from either map insertion; that must release the new allocation and return `nullptr` when the standard library supports throwing allocation failures. ```text duplicate RPC address P before: existing P -> first size second allocation returns P insertion fails -> free(P) -> first allocation now refers to freed memory after: existing P -> first size second allocation returns P reject with nullptr -> no free, no map changes first allocation's eventual FreeMem(P) frees it once ``` ## Solution One `memory_mutex_` protects allocation, freeing, alias bookkeeping and lookup operations. `GetAllocationInfo` obtains the allocated size and RPC file descriptor under one lock and returns `{0, -1}` for an allocation that is no longer recorded. `RegisterCustomMem` consumes that snapshot instead of taking independent size and descriptor snapshots. This makes the snapshot consistent; it does not retain memory after the lock is released. The public free API documents the client's lifetime obligation explicitly. Validate alignment and require `bytes <= INT32_MAX - alignment` before calling RPC. This avoids overflowing the sum and accepts the exact limit `bytes + alignment == INT32_MAX`. `AllocMem(0, alignment)` allocates only the padding. Reject a duplicate aligned address before taking ownership of it, because freeing it would destroy the existing live allocation. For a new address, insert the original-pointer and size mappings while holding the mutex. If an insertion throws `std::bad_alloc`, erase the partial pointer mapping, free the new RPC allocation, log the failure and return `nullptr`. There is no redundant assignment after insertion. Toolchains without throwing allocation support retain their allocation-failure behavior; the regression that injects throwing `operator new` is enabled only where those exceptions are supported. ```text AllocMem(bytes, alignment), size units: bytes validate -> lock -> RPC allocation -> align address | address already live? -> nullptr; preserve live owner | insert original-pointer mapping | insert allocated-size mapping | return aligned pointer either insertion throws bad_alloc -> erase partial bookkeeping -> RPC free -> nullptr ``` The host tests exercise offset alignment and freeing the original pointer, zero-byte compatibility, invalid requests without an RPC call, the exact signed-32-bit limit, duplicate live addresses, allocation snapshots before and after free, and concurrent bookkeeping on distinct allocations. Where supported, an injected allocation failure exercises each insertion separately and checks cleanup and successful retry. Test teardown uses the configured RPC free callback, so an offset allocator is cleaned up correctly even after a failed assertion. The OSS `test-qnn-buck-build-linux` job already installs the QNN SDK and builds the backend. It now runs `buck2 test //backends/qualcomm/runtime/tests:shared_buffer_test`. The suite uses fake RPC callbacks and needs no device or loaded RPC library; an actual green OSS run remains a separate verification gate. ## Compatibility and evidence limits No exported C function signature changes. Empty tensors keep their previous allocation behavior. Invalid alignment and requests exceeding the signed RPC allocation limit return `nullptr`; duplicate allocator addresses preserve the original allocation. The aggregate lookup is a C++ backend helper, not a change to model serialization. The mutex protects bookkeeping and metadata snapshots, not a registration-to-execution lease. Clients must coordinate `FreeMem` with all users of that allocation. Host tests verify allocator contracts; they do not establish device registration behavior or performance. Concrete results and earlier device evidence are recorded separately in the Test Plan. Authored with AI assistance from Codex and Claude Code. Differential Revision: D128567553 Signed-off-by: Ali Monfared <alimonfared@meta.com>

qti-horodnic

LGTM, thanks @alirezamonfared, only a minor comment.

qti-horodnic · 2026-11-09T19:51:59Z

@@ -61,9 +61,19 @@ class SharedBuffer final {

  size_t GetAllocatedSize(void* buf);


Minor: this method is only called by the tests, can it be removed, or use GetAllocationInfo in the tests?

Summary: ## TL;DR Serialize the QNN shared-buffer manager's bookkeeping, preserve zero-byte allocations, and return `nullptr` for invalid requests or recoverable bookkeeping failures. A duplicate pointer returned by the RPC allocator must leave the existing live allocation intact. Host regressions cover these contracts, and the existing OSS QNN Buck job now executes the suite instead of only building it. ## Context `QnnExecuTorchAllocCustomMem` allocates RPC-backed memory through the process-wide `SharedBuffer` singleton. Clients can associate tensor addresses with an allocation using `QnnExecuTorchAddCustomMemTensorAddr`; `QnnManager::RegisterMem` then selects direct ION registration or custom-memory registration during execution. `QnnExecuTorchFreeCustomMem` frees the original RPC pointer and removes the allocation and its tensor aliases. The manager maintains four maps: aligned pointer to original RPC pointer, aligned pointer to allocated size, tensor address to custom-memory base, and custom-memory base to tensor addresses. These maps describe ownership and registration metadata; they are not owners of the client's execution lifetime. The client must retain an allocation and its aliases until every registration and execution using them has finished. ```text client allocation client execution QnnExecuTorchAllocCustomMem QnnManager::RegisterMem -> SharedBuffer::AllocMem -> alias lookup -> RPC allocation [bytes + alignment] -> RegisterCustomMem / RegisterIonMem -> aligned address + bookkeeping -> QNN execution client teardown, after registration and execution have finished QnnExecuTorchFreeCustomMem -> FreeMem -> RPC free + removal of all aliases ``` This is allocator hardening for every QNN custom-memory client. It does not enable shared I/O, change exported model contents, or establish safe concurrent execution and freeing of the same allocation. ## Problem Map reads and writes previously had no common mutex. Calls on independent threads could mutate the same containers while another thread looked up registration metadata. Even locking each lookup separately does not make a series of size and file-descriptor reads one consistent observation. Allocation bookkeeping also had failure cases. The RPC allocator takes a signed 32-bit byte count, but the requested size and alignment are `size_t`. For example, `bytes = INT32_MAX - 63` and `alignment = 64` require `2,147,483,648` bytes, which exceeds its limit by one byte. Alignment zero and non-power-of-two alignment are invalid. Zero **bytes**, however, were historically valid: the alignment padding still makes the RPC request positive, and rejecting that request breaks callers with empty I/O tensors. The duplicate-address failure is distinct from allocation failure. If an allocator returns the same still-live pointer twice, failure to insert the second mapping must not free that pointer or overwrite the first allocation's size. A successful RPC allocation can also be followed by `std::bad_alloc` from either map insertion; that must release the new allocation and return `nullptr` when the standard library supports throwing allocation failures. ```text duplicate RPC address P before: existing P -> first size second allocation returns P insertion fails -> free(P) -> first allocation now refers to freed memory after: existing P -> first size second allocation returns P reject with nullptr -> no free, no map changes first allocation's eventual FreeMem(P) frees it once ``` ## Solution One `memory_mutex_` protects allocation, freeing, alias bookkeeping and lookup operations. `GetAllocationInfo` obtains the allocated size and RPC file descriptor under one lock and returns `{0, -1}` for an allocation that is no longer recorded. `RegisterCustomMem` consumes that snapshot instead of taking independent size and descriptor snapshots. This makes the snapshot consistent; it does not retain memory after the lock is released. The public free API documents the client's lifetime obligation explicitly. Remove the obsolete `GetAllocatedSize` accessor now that runtime registration uses `GetAllocationInfo`. Tests read `allocated_size` from that same aggregate API; the live-allocation test checks it directly against requested bytes plus alignment, preserving an independent expected value. Validate alignment and require `bytes <= INT32_MAX - alignment` before calling RPC. This avoids overflowing the sum and accepts the exact limit `bytes + alignment == INT32_MAX`. `AllocMem(0, alignment)` allocates only the padding. Reject a duplicate aligned address before taking ownership of it, because freeing it would destroy the existing live allocation. For a new address, insert the original-pointer and size mappings while holding the mutex. If an insertion throws `std::bad_alloc`, erase the partial pointer mapping, free the new RPC allocation, log the failure and return `nullptr`. There is no redundant assignment after insertion. Toolchains without throwing allocation support retain their allocation-failure behavior; the regression that injects throwing `operator new` is enabled only where those exceptions are supported. ```text AllocMem(bytes, alignment), size units: bytes validate -> lock -> RPC allocation -> align address | address already live? -> nullptr; preserve live owner | insert original-pointer mapping | insert allocated-size mapping | return aligned pointer either insertion throws bad_alloc -> erase partial bookkeeping -> RPC free -> nullptr ``` The host tests exercise offset alignment and freeing the original pointer, zero-byte compatibility, invalid requests without an RPC call, the exact signed-32-bit limit, duplicate live addresses, allocation snapshots before and after free, and concurrent bookkeeping on distinct allocations. Where supported, an injected allocation failure exercises each insertion separately and checks cleanup and successful retry. Test teardown uses the configured RPC free callback, so an offset allocator is cleaned up correctly even after a failed assertion. The OSS `test-qnn-buck-build-linux` job already installs the QNN SDK and builds the backend. It now runs `buck2 test //backends/qualcomm/runtime/tests:shared_buffer_test`. The suite uses fake RPC callbacks and needs no device or loaded RPC library; an actual green OSS run remains a separate verification gate. ## Compatibility and evidence limits No exported C function signature changes. Empty tensors keep their previous allocation behavior. Invalid alignment and requests exceeding the signed RPC allocation limit return `nullptr`; duplicate allocator addresses preserve the original allocation. The aggregate lookup is a C++ backend helper, not a change to model serialization. The mutex protects bookkeeping and metadata snapshots, not a registration-to-execution lease. Clients must coordinate `FreeMem` with all users of that allocation. Host tests verify allocator contracts; they do not establish device registration behavior or performance. Concrete results and earlier device evidence are recorded separately in the Test Plan. Authored with AI assistance from Codex and Claude Code. Differential Revision: D128567553 Signed-off-by: Ali Monfared <alimonfared@meta.com>

alirezamonfared requested a review from psiddh as a code owner November 2, 2026 20:35

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 2, 2026

meta-codesync Bot added the meta-exported label Nov 2, 2026

alirezamonfared force-pushed the export-D128567553 branch from f54891c to 87da567 Compare November 2, 2026 22:27

meta-codesync Bot changed the title ~~Validate and lock SharedBuffer allocations~~ Validate and lock SharedBuffer allocations (#24463) Nov 2, 2026

psiddh added the module: qnn Issues related to Qualcomm's QNN delegate and code under backends/qualcomm/ label Nov 3, 2026

alirezamonfared force-pushed the export-D128567553 branch from 87da567 to c0c300f Compare November 4, 2026 18:16

alirezamonfared force-pushed the export-D128567553 branch 2 times, most recently from f492733 to fb8307e Compare November 6, 2026 03:05

qti-horodnic suggested changes Nov 6, 2026

View reviewed changes

alirezamonfared force-pushed the export-D128567553 branch from fb8307e to 6727436 Compare November 9, 2026 17:32

qti-horodnic approved these changes Nov 9, 2026

View reviewed changes

alirezamonfared force-pushed the export-D128567553 branch from 6727436 to f1eb18b Compare November 9, 2026 21:11

psiddh approved these changes Nov 9, 2026

View reviewed changes

alirezamonfared merged commit fce7099 into pytorch:main Nov 9, 2026
222 of 223 checks passed

		@@ -61,9 +61,19 @@ class SharedBuffer final {

		size_t GetAllocatedSize(void* buf);

Conversation

alirezamonfared commented Nov 2, 2026 • edited by meta-codesync Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

TL;DR

Context

Problem

Solution

Compatibility and evidence limits

Uh oh!

pytorch-bot Bot commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24463

Uh oh!

linux-foundation-easycla Bot commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

meta-codesync Bot commented Nov 2, 2026

Uh oh!

github-actions Bot commented Nov 2, 2026

This PR needs a release notes: label

Uh oh!

psiddh commented Nov 4, 2026

Uh oh!

qti-horodnic left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

qti-horodnic left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

alirezamonfared commented Nov 2, 2026 •

edited by meta-codesync Bot

Loading

pytorch-bot Bot commented Nov 2, 2026 •

edited

Loading

linux-foundation-easycla Bot commented Nov 2, 2026 •

edited

Loading

This PR needs a `release notes:` label