diff --git a/runtime/core/exec_aten/util/tensor_util_aten.cpp b/runtime/core/exec_aten/util/tensor_util_aten.cpp index b8d8e266016..024991a2355 100644 --- a/runtime/core/exec_aten/util/tensor_util_aten.cpp +++ b/runtime/core/exec_aten/util/tensor_util_aten.cpp @@ -135,9 +135,17 @@ 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)); + // The destination keeps its own TensorImpl but adopts the source storage, so + // the devices must agree or TensorImpl and DataPtr would disagree. + ET_CHECK_OR_RETURN_ERROR( + t_dst.device() == t_src.device(), + InvalidArgument, + "Destination device %s does not match source device %s", + c10::str(t_dst.device()).c_str(), + c10::str(t_src.device()).c_str()); + // 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..5a2a4ef532a 100644 --- a/runtime/executor/tensor_parser_aten.cpp +++ b/runtime/executor/tensor_parser_aten.cpp @@ -30,6 +30,19 @@ namespace { void deleteNothing(void*); void deleteNothing(void*) {} +// Maps a serialized DeviceType to its ATen counterpart. Kept as a table so +// validation and conversion share one source of truth and a new device only +// needs a single entry here. +struct DeviceTypeMapping { + executorch_flatbuffer::DeviceType serialized; + c10::DeviceType aten; +}; + +constexpr DeviceTypeMapping kDeviceTypeMappings[] = { + {executorch_flatbuffer::DeviceType::CPU, c10::DeviceType::CPU}, + {executorch_flatbuffer::DeviceType::CUDA, c10::DeviceType::CUDA}, +}; + } // namespace Result parseTensor( @@ -53,6 +66,47 @@ 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; an unmapped value is rejected below so it + // cannot reach c10::Device as garbage. + const auto raw_device_type = s_tensor->extra_tensor_info()->device_type(); + bool valid_device_type = false; + for (const auto& mapping : kDeviceTypeMappings) { + if (mapping.serialized == raw_device_type) { + device_type = mapping.aten; + valid_device_type = true; + break; + } + } + ET_CHECK_OR_RETURN_ERROR( + valid_device_type, + InvalidProgram, + "Invalid DeviceType %" PRId8, + static_cast(raw_device_type)); + 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 +157,13 @@ 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 a CPU allocator so aten kernels can resize + // them. A device tensor cannot be honored here, so reject it rather than + // silently returning CPU storage that contradicts the serialized device. + ET_CHECK_OR_RETURN_ERROR( + device_type == c10::DeviceType::CPU, + NotSupported, + "DYNAMIC_UNBOUND tensors are only supported on CPU"); auto impl = tensor.unsafeGetTensorImpl(); at::StorageImpl* storage = impl->unsafe_storage().unsafeGetStorageImpl(); storage->set_allocator(at::getCPUAllocator()); @@ -128,8 +187,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; diff --git a/runtime/executor/test/targets.bzl b/runtime/executor/test/targets.bzl index 4a14285e381..8d397847ec1 100644 --- a/runtime/executor/test/targets.bzl +++ b/runtime/executor/test/targets.bzl @@ -319,6 +319,19 @@ def define_common_targets(is_fbcode = False): env = modules_env, ) + runtime.cxx_test( + name = "tensor_parser_aten_test", + srcs = [ + "tensor_parser_aten_test.cpp", + ], + deps = [ + ":managed_memory_manager", + "//executorch/runtime/executor:program_aten", + "//executorch/runtime/core/exec_aten:lib_aten", + "//executorch/schema:program", + ], + ) + runtime.cxx_test( name = "tensor_parser_device_test", srcs = [ diff --git a/runtime/executor/test/tensor_parser_aten_test.cpp b/runtime/executor/test/tensor_parser_aten_test.cpp new file mode 100644 index 00000000000..210f29fbd73 --- /dev/null +++ b/runtime/executor/test/tensor_parser_aten_test.cpp @@ -0,0 +1,186 @@ +/* + * 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. + */ + +/** + * Unit tests for the ATen-mode tensor parser's device handling. These exercise + * the validation and CPU-tagging paths that run without a real accelerator. The + * end-to-end CUDA device-pointer path is covered separately by + * tensor_parser_device_test (portable mode) with a mock allocator; ATen-mode + * CUDA execution needs a real device and is out of scope here. + */ + +#include + +#include +#include +#include +#include + +#include + +#include + +using executorch::aten::Tensor; +using executorch::runtime::Error; +using executorch::runtime::MemoryManager; +using executorch::runtime::Result; +using executorch::runtime::aten::deserialization::parseTensor; +using executorch::runtime::testing::ManagedMemoryManager; + +namespace { + +constexpr size_t kMemBytes = 32 * 1024U; + +// Builds a standalone serialized Tensor with no data_buffer and no +// allocation_info, i.e. an input/placeholder. parseTensor then resolves a null +// data pointer, so these tests reach the device validation and tagging logic +// without needing a Program or any device memory. +const executorch_flatbuffer::Tensor* buildTensor( + flatbuffers::FlatBufferBuilder& builder, + executorch_flatbuffer::TensorShapeDynamism dynamism, + bool with_extra_info, + executorch_flatbuffer::DeviceType device_type, + int8_t device_index) { + flatbuffers::Offset extra_info = 0; + if (with_extra_info) { + executorch_flatbuffer::ExtraTensorInfoBuilder eti(builder); + eti.add_device_type(device_type); + eti.add_device_index(device_index); + extra_info = eti.Finish(); + } + const std::vector sizes = {2, 2}; + const std::vector dim_order = {0, 1}; + const auto sizes_offset = builder.CreateVector(sizes); + const auto dim_order_offset = builder.CreateVector(dim_order); + + executorch_flatbuffer::TensorBuilder tb(builder); + tb.add_scalar_type(executorch_flatbuffer::ScalarType::FLOAT); + tb.add_storage_offset(0); + tb.add_sizes(sizes_offset); + tb.add_dim_order(dim_order_offset); + tb.add_data_buffer_idx(0); + tb.add_shape_dynamism(dynamism); + if (with_extra_info) { + tb.add_extra_tensor_info(extra_info); + } + builder.Finish(tb.Finish()); + return flatbuffers::GetRoot( + builder.GetBufferPointer()); +} + +} // namespace + +class TensorParserAtenTest : public ::testing::Test { + protected: + void SetUp() override { + executorch::runtime::runtime_init(); + } +}; + +// A CPU tensor must stay unindexed. An explicit cpu:0 would mismatch the +// graph's default cpu tensors and trip ATen's same-device check. +TEST_F(TensorParserAtenTest, CpuTensorIsUnindexed) { + flatbuffers::FlatBufferBuilder builder; + const auto* s_tensor = buildTensor( + builder, + executorch_flatbuffer::TensorShapeDynamism::STATIC, + /*with_extra_info=*/true, + executorch_flatbuffer::DeviceType::CPU, + /*device_index=*/0); + + ManagedMemoryManager mmm(kMemBytes, kMemBytes); + Result tensor = parseTensor(nullptr, &mmm.get(), s_tensor); + ASSERT_EQ(tensor.error(), Error::Ok); + EXPECT_TRUE(tensor->is_cpu()); + EXPECT_FALSE(tensor->device().has_index()); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); +} + +// A tensor with no extra_tensor_info (older PTE files) defaults to CPU. +TEST_F(TensorParserAtenTest, MissingExtraInfoDefaultsToCpu) { + flatbuffers::FlatBufferBuilder builder; + const auto* s_tensor = buildTensor( + builder, + executorch_flatbuffer::TensorShapeDynamism::STATIC, + /*with_extra_info=*/false, + executorch_flatbuffer::DeviceType::CPU, + /*device_index=*/0); + + ManagedMemoryManager mmm(kMemBytes, kMemBytes); + Result tensor = parseTensor(nullptr, &mmm.get(), s_tensor); + ASSERT_EQ(tensor.error(), Error::Ok); + EXPECT_TRUE(tensor->is_cpu()); + EXPECT_FALSE(tensor->device().has_index()); +} + +// An unmapped DeviceType from the untrusted PTE is rejected before it can reach +// c10::Device as garbage. +TEST_F(TensorParserAtenTest, InvalidDeviceTypeIsRejected) { + flatbuffers::FlatBufferBuilder builder; + const auto* s_tensor = buildTensor( + builder, + executorch_flatbuffer::TensorShapeDynamism::STATIC, + /*with_extra_info=*/true, + static_cast(7), + /*device_index=*/0); + + ManagedMemoryManager mmm(kMemBytes, kMemBytes); + EXPECT_EQ( + parseTensor(nullptr, &mmm.get(), s_tensor).error(), + Error::InvalidProgram); +} + +// A negative accelerator index is not a valid serialized placement. +TEST_F(TensorParserAtenTest, NegativeAcceleratorIndexIsRejected) { + flatbuffers::FlatBufferBuilder builder; + const auto* s_tensor = buildTensor( + builder, + executorch_flatbuffer::TensorShapeDynamism::STATIC, + /*with_extra_info=*/true, + executorch_flatbuffer::DeviceType::CUDA, + /*device_index=*/-1); + + ManagedMemoryManager mmm(kMemBytes, kMemBytes); + EXPECT_EQ( + parseTensor(nullptr, &mmm.get(), s_tensor).error(), + Error::InvalidProgram); +} + +// A fully dynamic tensor gets a CPU-resizable allocator, so a non-CPU device +// cannot be honored and must be rejected rather than silently tagged CPU. +TEST_F(TensorParserAtenTest, DynamicUnboundCudaIsRejected) { + flatbuffers::FlatBufferBuilder builder; + const auto* s_tensor = buildTensor( + builder, + executorch_flatbuffer::TensorShapeDynamism::DYNAMIC_UNBOUND, + /*with_extra_info=*/true, + executorch_flatbuffer::DeviceType::CUDA, + /*device_index=*/0); + + ManagedMemoryManager mmm(kMemBytes, kMemBytes); + EXPECT_EQ( + parseTensor(nullptr, &mmm.get(), s_tensor).error(), Error::NotSupported); +} + +// The DYNAMIC_UNBOUND guard only rejects non-CPU; CPU resizable tensors still +// parse so the guard does not regress the common case. +TEST_F(TensorParserAtenTest, DynamicUnboundCpuIsAllowed) { + flatbuffers::FlatBufferBuilder builder; + const auto* s_tensor = buildTensor( + builder, + executorch_flatbuffer::TensorShapeDynamism::DYNAMIC_UNBOUND, + /*with_extra_info=*/false, + executorch_flatbuffer::DeviceType::CPU, + /*device_index=*/0); + + ManagedMemoryManager mmm(kMemBytes, kMemBytes); + Result tensor = parseTensor(nullptr, &mmm.get(), s_tensor); + ASSERT_EQ(tensor.error(), Error::Ok); + EXPECT_TRUE(tensor->is_cpu()); + EXPECT_FALSE(tensor->device().has_index()); +}