From cbcf910f93e4bf45a7deba990755360dd357ede5 Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Sun, 9 Aug 2026 16:16:00 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../core/exec_aten/util/tensor_util_aten.cpp | 13 ++++ .../exec_aten/util/tensor_util_portable.cpp | 12 +++ .../core/exec_aten/util/test/CMakeLists.txt | 6 +- .../test/copy_tensor_data_device_test.cpp | 74 +++++++++++++++++++ runtime/core/exec_aten/util/test/targets.bzl | 9 +++ test/utils/OSSTestConfig.json | 1 + 6 files changed, 112 insertions(+), 3 deletions(-) create mode 100644 runtime/core/exec_aten/util/test/copy_tensor_data_device_test.cpp diff --git a/runtime/core/exec_aten/util/tensor_util_aten.cpp b/runtime/core/exec_aten/util/tensor_util_aten.cpp index 024991a2355..8a6605ee54f 100644 --- a/runtime/core/exec_aten/util/tensor_util_aten.cpp +++ b/runtime/core/exec_aten/util/tensor_util_aten.cpp @@ -177,6 +177,19 @@ 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, + "Cannot copy into a memory-planned tensor that does not live on the host: destination " + "device %s, source device %s. A program whose activations stay on a device must be exported " + "with MemoryPlanningPass(alloc_graph_input=False) so the runtime shares the caller's memory " + "instead of copying into its own buffer.", + 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..eb7b532858e 100644 --- a/runtime/core/exec_aten/util/tensor_util_portable.cpp +++ b/runtime/core/exec_aten/util/tensor_util_portable.cpp @@ -178,6 +178,18 @@ 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, + "Cannot copy into a memory-planned tensor that does not live on the host: destination " + "device type %d, source device type %d. A program whose activations stay on a device must " + "be exported with MemoryPlanningPass(alloc_graph_input=False) so the runtime shares the " + "caller's memory instead of copying into its own buffer.", + 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..d8f00c40a20 --- /dev/null +++ b/runtime/core/exec_aten/util/test/copy_tensor_data_device_test.cpp @@ -0,0 +1,74 @@ +/* + * 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::internal::copy_tensor_data; +using executorch::runtime::etensor::DeviceType; +using executorch::runtime::etensor::ScalarType; +using executorch::runtime::etensor::Tensor; +using executorch::runtime::etensor::TensorImpl; + +// 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(); + } + + // The data pointer is never dereferenced in the refusal cases below, because the copy is rejected + // before it happens. It has to be non-null so the earlier argument checks pass. + Tensor make(DeviceType device) { + impls_.push_back(std::make_unique( + ScalarType::Float, + static_cast(sizes_.size()), + sizes_.data(), + storage_.data(), + dim_order_.data(), + strides_.data(), + executorch::runtime::etensor::TensorShapeDynamism::STATIC, + device)); + return Tensor(impls_.back().get()); + } + + std::array sizes_{4}; + std::array dim_order_{0}; + std::array strides_{1}; + std::array storage_{1.0f, 2.0f, 3.0f, 4.0f}; + std::vector> impls_; +}; + +TEST_F(CopyTensorDataDeviceTest, HostToHostIsCopied) { + Tensor destination = make(DeviceType::CPU); + Tensor source = make(DeviceType::CPU); + EXPECT_EQ(copy_tensor_data(destination, source), Error::Ok); +} + +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",