tensor_parser_aten: tag planned tensor DataPtr with its real device (#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).
The serialized device_type is mapped to its ATen counterpart through a
small
table so validation and conversion share one source of truth and a new
device
only needs one entry. The untrusted device_type and device_index are
validated
(unmapped type or a negative accelerator index is rejected), and a
non-CPU
DYNAMIC_UNBOUND tensor is rejected 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.
Test Plan:
New ATen-mode unit test for the device parsing logic:
buck2 test
fbcode//executorch/runtime/executor/test:tensor_parser_aten_test
Pass 6, fail 0. Covers CPU tensors staying unindexed, missing
extra_tensor_info
defaulting to CPU, rejection of an unmapped device_type, rejection of a
negative
accelerator index, rejection of DYNAMIC_UNBOUND on a non-CPU device, and
DYNAMIC_UNBOUND still allowed on CPU.
ATen-mode load + execute path (the CPU device-index regression this
restores):
buck2 test fbcode//executorch/extension/module/test:module_test_aten
Pass 50, fail 0. This exercises parseTensor followed by aten kernels
(e.g.
aten::mul.out) on ModuleAddMul; a hardcoded CPU index 0 previously made
ATen see
cpu vs cpu:0 and abort the kernel, which this fix avoids by keeping CPU
unindexed.
Unblocks loading a coalesced TRT+CUDA .pte whose CUDA-delegate planned
buffers
are device memory. A
Anthony Shoumikhin committed
55d693b1655314526d25436754d83b9554049540
Parent: 16b780b
Committed by GitHub <noreply@github.com>
on 7/27/2026, 5:49:02 PM