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

runtime/core/exec_aten/util/tensor_util_aten.cpp

@@ -1,3 +1,3 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
@@ -177,6 +177,19 @@
         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());
   }
 

runtime/core/exec_aten/util/tensor_util_portable.cpp

@@ -1,3 +1,3 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
@@ -178,6 +178,18 @@
         "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<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,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 <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::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<TensorImpl>(
+        ScalarType::Float,
+        static_cast<ssize_t>(sizes_.size()),
+        sizes_.data(),
+        storage_.data(),
+        dim_order_.data(),
+        strides_.data(),
+        executorch::runtime::etensor::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::array<float, 4> storage_{1.0f, 2.0f, 3.0f, 4.0f};
+  std::vector<std::unique_ptr<TensorImpl>> 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);
+}

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