Cover caller-stream TLS semantics (#21291)

[shoumikhin]

Summary:

Add focused coverage for the caller-stream contract relied on by external CUDA
backends. Verify that an explicitly selected null/default stream remains
distinct from no guard, and that selections are isolated per execution thread.
Keep the fbcode and xplat mirrors identical.

Differential Revision: D113382078

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 9a00146 with merge base 8200b9e ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

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

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

[Copilot]

Pull request overview

This PR strengthens CUDA delegate correctness in coalesced / ATen-host scenarios by (1) adding focused tests for the caller-selected CUDA stream TLS contract and (2) ensuring tensor device metadata and AOTInductor shim symbol resolution are propagated consistently across runtime, backends, and build systems.

Changes:

• Add unit tests validating that CallerStreamGuard distinguishes “explicitly selected nullptr stream” from “no guard”, and that selection is thread-local.
• Propagate serialized device metadata into ATen tensor deserialization and preserve device tags when sharing tensor storage.
• Introduce/standardize an executorch_ AOTI shim symbol prefix (Buck + CMake + Python), and add an ATen-mode Buck target for the CUDA backend.

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
runtime/executor/tensor_parser_aten.cpp	Parse extra_tensor_info device fields and rebuild tensors with consistent device/dispatch metadata.
runtime/core/exec_aten/util/tensor_util_aten.cpp	Preserve source tensor device when sharing storage DataPtr.
backends/cuda/runtime/utils.h	Switch helper APIs to use mode-flexible executorch::aten::* tensor types.
backends/cuda/runtime/TARGETS	Add shim-prefix flags propagation and introduce cuda_backend_aten Buck target.
backends/cuda/runtime/shims/memory.h	Declare aoti_torch_empty_strided_pinned shim entrypoint.
backends/cuda/runtime/shims/memory.cpp	Implement pinned allocator shim (currently falls back to pageable allocation).
backends/cuda/runtime/cuda_backend.cpp	Register backend in ET_RUNTIME_NAMESPACE, avoid .so temp path collisions, and tighten CUDA IO device validation.
backends/cuda/cuda_backend.py	Set aot_inductor.shim_symbol_prefix and update custom-op shim symbol names to executorch_*.
backends/cuda/CMakeLists.txt	Add AOTI_SHIM_SYMBOL_PREFIX compile definition for CUDA shims.
backends/aoti/targets.bzl	Apply shim-prefix flags to AOTI common slim shims and force retention via link_whole.
backends/aoti/shim_symbol_prefix.bzl	Centralize the shim symbol prefix and corresponding preprocessor flag list.
backends/aoti/export.h	Add macro-based symbol renaming for aoti_torch_* shims when a prefix is defined.
backends/aoti/CMakeLists.txt	Apply AOTI_SHIM_SYMBOL_PREFIX to AOTI slim shims (incl. MSVC object-lib path).
backends/aoti/slim/cuda/test/test_cuda_stream_guard.cpp	Add TLS semantics tests for caller stream selection (including cross-thread isolation).
backends/aoti/aoti_delegate_handle.h	Use executorch::aten::Tensor as the mode-flexible host-tensor handle type.
backends/cuda/runtime/shims/tensor_attribute.h	Remove CUDA-specific tensor-attribute shim (now covered by common slim shims).
backends/cuda/runtime/shims/tensor_attribute.cpp	Remove CUDA-specific tensor-attribute shim implementation.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated 4 comments.

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated no new comments.

[Gasoonjia]

just curious why should be take care of the multi-cpu-thread scenerio? Is that something ET need to support?

[shoumikhin]

Not about multi-GPU or concurrent execution of one Method. getCallerStream()/CallerStreamGuard are thread_local by design so the CUDA and TRT delegates can each run on their own thread and stream. This is the only test that catches caller_stream_ accidentally losing thread_local; without it, concurrent inferences would silently stomp each other's stream. Keeping it.

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated no new comments.