Pass in correct backend id into data load from runtime (pytorch#4218)

Summary:
Pull Request resolved: pytorch#4218

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](https://www.internalfb.com/diff/D59399538?dst_version_fbid=2979141262226992)
- (  ) [[2/n][executorch] Switch to DataLoader::load in runtime](https://www.internalfb.com/diff/D59594142)
- (**x**) [[3/n][executorch] Pass in correct backend id into data load from runtime](https://www.internalfb.com/diff/D59606243)

Reviewed By: dbort

Differential Revision: D59606243

fbshipit-source-id: b4af5e1bb2e53f576ab75a65d2ef950cc051e8c9

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

@@ -414,8 +414,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;
@@ -431,12 +433,8 @@ 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.
-  // TODO(jackzhxng): "backend_segment" is a hardcode, pass in real backend id.
   return loader_->load(
-      segment_base_offset_ + segment->offset(),
-      segment->size(),
-      DataLoader::SegmentInfo(
-          DataLoader::SegmentInfo::Type::Backend, index, "backend_segment"));
+      segment_base_offset_ + segment->offset(), segment->size(), segment_info);
 }
 
 } // namespace executor

runtime/executor/program.h

@@ -239,8 +239,9 @@ 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[in] SegmentInfo Struct containing an index to load from the
+   * Program.segments list. The other fields of the struct, such as
+   * `segment_type` and `descriptor`, need to also be correct.
    *
    * @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 +250,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

@@ -545,8 +545,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);
   EXPECT_EQ(backend_load_was_called, 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);
 }
 
@@ -350,9 +355,13 @@ TEST_F(ProgramTest, LoadConstantSegment) {
   Result<Program> program = Program::load(&linear_loader.get());
   ASSERT_EQ(program.error(), Error::Ok);
 
-  // Load constant segment data.
+  // Load constant segment data, which is currently always in segment index
+  // zero.
+  const auto segment_info = DataLoader::SegmentInfo(
+      DataLoader::SegmentInfo::Type::Constant,
+      /*segment_index=*/0);
   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 =