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:
+  /**
+   * A struct to pass information about the segment to the data loader.
+   */
+  struct SegmentInfo {
+    /**
+     * Represents the function of the segment.
+     */
+    enum class Type {
+      /**
+       * Data for the actual program. These always comes before the rest of the
+       * segments.
+       */
+      Program,
+      /**
+       * Holds constant data, i.e. tensors.
+       */
+      Constant,
+      /**
+       * Data used for initializing a backend.
+       */
+      Backend
+    };
+
+    /// Type of the segment.
+    Type segment_type;
+
+    /// Index of the segment within the segment list. Irrelevant for program
+    /// segments.
+    size_t segment_index;
+
+    /// Some descriptor relevant to the segment type, i.e. this would be backend
+    /// id for backend segments and null otherwise.
+    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 `size` bytes at byte offset `offset` from the underlying data source
+   * into a `FreeableBuffer`, which owns the memory. This method also provides
+   * information about the segment being loaded in segment_info.
+   *
+   * NOTE: This must be thread-safe. If this call modifies common state, the
+   * implementation must do its own locking.
+   *
+   * @param offset The 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.
+   */
+  __ET_NODISCARD virtual inline Result<FreeableBuffer>
+  load(size_t offset, size_t size, const SegmentInfo& segment_info) {
+    (void)segment_info;
+    return Load(offset, size);
+  }
+
   /**
    * Returns the length of the underlying data source, typically the file size.
    */