Copy a tensor back from the device before ETDump writes it

devtools/etdump/CMakeLists.txt

@@ -92,3 +92,7 @@ install(
   INCLUDES
   DESTINATION ${_common_include_directories}
 )
+
+if(BUILD_TESTING)
+  add_subdirectory(tests)
+endif()

devtools/etdump/etdump_flatcc.cpp

@@ -15,11 +15,16 @@
 #include <executorch/devtools/etdump/etdump_schema_flatcc_builder.h>
 #include <executorch/devtools/etdump/etdump_schema_flatcc_reader.h>
 #include <executorch/devtools/etdump/utils.h>
+#include <executorch/runtime/core/device_allocator.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 #include <executorch/runtime/platform/assert.h>
 
+#ifdef USE_ATEN_LIB
+#include <executorch/extension/aten_util/aten_bridge.h>
+#endif
+
 #include <flatcc/flatcc_types.h>
 
 using ::executorch::aten::Tensor;
@@ -105,6 +110,75 @@ Result<etdump_Tensor_ref_t> add_tensor_entry(
   return etdump_Tensor_end(builder_);
 }
 
+// The allocator registry is keyed on the runtime's own device type. In ATen
+// mode a tensor carries the ATen enum instead, which names many more devices
+// than ExecuTorch can register an allocator for.
+Result<::executorch::runtime::etensor::DeviceType> to_runtime_device_type(
+    ::executorch::aten::DeviceType type) {
+#ifdef USE_ATEN_LIB
+  std::optional<::executorch::runtime::etensor::Device> device =
+      ::executorch::extension::torch_to_executorch_device(c10::Device(type));
+  if (!device.has_value()) {
+    return Error::NotSupported;
+  }
+  return device->type();
+#else
+  return type;
+#endif
+}
+
+// Stages a tensor that lives on an accelerator through host memory before
+// handing it to the data sink, which can only read host pointers. A data sink
+// stores the bytes with a plain host read, so handing it the accelerator
+// pointer straight would crash the process.
+Result<long> write_device_tensor(DataSinkBase* data_sink, Tensor tensor) {
+  Result<::executorch::runtime::etensor::DeviceType> device_type =
+      to_runtime_device_type(tensor.device().type());
+  if (!device_type.ok()) {
+    ET_LOG(
+        Error,
+        "ETDump cannot read a tensor on device type %d",
+        static_cast<int>(tensor.device().type()));
+    return device_type.error();
+  }
+
+  ::executorch::runtime::DeviceAllocator* allocator =
+      ::executorch::runtime::get_device_allocator(device_type.get());
+  if (allocator == nullptr) {
+    ET_LOG(
+        Error,
+        "No device allocator registered for device type %d, so a tensor on that device cannot be copied back to host memory",
+        static_cast<int>(device_type.get()));
+    return Error::NotFound;
+  }
+
+  const size_t nbytes = tensor.nbytes();
+  void* staging = malloc(nbytes);
+  if (staging == nullptr) {
+    ET_LOG(
+        Error, "Failed to allocate %zu bytes to stage a device tensor", nbytes);
+    return Error::MemoryAllocationFailed;
+  }
+
+  Error copy_error = allocator->copy_device_to_host(
+      staging,
+      tensor.const_data_ptr(),
+      nbytes,
+      static_cast<::executorch::runtime::etensor::DeviceIndex>(
+          tensor.device().index()));
+  if (copy_error != Error::Ok) {
+    free(staging);
+    return copy_error;
+  }
+
+  Result<size_t> ret = data_sink->write(staging, nbytes);
+  free(staging);
+  if (!ret.ok()) {
+    return ret.error();
+  }
+  return static_cast<long>(ret.get());
+}
+
 } // namespace
 
 // Constructor implementation
@@ -725,6 +799,11 @@ Result<long> ETDumpGen::write_tensor_or_return_error(Tensor tensor) {
   if (!data_sink_) {
     return Error::InvalidArgument;
   }
+
+  if (!tensor.device().is_cpu()) {
+    return write_device_tensor(data_sink_, tensor);
+  }
+
   Result<size_t> ret =
       data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
   if (!ret.ok()) {

devtools/etdump/targets.bzl

@@ -132,8 +132,11 @@ def define_common_targets():
                 "etdump_flatcc.h",
             ],
             deps = [
+                "//executorch/runtime/core:device_allocator",
                 "//executorch/runtime/platform:platform",
-            ],
+            ] + ([
+                "//executorch/extension/aten_util:aten_bridge",
+            ] if aten_mode else []),
             exported_deps = [
                 ":etdump_schema_flatcc",
                 ":utils",

devtools/etdump/tests/CMakeLists.txt

@@ -19,16 +19,51 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake)
 
 set(_test_srcs etdump_test.cpp)
 
