From 56182a8feab260ac19e72a0e3897619f7073a116 Mon Sep 17 00:00:00 2001
From: lucylq <lfq@meta.com>
Date: Tue, 4 Mar 2025 14:30:55 -0800
Subject: [PATCH] Use NDM in backend init context

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]
---
 runtime/backend/backend_init_context.h | 16 ++++++++++++++--
 runtime/backend/interface.h            |  1 +
 runtime/backend/targets.bzl            |  1 +
 runtime/executor/method.cpp            | 11 ++++++++++-
 4 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/runtime/backend/backend_init_context.h b/runtime/backend/backend_init_context.h
index 49bf4adbf9a..de1661c3af0 100644
--- a/runtime/backend/backend_init_context.h
+++ b/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
diff --git a/runtime/backend/interface.h b/runtime/backend/interface.h
index b74858a9d94..0a3c069a201 100644
--- a/runtime/backend/interface.h
+++ b/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>
 
diff --git a/runtime/backend/targets.bzl b/runtime/backend/targets.bzl
index 12d68396730..d2187afb5fc 100644
--- a/runtime/backend/targets.bzl
+++ b/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",
             ],
         )
diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp
index 7da7bafd3e5..41d44522a22 100644
--- a/runtime/executor/method.cpp
+++ b/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) {