Use NDM in backend init context (pytorch#9102)

Add NDM to backend init context, so backends have access to it at init time.

Add to backend init context instead of as an argument to backend.init, as it's easier to add/deprecate an arg there rather than in the init/execute interface.

TODO: add a merge/union method to named_data_map.h interface. If there is an external file (or multiple), and/or named_segments in the PTE, these should be merged into one.

Differential Revision: [D70279622](https://our.internmc.facebook.com/intern/diff/D70279622/)

[ghstack-poisoned]

Co-authored-by: lucylq <[email protected]>

Author: kirklandsign

runtime/backend/backend_init_context.h

@@ -8,6 +8,7 @@
 
 #pragma once
 #include <executorch/runtime/core/memory_allocator.h>
+#include <executorch/runtime/core/named_data_map.h>
 
 namespace executorch {
 namespace runtime {
@@ -21,8 +22,11 @@ class BackendInitContext final {
   explicit BackendInitContext(
       MemoryAllocator* runtime_allocator,
       EventTracer* event_tracer = nullptr,
-      const char* method_name = nullptr)
-      : runtime_allocator_(runtime_allocator), method_name_(method_name) {}
+      const char* method_name = nullptr,
+      const NamedDataMap* named_data_map = nullptr)
+      : runtime_allocator_(runtime_allocator),
+        method_name_(method_name),
+        named_data_map_(named_data_map) {}
 
   /** Get the runtime allocator passed from Method. It's the same runtime
    * executor used by the standard executor runtime and the life span is the
@@ -52,10 +56,18 @@ class BackendInitContext final {
     return method_name_;
   }
 
+  /** Get the named data map from ExecuTorch runtime.
+   * This provides a way for backends to retrieve data blobs by key.
+   */
+  const NamedDataMap* get_named_data_map() const {
+    return named_data_map_;
+  }
+
  private:
   MemoryAllocator* runtime_allocator_ = nullptr;
   EventTracer* event_tracer_ = nullptr;
   const char* method_name_ = nullptr;
+  const NamedDataMap* named_data_map_ = nullptr;
 };
 
 } // namespace runtime

runtime/backend/interface.h

@@ -17,6 +17,7 @@
 #include <executorch/runtime/core/evalue.h>
 #include <executorch/runtime/core/freeable_buffer.h>
 #include <executorch/runtime/core/memory_allocator.h>
+#include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/core/result.h>
 #include <executorch/runtime/platform/compiler.h>
 

runtime/backend/targets.bzl

@@ -29,5 +29,6 @@ def define_common_targets():
                 "//executorch/runtime/core:evalue" + aten_suffix,
                 "//executorch/runtime/core:event_tracer" + aten_suffix,
                 "//executorch/runtime/core:memory_allocator",
+                "//executorch/runtime/core:named_data_map",
             ],
         )

runtime/executor/method.cpp

@@ -798,6 +798,14 @@ Error Method::init(
       return Error::MemoryAllocationFailed;
     }
 
+    // Get NamedDataMap, if it exists.
+    const NamedDataMap* pte_data_map = nullptr;
+    Result<const NamedDataMap*> pte_data_map_res =
+        program_->get_named_data_map();
+    if (pte_data_map_res.ok()) {
+      pte_data_map = pte_data_map_res.get();
+    }
+
     // n_delegate_ counts the number of successfully-initialized delegates for
     // ~Method() to clean up, and is incremented at the bottom of the loop. This
     // makes it safe for errors to return without updating any state.
@@ -808,7 +816,8 @@ Error Method::init(
       BackendInitContext backend_init_context(
           method_allocator,
           /*event_tracer=*/event_tracer_,
-          /*method_name=*/serialization_plan_->name()->c_str());
+          /*method_name=*/serialization_plan_->name()->c_str(),
+          /*named_data_map=*/pte_data_map);
       Error err = BackendDelegate::Init(
           delegate, program_, backend_init_context, &delegates_[i]);
       if (err != Error::Ok) {