Introduce new DataLoader::load() with segment info

Summary: Adds a new DataLoader::load() function with an additonal parameter, SegmentInfo.

Differential Revision: D59399538

Author: jackzhxng

runtime/core/data_loader.h

@@ -9,6 +9,7 @@
 #pragma once
 
 #include <cstddef>
+#include <optional>
 
 #include <executorch/runtime/core/freeable_buffer.h>
 #include <executorch/runtime/core/result.h>
@@ -17,6 +18,27 @@
 namespace torch {
 namespace executor {
 
+enum class SegmentType { Program, Constant, Backend };
+
+struct SegmentInfo {
+  size_t segment_index = 0;
+  SegmentType segment_type = SegmentType::Program;
+  const char* descriptor =
+      nullptr; // This can be any descriptor to pass to the data loader. At the
+               // moment for backends this would be a backend id and for null
+               // otherwise.
+
+  SegmentInfo() = default;
+
+  SegmentInfo(
+      size_t segment_index,
+      SegmentType segment_type,
+      const char* descriptor)
+      : segment_index(segment_index),
+        segment_type(segment_type),
+        descriptor(descriptor) {}
+};
+
 /**
  * Loads from a data source.
  *
@@ -27,16 +49,33 @@ class DataLoader {
   virtual ~DataLoader() = default;
 
   /**
+   * DEPRECATED: Use `load()` going forward for access to segment info during
+   * the load.
+   *
    * Loads `size` bytes at byte offset `offset` from the underlying data source
    * into a `FreeableBuffer`, which owns the memory.
    *
    * NOTE: This must be thread-safe. If this call modifies common state, the
    * implementation must do its own locking.
    */
-  __ET_NODISCARD virtual Result<FreeableBuffer> Load(
+  __ET_DEPRECATED __ET_NODISCARD virtual Result<FreeableBuffer> Load(
       size_t offset,
       size_t size) = 0;
 
+  /**
+   * Loads `size` bytes at byte offset `offset` from the underlying data source
+   * into a `FreeableBuffer`, which owns the memory.
+   *
+   * NOTE: This must be thread-safe. If this call modifies common state, the
+   * implementation must do its own locking.
+   */
+  __ET_NODISCARD virtual inline Result<FreeableBuffer> load(
+      size_t _offset,
+      size_t _size,
+      const std::optional<SegmentInfo>& _segment_info) {
+    return Error::NotImplemented;
+  }
+
   /**
    * Returns the length of the underlying data source, typically the file size.
    */

runtime/executor/program.cpp

@@ -10,6 +10,8 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <iostream>
+#include <optional>
 
 #include <executorch/runtime/core/event_tracer_hooks.h>
 #include <executorch/runtime/executor/memory_manager.h>
@@ -423,8 +425,19 @@ Result<FreeableBuffer> Program::LoadSegment(size_t index) const {
   // Could fail if offset and size are out of bound for the data, or if this
   // is reading from a file and fails, or for many other reasons depending on
   // the implementation of the loader.
-  return loader_->Load(
-      segment_base_offset_ + segment->offset(), segment->size());
+  // TODO(jackzhxng): deprecate Load() and use load() instead.
+  Result<FreeableBuffer> segment_data = loader_->load(
+      segment_base_offset_ + segment->offset(),
+      segment->size(),
+      std::nullopt /* segment_info */);
+  if (segment_data.error() == Error::NotImplemented) {
+    std::cerr << "Using the old Load()" << std::endl; // TODO: remove debug
+    return loader_->Load(
+        segment_base_offset_ + segment->offset(), segment->size());
+  } else {
+    std::cerr << "Using the new load()" << std::endl; // TODO: remove debug
+    return segment_data;
+  }
 }
 
 } // namespace executor

runtime/executor/test/backend_integration_test.cpp

@@ -188,6 +188,24 @@ class DataLoaderSpy : public DataLoader {
         context->buffer.data(), context->buffer.size(), FreeBuffer, context);
   }
 
+  // TODO: figure out how to test that the segment_info is being used and
+  // passed in properly.
+  Result<FreeableBuffer> load(
+      size_t offset,
+      size_t size,
+      const std::optional<torch::executor::SegmentInfo>& segment_info)
+      override {
+    Result<FreeableBuffer> buf = delegate_->Load(offset, size);
+    if (!buf.ok()) {
+      return buf.error();
+    }
+    operations_.push_back({Operation::Load, offset, /*data=*/nullptr, size});
+    auto* context = new SpyContext(&operations_, std::move(buf.get()));
+    // Use context->buffer since buf has been moved.
+    return FreeableBuffer(
+        context->buffer.data(), context->buffer.size(), FreeBuffer, context);
+  }
+
   Result<size_t> size() const override {
     return delegate_->size();
   }