Cover caller-stream TLS semantics (#21291) by shoumikhin · Pull Request #21291 · pytorch/executorch

shoumikhin · 2026-07-23T17:50:02Z

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 · 2026-07-23T17:50:06Z

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

❌ 1 Cancelled Job, 1 Pending, 2 Unclassified Failures

As of commit b01f69a with merge base 8134bb2 ():

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

Test CUDA Builds / test-models-cuda / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
RuntimeError: Command docker exec -t c1980bf11e09a0ba49c701e2f457260f4b06fe6e985cd1d289b84d27eed8361f /exec failed with exit code 1
Test CUDA Builds / unittest-cuda / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
backends/cuda/tests/test_tq4_sdpa.py::TestTQ4Sdpa::test_export_cuda

CANCELLED JOB - The following job was cancelled. Please retry:

pull / unittest-editable / linux / linux-job (gh)
##[error]The operation was canceled.

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

meta-codesync · 2026-07-23T17:50:11Z

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

github-actions · 2026-07-23T17:50:59Z

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.

shoumikhin · 2026-07-23T20:04:33Z


 # Define CUDA_AVAILABLE to use SlimTensor on GPU in common_shims_slim.h
 target_compile_definitions(aoti_cuda_shims PRIVATE CUDA_AVAILABLE=1)
+# Private executorch_ prefix on the slim aoti_torch_* shims so they never collide
+# with libtorch's in a coalesced process (see backends/aoti/export.h). Must match
+# aot_inductor.shim_symbol_prefix set in backends/cuda/cuda_backend.py.
+target_compile_definitions(aoti_cuda_shims PRIVATE AOTI_SHIM_SYMBOL_PREFIX=executorch_)



Fixed: changed AOTI_SHIM_SYMBOL_PREFIX on aoti_cuda_shims from PRIVATE to PUBLIC. Direct consumers that include the shim headers (the shim gtests) now compile the declarations under the same executorch_ prefix and link the DLL's prefixed exports, matching the Buck target which already exports the define to header users.

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

Copilot

Pull request overview

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

Summary: The CUDA backend writes each compiled AOTInductor library to a temporary file, named from its content key plus the process id, before loading it. Two identical CUDA partitions can share the same content key, so both delegates computed the same path; loading the second library overwrote the first while it was still in use, and symbol lookup could then crash. Append a per-process atomic counter to the temporary file name so every delegate in a process gets a distinct path. The file is still removed when the delegate is destroyed, as before; only the name changes. Differential Revision: D113322459

…ytorch#21289) Summary: In ATen mode, parseTensor built every planned tensor's storage DataPtr with a hardcoded c10::DeviceType::CPU, and its options with at::CPU(type). A tensor whose planned buffer lives on an accelerator (e.g. a CUDA-delegate planned buffer) therefore got a CPU-tagged data pointer, so the runtime treated device memory as host memory and the delegate rejected it (Method::init fails in getDeviceFromPtr on a CPU-tagged device pointer). Fix: read the device from the serialized tensor's extra_tensor_info (device_type/device_index, defaulting to CPU when absent, matching the portable parser), then build the tensor with a single at::from_blob(ptr, sizes, strides, /*storage_offset=*/0, deleteNothing, at::TensorOptions().dtype(type).device(device), /*target_device=*/device) so the storage DataPtr, the TensorImpl device (device(), is_cuda()), and the dispatch key all agree. Passing target_device makes from_blob skip getDeviceFromPtr, so the same path works for a real device pointer and for a null runtime-bound pointer without inspecting the pointer. TRT-only / CPU programs are unchanged (device defaults to CPU). Also update internal_set_tensor_data (tensor_util_aten.cpp) to preserve the source tensor's device instead of hardcoding CPU, so sharing a device input's storage keeps its device tag. fbcode and xplat copies kept identical. Differential Revision: D113384858

Summary: Give ExecuTorch's SlimTensor aoti_torch_* shims a private executorch_ prefix (disjoint from libtorch's at::Tensor aoti_torch_*) and add an ATen-mode CUDA delegate variant, so a coalesced TensorRT[ATen] + CUDA .pte binds the delegate blob to the slim shims and registers CudaBackend in the runtime::aten registry. Landed atomically: the shim export rename, the runtime_shims prefix flag, and the generated blob's aot_inductor.shim_symbol_prefix all activate together, so there is no intermediate state where some shims are prefixed and others are not. aoti/cuda shim layer: - export.h renames the shim definitions to executorch_aoti_torch_* under -DAOTI_SHIM_SYMBOL_PREFIX (self-contained, keeps the slim layer libtorch-free); shim_symbol_prefix.bzl is the single source of the build flag, applied via preprocessor_flags (own defs) + exported_ (header includers) on the shim targets. - common_shims_slim: link_whole so the dlopen'd blob's shims survive --gc-sections. memory: add aoti_torch_empty_strided_pinned. - utils.h / aoti_delegate_handle.h: mode-agnostic executorch::aten::Tensor host tensor (etensor in portable, at::Tensor under USE_ATEN_LIB). Portable unchanged. - remove shims/tensor_attribute.{cpp,h}: deduped into common_shims_slim. CUDA backend: - cuda_backend.cpp: register through ET_RUNTIME_NAMESPACE (runtime in portable, runtime::aten under USE_ATEN_LIB). Mirrors XNNPACK / TensorRTBackend. - TARGETS: add cuda_backend_aten (-DUSE_ATEN_LIB, _aten runtime deps). - cuda_backend.py: set aot_inductor.shim_symbol_prefix so the blob imports the executorch_ names the shim libs export. Differential Revision: D113382057

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

