Cover caller-stream TLS semantics

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

⏳ 10 Pending, 1 Unrelated Failure

As of commit 81949e9 with merge base 6cd7ecd ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-binary-size-linux-gcc / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128

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 4 out of 4 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (2)

runtime/executor/tensor_parser_aten.cpp:145

• In the DYNAMIC_UNBOUND path, the tensor is left CPU-tagged (CPU options + CPU allocator), even if extra_tensor_info serialized a CUDA device. Since the PTE is untrusted, a malformed program could specify CUDA + DYNAMIC_UNBOUND and this would silently return a CPU tensor that contradicts the serialized device, leading to downstream device-mismatch failures. Reject DYNAMIC_UNBOUND for non-CPU devices (or implement dynamic CUDA tensors).

  if (s_tensor->shape_dynamism() ==
      executorch_flatbuffer::TensorShapeDynamism::DYNAMIC_UNBOUND) {
    // Fully dynamic tensors get an allocator so aten kernels can resize them.
    // Device-delegate planned buffers are statically bounded, so a device
    // tensor never reaches this CPU-tagged path.

runtime/core/exec_aten/util/tensor_util_aten.cpp:140

• share_tensor_data now tags Storage's DataPtr with t_src.device(), but it still doesn’t validate that t_dst.device() matches. If a method input is CUDA-tagged (via extra_tensor_info) and the caller forwards a CPU tensor (or vice versa), this will create an inconsistent tensor (TensorImpl device/dispatch keys disagree with Storage DataPtr device). Add a fast device equality check and fail with InvalidArgument before adopting the source pointer.

  // Preserve the source device; hardcoding CPU would mis-tag a device input's
  // storage as host and the backend would later reject it.
  storage->set_data_ptr(at::DataPtr(t_src.mutable_data_ptr(), t_src.device()));

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