Report a clear error instead of crashing on a device-planned copy

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

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<int>(t_dst.device().type()),
+        static_cast<int>(t_src.device().type()));
     std::memcpy(
         t_dst.mutable_data_ptr(), t_src.const_data_ptr(), t_src.nbytes());
   }

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)

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 <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
+#include <executorch/runtime/platform/runtime.h>
+
+#include <gtest/gtest.h>
+
+#include <array>
+#include <memory>
+#include <vector>
+
+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<float, 4> values = {0.0f, 0.0f, 0.0f, 0.0f}) {
+    storages_.push_back(std::make_unique<std::array<float, 4>>(values));
+    impls_.push_back(std::make_unique<TensorImpl>(
+        ScalarType::Float,
+        static_cast<ssize_t>(sizes_.size()),
+        sizes_.data(),
+        storages_.back()->data(),
+        dim_order_.data(),
+        strides_.data(),
+        executorch::runtime::TensorShapeDynamism::STATIC,
+        device));
+    return Tensor(impls_.back().get());
+  }
+
+  std::array<executorch::aten::SizesType, 1> sizes_{4};
+  std::array<executorch::aten::DimOrderType, 1> dim_order_{0};
+  std::array<executorch::aten::StridesType, 1> strides_{1};
+  std::vector<std::unique_ptr<std::array<float, 4>>> storages_;
+  std::vector<std::unique_ptr<TensorImpl>> 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<float>();
+  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);
+}

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"],

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",