private executorch_ aoti_torch_* ABI + ATen-mode CUDA delegate#21290
private executorch_ aoti_torch_* ABI + ATen-mode CUDA delegate#21290shoumikhin wants to merge 3 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21290
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 Cancelled Job, 2 Unclassified FailuresAs of commit e0178f1 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:
CANCELLED JOB - The following job was cancelled. Please retry:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
|
@shoumikhin has exported this pull request. If you are a Meta employee, you can view the originating Diff in D113382057. |
This PR needs a
|
There was a problem hiding this comment.
Pull request overview
This PR updates ExecuTorch’s AOTInductor (SlimTensor) shim ABI to use a private executorch_ symbol prefix (preventing collisions with libtorch’s aoti_torch_* ABI) and adds an ATen-mode CUDA delegate variant so ATen-hosted .pte loads register/resolve the CUDA backend correctly.
Changes:
- Prefix SlimTensor
aoti_torch_*shim exports toexecutorch_aoti_torch_*via a shared build flag, and plumb the matchingaot_inductor.shim_symbol_prefixinto CUDA export. - Add an ATen-mode CUDA backend library target and ensure backend registration happens in the correct runtime registry (
runtimevsruntime::aten). - Improve device propagation for ATen-mode tensor deserialization and fix ATen storage device tagging when sharing tensor data.
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| runtime/executor/tensor_parser_aten.cpp | Propagate serialized device info (CPU/CUDA) into at::Tensor construction during deserialization. |
| runtime/core/exec_aten/util/tensor_util_aten.cpp | Preserve source tensor device when sharing storage data pointers. |
| backends/cuda/runtime/utils.h | Switch CUDA runtime helper APIs to mode-flexible executorch::aten::* tensor/type spellings. |
| backends/cuda/runtime/TARGETS | Apply shim-symbol-prefix flags to CUDA shim builds; add cuda_backend_aten target. |
| backends/cuda/runtime/shims/tensor_attribute.h | Removed (deduped into common shim layer). |
| backends/cuda/runtime/shims/tensor_attribute.cpp | Removed (deduped into common shim layer). |
| backends/cuda/runtime/shims/memory.h | Add aoti_torch_empty_strided_pinned API surface. |
| backends/cuda/runtime/shims/memory.cpp | Implement aoti_torch_empty_strided_pinned (currently falls back to pageable host allocation). |
| backends/cuda/runtime/cuda_backend.cpp | Register backend via ET_RUNTIME_NAMESPACE; avoid temp .so path collisions; migrate to executorch::aten tensor/device types. |
| backends/cuda/cuda_backend.py | Set aot_inductor.shim_symbol_prefix="executorch_" and update custom op shim symbol spellings accordingly. |
| backends/cuda/CMakeLists.txt | Build CUDA shims with AOTI_SHIM_SYMBOL_PREFIX=executorch_. |
| backends/aoti/targets.bzl | Make common SlimTensor shim library link_whole and apply the shim symbol prefix flags consistently. |
| backends/aoti/shim_symbol_prefix.bzl | New single source of truth for the shim symbol prefix and corresponding preprocessor flags. |
| backends/aoti/export.h | Add preprocessor-based renaming of exported aoti_torch_* shims under AOTI_SHIM_SYMBOL_PREFIX. |
| backends/aoti/CMakeLists.txt | Apply AOTI_SHIM_SYMBOL_PREFIX=executorch_ for the common slim shims (and MSVC object lib variant). |
| backends/aoti/aoti_delegate_handle.h | Use mode-flexible executorch::aten::Tensor handle type for delegates. |
Comments suppressed due to low confidence (1)
runtime/executor/tensor_parser_aten.cpp:76
ExtraTensorInfo.device_indexis a FlatBuffersbyte(signed). It is treated as untrusted input here, but it’s cast toc10::DeviceIndexwithout validation; a negative value would construct an invalid CUDA device (and defeats the stated goal of validating untrusted fields). Validatedevice_index(>= 0) before casting/constructing thec10::Device.
device_type = raw_device_type == executorch_flatbuffer::DeviceType::CUDA
? c10::DeviceType::CUDA
: c10::DeviceType::CPU;
device_index = static_cast<c10::DeviceIndex>(
s_tensor->extra_tensor_info()->device_index());
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| using executorch::ET_RUNTIME_NAMESPACE::BackendOptionContext; | ||
| using executorch::ET_RUNTIME_NAMESPACE::CompileSpec; | ||
| using executorch::ET_RUNTIME_NAMESPACE::DelegateHandle; | ||
| using executorch::ET_RUNTIME_NAMESPACE::NamedDataMap; |
There was a problem hiding this comment.
NamedDataMap is declared inside namespace executorch::ET_RUNTIME_NAMESPACE (see runtime/core/named_data_map.h), not plain executorch::runtime. The class moves with the macro, so executorch::ET_RUNTIME_NAMESPACE::NamedDataMap resolves correctly in both portable (runtime) and ATen (runtime::aten) builds. The ATen-mode target compiles as-is, so this is intentional.
| // Convert sizes from int64_t to SizesType (int32_t) for resize | ||
| const size_t ndim = slim_tensor->dim(); | ||
| std::vector<executorch::runtime::etensor::TensorImpl::SizesType> new_sizes( | ||
| ndim); | ||
| std::vector<executorch::aten::SizesType> new_sizes(ndim); | ||
| auto slim_sizes = slim_tensor->sizes(); |
| // Defaults to CPU when extra_tensor_info is absent (older PTE files). A | ||
| // device-delegate planned buffer must be tagged with its real device or the | ||
| // runtime treats device memory as host memory. | ||
| c10::DeviceType device_type = c10::DeviceType::CPU; | ||
| c10::DeviceIndex device_index = 0; | ||
| if (s_tensor->extra_tensor_info() != nullptr) { | ||
| // Untrusted byte from the PTE; validate before the cast so a bogus value | ||
| // cannot reach c10::Device as garbage. | ||
| const auto raw_device_type = s_tensor->extra_tensor_info()->device_type(); | ||
| ET_CHECK_OR_RETURN_ERROR( | ||
| raw_device_type == executorch_flatbuffer::DeviceType::CPU || | ||
| raw_device_type == executorch_flatbuffer::DeviceType::CUDA, | ||
| InvalidProgram, | ||
| "Invalid DeviceType %" PRId8, | ||
| static_cast<int8_t>(raw_device_type)); | ||
| device_type = raw_device_type == executorch_flatbuffer::DeviceType::CUDA | ||
| ? c10::DeviceType::CUDA | ||
| : c10::DeviceType::CPU; | ||
| device_index = static_cast<c10::DeviceIndex>( | ||
| s_tensor->extra_tensor_info()->device_index()); | ||
| } | ||
| // CPU stays unindexed: an explicit cpu:0 would mismatch the graph's default | ||
| // cpu tensors and trip ATen's same-device check. Only accelerators carry an | ||
| // index. | ||
| const c10::Device device = device_type == c10::DeviceType::CPU | ||
| ? c10::Device(device_type) | ||
| : c10::Device(device_type, device_index); |
620cba8 to
232ff06
Compare
| const auto raw_device_type = s_tensor->extra_tensor_info()->device_type(); | ||
| ET_CHECK_OR_RETURN_ERROR( | ||
| raw_device_type == executorch_flatbuffer::DeviceType::CPU || | ||
| raw_device_type == executorch_flatbuffer::DeviceType::CUDA, | ||
| InvalidProgram, | ||
| "Invalid DeviceType %" PRId8, | ||
| static_cast<int8_t>(raw_device_type)); | ||
| device_type = raw_device_type == executorch_flatbuffer::DeviceType::CUDA | ||
| ? c10::DeviceType::CUDA | ||
| : c10::DeviceType::CPU; | ||
| 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)); |
There was a problem hiding this comment.
FlatBuffers byte maps to int8_t (signed), and the generated accessor returns int8_t, so there is no uint8->int8 wrap here (a stored 200 is already -56 on the wire, not produced by the cast). I've added a validation that rejects a negative accelerator index from the untrusted PTE (device_type == CPU || device_index >= 0) before it is used, which also covers the -1 "current device" case.
| try: | ||
| import torch | ||
|
|
There was a problem hiding this comment.
Good catch, fixed. The module already imports torch at the top, so the inner import torch here was redundant and made torch function-local, breaking the earlier hasattr(torch._inductor.config...) guard with UnboundLocalError. Removed the inner import so both references resolve to the module-level torch.
| // AOTInductor's PinnedStagingPool asks for page-locked host memory to overlap | ||
| // H2D copies; SlimTensor has no pinned (cudaHostAlloc) allocator yet, so this | ||
| // returns ordinary (pageable) host memory. Functionally correct -- the pool | ||
| // works and, on any error, falls back to a synchronous copy -- but without | ||
| // the async-overlap perf win. | ||
| // TODO: back this with cudaHostAlloc for true pinned semantics. | ||
| return aoti_torch_empty_strided( | ||
| ndim, | ||
| sizes_ptr, | ||
| strides_ptr, | ||
| dtype, | ||
| device_type, | ||
| device_index, | ||
| ret_new_tensor); |
There was a problem hiding this comment.
This is intentional and documented in the comment above: the only caller (AOTInductor's PinnedStagingPool) requests host staging memory, and SlimTensor has no pinned allocator yet, so this forwards to aoti_torch_empty_strided and returns ordinary pageable host memory (functionally correct; the pool falls back to a synchronous copy). The device_type=CUDA path is not exercised by that caller. Tracked by the TODO: back this with cudaHostAlloc note for true pinned semantics.
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
232ff06 to
e0178f1
Compare
| if not hasattr( | ||
| torch._inductor.config.aot_inductor, "shim_symbol_prefix" | ||
| ): | ||
| raise RuntimeError( | ||
| "The installed PyTorch does not support " | ||
| "aot_inductor.shim_symbol_prefix, which the ExecuTorch CUDA " | ||
| "backend requires to match its prefixed aoti_torch_* shims. " | ||
| "Update PyTorch to a version that provides it." | ||
| ) | ||
| options["aot_inductor.shim_symbol_prefix"] = "executorch_" |
|
CI note: the CUDA jobs ( This PR makes the ExecuTorch CUDA shims carry a private Required landing order:
The non-CUDA checks are green. This is a cross-repo sequencing dependency, not a defect in this change. |
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:
-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.
--gc-sections. memory: add aoti_torch_empty_strided_pinned.
tensor (etensor in portable, at::Tensor under USE_ATEN_LIB). Portable unchanged.
CUDA backend:
runtime::aten under USE_ATEN_LIB). Mirrors XNNPACK / TensorRTBackend.
executorch_ names the shim libs export.
Differential Revision: D113382057