tensor_parser_aten: tag planned tensor DataPtr with its real device

[shoumikhin]

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

The device_index read from the untrusted PTE is validated as non-negative for a
non-CPU device before use. A DYNAMIC_UNBOUND tensor annotated as non-CPU is
rejected (NotSupported) rather than silently returning CPU-resizable storage
that contradicts the serialized device.

Also update internal_set_tensor_data (tensor_util_aten.cpp) to preserve the
source tensor's device instead of hardcoding CPU, and verify the destination and
source devices match before sharing storage, so a shared tensor cannot end up
with a TensorImpl device that disagrees with its DataPtr device.

fbcode and xplat copies kept identical.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 3282ed6 with merge base 6cd7ecd ():
💚 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 D113384858.

[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 fixes incorrect device tagging for planned tensors in ATen-mode deserialization, ensuring that tensors backed by accelerator memory (e.g., CUDA planned buffers) are constructed with consistent device information across the storage DataPtr, TensorImpl device, and dispatch key. It also prevents CUDA backend delegate initialization from clobbering shared-library temp files when identical partitions are loaded multiple times.

Changes:

• Read device_type/device_index from extra_tensor_info during ATen tensor parsing and construct tensors via at::from_blob(..., TensorOptions().device(device), /*target_device=*/device) to keep device tagging consistent.
• Preserve the source tensor’s device when sharing tensor storage in internal::share_tensor_data.
• Ensure unique temp .so paths in the CUDA backend by adding an atomic counter suffix.

Reviewed changes

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

File	Description
runtime/executor/tensor_parser_aten.cpp	Parse device metadata and rebuild tensors with correct device tagging via from_blob(..., target_device=...).
runtime/core/exec_aten/util/tensor_util_aten.cpp	Preserve device when sharing tensor storage data pointers.
backends/cuda/runtime/cuda_backend.cpp	Avoid temp .so path collisions when loading identical CUDA partitions by adding a unique counter.

---

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

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

[Gasoonjia]

I feel like here we are too specific for cuda and cpu device and may have issue if in the future we will support more devices. Condier making it more general.

[shoumikhin]

The serialized DeviceType enum currently has only CPU and CUDA (schema program.fbs). The allowlist mirrors that exactly and fails loudly (InvalidProgram) on any value it does not recognize, so an unknown/future device can never be silently mis-mapped. When the schema adds a new device, its runtime mapping and tests should land together with it. Keeping it explicit for now avoids a speculative abstraction with no second real device to validate against.

[shoumikhin]

To add to the above: I kept it explicit intentionally, but your point about future devices is fair. The serialized DeviceType enum only defines CPU and CUDA today (schema program.fbs), so a generic mapping would have no second real device to validate against yet and would be speculative. When a new device is added to the schema, the right place to generalize is there, adding its c10 mapping and tests together. Happy to switch to a small enum->c10 mapping table now instead if you'd prefer that shape; let me know.

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (1)

runtime/executor/tensor_parser_aten.cpp:183

• The new at::from_blob(..., /*target_device=*/device) path is critical to avoid getDeviceFromPtr for runtime-bound (nullptr) planned buffers, but there’s no ATen-mode unit test exercising a CUDA-annotated tensor where getTensorDataPtr() returns nullptr and validating that (a) the tensor is tagged CUDA and (b) no pointer inspection occurs. Adding a focused gtest for the ATen tensor parser would help catch regressions across PyTorch upgrades.

    // Rebuild so storage DataPtr, TensorImpl device, and dispatch key agree.
    // target_device makes from_blob skip getDeviceFromPtr, so the same path
    // works for a real pointer and for a null runtime-bound one.
    tensor = at::from_blob(
        data_ptr.get(),
        sizes,
        strides,
        /*storage_offset=*/0,
        deleteNothing,
        at::TensorOptions().dtype(type).device(device),
        /*target_device=*/device);

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (1)

runtime/executor/tensor_parser_aten.cpp:177

• There is no automated test covering this new ATen-mode device-tagging path (extra_tensor_info->device_type CUDA + from_blob with target_device). The existing tensor_parser_device_test target is only built once (not in an aten_mode variant), so regressions in tensor_parser_aten.cpp (or future PyTorch API changes around from_blob/target_device) would go uncaught. Please add a gtest that builds with USE_ATEN_LIB and asserts that a CUDA-annotated serialized tensor produces an at::Tensor whose (a) storage DataPtr device and (b) tensor.device() both match CUDA (and that missing extra_tensor_info defaults to CPU).

    // Rebuild so storage DataPtr, TensorImpl device, and dispatch key agree.
    // target_device makes from_blob skip getDeviceFromPtr, so the same path
    // works for a real pointer and for a null runtime-bound one.
    tensor = at::from_blob(
        data_ptr.get(),

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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