[devtool] introduce datasink class to etdump (#8677)

Pull Request resolved: #8496

this diff introduce datasink class, the class for managing the customized debug data storage pipeline.

Detials can be found in https://docs.google.com/document/d/1y_m32mKdj-OgLcLUz9TKhBW3PC3bBDYSBbeAH544EfM/edit?tab=t.0
ghstack-source-id: 268137338

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

Co-authored-by: gasoonjia <[email protected]>
Co-authored-by: Scott Wolchok <[email protected]>

Author: 3 people

devtools/etdump/buffer_data_sink.cpp

@@ -0,0 +1,51 @@
+/*
+ * 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 <executorch/devtools/etdump/buffer_data_sink.h>
+#include <executorch/devtools/etdump/utils.h>
+
+using ::executorch::runtime::Error;
+using ::executorch::runtime::Result;
+
+namespace executorch {
+namespace etdump {
+
+Result<size_t> BufferDataSink::write(const void* ptr, size_t length) {
+  if (length == 0) {
+    return offset_;
+  }
+
+  uint8_t* last_data_end = debug_buffer_.data() + offset_;
+
+  // The beginning of the next data blob must be aligned to the alignment
+  uint8_t* cur_data_begin = internal::align_pointer(last_data_end, alignment_);
+  uint8_t* cur_data_end = cur_data_begin + length;
+
+  if (cur_data_end > debug_buffer_.data() + debug_buffer_.size()) {
+    ET_LOG(Error, "Ran out of space to store intermediate outputs.");
+    return Error::OutOfResources;
+  }
+
+  // Zero out the padding between data blobs
+  memset(last_data_end, 0, cur_data_begin - last_data_end);
+  memcpy(cur_data_begin, ptr, length);
+  offset_ = (size_t)(cur_data_end - debug_buffer_.data());
+
+  return (size_t)(cur_data_begin - debug_buffer_.data());
+}
+
+Result<size_t> BufferDataSink::get_storage_size() const {
+  return debug_buffer_.size();
+}
+
+size_t BufferDataSink::get_used_bytes() const {
+  return offset_;
+}
+
+} // namespace etdump
+} // namespace executorch

devtools/etdump/buffer_data_sink.h

@@ -0,0 +1,91 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <executorch/devtools/etdump/data_sink_base.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/span.h>
+
+namespace executorch {
+namespace etdump {
+
+/**
+ * BufferDataSink is a concrete implementation of the DataSinkBase class,
+ * designed to store debug data in a pre-allocated, user-owned buffer. This
+ * class provides methods to write raw data and tensor data into the buffer,
+ * ensuring proper alignment and managing padding as needed.
+ */
+class BufferDataSink : public DataSinkBase {
+ public:
+  /**
+   * Constructs a BufferDataSink with a given buffer.
+   *
+   * @param[in] buffer A Span object representing the buffer where data will be
+   * stored.
+   * @param[in] alignment The alignment requirement for the buffer. It must be
+   * a power of two. Default is 64.
+   */
+  explicit BufferDataSink(
+      ::executorch::runtime::Span<uint8_t> buffer,
+      size_t alignment = 64)
+      : debug_buffer_(buffer), offset_(0), alignment_(alignment) {}
+
+  // Uncopiable and unassignable to avoid double assignment and free of the
+  // internal buffer.
+  BufferDataSink(const BufferDataSink&) = delete;
+  BufferDataSink& operator=(const BufferDataSink&) = delete;
+
+  // Movable to be compatible with Result.
+  BufferDataSink(BufferDataSink&&) = default;
+  BufferDataSink& operator=(BufferDataSink&&) = default;
+
+  ~BufferDataSink() override = default;
+
+  /**
+   * Write data into the debug buffer and return the offset of the starting
+   * location of the data within the buffer.
+   *
+   * @param[in] ptr A pointer to the data to be written into the storage.
+   * @param[in] size The size of the data in bytes.
+   * @return A Result object containing either:
+   *         - The offset of the starting location of the data within the
+   *           debug buffer, or
+   *         - An error code indicating the failure reason, if any issue
+   *           occurs during the write process.
+   */
+  ::executorch::runtime::Result<size_t> write(const void* ptr, size_t size)
+      override;
+
+  /**
+   * Retrieves the total size of the buffer.
+   *
+   * @return A Result object containing the total size of the buffer in bytes.
+   */
+  ::executorch::runtime::Result<size_t> get_storage_size() const;
+
+  /**
+   * Retrieves the number of bytes currently used in the buffer.
+   *
+   * @return The amount of data currently stored in the buffer in bytes.
+   */
+  size_t get_used_bytes() const override;
+
+ private:
+  // A Span object representing the buffer used for storing debug data.
+  ::executorch::runtime::Span<uint8_t> debug_buffer_;
+
+  // The offset of the next available location in the buffer.
+  size_t offset_;
+
+  // The alignment of the buffer.
+  size_t alignment_;
+};
+
+} // namespace etdump
+} // namespace executorch

devtools/etdump/data_sink_base.h

