[executorch][runtime] Introduce CoreDataMap for weight sharing

CoreDataMap is the NamedDataMap that will live in the runtime. It is used to give delegates access to opaque named data stored in the PTE file. Open to alternative naming suggestions, maybe 'PTEDataMap' or 'ProgramDataMap'?

**Usage**
The CoreDataMap is owned by the program, and instantiated at program load time if named_data exists in the PTE file. We introduce usage of 'std::optional' here. I think we can also use executorch::aten::optional to avoid adding standard lib ?

When initializing delegates, the CoreDataMap is given to delegate_init. Delegates can retrieve opaque delegate data by key using 'get_data'.

**Testing**
This test uses the C++ flatbuffer API to build a fake program containing named data. Creates a temp file with sample data that the data loader can wrap around.

TODO: e2e test once delegate aot is ready and we can generate a file with named data.

**Note**
As the CoreDataMap wraps around flatbuffer constructs, the Program must outlive the CoreDataMap.

CoreDataMap does not implement
- get_metadata; currently, all data stored is opaque. Later, we can implement get_metadata if a backend stores plain tensor data.
- load_into; this is mostly used for the training case, and isn't used by delegates, at least not at the moment.

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

[ghstack-poisoned]

Author: lucylq

extension/testing_util/targets.bzl

@@ -16,5 +16,6 @@ def define_common_targets():
             "//executorch/extension/testing_util/test/...",
             "//executorch/extension/fb/ptez/decompression_methods/test/...",
             "//executorch/extension/fb/ptez/test/...",
+            "//executorch/runtime/executor/test/...",
         ],
     )

runtime/executor/core_data_map.cpp

@@ -0,0 +1,85 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include "executorch/runtime/executor/core_data_map.h"
+#include <executorch/schema/program_generated.h>
+
+namespace executorch {
+namespace runtime {
+
+/* static */ executorch::runtime::Result<CoreDataMap> CoreDataMap::load(
+    executorch::runtime::DataLoader* loader,
+    size_t segment_base_offset,
+    const flatbuffers::Vector<
+        flatbuffers::Offset<executorch_flatbuffer::NamedData>>* named_data,
+    const flatbuffers::Vector<
+        flatbuffers::Offset<executorch_flatbuffer::DataSegment>>* segments) {
+  ET_CHECK_OR_RETURN_ERROR(
+      loader != nullptr && named_data != nullptr && segments != nullptr,
+      InvalidArgument,
+      "CoreDataMap loader, named_data or segments is null; most likely the program does not have any named_data segments");
+  return CoreDataMap(loader, segment_base_offset, named_data, segments);
+}
+
+ET_NODISCARD
+executorch::runtime::Result<executorch::runtime::FreeableBuffer>
+CoreDataMap::get_data(const char* key) const {
+  for (size_t i = 0; i < named_data_->size(); i++) {
+    ET_CHECK_OR_RETURN_ERROR(
+        named_data_->Get(i) != nullptr && named_data_->Get(i)->key() != nullptr,
+        InvalidArgument,
+        "NamedData at index %zu is null",
+        i);
+    if (strcmp(named_data_->Get(i)->key()->c_str(), key) == 0) {
+      // Get the segment index.
+      size_t segment_index = named_data_->Get(i)->segment_index();
+
+      // Get the segment offset and size.
+      ET_CHECK_OR_RETURN_ERROR(
+          segment_index < segments_->size(),
+          InvalidArgument,
+          "Segment index %zu is out of range for segments size %u",
+          segment_index,
+          segments_->size());
+      size_t segment_offset = segments_->Get(segment_index)->offset();
+      size_t segment_size = segments_->Get(segment_index)->size();
+
+      return loader_->load(
+          /*offset=*/segment_base_offset_ + segment_offset,
+          segment_size,
+          DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External));
+    }
+  }
+  return Error::NotFound;
+}
+
+ET_NODISCARD executorch::runtime::Result<size_t> CoreDataMap::get_num_keys()
+    const {
+  return named_data_->size();
+}
+
+ET_NODISCARD executorch::runtime::Result<const char*> CoreDataMap::get_key(
+    size_t index) const {
+  ET_CHECK_OR_RETURN_ERROR(
+      index < named_data_->size(),
+      InvalidArgument,
+      "Index out of range: named_data size is %u, received index %zu",
+      named_data_->size(),
+      index);
+
+  ET_CHECK_OR_RETURN_ERROR(
+      named_data_->Get(index) != nullptr &&
+          named_data_->Get(index)->key() != nullptr,
+      InvalidArgument,
+      "NamedData at index %zu is null",
+      index);
+  return named_data_->Get(index)->key()->c_str();
+}
+
+} // namespace runtime
+} // namespace executorch

