[devtool] make ETDumpGen use datasink

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

Differential Revision: [D69647096](https://our.internmc.facebook.com/intern/diff/D69647096/)

ghstack-source-id: 266455642
Pull Request resolved: #8499

Author: Gasoonjia

devtools/etdump/data_sink.h

@@ -11,9 +11,16 @@ namespace etdump {
 
 class DataSink : public DataSinkBase {
  public:
-  explicit DataSink(::executorch::runtime::Span<uint8_t> buffer)
+  /// Construct an empty DataSink.
+  /* implicit */ constexpr DataSink() noexcept
+      : debug_buffer_({}), offset_(0) {}
+
+  DataSink(::executorch::runtime::Span<uint8_t> buffer)
       : debug_buffer_(buffer), offset_(0) {}
 
+  DataSink(void* ptr, size_t size)
+      : debug_buffer_((uint8_t*)ptr, size), offset_(0) {}
+
   size_t write_tensor(const executorch::aten::Tensor& tensor) override;
   size_t get_storage_size() const override;
   size_t get_used_bytes() const override;

devtools/etdump/etdump_flatcc.cpp

@@ -10,6 +10,7 @@
 
 #include <cstring>
 
+#include <executorch/devtools/etdump/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>
@@ -347,10 +348,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 +368,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 = data_sink_->write_tensor(output);
     etdump_Tensor_ref_t tensor_ref = add_tensor_entry(builder_, output, offset);
 
     etdump_Value_start(builder_);
@@ -377,7 +378,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 = data_sink_->write_tensor(output[i]);
       etdump_Tensor_vec_push(
           builder_, add_tensor_entry(builder_, output[i], offset));
     }
@@ -497,27 +498,15 @@ ETDumpResult ETDumpGen::get_etdump_data() {
 }
 
 void ETDumpGen::set_debug_buffer(Span<uint8_t> buffer) {
-  debug_buffer_ = buffer;
+  data_sink_ = std::make_shared<DataSink>(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> 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 +518,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 = data_sink_->write_tensor(tensor);
       etdump_Tensor_ref_t tensor_ref =
           add_tensor_entry(builder_, tensor, offset);
 
@@ -551,7 +540,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 = data_sink_->write_tensor(tensors[i]);
         etdump_Tensor_vec_push(
             builder_, add_tensor_entry(builder_, tensors[i], offset));
       }
@@ -636,7 +625,13 @@ bool ETDumpGen::is_static_etdump() {
 }
 
 size_t ETDumpGen::get_debug_buffer_size() const {
-  return debug_buffer_.size();
+  ET_CHECK_MSG(data_sink_, "Must set data sink before checking its size\n");
+  return data_sink_->get_storage_size();
+}
+
+size_t ETDumpGen::get_data_sink_size() const {
+  ET_CHECK_MSG(data_sink_, "Must set data sink before checking its size\n");
+  return data_sink_->get_storage_size();
 }
 
 } // namespace etdump

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> 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.
@@ -172,8 +175,7 @@ class ETDumpGen : public ::executorch::runtime::EventTracer {
 
   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_;

devtools/etdump/targets.bzl

@@ -154,6 +154,8 @@ def define_common_targets():
             exported_deps = [
                 ":etdump_schema_flatcc",
                 ":utils",
+                ":data_sink_base" + aten_suffix,
+                ":data_sink" + aten_suffix,
                 "//executorch/runtime/core:event_tracer" + aten_suffix,
                 "//executorch/runtime/core/exec_aten/util:scalar_type_util" + aten_suffix,
             ],