Copilot

Pull request overview

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

shoumikhin · 2026-07-23T21:49:13Z

+    device_index = static_cast<c10::DeviceIndex>(
+        s_tensor->extra_tensor_info()->device_index());
+    // Reject a negative accelerator index from the untrusted PTE; -1
+    // (any/current device) is not a valid serialized placement and would later
+    // confuse device matching.
+    ET_CHECK_OR_RETURN_ERROR(
+        device_type == c10::DeviceType::CPU || device_index >= 0,
+        InvalidProgram,
+        "Invalid device_index %" PRId8,
+        static_cast<int8_t>(device_index));


This is a false positive: in FlatBuffers, schema type byte maps to signed int8_t (it's ubyte that is uint8_t), and the generated ExtraTensorInfo.device_index accessor returns int8_t. So there is no uint8->int8 reinterpretation on the cast — a value like 200 cannot be stored in this field; the wire value is already a signed int8. The validation (device_type == CPU || device_index >= 0) correctly rejects a negative accelerator index (including the -1 "current device" sentinel) from an untrusted PTE, and the error prints the true int8 value. No change needed.

Copilot AI review requested due to automatic review settings July 23, 2026 17:50

shoumikhin requested review from JacobSzwejbka, kirklandsign and larryliu0820 as code owners July 23, 2026 17:50

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 Jul 23, 2026

meta-codesync Bot added the meta-exported label Jul 23, 2026

meta-codesync Bot had a problem deploying to cadence July 23, 2026 17:50 Failure

Copilot started reviewing on behalf of shoumikhin July 23, 2026 17:50 View session

Copilot AI reviewed Jul 23, 2026

View reviewed changes

meta-codesync Bot changed the title ~~Cover caller-stream TLS semantics~~ Cover caller-stream TLS semantics (#21291) Jul 23, 2026

Copilot AI review requested due to automatic review settings July 23, 2026 20:07

shoumikhin force-pushed the export-D113382078 branch from 74bb965 to b682d0f Compare July 23, 2026 20:07

Copilot started reviewing on behalf of shoumikhin July 23, 2026 20:08 View session

Copilot AI reviewed Jul 23, 2026

View reviewed changes

shoumikhin added 4 commits July 23, 2026 13:15

Copilot AI review requested due to automatic review settings July 23, 2026 20:16

shoumikhin force-pushed the export-D113382078 branch from b682d0f to b01f69a Compare July 23, 2026 20:16

Copilot started reviewing on behalf of shoumikhin July 23, 2026 20:16 View session

Copilot AI reviewed Jul 23, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Cover caller-stream TLS semantics (#21291)#21291

Cover caller-stream TLS semantics (#21291)#21291
shoumikhin wants to merge 4 commits into
pytorch:mainfrom
shoumikhin:export-D113382078

shoumikhin commented Jul 23, 2026 •

edited by meta-codesync Bot

Loading

Uh oh!

pytorch-bot Bot commented Jul 23, 2026 •

edited

Loading

Uh oh!

meta-codesync Bot commented Jul 23, 2026

Uh oh!

github-actions Bot commented Jul 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

shoumikhin Jul 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

shoumikhin Jul 23, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

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

Uh oh!

Uh oh!

pytorch-bot Bot commented Jul 23, 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/21291

❌ 1 Cancelled Job, 1 Pending, 2 Unclassified Failures

Uh oh!

meta-codesync Bot commented Jul 23, 2026

Uh oh!

github-actions Bot commented Jul 23, 2026

This PR needs a release notes: label

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

shoumikhin Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

shoumikhin Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

shoumikhin commented Jul 23, 2026 •

edited by meta-codesync Bot

Loading

pytorch-bot Bot commented Jul 23, 2026 •

edited

Loading

This PR needs a `release notes:` label