Introduce new DataLoader::load() with segment info (#4158)

Summary:
Pull Request resolved: #4158

Adds a new `DataLoader::load()` function with an additonal parameter, SegmentInfo. Makes the old `DataLoader::Load()` function deprecated.

Differential Revision: D59399538

Author: jackzhxng

runtime/core/data_loader.h

@@ -24,19 +24,85 @@ namespace executor {
  */
 class DataLoader {
  public:
+  /**
+   * Describes the content of the segment.
+   */
+  struct SegmentInfo {
+    /**
+     * Represents the purpose of the segment.
+     */
+    enum class Type {
+      /**
+       * Data for the actual program.
+       */
+      Program,
+      /**
+       * Holds constant tensor data.
+       */
+      Constant,
+      /**
+       * Data used for initializing a backend.
+       */
+      Backend,
+    };
+
+    /// Type of the segment.
+    Type segment_type;
+
+    /// Index of the segment within the segment list. Undefined for program
+    /// segments.
+    size_t segment_index;
+
+    /// An optional, null-terminated string describing the segment. For
+    /// `Backend` segments, this is the backend ID. Null for other segment
+    /// types.
+    const char* descriptor;
+
+    SegmentInfo() = default;
+
+    explicit SegmentInfo(
+        Type segment_type,
+        size_t segment_index = 0,
+        const char* descriptor = nullptr)
+        : segment_type(segment_type),
+          segment_index(segment_index),
+          descriptor(descriptor) {}
+  };
+
   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 data from the underlying data source.
+   *
+   * NOTE: This must be thread-safe. If this call modifies common state, the
+   * implementation must do its own locking.
+   *
+   * @param offset The byte offset in the data source to start loading from.
+   * @param size The number of bytes to load.
+   * @param segment_info Information about the segment being loaded.
+   *
+   * @returns a `FreeableBuffer` that owns the loaded data.
+   */
+  __ET_NODISCARD virtual Result<FreeableBuffer>
+  load(size_t offset, size_t size, const SegmentInfo& segment_info) {
+    (void)segment_info;
+    return Load(offset, size); // NOLINT
+  }
+
   /**
    * Returns the length of the underlying data source, typically the file size.
    */