Update on "[devtool] introduce datasink class to etdump"

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

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

[ghstack-poisoned]

Author: Gasoonjia

devtools/etdump/buffer_data_sink.cpp

@@ -0,0 +1,42 @@
+// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+#include <executorch/devtools/etdump/buffer_data_sink.h>
+#include <executorch/devtools/etdump/utils.h>
+
+namespace executorch {
+namespace etdump {
+
+Result<size_t> BufferDataSink::write(const void* ptr, size_t length) {
+  if (length == 0) {
+    return static_cast<size_t>(-1);
+  }
+  uint8_t* offset_ptr =
+      internal::alignPointer(debug_buffer_.data() + offset_, 64);
+
+  // Zero out the padding between data blobs.
+  size_t n_zero_pad = offset_ptr - debug_buffer_.data() - offset_;
+  memset(debug_buffer_.data() + offset_, 0, n_zero_pad);
+
+  offset_ = (offset_ptr - debug_buffer_.data()) + length;
+  ET_CHECK_MSG(
+      offset_ <= debug_buffer_.size(),
+      "Ran out of space to store tensor data.");
+  memcpy(offset_ptr, ptr, length);
+  return (size_t)(offset_ptr - debug_buffer_.data());
+}
+
+Result<size_t> BufferDataSink::write_tensor(
+    const executorch::aten::Tensor& tensor) {
+  return write(tensor.const_data_ptr(), tensor.nbytes());
+}
+
+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/data_sink.h renamed to devtools/etdump/buffer_data_sink.h

@@ -9,13 +9,14 @@
 namespace executorch {
 namespace etdump {
 
-class DataSink : public DataSinkBase {
+class BufferDataSink : public DataSinkBase {
  public:
-  explicit DataSink(::executorch::runtime::Span<uint8_t> buffer)
+  explicit BufferDataSink(::executorch::runtime::Span<uint8_t> buffer)
       : debug_buffer_(buffer), offset_(0) {}
 
-  size_t write_tensor(const executorch::aten::Tensor& tensor) override;
-  size_t get_storage_size() const override;
+  Result<size_t> write(const void* ptr, size_t length) override;
+  Result<size_t> write_tensor(const executorch::aten::Tensor& tensor);
+  Result<size_t> get_storage_size() const override;
   size_t get_used_bytes() const override;
 
  private:

devtools/etdump/data_sink.cpp

@@ -3,15 +3,17 @@
 #pragma once
 
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
-#include <cstdlib>
+#include <executorch/runtime/core/result.h>
+
+using ::executorch::runtime::Result;
 
 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 an basic and essential interface for writing tensor data to a
+ * class provides an basic and essential interface for writing datablob to a
  * user-defined storage, retrieving storage capacity, and tracking the amount of
  * data stored.
  */
@@ -22,21 +24,24 @@ class DataSinkBase {
    */
   virtual ~DataSinkBase() = default;
   /**
-   * Write tensor data into the debug storage. This method should be implemented
+   * Write data into the debug storage. This method should be implemented
    * by derived classes to handle the specifics of data storage.
    *
-   * @param[in] tensor The tensor data to be written into the storage.
-   * @return The offset of the starting location of the tensor data within the
-   *         debug storage, which will be recorded in corresponding tensor
-   * metadata of ETDump.
+   * @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 The offset of the starting location of the data within the
+   *         debug storage, which will be recorded in the corresponding
+   *         metadata of ETDump.
    */
-  virtual size_t write_tensor(const executorch::aten::Tensor& tensor) = 0;
+  virtual Result<size_t> write(const void* ptr, size_t length) = 0;
   /**
    * Get the maximum capacity of the debug storage in bytes.
+   * Should return Error::NotSupported if the capacity is not available
+   * (e.g. unbounded storage like internet or file)
    *
    * @return The total size of the debug storage.
    */
-  virtual size_t get_storage_size() const = 0;
+  virtual Result<size_t> get_storage_size() const = 0;
   /**
    * Get the number of bytes currently used in the debug storage.
    *

devtools/etdump/data_sink_base.h

@@ -116,12 +116,12 @@ def define_common_targets():
         )
 
         runtime.cxx_library(
-            name = "data_sink" + aten_suffix,
+            name = "buffer_data_sink" + aten_suffix,
             headers = [
-                "data_sink.h",
+                "buffer_data_sink.h",
             ],
             srcs = [
-                "data_sink.cpp",
+                "buffer_data_sink.cpp",
             ],
             deps = [
                 ":utils",
@@ -149,7 +149,6 @@ def define_common_targets():
             ],
             deps = [
                 "//executorch/runtime/platform:platform",
-                ":utils",
             ],
             exported_deps = [
                 ":etdump_schema_flatcc",

devtools/etdump/targets.bzl

@@ -6,7 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/devtools/etdump/data_sink.h>
+#include <executorch/devtools/etdump/buffer_data_sink.h>
 #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
 #include <executorch/runtime/core/span.h>
 #include <executorch/runtime/platform/runtime.h>
@@ -19,15 +19,15 @@ using executorch::aten::Tensor;
 using ::executorch::runtime::Span;
 using torch::executor::testing::TensorFactory;
 
-class DataSinkTest : public ::testing::Test {
+class BufferDataSinkTest : public ::testing::Test {
  protected:
   void SetUp() override {
     torch::executor::runtime_init();
     // Allocate a small buffer for testing
     buffer_size_ = 128; // Small size for testing
     buffer_ptr_ = malloc(buffer_size_);
     buffer_ = Span<uint8_t>(static_cast<uint8_t*>(buffer_ptr_), buffer_size_);
-    data_sink_ = std::make_unique<executorch::etdump::DataSink>(buffer_);
+    data_sink_ = std::make_unique<executorch::etdump::BufferDataSink>(buffer_);
   }
 
   void TearDown() override {
@@ -37,18 +37,35 @@ class DataSinkTest : public ::testing::Test {
   size_t buffer_size_;
   void* buffer_ptr_;
   Span<uint8_t> buffer_;
-  std::unique_ptr<executorch::etdump::DataSink> data_sink_;
+  std::unique_ptr<executorch::etdump::BufferDataSink> data_sink_;
 };
 
-TEST_F(DataSinkTest, StorageSizeCheck) {
-  EXPECT_EQ(data_sink_->get_storage_size(), buffer_size_);
+TEST_F(BufferDataSinkTest, StorageSizeCheck) {
+  Result<size_t> ret = data_sink_->get_storage_size();
+  if (!ret.ok()) {
+    ET_LOG(
+        Info,
+        "Get storage size failed with 0x%" PRIx32,
+        (unsigned int)ret.error());
+  }
+  size_t storage_size = ret.get();
+  EXPECT_EQ(storage_size, buffer_size_);
 }
 
-TEST_F(DataSinkTest, WriteOneTensorAndCheckData) {
+TEST_F(BufferDataSinkTest, WriteOneTensorAndCheckData) {
   TensorFactory<ScalarType::Float> tf;
   Tensor tensor = tf.make({1, 4}, {1.0, 2.0, 3.0, 4.0});
 
-  size_t offset = data_sink_->write_tensor(tensor);
+  Result<size_t> ret = data_sink_->write_tensor(tensor);
+  if (!ret.ok()) {
+    ET_LOG(
+        Info,
+        "Write tensor into debug storage failed with 0x%" PRIx32,
+        (unsigned int)ret.error());
+  }
+
+  size_t offset = ret.get();
+
   EXPECT_NE(offset, static_cast<size_t>(-1));
 
   // Check that the data in the buffer matches the tensor data
@@ -59,14 +76,21 @@ TEST_F(DataSinkTest, WriteOneTensorAndCheckData) {
   }
 }
 
-TEST_F(DataSinkTest, WriteMultiTensorsAndCheckData) {
+TEST_F(BufferDataSinkTest, WriteMultiTensorsAndCheckData) {
   TensorFactory<ScalarType::Float> tf;
   std::vector<Tensor> tensors = {
       tf.make({1, 4}, {1.0, 2.0, 3.0, 4.0}),
       tf.make({1, 4}, {5.0, 6.0, 7.0, 8.0})};
-  size_t offset = 0;
   for (const auto& tensor : tensors) {
-    offset = data_sink_->write_tensor(tensor);
+    Result<size_t> ret = data_sink_->write_tensor(tensor);
+    if (!ret.ok()) {
+      ET_LOG(
+          Info,
+          "Write tensor into debug storage failed with 0x%" PRIx32,
+          (unsigned int)ret.error());
+    }
+
+    size_t offset = ret.get();
     EXPECT_NE(offset, static_cast<size_t>(-1));
     // Check that the data in the buffer matches the tensor data
     const float* buffer_data =
@@ -77,17 +101,25 @@ TEST_F(DataSinkTest, WriteMultiTensorsAndCheckData) {
   }
 }
 
-TEST_F(DataSinkTest, PointerAlignmentCheck) {
+TEST_F(BufferDataSinkTest, PointerAlignmentCheck) {
   TensorFactory<ScalarType::Float> tf;
   Tensor tensor = tf.make({1, 4}, {1.0, 2.0, 3.0, 4.0});
-  size_t offset = data_sink_->write_tensor(tensor);
+  Result<size_t> ret = data_sink_->write_tensor(tensor);
+  if (!ret.ok()) {
+    ET_LOG(
+        Info,
+        "Write tensor into debug storage failed with 0x%" PRIx32,
+        (unsigned int)ret.error());
+  }
+
+  size_t offset = ret.get();
   EXPECT_NE(offset, static_cast<size_t>(-1));
   // Check that the offset pointer is 64-byte aligned
   const uint8_t* offset_ptr = buffer_.data() + offset;
   EXPECT_EQ(reinterpret_cast<uintptr_t>(offset_ptr) % 64, 0);
 }
 
-TEST_F(DataSinkTest, WriteUntilOverflow) {
+TEST_F(BufferDataSinkTest, WriteUntilOverflow) {
   TensorFactory<ScalarType::Float> tf;
   Tensor tensor = tf.zeros({1, 8}); // Large tensor to fill the buffer
 

devtools/etdump/tests/data_sink_test.cpp renamed to devtools/etdump/tests/buffer_data_sink_test.cpp

@@ -21,12 +21,12 @@ def define_common_targets():
     )
 
     runtime.cxx_test(
-        name = "data_sink_test",
+        name = "buffer_data_sink_test",
         srcs = [
-            "data_sink_test.cpp",
+            "buffer_data_sink_test.cpp",
         ],
         deps = [
-            "//executorch/devtools/etdump:data_sink",
+            "//executorch/devtools/etdump:buffer_data_sink",
             "//executorch/runtime/core/exec_aten/testing_util:tensor_util",
         ],
     )