[alirezamonfared]

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.

Authored with Claude Code and Codex.

Differential Revision: D128567553

[pytorch-bot]

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

❌ 1 New Failure, 1 Unrelated Failure

As of commit 87da567 with merge base 893d815 ():

NEW FAILURE - The following job has failed:

• pull / android / run-emulator (gh)
  adb: device offline

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / nxp-mcuxpresso-test / linux-job (gh) (trunk failure)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

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

[linux-foundation-easycla]

• ❌ The email address for the commit (87da567) is not linked to the GitHub account, preventing the EasyCLA check. Consult this Help Article and GitHub Help to resolve. (To view the commit's email address, add .patch at the end of this PR page's URL.) For further assistance with EasyCLA, please visit our EasyCLA portal and chat with our support bot.

[meta-codesync]

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

[github-actions]

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.