make ETDumpGen use bufferdatasink

Differential Revision: D69647096

Pull Request resolved: pytorch/executorch#8741

Author: Gasoonjia

devtools/CMakeLists.txt

@@ -176,6 +176,8 @@ add_custom_command(
 add_library(
   etdump ${CMAKE_CURRENT_SOURCE_DIR}/etdump/etdump_flatcc.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/etdump/emitter.cpp
+         ${CMAKE_CURRENT_SOURCE_DIR}/etdump/buffer_data_sink.cpp
+         ${CMAKE_CURRENT_SOURCE_DIR}/etdump/buffer_data_sink.h
 )
 
 target_link_libraries(

devtools/etdump/buffer_data_sink.cpp

@@ -27,6 +27,11 @@ Result<BufferDataSink> BufferDataSink::create(
   return BufferDataSink(buffer, alignment);
 }
 
+Result<BufferDataSink>
+BufferDataSink::create(void* ptr, size_t size, size_t alignment) noexcept {
+  return BufferDataSink::create({(uint8_t*)ptr, size}, alignment);
+}
+
 Result<size_t> BufferDataSink::write(const void* ptr, size_t length) {
   if (length == 0) {
     return offset_;

devtools/etdump/buffer_data_sink.h

@@ -39,6 +39,26 @@ class BufferDataSink : public DataSinkBase {
       ::executorch::runtime::Span<uint8_t> buffer,
       size_t alignment = 64) noexcept;
 
+  /**
+   * Creates a BufferDataSink with a given span buffer.
+   *
+   * @param[in] ptr A pointer to the data blob where data will be stored.
+   * @param[in] size The size of the data blob in bytes.
+   * @param[in] alignment The alignment requirement for the buffer. It must be
+   * a power of two and greater than zero. Default is 64.
+   * @return A Result object containing either:
+   *         - A BufferDataSink object if succees, or
+   *         - An error code indicating the failure reason, if any issue
+   *           occurs during the creation process.
+   */
+  static ::executorch::runtime::Result<BufferDataSink>
+  create(void* ptr, size_t size, size_t alignment = 64) noexcept;
+
+  /**
+   * Creates a empty BufferDataSink;
+   */
+  BufferDataSink() = default;
+
   // Uncopiable and unassignable to avoid double assignment and free of the
   // internal buffer.
   BufferDataSink(const BufferDataSink&) = delete;

devtools/etdump/etdump_flatcc.cpp

@@ -10,6 +10,7 @@
 
 #include <cstring>
 
+#include <executorch/devtools/etdump/buffer_data_sink.h>
 #include <executorch/devtools/etdump/emitter.h>
 #include <executorch/devtools/etdump/etdump_schema_flatcc_builder.h>
 #include <executorch/devtools/etdump/etdump_schema_flatcc_reader.h>
@@ -29,6 +30,7 @@ using ::executorch::runtime::DelegateDebugIdType;
 using ::executorch::runtime::EValue;
 using ::executorch::runtime::EventTracerEntry;
 using ::executorch::runtime::LoggedEValueType;
+using ::executorch::runtime::Result;
 using ::executorch::runtime::Span;
 using ::executorch::runtime::Tag;
 
@@ -142,6 +144,7 @@ ETDumpGen::~ETDumpGen() {
 void ETDumpGen::reset() {
   state_ = State::Init;
   num_blocks_ = 0;
+  data_sink_ = nullptr;
   flatcc_builder_reset(builder_);
   flatbuffers_buffer_start(builder_, etdump_ETDump_file_identifier);
   etdump_ETDump_start_as_root_with_size(builder_);
@@ -347,10 +350,10 @@ void ETDumpGen::log_intermediate_output_delegate_helper(
   ET_CHECK_MSG(
       (name == nullptr) ^ (delegate_debug_index == -1),
       "Only name or delegate_debug_index can be valid. Check DelegateMappingBuilder documentation for more details.");
-  if (debug_buffer_.empty()) {
-    ET_CHECK_MSG(0, "Must pre-set debug buffer with set_debug_buffer()\n");
-    return;
-  }
+
+  ET_CHECK_MSG(
+      data_sink_,
+      "Must pre-set data sink before logging evalue with set_data_sink() or set_debug_buffer()\n");
 
   check_ready_to_add_events();
   int64_t string_id = name != nullptr ? create_string_entry(name) : -1;
@@ -367,7 +370,7 @@ void ETDumpGen::log_intermediate_output_delegate_helper(
 
   // Check the type of `output` then call the corresponding logging functions
   if constexpr (std::is_same<T, Tensor>::value) {
-    long offset = copy_tensor_to_debug_buffer(output);
+    long offset = write_tensor_or_raise_error(output);
     etdump_Tensor_ref_t tensor_ref = add_tensor_entry(builder_, output, offset);
 
     etdump_Value_start(builder_);
@@ -377,7 +380,7 @@ void ETDumpGen::log_intermediate_output_delegate_helper(
   } else if constexpr (std::is_same<T, ArrayRef<Tensor>>::value) {
     etdump_Tensor_vec_start(builder_);
     for (size_t i = 0; i < output.size(); ++i) {
-      long offset = copy_tensor_to_debug_buffer(output[i]);
+      long offset = write_tensor_or_raise_error(output[i]);
       etdump_Tensor_vec_push(
           builder_, add_tensor_entry(builder_, output[i], offset));
     }
@@ -497,27 +500,22 @@ ETDumpResult ETDumpGen::get_etdump_data() {
 }
 
 void ETDumpGen::set_debug_buffer(Span<uint8_t> buffer) {
-  debug_buffer_ = buffer;
+  Result<BufferDataSink> bds_ret = BufferDataSink::create(buffer);
+  ET_CHECK_MSG(
+      bds_ret.ok(),
+      "Failed to write tensor with error 0x%" PRIx32,
+      static_cast<uint32_t>(bds_ret.error()));
+
+  buffer_data_sink_ = std::move(bds_ret.get());
+  data_sink_ = &buffer_data_sink_;
 }
 
-size_t ETDumpGen::copy_tensor_to_debug_buffer(executorch::aten::Tensor tensor) {
-  if (tensor.nbytes() == 0) {
-    return static_cast<size_t>(-1);
-  }
-  uint8_t* offset_ptr =
-      internal::align_pointer(debug_buffer_.data() + debug_buffer_offset_, 64);
-  debug_buffer_offset_ = (offset_ptr - debug_buffer_.data()) + tensor.nbytes();
-  ET_CHECK_MSG(
-      debug_buffer_offset_ <= debug_buffer_.size(),
-      "Ran out of space to store intermediate outputs.");
-  memcpy(offset_ptr, tensor.const_data_ptr(), tensor.nbytes());
-  return (size_t)(offset_ptr - debug_buffer_.data());
+void ETDumpGen::set_data_sink(DataSinkBase* data_sink) {
+  data_sink_ = data_sink;
 }
 
 void ETDumpGen::log_evalue(const EValue& evalue, LoggedEValueType evalue_type) {
-  if (debug_buffer_.empty()) {
-    return;
-  }
+  ET_CHECK_MSG(data_sink_, "Must set data sink before logging evalue\n");
 
   check_ready_to_add_events();
 
@@ -529,7 +527,7 @@ void ETDumpGen::log_evalue(const EValue& evalue, LoggedEValueType evalue_type) {
   switch (evalue.tag) {
     case Tag::Tensor: {
       executorch::aten::Tensor tensor = evalue.toTensor();
-      long offset = copy_tensor_to_debug_buffer(tensor);
+      long offset = write_tensor_or_raise_error(tensor);
       etdump_Tensor_ref_t tensor_ref =
           add_tensor_entry(builder_, tensor, offset);
 
@@ -551,7 +549,7 @@ void ETDumpGen::log_evalue(const EValue& evalue, LoggedEValueType evalue_type) {
           evalue.toTensorList();
       etdump_Tensor_vec_start(builder_);
       for (size_t i = 0; i < tensors.size(); ++i) {
-        long offset = copy_tensor_to_debug_buffer(tensors[i]);
+        long offset = write_tensor_or_raise_error(tensors[i]);
         etdump_Tensor_vec_push(
             builder_, add_tensor_entry(builder_, tensors[i], offset));
       }
@@ -635,8 +633,31 @@ bool ETDumpGen::is_static_etdump() {
   return alloc_.data != nullptr;
 }
 
-size_t ETDumpGen::get_debug_buffer_size() const {
-  return debug_buffer_.size();
+DataSinkBase* ETDumpGen::get_data_sink() {
+  return data_sink_;
+}
+
+long ETDumpGen::write_tensor_or_raise_error(Tensor tensor) {
+  // Previously, the function copy_tensor_to_debug_buffer returned 0xFF..F when
+  // given an empty tensor, which is an invalid offset for most buffers. In our
+  // data sink, we will return the current debug_buffer_offset for better
+  // clarity. We are isolating the empty tensor case here using the old logic to
+  // avoid any backward compatibility issues while introducing the data sink.
+  // Once the data sink is fully implemented, we can remove this check and apply
+  // the new logic to all cases.
+  // TODO(gasoonjia): remove this check after datasink is fully rolled out.
+  if (tensor.nbytes() == 0) {
+    return static_cast<size_t>(-1);
+  }
+
+  ET_CHECK_MSG(data_sink_, "Must set data sink before writing data");
+  Result<size_t> ret =
+      data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
+  ET_CHECK_MSG(
+      ret.ok(),
+      "Failed to write tensor with error 0x%" PRIx32,
+      static_cast<uint32_t>(ret.error()));
+  return static_cast<long>(ret.get());
 }
 
 } // namespace etdump

devtools/etdump/etdump_flatcc.h

@@ -9,7 +9,10 @@
 #pragma once
 
 #include <cstdint>
+#include <memory>
 
+#include <executorch/devtools/etdump/buffer_data_sink.h>
+#include <executorch/devtools/etdump/data_sink_base.h>
 #include <executorch/runtime/core/event_tracer.h>
 #include <executorch/runtime/core/span.h>
 #include <executorch/runtime/platform/platform.h>
@@ -141,9 +144,10 @@ class ETDumpGen : public ::executorch::runtime::EventTracer {
       ::executorch::runtime::DebugHandle delegate_debug_index,
       const double& output) override;
   void set_debug_buffer(::executorch::runtime::Span<uint8_t> buffer);
+  void set_data_sink(DataSinkBase* data_sink);
   ETDumpResult get_etdump_data();
-  size_t get_debug_buffer_size() const;
   size_t get_num_blocks();
+  DataSinkBase* get_data_sink();
   bool is_static_etdump();
   void reset();
 
@@ -158,7 +162,6 @@ class ETDumpGen : public ::executorch::runtime::EventTracer {
 
   void check_ready_to_add_events();
   int64_t create_string_entry(const char* name);
-  size_t copy_tensor_to_debug_buffer(executorch::aten::Tensor tensor);
 
   /**
    * Templated helper function used to log various types of intermediate output.
@@ -170,10 +173,15 @@ class ETDumpGen : public ::executorch::runtime::EventTracer {
       ::executorch::runtime::DebugHandle delegate_debug_index,
       const T& output);
 
+  long write_tensor_or_raise_error(executorch::aten::Tensor tensor);
+
   struct flatcc_builder* builder_;
   size_t num_blocks_ = 0;
-  ::executorch::runtime::Span<uint8_t> debug_buffer_;
-  size_t debug_buffer_offset_ = 0;
+  DataSinkBase* data_sink_;
+
+  // It is only for set_debug_buffer function.
+  BufferDataSink buffer_data_sink_;
+
   int bundled_input_index_ = -1;
   State state_ = State::Init;
   struct internal::ETDumpStaticAllocator alloc_;

devtools/etdump/targets.bzl

@@ -117,7 +117,7 @@ def define_common_targets():
 
         runtime.cxx_library(
             name = "buffer_data_sink" + aten_suffix,
-            headers = [
+            exported_headers = [
                 "buffer_data_sink.h",
             ],
             srcs = [
@@ -153,6 +153,8 @@ def define_common_targets():
             exported_deps = [
                 ":etdump_schema_flatcc",
                 ":utils",
+                ":data_sink_base" + aten_suffix,
+                ":buffer_data_sink" + aten_suffix,
                 "//executorch/runtime/core:event_tracer" + aten_suffix,
                 "//executorch/runtime/core/exec_aten/util:scalar_type_util" + aten_suffix,
             ],