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); 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); 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;