runtime/executor/core_data_map.h

@@ -0,0 +1,138 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <executorch/runtime/core/data_loader.h>
+#include <executorch/runtime/core/named_data_map.h>
+
+// Forward declare flatbuffer types. This is a public header and must not
+// include the generated flatbuffer header.
+namespace flatbuffers {
+template <typename T>
+class Vector;
+template <typename T>
+struct Offset;
+} // namespace flatbuffers
+
+namespace executorch_flatbuffer {
+struct NamedData;
+struct DataSegment;
+} // namespace executorch_flatbuffer
+
+namespace executorch {
+namespace runtime {
+
+/**
+ * A NamedDataMap implementation for Flatbuffer-serialized named data
+ * originating from a PTE file.
+ */
+class CoreDataMap final : public executorch::runtime::NamedDataMap {
+ public:
+  /**
+   * Creates a new DataMap that wraps named_data from the PTE file.
+   *
+   * @param[in] loader The DataLoader that accesses the PTE file.
+   * Note: the loader must outlive the CoreDataMap instance.
+   * @param[in] segment_base_offset The offset to the first segment in the PTE
+   * file, in bytes.
+   * @param[in] named_data The named_data from the PTE file. Note: the pointer
+   * passed here must outlive the CoreDataMap instance.
+   * @param[in] segments The segments from the PTE file. Note: the pointer
+   * passed here must outlive the CoreDataMap instance.
+   */
+  static executorch::runtime::Result<CoreDataMap> load(
+      executorch::runtime::DataLoader* loader,
+      size_t segment_base_offset,
+      const flatbuffers::Vector<
+          flatbuffers::Offset<executorch_flatbuffer::NamedData>>* named_data,
+      const flatbuffers::Vector<
+          flatbuffers::Offset<executorch_flatbuffer::DataSegment>>* segments);
+
+  /**
+   * The CoreDataMap currently only handles opaque data that does not contain
+   * tensor-specific metadata.
+   */
+  ET_NODISCARD
+  executorch::runtime::Result<const executorch::runtime::TensorLayout>
+  get_metadata(ET_UNUSED const char* key) const override {
+    return Error::NotImplemented;
+  }
+
+  /**
+   * Retrieve read-only data for the specified key.
+   *
+   * @param[in] key The name of the blob to get data on.
+   *
+   * @return error if the key is not present or data cannot be loaded.
+   */
+  ET_NODISCARD
+  executorch::runtime::Result<executorch::runtime::FreeableBuffer> get_data(
+      const char* key) const override;
+
+  /**
+   * The CoreDataMap currently does not implement load_into.
+   */
+  ET_NODISCARD executorch::runtime::Error load_data_into(
+      ET_UNUSED const char* key,
+      ET_UNUSED void* buffer,
+      ET_UNUSED size_t size) const override {
+    return Error::NotImplemented;
+  }
+
+  /**
+   * @returns The number of keys in the map.
+   */
+  ET_NODISCARD executorch::runtime::Result<size_t> get_num_keys()
+      const override;
+
+  /**
+   * @returns The key at the specified index, error if index out of bounds.
+   */
+  ET_NODISCARD executorch::runtime::Result<const char*> get_key(
+      size_t index) const override;
+
+  // Moveable, to be compatible with Result.
+  CoreDataMap(CoreDataMap&&) noexcept = default;
+  ~CoreDataMap() override = default;
+
+ private:
+  CoreDataMap(
+      executorch::runtime::DataLoader* loader,
+      size_t segment_base_offset,
+      const flatbuffers::Vector<
+          flatbuffers::Offset<executorch_flatbuffer::NamedData>>* named_data,
+      const flatbuffers::Vector<
+          flatbuffers::Offset<executorch_flatbuffer::DataSegment>>* segments)
+      : loader_(loader),
+        segment_base_offset_(segment_base_offset),
+        named_data_(named_data),
+        segments_(segments) {}
+
+  // Not copyable or assignable.
+  CoreDataMap(const CoreDataMap& rhs) = delete;
+  CoreDataMap& operator=(CoreDataMap&& rhs) noexcept = delete;
+  CoreDataMap& operator=(const CoreDataMap& rhs) = delete;
+
+  // Data loader, used to load segment data.
+  executorch::runtime::DataLoader* loader_;
+
+  // Segment base offset.
+  size_t segment_base_offset_;
+
+  // Named data, containing name and segment index.
+  const flatbuffers::Vector<
+      flatbuffers::Offset<executorch_flatbuffer::NamedData>>* named_data_;
+
+  // Segments, to retrieve offset and size for the loader.
+  const flatbuffers::Vector<
+      flatbuffers::Offset<executorch_flatbuffer::DataSegment>>* segments_;
+};
+
+} // namespace runtime
+} // namespace executorch

