From 8b3f986561396212e0a922ce8552f300a98891e7 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 23 Jul 2026 18:41:50 -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 | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 0fcc2a59404..8666fc9c098 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -330,10 +330,18 @@ 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. so_blob_key already selected the + // blob above and is deliberately kept out of the filename: it can be an + // untrusted, variable-length payload key, so embedding it risks path + // traversal and cross-key collisions. Uniqueness comes from the pid + // (across processes) and an atomic counter (within a process, since two + // identical CUDA partitions would otherwise clobber each other's .so). + 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 / + ("executorch_cuda_" + 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 f6f99ce4a1b85c4e4ff613b06bc08398851f851f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 23 Jul 2026 18:41:50 -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 6637f8149354ceef340e95a6add296f31a5fa9ba Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 23 Jul 2026 18:41:50 -0700 Subject: [PATCH 3/3] Cover caller-stream TLS semantics (#21291) Summary: Add focused coverage for the caller-stream contract relied on by external CUDA backends. Verify that an explicitly selected null/default stream remains distinct from no guard, and that selections are isolated per execution thread. Keep the fbcode and xplat mirrors identical. Differential Revision: D113382078 --- .../slim/cuda/test/test_cuda_stream_guard.cpp | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/backends/aoti/slim/cuda/test/test_cuda_stream_guard.cpp b/backends/aoti/slim/cuda/test/test_cuda_stream_guard.cpp index df618a7b8e9..d06a22e4594 100644 --- a/backends/aoti/slim/cuda/test/test_cuda_stream_guard.cpp +++ b/backends/aoti/slim/cuda/test/test_cuda_stream_guard.cpp @@ -12,6 +12,7 @@ #include #include +#include #include using namespace executorch::backends::cuda; @@ -296,6 +297,30 @@ TEST(CallerStreamGuardTest, GuardSelectsThenRestores) { EXPECT_FALSE(getCallerStream().has_value()); } +TEST(CallerStreamGuardTest, ExplicitNullStreamIsStillSelected) { + CallerStreamGuard guard(nullptr); + ASSERT_TRUE(getCallerStream().has_value()); + EXPECT_EQ(*getCallerStream(), nullptr); +} + +TEST(CallerStreamGuardTest, SelectionIsThreadLocal) { + const cudaStream_t selected = fake_stream(0); + CallerStreamGuard guard(selected); + + bool child_started_without_stream = false; + bool child_selected_own_stream = false; + std::thread child([&]() { + child_started_without_stream = !getCallerStream().has_value(); + CallerStreamGuard child_guard(fake_stream(1)); + child_selected_own_stream = getCallerStream() == fake_stream(1); + }); + child.join(); + + EXPECT_TRUE(child_started_without_stream); + EXPECT_TRUE(child_selected_own_stream); + EXPECT_EQ(getCallerStream(), selected); +} + TEST(CallerStreamGuardTest, NestedGuardsRestoreOuter) { const cudaStream_t outer = fake_stream(1); const cudaStream_t inner = fake_stream(2);