+# etdump_test.cpp includes etdump_filter.h, which includes <re2/re2.h>. re2 only
+# reaches CMake through the tokenizers submodule, which a default checkout does
+# not populate, so re2::re2 usually does not exist and this suite is skipped
+# here. Buck builds it unconditionally, which is where it runs today.
+#
+# The guard is not sufficient on its own: this target links neither re2 nor
+# etdump_filter.cpp (the etdump library does not compile that file), so where
+# re2::re2 does exist the suite fails to link. That predates this file's change,
+# which only narrows when it is attempted. Fixing it properly means adding
+# etdump_filter.cpp to the etdump library and re2 to EXTRA_LIBS.
+#
+# The two tests below avoid etdump_filter.h so they always build.
+if(TARGET re2::re2)
+  et_cxx_test(
+    sdk_etdump_tests
+    SOURCES
+    ${_test_srcs}
+    EXTRA_LIBS
+    bundled_program
+    etdump
+    flatccrt
+  )
+  target_include_directories(
+    sdk_etdump_tests PRIVATE ${CMAKE_INSTALL_PREFIX}/sdk/include
+                             ${EXECUTORCH_ROOT}/third-party/flatcc/include
+  )
+endif()
+
 et_cxx_test(
-  sdk_etdump_tests
-  SOURCES
-  ${_test_srcs}
-  EXTRA_LIBS
-  bundled_program
-  etdump
-  flatccrt
+  etdump_device_test SOURCES etdump_device_test.cpp EXTRA_LIBS etdump flatccrt
 )
-target_include_directories(
-  sdk_etdump_tests PRIVATE ${CMAKE_INSTALL_PREFIX}/sdk/include
-                           ${EXECUTORCH_ROOT}/third-party/flatcc/include
+
+et_cxx_test(
+  etdump_device_no_allocator_test SOURCES etdump_device_no_allocator_test.cpp
+  EXTRA_LIBS etdump flatccrt
 )
+
+# These two register a device allocator and then call into etdump, which reads
+# it back out of the registry. et_cxx_test links the static executorch_core, so
+# in a shared build, where etdump resolves the runtime from libexecutorch.so,
+# the archive would win the runtime symbols and the test would hold a private
+# registry (Utils.cmake documents that ordering). The usual remedy,
+# executorch_target_link_shared_runtime, cannot be called from here: the root
+# CMakeLists adds devtools well before it defines executorch_shared, so the
+# target does not exist yet at this point. These tests are run by the Buck
+# targets and by the default static CMake build, both of which have one
+# registry; a shared-build CMake run of them needs that ordering addressed
+# first.

devtools/etdump/tests/etdump_device_no_allocator_test.cpp

@@ -0,0 +1,67 @@
+/*
+ * 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.
+ */
+
+// The allocator registry is a process wide static with no way to remove an
+// entry, so the "nothing registered for this device" path needs a binary that
+// never registers anything.
+
+#include <gtest/gtest.h>
+#include <cstdint>
+#include <cstdlib>
+
+#include <executorch/devtools/etdump/data_sinks/buffer_data_sink.h>
+#include <executorch/devtools/etdump/etdump_flatcc.h>
+#include <executorch/runtime/core/device_allocator.h>
+#include <executorch/runtime/core/portable_type/tensor_impl.h>
+#include <executorch/runtime/platform/runtime.h>
+#include <executorch/test/utils/DeathTest.h>
+
+using namespace ::executorch::etdump;
+using namespace ::executorch::runtime;
+using namespace ::executorch::runtime::etensor;
+
+namespace {
+
+TEST(ETDumpNoDeviceAllocatorTest, LogTensorOnUnregisteredDeviceAborts) {
+  runtime_init();
+  ASSERT_EQ(get_device_allocator(DeviceType::CUDA), nullptr);
+
+  // Host memory, so the test fails by not dying rather than by crashing if
+  // ETDump ever reads the pointer instead of reporting the missing allocator.
+  float data[] = {1.5f, 2.5f, 3.5f, 4.5f};
+  int32_t sizes[] = {4};
+  uint8_t dim_order[] = {0};
+  int32_t strides[] = {1};
+  TensorImpl impl(
+      ScalarType::Float,
+      1,
+      sizes,
+      data,
+      dim_order,
+      strides,
+      TensorShapeDynamism::STATIC,
+      DeviceType::CUDA,
+      0);
+  Tensor tensor(&impl);
+
+  const size_t debug_buf_size = 2048;
+  void* debug_buf = malloc(debug_buf_size);
+  auto buffer_data_sink = BufferDataSink::create(debug_buf, debug_buf_size);
+  ASSERT_TRUE(buffer_data_sink.ok());
+
+  ETDumpGen etdump_gen;
+  etdump_gen.create_event_block("test_block");
+  etdump_gen.set_data_sink(&buffer_data_sink.get());
+
+  ET_EXPECT_DEATH(
+      etdump_gen.log_evalue(EValue(tensor)), "No device allocator registered");
+
+  free(debug_buf);
+}
+
+} // namespace