runtime/executor/program.cpp

@@ -150,6 +150,21 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
   const executorch_flatbuffer::Program* flatbuffer_program =
       executorch_flatbuffer::GetProgram(program_data->data());
 
+  // Instantiate CoreDataMap if named_data is present.
+  const auto named_data = flatbuffer_program->named_data();
+  std::optional<CoreDataMap> core_data_map = std::nullopt;
+  if (named_data != nullptr) {
+    Result<CoreDataMap> core_data_map_result = CoreDataMap::load(
+        loader,
+        segment_base_offset,
+        named_data,
+        flatbuffer_program->segments());
+    if (!core_data_map_result.ok()) {
+      return core_data_map_result.error();
+    }
+    core_data_map.emplace(std::move(core_data_map_result.get()));
+  }
+
   // Constant data may live inside the flatbuffer data (constant_buffer) or in a
   // separate segment (constant_segment). It should not be in both.
   // Check constant_segment->offsets()->size() > 1, as the offsets list will
@@ -199,7 +214,8 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
         segment_base_offset,
         std::move(program_data.get()),
         flatbuffer_program,
-        std::move(constant_segment_data.get()));
+        std::move(constant_segment_data.get()),
+        std::move(core_data_map));
   } else {
     // The constant data is stored inside the flatbuffer, so this program does
     // not contain a separate segment for it.
@@ -208,7 +224,8 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
         segment_base_offset,
         std::move(program_data.get()),
         flatbuffer_program,
-        /*constant_segment_data=*/FreeableBuffer{});
+        /*constant_segment_data=*/FreeableBuffer{},
+        std::move(core_data_map));
   }
 }
 

runtime/executor/program.h

@@ -13,12 +13,14 @@
 
 #include <cinttypes>
 #include <cstdint>
+#include <optional>
 
 #include <executorch/runtime/core/data_loader.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/event_tracer.h>
 #include <executorch/runtime/core/freeable_buffer.h>
 #include <executorch/runtime/core/result.h>
+#include <executorch/runtime/executor/core_data_map.h>
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/method.h>
 #include <executorch/runtime/executor/method_meta.h>
@@ -266,13 +268,15 @@ class Program final {
       size_t segment_base_offset,
       FreeableBuffer&& program_data,
       const executorch_flatbuffer::Program* internal_program,
-      FreeableBuffer&& constant_segment_data)
+      FreeableBuffer&& constant_segment_data,
+      std::optional<CoreDataMap>&& core_data_map)
       : program_data_(std::move(program_data)),
         // Don't need the loader if there are no segments.
         loader_(segment_base_offset > 0 ? loader : nullptr),
         internal_program_(internal_program),
         segment_base_offset_(segment_base_offset),
-        constant_segment_data_(std::move(constant_segment_data)) {}
+        constant_segment_data_(std::move(constant_segment_data)),
+        core_data_map_(std::move(core_data_map)) {}
 
   // Not copyable or assignable.
   Program(const Program& rhs) = delete;
@@ -295,6 +299,9 @@ class Program final {
 
   /// Constant segment data.
   FreeableBuffer constant_segment_data_;
+
+  /// NamedDataMap holding named data from the program.
+  std::optional<CoreDataMap> core_data_map_;
 };
 
 } // namespace runtime

runtime/executor/targets.bzl

@@ -42,6 +42,25 @@ def define_common_targets():
         ],
     )
 
+    runtime.cxx_library(
+        name = "core_data_map",
+        srcs = [
+            "core_data_map.cpp",
+        ],
+        exported_headers = [
+            "core_data_map.h",
+        ],
+        visibility = [
+            "//executorch/runtime/executor/...",
+            "@EXECUTORCH_CLIENTS",
+        ],
+        deps = [
+            "//executorch/runtime/core:core",
+            "//executorch/runtime/core:named_data_map",
+            "//executorch/schema:program",
+        ],
+    )
+
     for aten_mode in get_aten_mode_options():
         aten_suffix = "_aten" if aten_mode else ""
         runtime.cxx_library(
@@ -59,6 +78,7 @@ def define_common_targets():
         runtime.cxx_library(
             name = "program_no_prim_ops" + aten_suffix,
             srcs = [
+                "core_data_map.cpp",
                 "method.cpp",
                 "method_meta.cpp",
                 "program.cpp",
@@ -69,6 +89,7 @@ def define_common_targets():
                 "platform_memory_allocator.h",
             ],
             exported_headers = [
+                "core_data_map.h",
                 "method.h",
                 "method_meta.h",
                 "program.h",