Pass in correct backend id into data load from runtime

Summary:
The runtime now passes in the backend id into the SegmentInfo for its load() call.

Diff stack:
- (  ) [1/n][executorch] Introduce new DataLoader::load() with segment info
- (  ) [2/n][executorch] Switch to DataLoader::load in runtime
- (**x**) [3/n][executorch] Pass in correct backend id into data load from runtime

Differential Revision: D59606243

Author: jackzhxng

runtime/executor/method.cpp

@@ -182,7 +182,11 @@ class BackendDelegate final {
             /*free_fn=*/nullptr);
       }
       case executorch_flatbuffer::DataLocation::SEGMENT: {
-        return program->LoadSegment(processed->index());
+        const char* backend_id = delegate.id()->c_str();
+        return program->LoadSegment(DataLoader::SegmentInfo(
+            DataLoader::SegmentInfo::Type::Backend,
+            processed->index(),
+            backend_id));
       }
       default:
         ET_LOG(

runtime/executor/program.cpp

@@ -415,8 +415,10 @@ Error Program::get_backend_delegate_data(
   return HeaderStatus::NotPresent;
 }
 
-Result<FreeableBuffer> Program::LoadSegment(size_t index) const {
+Result<FreeableBuffer> Program::LoadSegment(
+    const DataLoader::SegmentInfo& segment_info) const {
   EXECUTORCH_SCOPE_PROF("Program::LoadSegment");
+  size_t index = segment_info.segment_index;
   if (loader_ == nullptr || segment_base_offset_ == 0) {
     ET_LOG(Error, "No segments in program: requested index %zu", index);
     return Error::NotFound;
@@ -432,8 +434,6 @@ 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.
-  auto segment_info = DataLoader::SegmentInfo(
-      DataLoader::SegmentInfo::Type::Backend, index, "backend_segment");
   return loader_->load(
       segment_base_offset_ + segment->offset(), segment->size(), segment_info);
 }

runtime/executor/program.h

@@ -239,8 +239,8 @@ class Program final {
   /**
    * Loads a segment by index.
    *
-   * @param[in] index The sement index to load. This should be an index into
-   *     the Program.segments list.
+   * @param SegmentInfo Struct containing an index to load from the
+   * Program.segments list.
    *
    * @returns The data as a FreeableBuffer, if the index is valid.
    * @retval Error::NotFound The program does not contain any segments or the
@@ -249,7 +249,8 @@ class Program final {
    *     DataLoader: The Program.segment table is inconsistent, or the
    *     data cannot be accessed.
    */
-  __ET_NODISCARD Result<FreeableBuffer> LoadSegment(size_t index) const;
+  __ET_NODISCARD Result<FreeableBuffer> LoadSegment(
+      const DataLoader::SegmentInfo& segment_info) const;
 
  private:
   Program(

runtime/executor/test/backend_integration_test.cpp

@@ -542,8 +542,8 @@ TEST_P(BackendIntegrationTest, SegmentInfoIsPassedIntoDataLoader) {
   // Expect that load was called correctly on a backend segment.
   bool backend_load_was_called = spy_loader.UsedLoad(
       DataLoader::SegmentInfo::Type::Backend,
-      "backend_segment"); // TODO(jackzhxng): replace with actual mock PTE
-                          // file's backend_id in next chained PR.
+      "StubBackend"); // This backend id is taken from the StubBackend defined
+                      // in export_delegated_program.py.
 
   EXPECT_TRUE(program_load_was_called);
   if (using_segments()) {

runtime/executor/test/program_test.cpp

@@ -24,6 +24,7 @@
 #include <gtest/gtest.h>
 
 using namespace ::testing;
+using torch::executor::DataLoader;
 using torch::executor::Error;
 using torch::executor::FreeableBuffer;
 using torch::executor::Program;
@@ -87,8 +88,8 @@ class ProgramTestFriend final {
  public:
   __ET_NODISCARD static Result<FreeableBuffer> LoadSegment(
       const Program* program,
-      size_t index) {
-    return program->LoadSegment(index);
+      const DataLoader::SegmentInfo& segment_info) {
+    return program->LoadSegment(segment_info);
   }
 
   const static executorch_flatbuffer::Program* GetInternalProgram(
@@ -227,14 +228,18 @@ TEST_F(ProgramTest, UnalignedProgramDataFails) {
 }
 
 TEST_F(ProgramTest, LoadSegmentWithNoSegments) {
-  // Load a program with no segments.
+  // Load a program with no appended segments.
   Result<Program> program =
       Program::load(add_loader_.get(), kDefaultVerification);
   EXPECT_EQ(program.error(), Error::Ok);
 
-  // Loading a segment should fail.
+  // Loading a non-program segment should fail.
+  const auto segment_info = DataLoader::SegmentInfo(
+      DataLoader::SegmentInfo::Type::Backend,
+      /*segment_index=*/0,
+      "some-backend");
   Result<FreeableBuffer> segment =
-      ProgramTestFriend::LoadSegment(&program.get(), 0);
+      ProgramTestFriend::LoadSegment(&program.get(), segment_info);
   EXPECT_NE(segment.error(), Error::Ok);
 }
 
@@ -351,8 +356,12 @@ TEST_F(ProgramTest, LoadConstantSegment) {
   ASSERT_EQ(program.error(), Error::Ok);
 
   // Load constant segment data.
+  const auto segment_info = DataLoader::SegmentInfo(
+      DataLoader::SegmentInfo::Type::Constant,
+      /*segment_index=*/0,
+      /*descriptor=*/nullptr);
   Result<FreeableBuffer> segment =
-      ProgramTestFriend::LoadSegment(&program.get(), 0);
+      ProgramTestFriend::LoadSegment(&program.get(), segment_info);
   EXPECT_EQ(segment.error(), Error::Ok);
 
   const executorch_flatbuffer::Program* flatbuffer_program =