@@ -0,0 +1,61 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/result.h>
+
+namespace executorch {
+namespace etdump {
+
+/**
+ * DataSinkBase is an abstract class that users can inherit and implement
+ * to customize the storage and management of debug data in ETDumpGen. This
+ * class provides a basic and essential interface for writing datablob to a
+ * user-defined storage, retrieving storage capacity, and tracking the amount of
+ * data stored.
+ */
+class DataSinkBase {
+ public:
+  /**
+   * Virtual destructor to ensure proper cleanup of derived classes.
+   */
+
+  virtual ~DataSinkBase() = default;
+  /**
+   * Write data into the debug storage. This method should be implemented
+   * by derived classes to handle the specifics of data storage.
+   *
+   * This function should return the offset of the starting location of the
+   * data within the debug storage if the write operation succeeds, or an
+   * Error code if any issue occurs during the write process.
+   *
+   * @param[in] ptr A pointer to the data to be written into the storage.
+   * @param[in] length The size of the data in bytes.
+   * @return A Result object containing either:
+   *         - The offset of the starting location of the data within the
+   *           debug storage, which will be recorded in the corresponding
+   *           metadata of ETDump, or
+   *         - An error code indicating the failure reason, if any issue
+   *           occurs during the write process.
+   */
+  virtual ::executorch::runtime::Result<size_t> write(
+      const void* ptr,
+      size_t length) = 0;
+
+  /**
+   * Get the number of bytes currently used in the debug storage.
+   *
+   * @return The amount of data currently stored in bytes.
+   */
+  virtual size_t get_used_bytes() const = 0;
+};
+
+} // namespace etdump
+} // namespace executorch

devtools/etdump/etdump_flatcc.cpp

@@ -13,6 +13,7 @@
 #include <executorch/devtools/etdump/emitter.h>
 #include <executorch/devtools/etdump/etdump_schema_flatcc_builder.h>
 #include <executorch/devtools/etdump/etdump_schema_flatcc_reader.h>
+#include <executorch/devtools/etdump/utils.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 #include <executorch/runtime/platform/assert.h>
@@ -94,16 +95,6 @@ etdump_Tensor_ref_t add_tensor_entry(
   return etdump_Tensor_end(builder_);
 }
 
-static uint8_t* alignPointer(void* ptr, size_t alignment) {
-  intptr_t addr = reinterpret_cast<intptr_t>(ptr);
-  if ((addr & (alignment - 1)) == 0) {
-    // Already aligned.
-    return reinterpret_cast<uint8_t*>(ptr);
-  }
-  addr = (addr | (alignment - 1)) + 1;
-  return reinterpret_cast<uint8_t*>(addr);
-}
-
 } // namespace
 
 // Constructor implementation
@@ -113,9 +104,10 @@ ETDumpGen::ETDumpGen(Span<uint8_t> buffer) {
   // Initialize the flatcc builder_ using the buffer and buffer size.
 
   if (buffer.data() != nullptr) {
-    builder_ = (struct flatcc_builder*)alignPointer(buffer.data(), 64);
-    uintptr_t buffer_with_builder =
-        (uintptr_t)alignPointer(builder_ + sizeof(struct flatcc_builder), 64);
+    builder_ =
+        (struct flatcc_builder*)internal::align_pointer(buffer.data(), 64);
+    uintptr_t buffer_with_builder = (uintptr_t)internal::align_pointer(
+        builder_ + sizeof(struct flatcc_builder), 64);
     size_t builder_size =
         (size_t)(buffer_with_builder - (uintptr_t)buffer.data());
     size_t min_buf_size = max_alloc_buf_size + builder_size;
@@ -513,7 +505,7 @@ size_t ETDumpGen::copy_tensor_to_debug_buffer(executorch::aten::Tensor tensor) {
     return static_cast<size_t>(-1);
   }
   uint8_t* offset_ptr =
-      alignPointer(debug_buffer_.data() + debug_buffer_offset_, 64);
+      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(),

devtools/etdump/targets.bzl

@@ -87,8 +87,54 @@ def define_common_targets():
         exported_external_deps = ["flatccrt"],
     )
 
+    runtime.cxx_library(
+        name = "utils",
+        srcs = [],
+        exported_headers = [
+            "utils.h",
+        ],
+        visibility = [
+
+        ],
+    )
+
     for aten_mode in get_aten_mode_options():
         aten_suffix = "_aten" if aten_mode else ""
+
+        runtime.cxx_library(
+            name = "data_sink_base" + aten_suffix,
+            exported_headers = [
+                "data_sink_base.h",
+            ],
+            exported_deps = [
+                "//executorch/runtime/core/exec_aten/util:scalar_type_util" + aten_suffix,
+            ],
+            visibility = [
+                "//executorch/...",
+                "@EXECUTORCH_CLIENTS",
+            ],
+        )
+
+        runtime.cxx_library(
+            name = "buffer_data_sink" + aten_suffix,
+            headers = [
+                "buffer_data_sink.h",
+            ],
+            srcs = [
+                "buffer_data_sink.cpp",
+            ],
+            deps = [
+                ":utils",
+            ],
+            exported_deps = [
+                "//executorch/runtime/core/exec_aten:lib" + aten_suffix,
+                ":data_sink_base" + aten_suffix,
+            ],
+            visibility = [
+                "//executorch/...",
+                "@EXECUTORCH_CLIENTS",
+            ],
+        )
         runtime.cxx_library(
             name = "etdump_flatcc" + aten_suffix,
             srcs = [
@@ -106,6 +152,7 @@ def define_common_targets():
             ],
             exported_deps = [
                 ":etdump_schema_flatcc",
+                ":utils",
                 "//executorch/runtime/core:event_tracer" + aten_suffix,
                 "//executorch/runtime/core/exec_aten/util:scalar_type_util" + aten_suffix,
             ],