diff --git a/runtime/core/exec_aten/util/tensor_util_aten.cpp b/runtime/core/exec_aten/util/tensor_util_aten.cpp index 024991a2355..cedb18ae2bb 100644 --- a/runtime/core/exec_aten/util/tensor_util_aten.cpp +++ b/runtime/core/exec_aten/util/tensor_util_aten.cpp @@ -177,6 +177,21 @@ Error copy_tensor_data(const at::Tensor& t_dst, const at::Tensor& t_src) { t_src.nbytes()); // Copy the source data to the preallocated memory of the destination, which // must be the same size as the source. + // + // Both sides have to be host memory. Reaching here with device memory means + // a program planned a buffer for a tensor that lives on an accelerator, and + // a host memcpy into it is undefined. Reported because the alternative is a + // crash with no message. + ET_CHECK_OR_RETURN_ERROR( + t_dst.device().is_cpu() && t_src.device().is_cpu(), + NotSupported, + // Kept under the 256-char log buffer (runtime/platform/log.cpp) so the + // MemoryPlanningPass hint, which is the actionable part, is not + // truncated away. + "Planned-buffer copy needs host memory on both sides: dst device %s, src device %s. " + "Export with MemoryPlanningPass(alloc_graph_input=False) to share the caller's memory.", + c10::DeviceTypeName(t_dst.device().type()).c_str(), + c10::DeviceTypeName(t_src.device().type()).c_str()); std::memcpy(dst_data_ptr, t_src.const_data_ptr(), t_src.nbytes()); } diff --git a/runtime/core/exec_aten/util/tensor_util_portable.cpp b/runtime/core/exec_aten/util/tensor_util_portable.cpp index 9626974ad7d..e3c258af1d5 100644 --- a/runtime/core/exec_aten/util/tensor_util_portable.cpp +++ b/runtime/core/exec_aten/util/tensor_util_portable.cpp @@ -178,6 +178,20 @@ Error copy_tensor_data( "t_dst.nbytes() %zu != t_src.nbytes(). %zu", t_dst.nbytes(), t_src.nbytes()); + // The copy below assumes host memory on both sides. Reaching it with device + // memory means a program planned a buffer for a tensor that lives on an + // accelerator, and copying into it with a host memcpy is undefined. + // Reported here because the alternative is a crash with no message. + ET_CHECK_OR_RETURN_ERROR( + t_dst.device().is_cpu() && t_src.device().is_cpu(), + NotSupported, + // Kept under the 256-char log buffer (runtime/platform/log.cpp) so the + // MemoryPlanningPass hint, which is the actionable part, is not + // truncated away. + "Planned-buffer copy needs host memory on both sides: dst device %d, src device %d. " + "Export with MemoryPlanningPass(alloc_graph_input=False) to share the caller's memory.", + static_cast(t_dst.device().type()), + static_cast(t_src.device().type())); std::memcpy( t_dst.mutable_data_ptr(), t_src.const_data_ptr(), t_src.nbytes()); } diff --git a/runtime/core/exec_aten/util/test/CMakeLists.txt b/runtime/core/exec_aten/util/test/CMakeLists.txt index e806419e21e..cf7dba2dbc7 100644 --- a/runtime/core/exec_aten/util/test/CMakeLists.txt +++ b/runtime/core/exec_aten/util/test/CMakeLists.txt @@ -20,9 +20,9 @@ set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..) include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) set(_test_srcs - dim_order_util_test.cpp operator_impl_example_test.cpp - scalar_type_util_test.cpp tensor_shape_to_c_string_test.cpp - tensor_util_test.cpp + copy_tensor_data_device_test.cpp dim_order_util_test.cpp + operator_impl_example_test.cpp scalar_type_util_test.cpp + tensor_shape_to_c_string_test.cpp tensor_util_test.cpp ) et_cxx_test(runtime_core_exec_aten_util_test SOURCES ${_test_srcs} EXTRA_LIBS) diff --git a/runtime/core/exec_aten/util/test/copy_tensor_data_device_test.cpp b/runtime/core/exec_aten/util/test/copy_tensor_data_device_test.cpp new file mode 100644 index 00000000000..8b4983e07d7 --- /dev/null +++ b/runtime/core/exec_aten/util/test/copy_tensor_data_device_test.cpp @@ -0,0 +1,88 @@ +/* + * 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 +#include +#include + +#include + +#include +#include +#include + +using executorch::runtime::Error; +using executorch::runtime::etensor::DeviceType; +using executorch::runtime::etensor::ScalarType; +using executorch::runtime::etensor::Tensor; +using executorch::runtime::etensor::TensorImpl; +using executorch::runtime::internal::copy_tensor_data; + +// Copying into a memory-planned tensor uses a host memcpy. A tensor that lives +// on an accelerator cannot be filled that way, and doing it anyway crashed the +// process with no message, so the runtime reports it instead. +class CopyTensorDataDeviceTest : public ::testing::Test { + protected: + void SetUp() override { + executorch::runtime::runtime_init(); + } + + // Each tensor gets its own storage, so a host-to-host copy has a distinct + // destination to land in and the data can be checked afterwards. Sharing one + // buffer would make the copy a self-memcpy, which is undefined and which no + // assertion could distinguish from the copy being skipped. + Tensor make( + DeviceType device, + std::array values = {0.0f, 0.0f, 0.0f, 0.0f}) { + storages_.push_back(std::make_unique>(values)); + impls_.push_back(std::make_unique( + ScalarType::Float, + static_cast(sizes_.size()), + sizes_.data(), + storages_.back()->data(), + dim_order_.data(), + strides_.data(), + executorch::runtime::TensorShapeDynamism::STATIC, + device)); + return Tensor(impls_.back().get()); + } + + std::array sizes_{4}; + std::array dim_order_{0}; + std::array strides_{1}; + std::vector>> storages_; + std::vector> impls_; +}; + +TEST_F(CopyTensorDataDeviceTest, HostToHostIsCopied) { + Tensor destination = make(DeviceType::CPU); + Tensor source = make(DeviceType::CPU, {1.0f, 2.0f, 3.0f, 4.0f}); + EXPECT_EQ(copy_tensor_data(destination, source), Error::Ok); + // Assert the data actually moved. Checking only the returned Error would + // still pass if the copy were removed entirely. + const float* copied = destination.const_data_ptr(); + ASSERT_NE(copied, nullptr); + EXPECT_EQ(copied[0], 1.0f); + EXPECT_EQ(copied[1], 2.0f); + EXPECT_EQ(copied[2], 3.0f); + EXPECT_EQ(copied[3], 4.0f); +} + +TEST_F(CopyTensorDataDeviceTest, ADeviceDestinationIsRefused) { + // This is the case that used to crash: a planned buffer on an accelerator, + // filled by a host memcpy. + Tensor destination = make(DeviceType::CUDA); + Tensor source = make(DeviceType::CPU); + EXPECT_EQ(copy_tensor_data(destination, source), Error::NotSupported); +} + +TEST_F(CopyTensorDataDeviceTest, ADeviceSourceIsRefused) { + Tensor destination = make(DeviceType::CPU); + Tensor source = make(DeviceType::CUDA); + EXPECT_EQ(copy_tensor_data(destination, source), Error::NotSupported); +} diff --git a/runtime/core/exec_aten/util/test/targets.bzl b/runtime/core/exec_aten/util/test/targets.bzl index 25077550e6c..dfc770da568 100644 --- a/runtime/core/exec_aten/util/test/targets.bzl +++ b/runtime/core/exec_aten/util/test/targets.bzl @@ -50,6 +50,15 @@ def define_common_targets(): ], ) + runtime.cxx_test( + name = "copy_tensor_data_device_test", + srcs = ["copy_tensor_data_device_test.cpp"], + deps = [ + "//executorch/runtime/core/exec_aten/util:tensor_util", + "//executorch/runtime/core/portable_type:portable_type", + ], + ) + runtime.cxx_test( name = "tensor_shape_to_c_string_test", srcs = ["tensor_shape_to_c_string_test.cpp"], diff --git a/test/utils/OSSTestConfig.json b/test/utils/OSSTestConfig.json index d7d0cb08567..97597db543e 100644 --- a/test/utils/OSSTestConfig.json +++ b/test/utils/OSSTestConfig.json @@ -97,6 +97,7 @@ { "directory": "runtime/core/exec_aten/util/test", "sources": [ + "copy_tensor_data_device_test.cpp", "dim_order_util_test.cpp", "operator_impl_example_test.cpp", "scalar_type_util_test.cpp",