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..194a6d5c2d3 100644 --- a/devtools/etdump/etdump_flatcc.cpp +++ b/devtools/etdump/etdump_flatcc.cpp @@ -15,11 +15,16 @@ #include #include #include +#include #include #include #include #include +#ifdef USE_ATEN_LIB +#include +#endif + #include using ::executorch::aten::Tensor; @@ -105,6 +110,75 @@ 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 + 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 @@ -725,6 +799,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(data_sink_, tensor); + } + Result ret = data_sink_->write(tensor.const_data_ptr(), tensor.nbytes()); if (!ret.ok()) { diff --git a/devtools/etdump/targets.bzl b/devtools/etdump/targets.bzl index 44c8c42c1b9..3759206f5d8 100644 --- a/devtools/etdump/targets.bzl +++ b/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", diff --git a/devtools/etdump/tests/CMakeLists.txt b/devtools/etdump/tests/CMakeLists.txt index 1443457932e..8d565af0a7e 100644 --- a/devtools/etdump/tests/CMakeLists.txt +++ b/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 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. 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 = [