From 7f7c703fdb5cbc6d60f494c6fdb22d7547c45112 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 24 Aug 2026 11:59:45 -0700 Subject: [PATCH 1/3] Copy a tensor back from the device before ETDump writes it 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: 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, 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. Not covered: The ATen mode branch of the device type conversion only compiles. There is no ATen mode CMake build to run it in, and an ATen mode tensor in ExecuTorch carries no device metadata today, so the branch has no caller that can reach it. --- devtools/etdump/CMakeLists.txt | 4 + devtools/etdump/etdump_flatcc.cpp | 77 +++++++++ devtools/etdump/etdump_flatcc.h | 5 + devtools/etdump/targets.bzl | 1 + devtools/etdump/tests/CMakeLists.txt | 33 ++-- .../tests/etdump_device_no_allocator_test.cpp | 67 ++++++++ devtools/etdump/tests/etdump_device_test.cpp | 148 ++++++++++++++++++ devtools/etdump/tests/targets.bzl | 28 ++++ 8 files changed, 353 insertions(+), 10 deletions(-) create mode 100644 devtools/etdump/tests/etdump_device_no_allocator_test.cpp create mode 100644 devtools/etdump/tests/etdump_device_test.cpp diff --git a/devtools/etdump/CMakeLists.txt b/devtools/etdump/CMakeLists.txt index 482bd03023f..85d73143455 100644 --- a/devtools/etdump/CMakeLists.txt +++ b/devtools/etdump/CMakeLists.txt @@ -92,3 +92,7 @@ install( INCLUDES DESTINATION ${_common_include_directories} ) + +if(BUILD_TESTING) + add_subdirectory(tests) +endif() diff --git a/devtools/etdump/etdump_flatcc.cpp b/devtools/etdump/etdump_flatcc.cpp index 92ce2070fc8..200ea18d869 100644 --- a/devtools/etdump/etdump_flatcc.cpp +++ b/devtools/etdump/etdump_flatcc.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -105,6 +106,25 @@ Result 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 + switch (type) { + case c10::DeviceType::CPU: + return ::executorch::runtime::etensor::DeviceType::CPU; + case c10::DeviceType::CUDA: + return ::executorch::runtime::etensor::DeviceType::CUDA; + default: + return Error::NotSupported; + } +#else + return type; +#endif +} + } // namespace // Constructor implementation @@ -725,6 +745,11 @@ Result ETDumpGen::write_tensor_or_return_error(Tensor tensor) { if (!data_sink_) { return Error::InvalidArgument; } + + if (!tensor.device().is_cpu()) { + return write_device_tensor_or_return_error(tensor); + } + Result ret = data_sink_->write(tensor.const_data_ptr(), tensor.nbytes()); if (!ret.ok()) { @@ -733,5 +758,57 @@ Result ETDumpGen::write_tensor_or_return_error(Tensor tensor) { return static_cast(ret.get()); } +Result ETDumpGen::write_device_tensor_or_return_error(Tensor tensor) { + // A data sink stores the bytes with a plain host read, so the data of a + // tensor that lives on an accelerator has to be brought back to host memory + // first. Handing the accelerator pointer straight to the sink crashes the + // process. + 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(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(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 ret = data_sink_->write(staging, nbytes); + free(staging); + if (!ret.ok()) { + return ret.error(); + } + return static_cast(ret.get()); +} + } // namespace etdump } // namespace executorch diff --git a/devtools/etdump/etdump_flatcc.h b/devtools/etdump/etdump_flatcc.h index 8b39b243165..217de4155e4 100644 --- a/devtools/etdump/etdump_flatcc.h +++ b/devtools/etdump/etdump_flatcc.h @@ -186,6 +186,11 @@ class ETDumpGen : public ::executorch::runtime::EventTracer { Result write_tensor_or_return_error(executorch::aten::Tensor tensor); + /// Stages a tensor that lives on an accelerator through host memory before + /// handing it to the data sink, which can only read host pointers. + Result write_device_tensor_or_return_error( + executorch::aten::Tensor tensor); + struct flatcc_builder* builder_; size_t num_blocks_ = 0; DataSinkBase* data_sink_; diff --git a/devtools/etdump/targets.bzl b/devtools/etdump/targets.bzl index 44c8c42c1b9..431b4898d6b 100644 --- a/devtools/etdump/targets.bzl +++ b/devtools/etdump/targets.bzl @@ -132,6 +132,7 @@ def define_common_targets(): "etdump_flatcc.h", ], deps = [ + "//executorch/runtime/core:device_allocator", "//executorch/runtime/platform:platform", ], exported_deps = [ diff --git a/devtools/etdump/tests/CMakeLists.txt b/devtools/etdump/tests/CMakeLists.txt index 1443457932e..e4be4e7320a 100644 --- a/devtools/etdump/tests/CMakeLists.txt +++ b/devtools/etdump/tests/CMakeLists.txt @@ -19,16 +19,29 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) set(_test_srcs etdump_test.cpp) +# etdump_test.cpp includes etdump_filter.h, which needs re2, and the devtools +# build does not pull re2 in. +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 ) diff --git a/devtools/etdump/tests/etdump_device_no_allocator_test.cpp b/devtools/etdump/tests/etdump_device_no_allocator_test.cpp new file mode 100644 index 00000000000..e2fb5eb526c --- /dev/null +++ b/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 +#include +#include + +#include +#include +#include +#include +#include +#include + +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 diff --git a/devtools/etdump/tests/etdump_device_test.cpp b/devtools/etdump/tests/etdump_device_test.cpp new file mode 100644 index 00000000000..ed16e605908 --- /dev/null +++ b/devtools/etdump/tests/etdump_device_test.cpp @@ -0,0 +1,148 @@ +/* + * 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 +#include +#include +#include +#include +#include + +using namespace ::executorch::etdump; +using namespace ::executorch::runtime; +using namespace ::executorch::runtime::etensor; +using namespace ::executorch::runtime::testing; + +namespace { + +// Backs its device memory with host memory, so the copy path runs without a +// GPU. +MockCudaAllocator g_mock_cuda; + +class ETDumpDeviceTest : public ::testing::Test { + protected: + // The registry only accepts allocators with static lifetime and aborts on a + // second registration for the same device type, so register once per binary. + static void SetUpTestSuite() { + register_device_allocator(&g_mock_cuda); + } + + void SetUp() override { + runtime_init(); + g_mock_cuda.reset(); + etdump_gen_ = new ETDumpGen(); + debug_buf_ = malloc(kDebugBufSize); + } + + void TearDown() override { + delete etdump_gen_; + free(debug_buf_); + } + + static constexpr size_t kDebugBufSize = 2048; + + ETDumpGen* etdump_gen_ = nullptr; + void* debug_buf_ = nullptr; + int32_t sizes_[1] = {4}; + uint8_t dim_order_[1] = {0}; + int32_t strides_[1] = {1}; +}; + +TEST_F(ETDumpDeviceTest, LogTensorOnDeviceCopiesItBackToHost) { + ASSERT_EQ(get_device_allocator(DeviceType::CUDA), &g_mock_cuda); + + float device_data[] = {1.5f, 2.5f, 3.5f, 4.5f}; + TensorImpl impl( + ScalarType::Float, + 1, + sizes_, + device_data, + dim_order_, + strides_, + TensorShapeDynamism::STATIC, + DeviceType::CUDA, + 0); + Tensor tensor(&impl); + + auto buffer_data_sink = BufferDataSink::create(debug_buf_, kDebugBufSize); + ASSERT_TRUE(buffer_data_sink.ok()); + + etdump_gen_->create_event_block("test_block"); + etdump_gen_->set_data_sink(&buffer_data_sink.get()); + etdump_gen_->log_evalue(EValue(tensor)); + + EXPECT_EQ(g_mock_cuda.d2h_count_, 1); + EXPECT_EQ(g_mock_cuda.last_d2h_size_, sizeof(device_data)); + EXPECT_EQ(g_mock_cuda.last_d2h_src_, device_data); + EXPECT_EQ(g_mock_cuda.last_d2h_index_, 0); + + ETDumpResult result = etdump_gen_->get_etdump_data(); + ASSERT_TRUE(result.buf != nullptr); + ASSERT_TRUE(result.size != 0); + + size_t size = 0; + void* buf = flatbuffers_read_size_prefix(result.buf, &size); + etdump_ETDump_table_t etdump = + etdump_ETDump_as_root_with_identifier(buf, etdump_ETDump_file_identifier); + etdump_RunData_vec_t run_data_vec = etdump_ETDump_run_data(etdump); + ASSERT_EQ(etdump_RunData_vec_len(run_data_vec), 1); + etdump_Event_vec_t events = + etdump_RunData_events(etdump_RunData_vec_at(run_data_vec, 0)); + ASSERT_EQ(etdump_Event_vec_len(events), 1); + etdump_Tensor_table_t logged = + etdump_Value_tensor(etdump_DebugEvent_debug_entry( + etdump_Event_debug_event(etdump_Event_vec_at(events, 0)))); + + // The sink aligns each blob, so the bytes start at the recorded offset rather + // than at the front of the buffer. + const long offset = etdump_Tensor_offset(logged); + ASSERT_GE(offset, 0); + EXPECT_EQ( + memcmp((uint8_t*)debug_buf_ + offset, device_data, sizeof(device_data)), + 0); + + free(result.buf); +} + +TEST_F(ETDumpDeviceTest, LogTensorOnCpuDoesNotStageThroughTheAllocator) { + float host_data[] = {1.5f, 2.5f, 3.5f, 4.5f}; + TensorImpl impl( + ScalarType::Float, + 1, + sizes_, + host_data, + dim_order_, + strides_, + TensorShapeDynamism::STATIC, + DeviceType::CPU, + 0); + Tensor tensor(&impl); + + auto buffer_data_sink = BufferDataSink::create(debug_buf_, kDebugBufSize); + ASSERT_TRUE(buffer_data_sink.ok()); + + etdump_gen_->create_event_block("test_block"); + etdump_gen_->set_data_sink(&buffer_data_sink.get()); + etdump_gen_->log_evalue(EValue(tensor)); + + EXPECT_EQ(g_mock_cuda.d2h_count_, 0); + + ETDumpResult result = etdump_gen_->get_etdump_data(); + ASSERT_TRUE(result.buf != nullptr); + ASSERT_TRUE(result.size != 0); + free(result.buf); +} + +} // namespace diff --git a/devtools/etdump/tests/targets.bzl b/devtools/etdump/tests/targets.bzl index f1ab3025955..2044b54c73f 100644 --- a/devtools/etdump/tests/targets.bzl +++ b/devtools/etdump/tests/targets.bzl @@ -24,6 +24,34 @@ def define_common_targets(is_fbcode = False): ], ) + runtime.cxx_test( + name = "etdump_device_test", + srcs = [ + "etdump_device_test.cpp", + ], + deps = [ + "//executorch/devtools/etdump:etdump_flatcc", + "//executorch/devtools/etdump:etdump_schema_flatcc", + "//executorch/devtools/etdump/data_sinks:buffer_data_sink", + "//executorch/runtime/core:device_allocator", + "//executorch/runtime/core/test:mock_cuda_allocator", + "//executorch/runtime/platform:platform", + ], + ) + + runtime.cxx_test( + name = "etdump_device_no_allocator_test", + srcs = [ + "etdump_device_no_allocator_test.cpp", + ], + deps = [ + "//executorch/devtools/etdump:etdump_flatcc", + "//executorch/devtools/etdump/data_sinks:buffer_data_sink", + "//executorch/runtime/core:device_allocator", + "//executorch/runtime/platform:platform", + ], + ) + runtime.cxx_test( name = "etdump_filter_test", srcs = [ From 80bda2942f1efc2e0a8245e5ada5270b272bf32f Mon Sep 17 00:00:00 2001 From: r Date: Mon, 24 Aug 2026 13:05:16 -0700 Subject: [PATCH 2/3] Document the re2 guard and the shared-build registry ordering The re2-guarded suite links neither re2 nor etdump_filter.cpp, so it cannot link where re2::re2 does exist. That predates this change, which only narrows when it is attempted; say so rather than implying the branch is merely rare. Also record why the two new tests do not call executorch_target_link_shared_runtime: the root CMakeLists adds devtools before it defines executorch_shared, so the target does not exist at that point. They run under Buck and the default static CMake build, which have a single registry. --- devtools/etdump/tests/CMakeLists.txt | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/devtools/etdump/tests/CMakeLists.txt b/devtools/etdump/tests/CMakeLists.txt index e4be4e7320a..8d565af0a7e 100644 --- a/devtools/etdump/tests/CMakeLists.txt +++ b/devtools/etdump/tests/CMakeLists.txt @@ -19,8 +19,18 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) set(_test_srcs etdump_test.cpp) -# etdump_test.cpp includes etdump_filter.h, which needs re2, and the devtools -# build does not pull re2 in. +# etdump_test.cpp includes etdump_filter.h, which includes . 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 @@ -45,3 +55,15 @@ 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. From 74995b78348adf58fc864c4a470a9bec67475133 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 24 Aug 2026 17:42:25 -0700 Subject: [PATCH 3/3] Address review: reuse aten_bridge device mapping and hide the device-copy helper Two changes from review feedback, no behavior change. Reuse the existing device mapping. to_runtime_device_type had its own CPU/CUDA switch in the USE_ATEN_LIB branch; that is exactly what extension/aten_util/aten_bridge.h's torch_to_executorch_device already does, so call it instead of duplicating the mapping. This pulls aten_bridge in as an aten-only dependency of etdump_flatcc; the non-ATen build and the CMake build, which never compile this branch, are unaffected. Stop exposing the device-copy path. write_device_tensor_or_return_error was a private member declared in the header. Nothing outside the class uses it, so move it into the anonymous namespace as a free function that takes the DataSinkBase it needs, and drop the declaration from etdump_flatcc.h. Only write_tensor_or_return_error remains, which is the one entry point callers use; it routes device tensors to the helper as before. --- devtools/etdump/etdump_flatcc.cpp | 122 +++++++++++++++--------------- devtools/etdump/etdump_flatcc.h | 5 -- devtools/etdump/targets.bzl | 4 +- 3 files changed, 65 insertions(+), 66 deletions(-) diff --git a/devtools/etdump/etdump_flatcc.cpp b/devtools/etdump/etdump_flatcc.cpp index 200ea18d869..194a6d5c2d3 100644 --- a/devtools/etdump/etdump_flatcc.cpp +++ b/devtools/etdump/etdump_flatcc.cpp @@ -21,6 +21,10 @@ #include #include +#ifdef USE_ATEN_LIB +#include +#endif + #include using ::executorch::aten::Tensor; @@ -112,19 +116,69 @@ Result add_tensor_entry( Result<::executorch::runtime::etensor::DeviceType> to_runtime_device_type( ::executorch::aten::DeviceType type) { #ifdef USE_ATEN_LIB - switch (type) { - case c10::DeviceType::CPU: - return ::executorch::runtime::etensor::DeviceType::CPU; - case c10::DeviceType::CUDA: - return ::executorch::runtime::etensor::DeviceType::CUDA; - default: - return Error::NotSupported; + 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 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(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(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 ret = data_sink->write(staging, nbytes); + free(staging); + if (!ret.ok()) { + return ret.error(); + } + return static_cast(ret.get()); +} + } // namespace // Constructor implementation @@ -747,7 +801,7 @@ Result ETDumpGen::write_tensor_or_return_error(Tensor tensor) { } if (!tensor.device().is_cpu()) { - return write_device_tensor_or_return_error(tensor); + return write_device_tensor(data_sink_, tensor); } Result ret = @@ -758,57 +812,5 @@ Result ETDumpGen::write_tensor_or_return_error(Tensor tensor) { return static_cast(ret.get()); } -Result ETDumpGen::write_device_tensor_or_return_error(Tensor tensor) { - // A data sink stores the bytes with a plain host read, so the data of a - // tensor that lives on an accelerator has to be brought back to host memory - // first. Handing the accelerator pointer straight to the sink crashes the - // process. - 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(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(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 ret = data_sink_->write(staging, nbytes); - free(staging); - if (!ret.ok()) { - return ret.error(); - } - return static_cast(ret.get()); -} - } // namespace etdump } // namespace executorch diff --git a/devtools/etdump/etdump_flatcc.h b/devtools/etdump/etdump_flatcc.h index 217de4155e4..8b39b243165 100644 --- a/devtools/etdump/etdump_flatcc.h +++ b/devtools/etdump/etdump_flatcc.h @@ -186,11 +186,6 @@ class ETDumpGen : public ::executorch::runtime::EventTracer { Result write_tensor_or_return_error(executorch::aten::Tensor tensor); - /// Stages a tensor that lives on an accelerator through host memory before - /// handing it to the data sink, which can only read host pointers. - Result write_device_tensor_or_return_error( - executorch::aten::Tensor tensor); - struct flatcc_builder* builder_; size_t num_blocks_ = 0; DataSinkBase* data_sink_; diff --git a/devtools/etdump/targets.bzl b/devtools/etdump/targets.bzl index 431b4898d6b..3759206f5d8 100644 --- a/devtools/etdump/targets.bzl +++ b/devtools/etdump/targets.bzl @@ -134,7 +134,9 @@ def define_common_targets(): 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",