Copy a tensor back from the device before ETDump writes it (pytorch#22095)

### Summary

ETDump records an intermediate tensor by handing the tensor's data
pointer to a
data sink, and every sink reads those bytes with a plain host read:

```cpp
memcpy(cur_data_begin, ptr, length);
```

When the tensor lives on an accelerator that pointer is not host memory,
so the
read segfaults. A program placed on CUDA crashes as soon as tracing is
turned
on, which is exactly when someone is trying to debug it.

Returning an error instead would not help. All four callers wrap the
result in
`ET_CHECK_MSG`, so an error aborts the process rather than skipping the
tensor.

This change brings the data back to host memory first. When the tensor
is not on
CPU, ETDump looks up the allocator registered for that device type,
stages the
bytes into a temporary host buffer with `copy_device_to_host`, writes
that buffer
to the sink and frees it. A tensor on CPU keeps the old path and copies
nothing
extra.

If no allocator is registered for the device, ETDump now reports
`NotFound` and
logs the device type instead of reading the pointer anyway.

### Test plan

Two new test files, each with a CMake target and a Buck target.

`devtools/etdump/tests/etdump_device_test.cpp` registers the existing
`MockCudaAllocator`, which backs its device memory with host memory, and
has two
tests. One logs a tensor tagged as CUDA and checks both that ETDump went
through
the allocator and that the bytes reached the debug buffer. The other
logs a
tensor on CPU and checks that the allocator was not used at all, so the
CPU path
is unchanged.

`devtools/etdump/tests/etdump_device_no_allocator_test.cpp` covers the
case where
nothing is registered for the device. The registry is a process wide
static with
no way to remove an entry, so that case needs a binary that never
registers
anything, which is why it is a second file.

`devtools/etdump/tests/CMakeLists.txt` was not referenced by any parent
`CMakeLists.txt`, so nothing in that directory was built by CMake. This
adds
`add_subdirectory(tests)` to `devtools/etdump/CMakeLists.txt` under
`BUILD_TESTING`, so both new tests are picked up by `ctest`, which is
how the C++
tests run.

The pre-existing `sdk_etdump_tests` target stays out of the CMake build.
It
compiles `etdump_test.cpp`, which includes `etdump_filter.h`, which
needs re2,
and the devtools build does not pull re2 in. It is now guarded on re2
being
available rather than being silently unreachable.

With this change both new binaries pass under `ctest`:

```
1/2 Test #1: etdump_device_test ................   Passed
2/2 Test #2: etdump_device_no_allocator_test ...   Passed
```

With `etdump_flatcc.cpp` reverted to the old code and everything
rebuilt, both
fail:

```
Expected equality of these values:
  g_mock_cuda.d2h_count_
    Which is: 0
  1

Death test: etdump_gen.log_evalue(EValue(tensor))
    Result: failed to die.
```

Also reproduced the real crash on one NVIDIA H100, with a small program
that
allocates through `cudaMalloc`, tags a tensor as CUDA and logs it:

```
before: Segmentation fault (core dumped)
after:  debug buffer holds 1.5 2.5 3.5 4.5
```

Checked that `etdump_flatcc.cpp` still compiles with `-DUSE_ATEN_LIB`.

`clang-format` reports no changes needed on the four touched C++ files.

### Landing order

pytorch#22058 adds device-planned arenas to the Python bindings that ETDump can
be handed
pointers into, and without this change `BufferDataSink::write` would
`memcpy` device
memory from the host. This should land before or together with it.

### Not covered

The ATen mode branch of the device type conversion only compiles. There
is no
ATen mode CMake build to run it in, so nothing here executes it. It is
reachable
in principle: `runtime/executor/tensor_parser_aten.cpp` reads the
serialized
device type and index, including CUDA, and builds the ATen tensor on
that device.
An earlier version of this description said ATen tensors carry no device
metadata,
which is wrong.

`LogTensorOnCpuDoesNotStageThroughTheAllocator` passes with the
production change
reverted, since CPU tensors already went straight to the data sink. It
documents
the CPU path rather than locking the fix; the other two tests are the
ones that
require the new branch.

---------

Co-authored-by: r <r@e>

Author: shoumikhin

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