[executorch][runtime] Add NamedDataMap to method load (#8383)

Pull Request resolved: #8200

Add NamedDataMap as an arg to:
- Method
- load_method
- parseTensor

Use NamedDataMap to resolve external tensors in parseTensor.

Test that the PTE + PTD file run well inside method_test.

TODO: additional test cases. also, tests for mutable case with load_data_into when that is landed, see D69148652.

Land after: #8349, which adds CMake support for FlatTensor.

Note: dependency on D66580784, which fixes some failing RL tests exposed by this diff from inclusion of portable and aten headers.
https://www.internalfb.com/code/fbsource/[2b74908ed03c34fbfd53fb57575eba0964fe0136][history]/arvr/projects/depthsensing/products/depthmanager/depth/stereo/MlStereoDisparityModelExecutorchBackend.h?lines=13-15


ghstack-source-id: 265828564
@exported-using-ghexport

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

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

Author: pytorchbot

extension/flat_tensor/test/targets.bzl

@@ -35,8 +35,8 @@ def define_common_targets(is_fbcode=False):
             # The tests use this var to find the program file to load. This uses
             # an fbcode target path because the authoring/export tools
             # intentionally don't work in xplat (since they're host-only tools).
-            "ET_MODULE_LINEAR_PROGRAM_PATH": "$(location fbcode//executorch/test/models:exported_programs_with_data_separated[ModuleLinear.pte])",
-            "ET_MODULE_LINEAR_DATA_PATH": "$(location fbcode//executorch/test/models:exported_programs_with_data_separated[ModuleLinear.ptd])",
+            "ET_MODULE_LINEAR_PROGRAM_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleLinear.pte])",
+            "ET_MODULE_LINEAR_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleLinear.ptd])",
         }
 
         runtime.cxx_test(

runtime/executor/method.cpp

@@ -16,6 +16,7 @@
 #include <executorch/runtime/backend/interface.h>
 #include <executorch/runtime/core/event_tracer_hooks.h>
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
+#include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/core/span.h>
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/platform_memory_allocator.h>
@@ -288,7 +289,7 @@ Result<bool> parse_cond_value(const EValue& cond_value) {
 
 } // namespace
 
-Error Method::parse_values() {
+Error Method::parse_values(const NamedDataMap* named_data_map) {
   auto flatbuffer_values = serialization_plan_->values();
   ET_CHECK_OR_RETURN_ERROR(
       flatbuffer_values != nullptr, InvalidProgram, "Missing values");
@@ -414,7 +415,8 @@ Error Method::parse_values() {
         auto t = deserialization::parseTensor(
             program_,
             memory_manager_,
-            static_cast<const executorch_flatbuffer::Tensor*>(val));
+            static_cast<const executorch_flatbuffer::Tensor*>(val),
+            named_data_map);
         if (!t.ok()) {
           ET_LOG(
               Error,
@@ -607,7 +609,8 @@ Result<Method> Method::load(
     executorch_flatbuffer::ExecutionPlan* s_plan,
     const Program* program,
     MemoryManager* memory_manager,
-    EventTracer* event_tracer) {
+    EventTracer* event_tracer,
+    const NamedDataMap* named_data_map) {
   MemoryAllocator* temp_allocator = memory_manager->temp_allocator();
   if (temp_allocator == nullptr) {
     PlatformMemoryAllocator* platform_allocator =
@@ -621,7 +624,7 @@ Result<Method> Method::load(
   }
   Method method(program, memory_manager, event_tracer, temp_allocator);
 
-  Error err = method.init(s_plan);
+  Error err = method.init(s_plan, named_data_map);
   if (err != Error::Ok) {
     return err;
   } else {
@@ -630,7 +633,9 @@ Result<Method> Method::load(
   }
 }
 
-Error Method::init(executorch_flatbuffer::ExecutionPlan* s_plan) {
+Error Method::init(
+    executorch_flatbuffer::ExecutionPlan* s_plan,
+    const NamedDataMap* named_data_map) {
   EXECUTORCH_SCOPE_PROF("Method::init");
   internal::EventTracerProfileMethodScope event_tracer_profile_scope =
       internal::EventTracerProfileMethodScope(event_tracer_, "Method::init");
@@ -647,7 +652,7 @@ Error Method::init(executorch_flatbuffer::ExecutionPlan* s_plan) {
 
   {
     // Parse the elements of the values_ array.
-    Error err = parse_values();
+    Error err = parse_values(named_data_map);
     if (err != Error::Ok) {
       return err;
     }

runtime/executor/method.h

@@ -14,6 +14,7 @@
 #include <executorch/runtime/core/evalue.h>
 #include <executorch/runtime/core/event_tracer.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/core/span.h>
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/method_meta.h>
@@ -294,14 +295,17 @@ class Method final {
       executorch_flatbuffer::ExecutionPlan* s_plan,
       const Program* program,
       MemoryManager* memory_manager,
-      EventTracer* event_tracer);
+      EventTracer* event_tracer,
+      const NamedDataMap* named_data_map);
 
   /**
    * Initialize the method from its serialized representation.
    *
    * @returns Error::Ok on success, non-Ok on failure.
    */
-  ET_NODISCARD Error init(executorch_flatbuffer::ExecutionPlan* s_plan);
+  ET_NODISCARD Error init(
+      executorch_flatbuffer::ExecutionPlan* s_plan,
+      const NamedDataMap* named_data_map);
 
   /// Returns true if the Method was successfully initialized.
   inline bool initialized() const {
@@ -339,7 +343,7 @@ class Method final {
    * the number of successfully-initialized entries so that ~Method doesn't try
    * to clean up uninitialized entries.
    */
-  ET_NODISCARD Error parse_values();
+  ET_NODISCARD Error parse_values(const NamedDataMap* named_data_map);
 
   ET_NODISCARD Error resolve_operator(
       int32_t op_index,

runtime/executor/program.cpp

@@ -240,7 +240,8 @@ Result<const char*> Program::get_method_name(size_t plan_index) const {
 Result<Method> Program::load_method(
     const char* method_name,
     MemoryManager* memory_manager,
-    EventTracer* event_tracer) const {
+    EventTracer* event_tracer,
+    const NamedDataMap* named_data_map) const {
   EXECUTORCH_SCOPE_PROF("Program::load_method");
   internal::event_tracer_create_event_block(event_tracer, "Default");
   internal::EventTracerProfileMethodScope event_tracer_scope =
@@ -257,7 +258,8 @@ Result<Method> Program::load_method(
   if (!plan.ok()) {
     return plan.error();
   }
-  return Method::load(plan.get(), this, memory_manager, event_tracer);
+  return Method::load(
+      plan.get(), this, memory_manager, event_tracer, named_data_map);
 }
 
 Result<MethodMeta> Program::method_meta(const char* method_name) const {

runtime/executor/program.h

@@ -129,13 +129,16 @@ class Program final {
    *     execution of the loaded method. If `memory_manager.temp_allocator()` is
    *     null, the runtime will allocate temp memory using `et_pal_allocate()`.
    * @param[in] event_tracer The event tracer to use for this method run.
+   * @param[in] named_data_map An optional map of {name, blob} used to resolve
+   *     data that is external to the PTE, if any.
    *
    * @returns The loaded method on success, or an error on failure.
    */
   Result<Method> load_method(
       const char* method_name,
       MemoryManager* memory_manager,
-      EventTracer* event_tracer = nullptr) const;
+      EventTracer* event_tracer = nullptr,
+      const NamedDataMap* named_data_map = nullptr) const;
 
   /**
    * Gathers metadata for the named method.

runtime/executor/targets.bzl

@@ -79,6 +79,7 @@ def define_common_targets():
                 ":memory_manager",
                 "//executorch/runtime/backend:interface",
                 "//executorch/runtime/core:core",
+                "//executorch/runtime/core:named_data_map",
                 "//executorch/runtime/core:evalue" + aten_suffix,
                 "//executorch/runtime/core:event_tracer" + aten_suffix,
                 "//executorch/runtime/core/exec_aten:lib" + aten_suffix,

runtime/executor/tensor_parser.h

@@ -7,6 +7,9 @@
  */
 
 #pragma once
+// Disable -Wdeprecated-declarations, as some builds use 'Werror'.
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
 
 #include <executorch/runtime/core/evalue.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
@@ -21,7 +24,8 @@ namespace deserialization {
 ET_NODISCARD Result<executorch::aten::Tensor> parseTensor(
     const Program* program,
     MemoryManager* memory_manager,
-    const executorch_flatbuffer::Tensor* s_tensor);
+    const executorch_flatbuffer::Tensor* s_tensor,
+    const NamedDataMap* named_data_map = nullptr);
 
 ET_NODISCARD Result<BoxedEvalueList<executorch::aten::Tensor>> parseTensorList(
     const flatbuffers::Vector<int32_t>* tensor_indices,
@@ -100,6 +104,8 @@ parseListOptionalType(
  * @param[in] program The Program to use for constant buffer data.
  * @param[in] nbytes The amount of memory to get from the allocator.
  * @param[in] allocator The source of memory for non-constant tensors.
+ * @param[in] named_data_map An optional map of {name, blob} used to resolve
+ *     data that is external to the PTE, if any.
  *
  * @returns On success, the data pointer to use for the tensor. On failure, a
  *     non-Ok Error.
@@ -108,7 +114,8 @@ ET_NODISCARD Result<void*> getTensorDataPtr(
     const executorch_flatbuffer::Tensor* s_tensor,
     const Program* program,
     size_t nbytes,
-    HierarchicalAllocator* allocator);
+    HierarchicalAllocator* allocator,
+    const NamedDataMap* named_data_map = nullptr);
 
 } // namespace deserialization
 } // namespace runtime
@@ -126,3 +133,4 @@ using ::executorch::runtime::deserialization::parseTensorList;
 } // namespace deserialization
 } // namespace executor
 } // namespace torch
+#pragma GCC diagnostic pop

runtime/executor/tensor_parser_aten.cpp

@@ -10,6 +10,7 @@
 
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
+#include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/program.h>
 #include <executorch/runtime/platform/profiler.h>
@@ -31,7 +32,8 @@ void deleteNothing(void*) {}
 Result<at::Tensor> parseTensor(
     const Program* program,
     MemoryManager* memory_manager,
-    const executorch_flatbuffer::Tensor* s_tensor) {
+    const executorch_flatbuffer::Tensor* s_tensor,
+    const NamedDataMap* named_data_map) {
   EXECUTORCH_SCOPE_PROF("TensorParser::parseTensor");
 
   ET_CHECK_OR_RETURN_ERROR(
@@ -102,7 +104,11 @@ Result<at::Tensor> parseTensor(
   } else {
     // Now that we know how big the tensor is, find and assign its memory.
     Result<void*> data_ptr = getTensorDataPtr(
-        s_tensor, program, tensor.nbytes(), memory_manager->planned_memory());
+        s_tensor,
+        program,
+        tensor.nbytes(),
+        memory_manager->planned_memory(),
+        named_data_map);
     if (!data_ptr.ok()) {
       ET_LOG(
           Error,

runtime/executor/tensor_parser_exec_aten.cpp

@@ -19,6 +19,8 @@ namespace executorch {
 namespace runtime {
 namespace deserialization {
 
+using executorch::aten::ScalarType;
+using executorch::runtime::TensorLayout;
 // Provides access to private Program methods.
 class TensorParser final {
  public:
@@ -113,7 +115,8 @@ ET_NODISCARD Result<void*> getTensorDataPtr(
     const executorch_flatbuffer::Tensor* s_tensor,
     const Program* program,
     size_t nbytes,
-    HierarchicalAllocator* allocator) {
+    HierarchicalAllocator* allocator,
+    const NamedDataMap* named_data_map) {
   auto data_buffer_idx = s_tensor->data_buffer_idx();
   const executorch_flatbuffer::AllocationDetails* allocation_info =
       s_tensor->allocation_info();
@@ -131,9 +134,103 @@ ET_NODISCARD Result<void*> getTensorDataPtr(
       return err;
     }
     return planned_ptr;
+  }
+
+  // External tensors.
+  else if (
+      s_tensor->extra_tensor_info() != nullptr &&
+      s_tensor->extra_tensor_info()->location() ==
+          executorch_flatbuffer::TensorDataLocation::EXTERNAL) {
+    // Check that fqn is not null.
+    ET_CHECK_OR_RETURN_ERROR(
+        s_tensor->extra_tensor_info()->fully_qualified_name() != nullptr,
+        InvalidExternalData,
+        "Fully qualified name of external tensor is null");
+    // Look up tensor in named data map.
+    Result<const TensorLayout> tensor_layout_res = named_data_map->get_metadata(
+        s_tensor->extra_tensor_info()->fully_qualified_name()->c_str());
+    if (!tensor_layout_res.ok()) {
+      return tensor_layout_res.error();
+    }
+    const TensorLayout& tensor_layout = tensor_layout_res.get();
+
+    // Compatibility checking.
+    ET_CHECK_OR_RETURN_ERROR(
+        static_cast<ScalarType>(s_tensor->scalar_type()) ==
+            tensor_layout.scalar_type(),
+        InvalidExternalData,
+        "Scalar type mismatch. Expected %hhd, got %hhd.",
+        static_cast<int8_t>(s_tensor->scalar_type()),
+        static_cast<int8_t>(tensor_layout.scalar_type()));
+    ET_CHECK_OR_RETURN_ERROR(
+        nbytes == tensor_layout.nbytes(),
+        InvalidExternalData,
+        "Nbytes mismatch. Expected %zu, got %zu.",
+        nbytes,
+        tensor_layout.nbytes());
+    int dim = s_tensor->sizes()->size();
+    ET_CHECK_OR_RETURN_ERROR(
+        dim == tensor_layout.sizes().size(),
+        InvalidExternalData,
+        "Dim mismatch. Expected %d, got %zu.",
+        dim,
+        tensor_layout.sizes().size());
+    for (int i = 0; i < dim; i++) {
+      ET_CHECK_OR_RETURN_ERROR(
+          s_tensor->sizes()->Get(i) == tensor_layout.sizes()[i],
+          InvalidExternalData,
+          "Sizes mismatch. Expected %d, got %d for size at index %d.",
+          s_tensor->sizes()->Get(i),
+          tensor_layout.sizes()[i],
+          i);
+      ET_CHECK_OR_RETURN_ERROR(
+          s_tensor->dim_order()->Get(i) == tensor_layout.dim_order()[i],
+          InvalidExternalData,
+          "Dim order mismatch. Expected %d, got %d for dim at index %d.",
+          s_tensor->dim_order()->Get(i),
+          tensor_layout.dim_order()[i],
+          i);
+    }
+
+    // Constant value.
+    if (allocation_info == nullptr) {
+      Result<FreeableBuffer> data_res = named_data_map->get_data(
+          s_tensor->extra_tensor_info()->fully_qualified_name()->c_str());
+      if (!data_res.ok()) {
+        return data_res.error();
+      }
+      // The const_cast is 'ok' here because program and runtime should
+      // guarantee that this data is never modified. Temporary until runtime
+      // takes ownership of FreeableBuffers in TODO(T214294528).
+      return const_cast<void*>(data_res.get().data());
+    }
+
+    // Mutable value.
+    else {
+      // Call load_into.
+      auto planned_ptr = getMemPlannedPtr(allocation_info, nbytes, allocator);
+      if (!planned_ptr.ok()) {
+        return planned_ptr.error();
+      }
+      auto size = named_data_map->load_data_into(
+          s_tensor->extra_tensor_info()->fully_qualified_name()->c_str(),
+          planned_ptr.get(),
+          nbytes);
+      if (size.error() != Error::Ok) {
+        return size.error();
+      }
+      ET_CHECK_OR_RETURN_ERROR(
+          size.get() == nbytes,
+          InvalidExternalData,
+          "Expected to load %zu bytes, actually loaded %u bytes",
+          nbytes,
+          static_cast<unsigned int>(size.get()));
+      return planned_ptr;
+    }
+  }
 
-    // Constant
-  } else if (data_buffer_idx > 0 && allocation_info == nullptr) {
+  // Constant, stored in PTE file.
+  else if (data_buffer_idx > 0 && allocation_info == nullptr) {
     auto const_data =
         program->get_constant_buffer_data(data_buffer_idx, nbytes);
     if (!const_data.ok()) {

runtime/executor/tensor_parser_portable.cpp

@@ -11,6 +11,7 @@
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
+#include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/program.h>
 #include <executorch/runtime/platform/profiler.h>
@@ -27,7 +28,8 @@ using torch::executor::TensorImpl;
 Result<Tensor> parseTensor(
     const Program* program,
     MemoryManager* memory_manager,
-    const executorch_flatbuffer::Tensor* s_tensor) {
+    const executorch_flatbuffer::Tensor* s_tensor,
+    const NamedDataMap* named_data_map) {
   EXECUTORCH_SCOPE_PROF("TensorParser::parseTensor");
   auto method_allocator = memory_manager->method_allocator();
 
@@ -146,7 +148,8 @@ Result<Tensor> parseTensor(
       s_tensor,
       program,
       tensor_impl->nbytes(),
-      memory_manager->planned_memory());
+      memory_manager->planned_memory(),
+      named_data_map);
   if (!data_ptr.ok()) {
     ET_LOG(
         Error,