From 72b20cc876e9507cfb10f44c049396efbe3bcea7 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 23 Jul 2026 13:15:03 -0700 Subject: [PATCH 1/3] prevent AOTI library collisions between delegates (#21287) 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 --- backends/cuda/runtime/cuda_backend.cpp | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 0fcc2a59404..1c75955e7ff 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -330,10 +330,16 @@ class ET_EXPERIMENTAL CudaBackend final so_blob_key.c_str(), static_cast(aoti_dso_buffer.error())); - // Generate dynamic temporary file path + // Generate a unique temporary file path. Two identical CUDA partitions can + // share the same so_blob_key, so a per-delegate counter is needed to keep + // their paths distinct; otherwise loading the second delegate would + // overwrite the first delegate's library while it is still in use. + static std::atomic so_file_counter{0}; filesystem::path temp_dir = filesystem::temp_directory_path(); - filesystem::path so_path = - temp_dir / (so_blob_key + to_string(get_process_id()) + ".so"); + filesystem::path so_path = temp_dir / + (so_blob_key + to_string(get_process_id()) + "_" + + to_string(so_file_counter.fetch_add(1, std::memory_order_relaxed)) + + ".so"); // Create a temporary file ofstream outfile(so_path, ios::binary); From 65e6ad5fe731657863d07569d80ac5416236f06d Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 23 Jul 2026 13:15:03 -0700 Subject: [PATCH 2/3] 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). 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 --- .../core/exec_aten/util/tensor_util_aten.cpp | 6 +- runtime/executor/tensor_parser_aten.cpp | 55 +++++++++++++++++-- 2 files changed, 54 insertions(+), 7 deletions(-) diff --git a/runtime/core/exec_aten/util/tensor_util_aten.cpp b/runtime/core/exec_aten/util/tensor_util_aten.cpp index b8d8e266016..d435dc21fa7 100644 --- a/runtime/core/exec_aten/util/tensor_util_aten.cpp +++ b/runtime/core/exec_aten/util/tensor_util_aten.cpp @@ -135,9 +135,9 @@ Error share_tensor_data(const at::Tensor& t_dst, const at::Tensor& t_src) { t_src.mutable_data_ptr() != nullptr, InvalidArgument, "Source tensor should have data_ptr not being nullptr."); - // Assign the dataptr as the input tensor dataptr - storage->set_data_ptr( - at::DataPtr(t_src.mutable_data_ptr(), at::DeviceType::CPU)); + // 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())); storage->set_nbytes(t_src.nbytes()); return Error::Ok; diff --git a/runtime/executor/tensor_parser_aten.cpp b/runtime/executor/tensor_parser_aten.cpp index ad980177cf1..4ff7761f8b6 100644 --- a/runtime/executor/tensor_parser_aten.cpp +++ b/runtime/executor/tensor_parser_aten.cpp @@ -53,6 +53,43 @@ Result parseTensor( InvalidProgram, "Invalid ScalarType %" PRId8, static_cast(type)); + + // 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(raw_device_type)); + device_type = raw_device_type == executorch_flatbuffer::DeviceType::CUDA + ? c10::DeviceType::CUDA + : c10::DeviceType::CPU; + device_index = static_cast( + 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(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); + // Sized with null data to compute nbytes; real device is applied below. auto options = at::CPU(type).options(); ET_CHECK_OR_RETURN_ERROR( @@ -103,8 +140,9 @@ Result parseTensor( if (s_tensor->shape_dynamism() == executorch_flatbuffer::TensorShapeDynamism::DYNAMIC_UNBOUND) { - // Provide fully dynamic tensors with an allocator so they can be resized - // within aten kernels. + // 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. auto impl = tensor.unsafeGetTensorImpl(); at::StorageImpl* storage = impl->unsafe_storage().unsafeGetStorageImpl(); storage->set_allocator(at::getCPUAllocator()); @@ -128,8 +166,17 @@ Result parseTensor( static_cast(data_ptr.error())); return data_ptr.error(); } - tensor.unsafeGetTensorImpl()->unsafe_storage().set_data_ptr( - at::DataPtr(data_ptr.get(), c10::DeviceType::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(), + sizes, + strides, + /*storage_offset=*/0, + deleteNothing, + at::TensorOptions().dtype(type).device(device), + /*target_device=*/device); } return tensor; From e0178f180e9427232e50bd3ef81af172b44b7873 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 23 Jul 2026 13:15:03 -0700 Subject: [PATCH 3/3] private executorch_ aoti_torch_* ABI + ATen-mode CUDA delegate 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 --- backends/aoti/CMakeLists.txt | 10 ++ backends/aoti/aoti_delegate_handle.h | 7 +- backends/aoti/export.h | 101 ++++++++++++++++++ backends/aoti/shim_symbol_prefix.bzl | 21 ++++ backends/aoti/targets.bzl | 13 +++ backends/cuda/CMakeLists.txt | 9 ++ backends/cuda/cuda_backend.py | 28 +++-- backends/cuda/runtime/TARGETS | 54 +++++++++- backends/cuda/runtime/cuda_backend.cpp | 29 +++-- backends/cuda/runtime/shims/memory.cpp | 24 +++++ backends/cuda/runtime/shims/memory.h | 9 ++ .../cuda/runtime/shims/tensor_attribute.cpp | 32 ------ .../cuda/runtime/shims/tensor_attribute.h | 36 ------- backends/cuda/runtime/utils.h | 55 ++-------- 14 files changed, 295 insertions(+), 133 deletions(-) create mode 100644 backends/aoti/shim_symbol_prefix.bzl delete mode 100644 backends/cuda/runtime/shims/tensor_attribute.cpp delete mode 100644 backends/cuda/runtime/shims/tensor_attribute.h diff --git a/backends/aoti/CMakeLists.txt b/backends/aoti/CMakeLists.txt index 4634f36eb9d..8d65156f19c 100644 --- a/backends/aoti/CMakeLists.txt +++ b/backends/aoti/CMakeLists.txt @@ -117,6 +117,13 @@ target_compile_options( target_compile_definitions( aoti_common_shims_slim PUBLIC $<$:EXPORT_AOTI_FUNCTIONS> ) +# 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_common_shims_slim PUBLIC AOTI_SHIM_SYMBOL_PREFIX=executorch_ +) target_link_libraries(aoti_common_shims_slim PUBLIC slimtensor ${CMAKE_DL_LIBS}) @@ -139,6 +146,9 @@ if(MSVC) aoti_common_shims_slim_obj PUBLIC $<$:EXPORT_AOTI_FUNCTIONS> ) + target_compile_definitions( + aoti_common_shims_slim_obj PUBLIC AOTI_SHIM_SYMBOL_PREFIX=executorch_ + ) target_link_libraries( aoti_common_shims_slim_obj PUBLIC slimtensor ${CMAKE_DL_LIBS} ) diff --git a/backends/aoti/aoti_delegate_handle.h b/backends/aoti/aoti_delegate_handle.h index fbd748306cc..d627a8790cd 100644 --- a/backends/aoti/aoti_delegate_handle.h +++ b/backends/aoti/aoti_delegate_handle.h @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -19,7 +20,11 @@ namespace aoti { using executorch::runtime::Error; using executorch::runtime::FreeableBuffer; -using executorch::runtime::etensor::Tensor; +// Mode-flexible host-tensor spelling: `etensor::Tensor` in portable mode (the +// slim ETensor the AOTI blob passes around), `at::Tensor` under USE_ATEN_LIB. +// Used here only as an opaque handle type; the CUDA delegate reinterpret_casts +// its SlimTensor* through it, and Metal passes its own portable handle. +using executorch::aten::Tensor; extern "C" { diff --git a/backends/aoti/export.h b/backends/aoti/export.h index 7c945f405b0..bf369004199 100644 --- a/backends/aoti/export.h +++ b/backends/aoti/export.h @@ -24,3 +24,104 @@ #else #define AOTI_SHIM_EXPORT #endif + +// Keep ExecuTorch's aoti_torch_* shim *definitions* in lockstep with the symbol +// names an AOTInductor blob imports. The blob and these (libtorch-free) shim +// libraries are both compiled with -DAOTI_SHIM_SYMBOL_PREFIX=executorch_ (see +// cuda_backend.py get_aoti_compile_options and the shim targets' +// exported_preprocessor_flags), so this block renames the exported definitions +// to executorch_aoti_torch_*. The blob (whose imports torch's +// aoti_torch/c/macros.h renames identically) then binds to ExecuTorch's +// SlimTensor shims and never to libtorch's at::Tensor aoti_torch_* in a +// coalesced (TensorRT[ATen] + CUDA) process. +// +// This block is kept self-contained (no torch include) to preserve the slim +// layer's libtorch-free property. It MUST stay in sync with the identical list +// in torch's aoti_torch/c/shim_symbol_prefix.h; a mismatch is a loud link error +// (an unresolved executorch_aoti_torch_* symbol), never a silent crash. With +// the define unset (portable builds) it expands to nothing. +#ifdef AOTI_SHIM_SYMBOL_PREFIX +#define AOTI_SHIM_CONCAT2_(a, b) a##b +#define AOTI_SHIM_CONCAT_(a, b) AOTI_SHIM_CONCAT2_(a, b) +#define AOTI_SHIM_RENAME_(name) AOTI_SHIM_CONCAT_(AOTI_SHIM_SYMBOL_PREFIX, name) + +#define aoti_torch_assign_tensors_out \ + AOTI_SHIM_RENAME_(aoti_torch_assign_tensors_out) +#define aoti_torch_check AOTI_SHIM_RENAME_(aoti_torch_check) +#define aoti_torch_clone AOTI_SHIM_RENAME_(aoti_torch_clone) +#define aoti_torch_clone_preserve_strides \ + AOTI_SHIM_RENAME_(aoti_torch_clone_preserve_strides) +#define aoti_torch_copy_ AOTI_SHIM_RENAME_(aoti_torch_copy_) +#define aoti_torch_create_cuda_guard \ + AOTI_SHIM_RENAME_(aoti_torch_create_cuda_guard) +#define aoti_torch_create_cuda_stream_guard \ + AOTI_SHIM_RENAME_(aoti_torch_create_cuda_stream_guard) +#define aoti_torch_create_tensor_from_blob \ + AOTI_SHIM_RENAME_(aoti_torch_create_tensor_from_blob) +#define aoti_torch_create_tensor_from_blob_v2 \ + AOTI_SHIM_RENAME_(aoti_torch_create_tensor_from_blob_v2) +#define aoti_torch_cuda_guard_set_index \ + AOTI_SHIM_RENAME_(aoti_torch_cuda_guard_set_index) +#define aoti_torch_cuda_int4_plain_mm \ + AOTI_SHIM_RENAME_(aoti_torch_cuda_int4_plain_mm) +#define aoti_torch_cuda_int5_plain_mm \ + AOTI_SHIM_RENAME_(aoti_torch_cuda_int5_plain_mm) +#define aoti_torch_cuda_int6_plain_mm \ + AOTI_SHIM_RENAME_(aoti_torch_cuda_int6_plain_mm) +#define aoti_torch_cuda_int8_plain_mm \ + AOTI_SHIM_RENAME_(aoti_torch_cuda_int8_plain_mm) +#define aoti_torch_cuda_rand AOTI_SHIM_RENAME_(aoti_torch_cuda_rand) +#define aoti_torch_cuda_randint_low_out \ + AOTI_SHIM_RENAME_(aoti_torch_cuda_randint_low_out) +#define aoti_torch_cuda_sort_stable \ + AOTI_SHIM_RENAME_(aoti_torch_cuda_sort_stable) +#define aoti_torch_cuda__weight_int4pack_mm \ + AOTI_SHIM_RENAME_(aoti_torch_cuda__weight_int4pack_mm) +#define aoti_torch_delete_cuda_guard \ + AOTI_SHIM_RENAME_(aoti_torch_delete_cuda_guard) +#define aoti_torch_delete_cuda_stream_guard \ + AOTI_SHIM_RENAME_(aoti_torch_delete_cuda_stream_guard) +#define aoti_torch_delete_tensor_object \ + AOTI_SHIM_RENAME_(aoti_torch_delete_tensor_object) +#define aoti_torch_device_type_cpu AOTI_SHIM_RENAME_(aoti_torch_device_type_cpu) +#define aoti_torch_device_type_cuda \ + AOTI_SHIM_RENAME_(aoti_torch_device_type_cuda) +#define aoti_torch_dtype_bfloat16 AOTI_SHIM_RENAME_(aoti_torch_dtype_bfloat16) +#define aoti_torch_dtype_bool AOTI_SHIM_RENAME_(aoti_torch_dtype_bool) +#define aoti_torch_dtype_float16 AOTI_SHIM_RENAME_(aoti_torch_dtype_float16) +#define aoti_torch_dtype_float32 AOTI_SHIM_RENAME_(aoti_torch_dtype_float32) +#define aoti_torch_dtype_int16 AOTI_SHIM_RENAME_(aoti_torch_dtype_int16) +#define aoti_torch_dtype_int32 AOTI_SHIM_RENAME_(aoti_torch_dtype_int32) +#define aoti_torch_dtype_int64 AOTI_SHIM_RENAME_(aoti_torch_dtype_int64) +#define aoti_torch_dtype_int8 AOTI_SHIM_RENAME_(aoti_torch_dtype_int8) +#define aoti_torch_dtype_uint8 AOTI_SHIM_RENAME_(aoti_torch_dtype_uint8) +#define aoti_torch_empty_strided AOTI_SHIM_RENAME_(aoti_torch_empty_strided) +#define aoti_torch_empty_strided_pinned \ + AOTI_SHIM_RENAME_(aoti_torch_empty_strided_pinned) +#define aoti_torch_get_current_cuda_stream \ + AOTI_SHIM_RENAME_(aoti_torch_get_current_cuda_stream) +#define aoti_torch_get_data_ptr AOTI_SHIM_RENAME_(aoti_torch_get_data_ptr) +#define aoti_torch_get_device_index \ + AOTI_SHIM_RENAME_(aoti_torch_get_device_index) +#define aoti_torch_get_device_type AOTI_SHIM_RENAME_(aoti_torch_get_device_type) +#define aoti_torch_get_dim AOTI_SHIM_RENAME_(aoti_torch_get_dim) +#define aoti_torch_get_dtype AOTI_SHIM_RENAME_(aoti_torch_get_dtype) +#define aoti_torch_get_numel AOTI_SHIM_RENAME_(aoti_torch_get_numel) +#define aoti_torch_get_sizes AOTI_SHIM_RENAME_(aoti_torch_get_sizes) +#define aoti_torch_get_storage_offset \ + AOTI_SHIM_RENAME_(aoti_torch_get_storage_offset) +#define aoti_torch_get_storage_size \ + AOTI_SHIM_RENAME_(aoti_torch_get_storage_size) +#define aoti_torch_get_strides AOTI_SHIM_RENAME_(aoti_torch_get_strides) +#define aoti_torch_grad_mode_is_enabled \ + AOTI_SHIM_RENAME_(aoti_torch_grad_mode_is_enabled) +#define aoti_torch_grad_mode_set_enabled \ + AOTI_SHIM_RENAME_(aoti_torch_grad_mode_set_enabled) +#define aoti_torch_item_bool AOTI_SHIM_RENAME_(aoti_torch_item_bool) +#define aoti_torch_layout_strided AOTI_SHIM_RENAME_(aoti_torch_layout_strided) +#define aoti_torch_new_tensor_handle \ + AOTI_SHIM_RENAME_(aoti_torch_new_tensor_handle) +#define aoti_torch__reinterpret_tensor \ + AOTI_SHIM_RENAME_(aoti_torch__reinterpret_tensor) +#define aoti_torch_warn AOTI_SHIM_RENAME_(aoti_torch_warn) +#endif // AOTI_SHIM_SYMBOL_PREFIX diff --git a/backends/aoti/shim_symbol_prefix.bzl b/backends/aoti/shim_symbol_prefix.bzl new file mode 100644 index 00000000000..49bf1d8e2e5 --- /dev/null +++ b/backends/aoti/shim_symbol_prefix.bzl @@ -0,0 +1,21 @@ +# Single source of truth for the ExecuTorch AOTI shim symbol prefix. +# +# The AOTInductor blob generated by the CUDA delegate is compiled with +# -DAOTI_SHIM_SYMBOL_PREFIX= (set from Python via +# cuda_backend.get_aoti_compile_options -> aot_inductor.shim_symbol_prefix), and +# ExecuTorch's shim libraries must be compiled with the SAME define so their +# exported aoti_torch_* definitions rename to aoti_torch_* and satisfy +# the blob's now-prefixed imports. Keeping the value here (and mirroring it in +# the Python default) prevents blob-import / shim-export drift, which would +# surface as a loud unresolved-symbol link error rather than a runtime crash. +# +# See: +# - fbcode/executorch/backends/aoti/export.h (shim-export rename block) +# - fbcode/caffe2/torch/csrc/inductor/aoti_torch/c/shim_symbol_prefix.h +# - fbcode/executorch/backends/cuda/cuda_backend.py (Python side) + +AOTI_SHIM_SYMBOL_PREFIX = "executorch_" + +AOTI_SHIM_PREFIX_PREPROCESSOR_FLAGS = [ + "-DAOTI_SHIM_SYMBOL_PREFIX=" + AOTI_SHIM_SYMBOL_PREFIX, +] diff --git a/backends/aoti/targets.bzl b/backends/aoti/targets.bzl index 51d47e01541..4442d5ecb4c 100644 --- a/backends/aoti/targets.bzl +++ b/backends/aoti/targets.bzl @@ -1,5 +1,6 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbsource//tools/build_defs:fbsource_utils.bzl", "is_fbcode") +load(":shim_symbol_prefix.bzl", "AOTI_SHIM_PREFIX_PREPROCESSOR_FLAGS") def define_common_targets(): if not is_fbcode(): @@ -105,6 +106,18 @@ def define_common_targets(): "export.h", "utils.h", ], + # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole) + # These extern "C" aoti_torch_* shims are only referenced by the CUDA + # delegate's dlopen'd AOTInductor blob, which is invisible to the static + # linker; without link_whole, --gc-sections drops the whole TU and the + # blob fails to resolve them at load. + link_whole = True, + # Export aoti_torch_* under the executorch_ prefix (see export.h) so the + # blob binds to these SlimTensor shims, not libtorch's, in a coalesced + # process. In BOTH lists: preprocessor_flags renames this target's own + # definitions (common_shims_slim.cpp); exported_ renames header includers. + preprocessor_flags = AOTI_SHIM_PREFIX_PREPROCESSOR_FLAGS, + exported_preprocessor_flags = AOTI_SHIM_PREFIX_PREPROCESSOR_FLAGS, visibility = ["@EXECUTORCH_CLIENTS"], exported_deps = [ "//executorch/runtime/core:core", diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 06990692428..ec5c2074b95 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -126,6 +126,15 @@ add_library(aoti_cuda_shims SHARED ${_aoti_cuda_shim_sources}) # 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. PUBLIC so direct consumers that include the +# shim headers (e.g. the shim gtests) compile the declarations under the same +# prefix and link the DLL's prefixed exports. +target_compile_definitions( + aoti_cuda_shims PUBLIC AOTI_SHIM_SYMBOL_PREFIX=executorch_ +) # Define export macros for shared library. Use WIN32 (not just MSVC) so MinGW # cross-compiles also emit dllexport symbols for AOTI shims. diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index bf28077c62c..f4d815bea75 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -469,29 +469,45 @@ def get_aoti_compile_options( "aot_inductor.emit_multi_arch_kernel": emit_multi_arch_kernel, } - try: - import torch + # Give the generated blob a private aoti_torch_* ABI so it binds to + # ExecuTorch's SlimTensor shims, never libtorch's at::Tensor shims, in a + # coalesced process (e.g. a TensorRT[ATen] + CUDA .pte). Must match the + # -DAOTI_SHIM_SYMBOL_PREFIX the ExecuTorch shim libs are built with + # (backends/aoti/shim_symbol_prefix.bzl). The shim libs export only the + # prefixed names, so this PyTorch capability is required, not optional; a + # missing key means the installed torch predates it. + 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_" + try: options["aot_inductor.custom_ops_to_c_shims"] = { torch.ops.executorch_cuda.int4_plain_mm.default: [ - "AOTITorchError aoti_torch_cuda_int4_plain_mm(" + "AOTITorchError executorch_aoti_torch_cuda_int4_plain_mm(" "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, " "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, " "int64_t, AtenTensorHandle*)" ], torch.ops.executorch_cuda.int5_plain_mm.default: [ - "AOTITorchError aoti_torch_cuda_int5_plain_mm(" + "AOTITorchError executorch_aoti_torch_cuda_int5_plain_mm(" "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, " "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, " "AtenTensorHandle, int64_t, AtenTensorHandle*)" ], torch.ops.executorch_cuda.int6_plain_mm.default: [ - "AOTITorchError aoti_torch_cuda_int6_plain_mm(" + "AOTITorchError executorch_aoti_torch_cuda_int6_plain_mm(" "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, " "AtenTensorHandle, AtenTensorHandle, int64_t, AtenTensorHandle*)" ], torch.ops.executorch_cuda.int8_plain_mm.default: [ - "AOTITorchError aoti_torch_cuda_int8_plain_mm(" + "AOTITorchError executorch_aoti_torch_cuda_int8_plain_mm(" "AtenTensorHandle, AtenTensorHandle, AtenTensorHandle, " "AtenTensorHandle, int64_t, AtenTensorHandle*)" ], diff --git a/backends/cuda/runtime/TARGETS b/backends/cuda/runtime/TARGETS index 122560e98ec..01a27a1bc2b 100644 --- a/backends/cuda/runtime/TARGETS +++ b/backends/cuda/runtime/TARGETS @@ -1,6 +1,7 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:cpp_unittest.bzl", "cpp_unittest") load("@fbcode_macros//build_defs/lib:re_test_utils.bzl", "re_test_utils") +load("//executorch/backends/aoti:shim_symbol_prefix.bzl", "AOTI_SHIM_PREFIX_PREPROCESSOR_FLAGS") load("//tools/build/buck:nvcc_flags.bzl", "get_nvcc_arch_args") oncall("executorch") @@ -37,7 +38,6 @@ runtime.cxx_library( "shims/memory.cpp", "shims/rand.cu", "shims/sort.cu", - "shims/tensor_attribute.cpp", ], headers = [ "shims/cuda_guard.h", @@ -46,7 +46,6 @@ runtime.cxx_library( "shims/memory.h", "shims/rand.h", "shims/sort.h", - "shims/tensor_attribute.h", "utils.h", ], # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole) @@ -54,7 +53,13 @@ runtime.cxx_library( supports_python_dlopen = True, # Constructor needed for backend registration. compiler_flags = ["-Wno-global-constructors"], - preprocessor_flags = ["-DCUDA_AVAILABLE=1"], + # AOTI_SHIM_PREFIX flags in BOTH lists: preprocessor_flags renames THIS + # target's own shim definitions (memory.cpp etc.); exported_preprocessor_flags + # renames references in downstream header includers so both agree. + preprocessor_flags = ["-DCUDA_AVAILABLE=1"] + AOTI_SHIM_PREFIX_PREPROCESSOR_FLAGS, + # Export aoti_torch_* under the executorch_ prefix (see aoti/export.h) so a + # coalesced blob binds to these SlimTensor shims, not libtorch's. + exported_preprocessor_flags = AOTI_SHIM_PREFIX_PREPROCESSOR_FLAGS, visibility = ["PUBLIC"], deps = [ "//executorch/backends/aoti:aoti_common_slim", @@ -141,6 +146,49 @@ runtime.cxx_library( ], ) +# ATen-mode variant of the CUDA delegate. -DUSE_ATEN_LIB routes the backend +# registration + interface types to the runtime::aten BackendRegistry (the one an +# _aten Module/Method reads at .pte load); without it an ATen host reports +# "CudaBackend is not registered". The slim/AOTI tensor layer stays portable +# (SlimTensor + mode-flexible executorch::aten:: host-tensor spelling), so this +# builds from the same sources. Mirrors XNNPACK / TensorRTBackend. Used by a +# coalesced TensorRT (ATen) + CUDA .pte runner. +runtime.cxx_library( + name = "cuda_backend_aten", + srcs = [ + "cuda_backend.cpp", + "cuda_mutable_state.cpp", + ], + headers = [ + "cuda_delegate_handle.h", + "cuda_mutable_state.h", + ], + # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole) + link_whole = True, + supports_python_dlopen = True, + # Constructor needed for backend registration. + compiler_flags = ["-Wno-global-constructors"], + preprocessor_flags = ["-DCUDA_AVAILABLE=1", "-DUSE_ATEN_LIB"], + visibility = ["PUBLIC"], + deps = [ + ":cuda_platform", + ":runtime_shims", + ":cuda_allocator", + "//executorch/backends/aoti:aoti_common_slim", + "//executorch/backends/aoti/slim/core:slimtensor", + "//executorch/backends/aoti/slim/factory:empty", + "//executorch/backends/aoti/slim/factory:from_blob", + "//executorch/backends/aoti/slim/factory:from_etensor", + "//executorch/extension/cuda:caller_stream", + "//executorch/extension/tensor:tensor_aten", + "//executorch/runtime/backend:interface_aten", + "//executorch/runtime/core/exec_aten/util:tensor_util_aten", + ], + external_deps = [ + ("cuda", None, "cuda-lazy"), + ], +) + cpp_unittest( name = "test_cuda_mutable_state", srcs = [ diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 1c75955e7ff..ac740f8a194 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -54,21 +54,23 @@ namespace executorch::backends::cuda { using namespace std; using namespace aoti; +using executorch::aten::ScalarType; +using executorch::aten::Tensor; +using executorch::ET_RUNTIME_NAMESPACE::Backend; +using executorch::ET_RUNTIME_NAMESPACE::BackendExecutionContext; +using executorch::ET_RUNTIME_NAMESPACE::BackendInitContext; +using executorch::ET_RUNTIME_NAMESPACE::BackendOptionContext; +using executorch::ET_RUNTIME_NAMESPACE::CompileSpec; +using executorch::ET_RUNTIME_NAMESPACE::DelegateHandle; +using executorch::ET_RUNTIME_NAMESPACE::NamedDataMap; using executorch::runtime::ArrayRef; -using executorch::runtime::BackendExecutionContext; -using executorch::runtime::BackendInitContext; using executorch::runtime::BackendOption; -using executorch::runtime::BackendOptionContext; -using executorch::runtime::CompileSpec; -using executorch::runtime::DelegateHandle; using executorch::runtime::Error; using executorch::runtime::EValue; using executorch::runtime::FreeableBuffer; using executorch::runtime::kMaxOptionValueLength; -using executorch::runtime::NamedDataMap; using executorch::runtime::Result; using executorch::runtime::Span; -using executorch::runtime::etensor::Tensor; // SlimTensor type aliases using cuda::CudaGraphPhase; @@ -86,7 +88,7 @@ constexpr char kWeightSharingAcrossMethods[] = "weight_sharing_across_methods"; } // anonymous namespace class ET_EXPERIMENTAL CudaBackend final - : public ::executorch::runtime::BackendInterface { + : public ::executorch::ET_RUNTIME_NAMESPACE::BackendInterface { private: // Trim leading/trailing whitespace from a view of the string. static std::string_view trim(std::string_view s) { @@ -517,7 +519,7 @@ class ET_EXPERIMENTAL CudaBackend final auto* tensor = &(args[i]->toTensor()); auto device_type = tensor->unsafeGetTensorImpl()->device_type(); ET_CHECK_OR_RETURN_ERROR( - device_type == executorch::runtime::etensor::DeviceType::CUDA, + device_type == executorch::aten::DeviceType::CUDA, InvalidArgument, "Tensor %zu expected device_type=CUDA (1), got %d. " "Device info may not be properly propagated from CudaPartitioner.", @@ -1242,9 +1244,14 @@ class ET_EXPERIMENTAL CudaBackend final namespace executorch::backends { namespace { auto cls = cuda::CudaBackend(); -executorch::runtime::Backend backend{"CudaBackend", &cls}; +// Register into ET_RUNTIME_NAMESPACE so the delegate resolves in whichever +// BackendRegistry the host runtime reads: `runtime` in portable builds, +// `runtime::aten` under -DUSE_ATEN_LIB (the _aten Module/Method look up the +// latter). Hardcoding `runtime::` would leave an ATen host with "CudaBackend is +// not registered". Mirrors XNNPACK / TensorRTBackend. +executorch::ET_RUNTIME_NAMESPACE::Backend backend{"CudaBackend", &cls}; static executorch::runtime::Error success_with_compiler = - register_backend(backend); + executorch::ET_RUNTIME_NAMESPACE::register_backend(backend); // Auto-register the CudaAllocator so that DeviceMemoryBuffer::create(CUDA) // works whenever the CUDA backend library is linked. diff --git a/backends/cuda/runtime/shims/memory.cpp b/backends/cuda/runtime/shims/memory.cpp index 8a81916ab6c..12a9729b93a 100644 --- a/backends/cuda/runtime/shims/memory.cpp +++ b/backends/cuda/runtime/shims/memory.cpp @@ -195,6 +195,30 @@ AOTITorchError aoti_torch_empty_strided( return Error::Ok; } +AOTITorchError aoti_torch_empty_strided_pinned( + int64_t ndim, + const int64_t* sizes_ptr, + const int64_t* strides_ptr, + int32_t dtype, + int32_t device_type, + int32_t device_index, + SlimTensor** ret_new_tensor) { + // 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); +} + AOTITorchError aoti_torch_delete_tensor_object(SlimTensor* tensor) { ET_CHECK_OR_RETURN_ERROR( tensor != nullptr, diff --git a/backends/cuda/runtime/shims/memory.h b/backends/cuda/runtime/shims/memory.h index ca464a9acf5..ec8398daaa9 100644 --- a/backends/cuda/runtime/shims/memory.h +++ b/backends/cuda/runtime/shims/memory.h @@ -95,6 +95,15 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_empty_strided( int32_t device_index, SlimTensor** ret_new_tensor); +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_empty_strided_pinned( + int64_t ndim, + const int64_t* sizes_ptr, + const int64_t* strides_ptr, + int32_t dtype, + int32_t device_type, + int32_t device_index, + SlimTensor** ret_new_tensor); + /** * Deletes a tensor object and frees associated resources. * diff --git a/backends/cuda/runtime/shims/tensor_attribute.cpp b/backends/cuda/runtime/shims/tensor_attribute.cpp deleted file mode 100644 index 1a14c79f9f2..00000000000 --- a/backends/cuda/runtime/shims/tensor_attribute.cpp +++ /dev/null @@ -1,32 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#include - -namespace executorch::backends::cuda { - -extern "C" { - -// Device type functions for tensor attributes -AOTITorchError aoti_torch_get_device_type( - Tensor* tensor, - int32_t* ret_device_type) { - // All tensors in aoti-cuda delegate are on CUDA - *ret_device_type = aoti_torch_device_type_cuda(); - return Error::Ok; -} - -// Device type constants -int32_t aoti_torch_device_type_cuda() { - // Let's say cuda is 1 for ET as well - return 1; -} - -} // extern "C" - -} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/shims/tensor_attribute.h b/backends/cuda/runtime/shims/tensor_attribute.h deleted file mode 100644 index efa15ca072a..00000000000 --- a/backends/cuda/runtime/shims/tensor_attribute.h +++ /dev/null @@ -1,36 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#pragma once - -#include -#include -#include -#include - -namespace executorch::backends::cuda { - -// Common using declarations for ExecuTorch types -using executorch::runtime::Error; -using executorch::runtime::etensor::Tensor; - -extern "C" { - -// Common AOTI type aliases -using AOTITorchError = Error; - -// Device type functions for tensor attributes -AOTI_SHIM_EXPORT AOTITorchError -aoti_torch_get_device_type(Tensor* tensor, int32_t* ret_device_type); - -// Device type constants -AOTI_SHIM_EXPORT int32_t aoti_torch_device_type_cuda(); - -} // extern "C" - -} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/utils.h b/backends/cuda/runtime/utils.h index aaed5108d4f..c3470c53db6 100644 --- a/backends/cuda/runtime/utils.h +++ b/backends/cuda/runtime/utils.h @@ -25,7 +25,7 @@ namespace executorch::backends::cuda { namespace { inline executorch::runtime::Error _check_tensor_metadata( const executorch::backends::aoti::slim::SlimTensor* slim_tensor, - executorch::runtime::etensor::Tensor* etensor) { + executorch::aten::Tensor* etensor) { ET_CHECK_OR_RETURN_ERROR( slim_tensor != nullptr, InvalidArgument, @@ -50,8 +50,7 @@ inline executorch::runtime::Error _check_tensor_metadata( // Check dtype matches executorch::backends::aoti::slim::c10::ScalarType slim_dtype = slim_tensor->dtype(); - executorch::runtime::etensor::ScalarType etensor_dtype = - etensor->scalar_type(); + executorch::aten::ScalarType etensor_dtype = etensor->scalar_type(); ET_CHECK_OR_RETURN_ERROR( static_cast(slim_dtype) == static_cast(etensor_dtype), InvalidArgument, @@ -67,23 +66,18 @@ inline executorch::runtime::Error _check_tensor_metadata( slim_tensor->dim(), etensor->dim()); - // Convert sizes from int64_t to SizesType (int32_t) for resize const size_t ndim = slim_tensor->dim(); - std::vector new_sizes( - ndim); + std::vector new_sizes(ndim); auto slim_sizes = slim_tensor->sizes(); for (size_t i = 0; i < ndim; ++i) { - new_sizes[i] = - static_cast( - slim_sizes[i]); + new_sizes[i] = static_cast(slim_sizes[i]); } // Resize ETensor to match SlimTensor sizes executorch::runtime::Error resize_err = executorch::ET_RUNTIME_NAMESPACE::resize_tensor( *etensor, - executorch::runtime::ArrayRef< - executorch::runtime::etensor::TensorImpl::SizesType>( + executorch::aten::ArrayRef( new_sizes.data(), new_sizes.size())); ET_CHECK_OK_OR_RETURN_ERROR(resize_err, "failed to resize ETensor"); @@ -92,7 +86,7 @@ inline executorch::runtime::Error _check_tensor_metadata( // Check if src and dst strides match (same layout, no rearrangement needed). inline bool _strides_match( const executorch::backends::aoti::slim::SlimTensor* slim_tensor, - const executorch::runtime::etensor::Tensor* etensor) { + const executorch::aten::Tensor* etensor) { const size_t ndim = slim_tensor->dim(); auto slim_strides = slim_tensor->strides(); auto et_strides = etensor->strides(); @@ -152,7 +146,7 @@ inline void _strided_copy( // path). When stream is null, GPU copies are synchronous. inline executorch::runtime::Error _copy_slimtensor_to_etensor_impl( const executorch::backends::aoti::slim::SlimTensor* slim_tensor, - executorch::runtime::etensor::Tensor* etensor, + executorch::aten::Tensor* etensor, const executorch::backends::aoti::slim::c10::Device& dst_device, cudaStream_t stream) { ET_CHECK_OK_OR_RETURN_ERROR(_check_tensor_metadata(slim_tensor, etensor)); @@ -272,7 +266,7 @@ inline executorch::runtime::Error _copy_slimtensor_to_etensor_impl( */ inline executorch::runtime::Error copy_slimtensor_to_etensor_async( const executorch::backends::aoti::slim::SlimTensor* slim_tensor, - executorch::runtime::etensor::Tensor* etensor, + executorch::aten::Tensor* etensor, cudaStream_t stream) { return _copy_slimtensor_to_etensor_impl( slim_tensor, @@ -303,7 +297,7 @@ inline executorch::runtime::Error copy_slimtensor_to_etensor_async( */ inline executorch::runtime::Error copy_slimtensor_to_device_etensor_async( const executorch::backends::aoti::slim::SlimTensor* slim_tensor, - executorch::runtime::etensor::Tensor* etensor, + executorch::aten::Tensor* etensor, cudaStream_t stream) { return _copy_slimtensor_to_etensor_impl( slim_tensor, etensor, slim_tensor->device(), stream); @@ -321,7 +315,7 @@ inline executorch::runtime::Error copy_slimtensor_to_device_etensor_async( */ inline executorch::runtime::Error copy_slimtensor_to_etensor( const executorch::backends::aoti::slim::SlimTensor* slim_tensor, - executorch::runtime::etensor::Tensor* etensor) { + executorch::aten::Tensor* etensor) { return _copy_slimtensor_to_etensor_impl( slim_tensor, etensor, @@ -329,33 +323,6 @@ inline executorch::runtime::Error copy_slimtensor_to_etensor( nullptr); } -/** - * Wraps a SlimTensor's data into an existing ETensor (zero-copy). - * - * This function resizes the ETensor to match the SlimTensor's shape and - * sets its data pointer to point directly to the SlimTensor's data buffer. - * No data is copied - the ETensor becomes a view of the SlimTensor's data. - * - * IMPORTANT: The caller must ensure the SlimTensor remains alive as long - * as the ETensor is in use, since the ETensor will reference the SlimTensor's - * data directly. - * - * @param slim_tensor Pointer to the source SlimTensor (must not be null). - * @param etensor Pointer to the destination ETensor (must not be null). - * @return Error::Ok on success, or an appropriate error code on failure. - */ -inline executorch::runtime::Error wrap_slimtensor_to_etensor( - const executorch::backends::aoti::slim::SlimTensor* slim_tensor, - executorch::runtime::etensor::Tensor* etensor) { - ET_CHECK_OK_OR_RETURN_ERROR(_check_tensor_metadata(slim_tensor, etensor)); - - // Set data pointer to point directly to SlimTensor's data (zero-copy) - etensor->unsafeGetTensorImpl()->set_data( - const_cast(slim_tensor->data_ptr())); - - return executorch::runtime::Error::Ok; -} - /** * Deletes all SlimTensor pointers in a vector and clears the vector. * @@ -394,7 +361,7 @@ inline void delete_slimtensor_vector( inline executorch::backends::aoti::slim::SlimTensor* make_slimtensor_from_blob_with_etensor_metadata( void* data_ptr, - const executorch::runtime::etensor::Tensor* etensor, + const executorch::aten::Tensor* etensor, const executorch::backends::aoti::slim::c10::Device& device = executorch::backends::aoti::slim::DEFAULT_CUDA_DEVICE) { auto sizes = etensor->sizes();