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 0fcc2a59404..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) { @@ -330,10 +332,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); @@ -511,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.", @@ -1236,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(); 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;