[devtool] make ETDumpGen use bufferdatasink

Gasoonjia · Gasoonjia · commit 87a9c46fd67b · 2025-02-16T00:34:57.000-08:00
Pull Request resolved: #8499 This diff enables customized debug data pipeline by making ETDumpGen leverage user-provided datasink. Details can be found in https://docs.google.com/document/d/1y_m32mKdj-OgLcLUz9TKhBW3PC3bBDYSBbeAH544EfM/edit?tab=t.0#heading=h.jlkcrurw482r ghstack-source-id: 266706035 @exported-using-ghexport Differential Revision: [D69647096](https://our.internmc.facebook.com/intern/diff/D69647096/)
diff --git a/devtools/etdump/buffer_data_sink.h b/devtools/etdump/buffer_data_sink.h
@@ -25,19 +25,29 @@ namespace etdump {
 class BufferDataSink : public DataSinkBase {
  public:
   /**
-   * Constructs a BufferDataSink with a given buffer.
+   * Constructs a BufferDataSink with a given span buffer.
    *
    * @param[in] buffer A Span object representing the buffer where data will be
    * stored.
    */
   explicit BufferDataSink(::executorch::runtime::Span<uint8_t> buffer)
       : debug_buffer_(buffer), offset_(0) {}
 
+/**
+   * Constructs a BufferDataSink with a given ptr to data blob, and the size of data blob.
+   *
+   * @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.
+   */
+  BufferDataSink(void* ptr, size_t size)
+      : debug_buffer_((uint8_t*)ptr, size), offset_(0) {}
+
   BufferDataSink(const BufferDataSink&) = delete;
   BufferDataSink& operator=(const BufferDataSink&) = delete;
   BufferDataSink(BufferDataSink&&) = default;
   BufferDataSink& operator=(BufferDataSink&&) = default;
 
+
   /**
    * Write data into the debug buffer and return the offset of the starting
    * location of the data within the buffer.
diff --git a/devtools/etdump/etdump_flatcc.cpp b/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>
@@ -28,6 +29,7 @@ using ::executorch::runtime::DebugHandle;
 using ::executorch::runtime::DelegateDebugIdType;
 using ::executorch::runtime::EValue;
 using ::executorch::runtime::EventTracerEntry;
+using ::executorch::runtime::Result;
 using ::executorch::runtime::LoggedEValueType;
 using ::executorch::runtime::Span;
 using ::executorch::runtime::Tag;
@@ -347,10 +349,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 +369,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 +379,8 @@ 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,15 @@ ETDumpResult ETDumpGen::get_etdump_data() {
 }
 
 void ETDumpGen::set_debug_buffer(Span<uint8_t> buffer) {
-  debug_buffer_ = buffer;
+  data_sink_ = std::make_shared<BufferDataSink>(buffer);
 }
 
-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::alignPointer(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(std::shared_ptr<DataSinkBase> buffer_data_sink) {
+  data_sink_ = buffer_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 +520,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 +542,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));
       }
@@ -636,7 +627,21 @@ bool ETDumpGen::is_static_etdump() {
 }
 
 size_t ETDumpGen::get_debug_buffer_size() const {
-  return debug_buffer_.size();
+  return ETDumpGen::get_data_sink_size();
+}
+
+size_t ETDumpGen::get_data_sink_size() const {
+  ET_CHECK_MSG(data_sink_, "Must set data sink before checking its size\n");
+  Result<size_t> ret = data_sink_->get_storage_size();
+  ET_CHECK_MSG(ret.ok(), "Failed to get storage size with error 0x%" PRIx32, static_cast<uint32_t>(ret.error()));
+  return ret.get();
+}
+
+long ETDumpGen::write_tensor_or_raise_error(Tensor tensor) {
+  ET_CHECK_MSG(data_sink_, "Must set data sink before writing data\n");
+  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
diff --git a/devtools/etdump/etdump_flatcc.h b/devtools/etdump/etdump_flatcc.h
@@ -9,7 +9,9 @@
 #pragma once
 
 #include <cstdint>
+#include <memory>
 
+#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,8 +143,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(std::shared_ptr<DataSinkBase> buffer_data_sink);
   ETDumpResult get_etdump_data();
   size_t get_debug_buffer_size() const;
+  size_t get_data_sink_size() const;
   size_t get_num_blocks();
   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,11 @@ 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;
+  std::shared_ptr<DataSinkBase> data_sink_;
   int bundled_input_index_ = -1;
   State state_ = State::Init;
   struct internal::ETDumpStaticAllocator alloc_;
diff --git a/devtools/etdump/targets.bzl b/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,
             ],
diff --git a/devtools/etdump/tests/etdump_test.cpp b/devtools/etdump/tests/etdump_test.cpp
