[wip] Add named data map merge

backends/xnnpack/test/runtime/test_xnn_data_separation.cpp

@@ -87,15 +87,15 @@ TEST_F(DataSeparationTest, TestExternalData) {
   // Check that accessing keys out of bounds fails.
   EXPECT_EQ(data_map->get_key(2).error(), Error::InvalidArgument);
 
-  // Linear.weight
+  // Linear.bias
   Result<FreeableBuffer> data0 = data_map->get_data(key0.get());
   EXPECT_EQ(data0.error(), Error::Ok);
-  EXPECT_EQ(data0.get().size(), 36); // 3*3*4 (3*3 matrix, 4 bytes per float)
+  EXPECT_EQ(data0.get().size(), 12); // 3*4 (3 vector, 4 bytes per float)
 
-  // Linear.bias
+  // Linear.weight
   Result<FreeableBuffer> data1 = data_map->get_data(key1.get());
   EXPECT_EQ(data1.error(), Error::Ok);
-  EXPECT_EQ(data1.get().size(), 12); // 3*4 (3 vector, 4 bytes per float)
+  EXPECT_EQ(data1.get().size(), 36); // 3*3*4 (3*3 matrix, 4 bytes per float)
 
   // Check that accessing non-existent data fails.
   Result<FreeableBuffer> data2 = data_map->get_data("nonexistent");

extension/flat_tensor/flat_tensor_data_map.cpp

@@ -19,15 +19,15 @@
 #include <executorch/runtime/core/span.h>
 #include <executorch/runtime/platform/compiler.h>
 
+using executorch::aten::ScalarType;
+using executorch::ET_RUNTIME_NAMESPACE::NamedDataMap;
+using executorch::ET_RUNTIME_NAMESPACE::TensorLayout;
+using executorch::runtime::DataLoader;
 using executorch::runtime::Error;
 using executorch::runtime::FreeableBuffer;
 using executorch::runtime::Result;
 using executorch::runtime::Span;
 
-using executorch::aten::ScalarType;
-using executorch::ET_RUNTIME_NAMESPACE::TensorLayout;
-using executorch::runtime::DataLoader;
-
 namespace executorch {
 namespace extension {
 
@@ -103,82 +103,111 @@ Result<const TensorLayout> create_tensor_layout(
 
 ET_NODISCARD Result<const TensorLayout> FlatTensorDataMap::get_tensor_layout(
     executorch::aten::string_view key) const {
-  Result<const flat_tensor_flatbuffer::NamedData*> named_data = get_named_data(
-      key,
-      flat_tensor_->named_data(),
-      flat_tensor_->segments(),
-      header_.segment_base_offset + header_.segment_data_size);
-  if (!named_data.ok()) {
+  if (key_to_map_index_.find(key.data()) == key_to_map_index_.end()) {
+    return Error::NotFound;
+  }
+  auto index = key_to_map_index_.at(key.data());
+  if (index == -1) {
+    Result<const flat_tensor_flatbuffer::NamedData*> named_data =
+        get_named_data(
+            key,
+            flat_tensor_->named_data(),
+            flat_tensor_->segments(),
+            header_.segment_base_offset + header_.segment_data_size);
+    if (named_data.ok()) {
+      return create_tensor_layout(named_data.get()->tensor_layout());
+    }
     return named_data.error();
+  } else {
+    return merged_maps_[index]->get_tensor_layout(key);
   }
-  return create_tensor_layout(named_data.get()->tensor_layout());
 }
 
 ET_NODISCARD Result<FreeableBuffer> FlatTensorDataMap::get_data(
     executorch::aten::string_view key) const {
-  Result<const flat_tensor_flatbuffer::NamedData*> named_data = get_named_data(
-      key,
-      flat_tensor_->named_data(),
-      flat_tensor_->segments(),
-      header_.segment_base_offset + header_.segment_data_size);
-  if (!named_data.ok()) {
+  if (key_to_map_index_.find(key.data()) == key_to_map_index_.end()) {
+    return Error::NotFound;
+  }
+  auto index = key_to_map_index_.at(key.data());
+  if (index == -1) {
+    Result<const flat_tensor_flatbuffer::NamedData*> named_data =
+        get_named_data(
+            key,
+            flat_tensor_->named_data(),
+            flat_tensor_->segments(),
+            header_.segment_base_offset + header_.segment_data_size);
+    if (named_data.ok()) {
+      uint32_t segment_index = named_data.get()->segment_index();
+      uint64_t segment_offset =
+          flat_tensor_->segments()->Get(segment_index)->offset();
+      uint64_t segment_size =
+          flat_tensor_->segments()->Get(segment_index)->size();
+
+      return loader_->load(
+          /*offset=*/header_.segment_base_offset + segment_offset,
+          segment_size,
+          DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External));
+    }
     return named_data.error();
+  } else {
+    return merged_maps_[index]->get_data(key);
   }
-
-  uint32_t segment_index = named_data.get()->segment_index();
-  uint64_t segment_offset =
-      flat_tensor_->segments()->Get(segment_index)->offset();
-  uint64_t segment_size = flat_tensor_->segments()->Get(segment_index)->size();
-
-  return loader_->load(
-      /*offset=*/header_.segment_base_offset + segment_offset,
-      segment_size,
-      DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External));
 }
 
 ET_NODISCARD Error FlatTensorDataMap::load_data_into(
     ET_UNUSED executorch::aten::string_view key,
     ET_UNUSED void* buffer,
     ET_UNUSED size_t size) const {
-  Result<const flat_tensor_flatbuffer::NamedData*> named_data = get_named_data(
-      key,
-      flat_tensor_->named_data(),
-      flat_tensor_->segments(),
-      header_.segment_base_offset + header_.segment_data_size);
-  if (!named_data.ok()) {
-    return named_data.error();
+  if (key_to_map_index_.find(key.data()) == key_to_map_index_.end()) {
+    return Error::NotFound;
   }
+  auto index = key_to_map_index_.at(key.data());
+  if (index == -1) {
+    Result<const flat_tensor_flatbuffer::NamedData*> named_data =
+        get_named_data(
+            key,
+            flat_tensor_->named_data(),
+            flat_tensor_->segments(),
+            header_.segment_base_offset + header_.segment_data_size);
+    if (!named_data.ok()) {
+      return named_data.error();
+    }
 
-  uint32_t segment_index = named_data.get()->segment_index();
-  uint64_t segment_offset =
-      flat_tensor_->segments()->Get(segment_index)->offset();
+    uint32_t segment_index = named_data.get()->segment_index();
+    uint64_t segment_offset =
+        flat_tensor_->segments()->Get(segment_index)->offset();
 
-  Result<const TensorLayout> tensor_layout =
-      create_tensor_layout(named_data.get()->tensor_layout());
+    Result<const TensorLayout> tensor_layout =
+        create_tensor_layout(named_data.get()->tensor_layout());
 
-  if (!tensor_layout.ok()) {
-    return tensor_layout.error();
-  }
+    if (!tensor_layout.ok()) {
+      return tensor_layout.error();
+    }
 
-  ET_CHECK_OR_RETURN_ERROR(
-      size <= tensor_layout.get().nbytes(),
-      InvalidArgument,
-      "Buffer size %zu is smaller than tensor size %zu",
-      size,
-      tensor_layout.get().nbytes());
-
-  // Load mutable data.
-  DataLoader::SegmentInfo info = DataLoader::SegmentInfo(
-      DataLoader::SegmentInfo::Type::Mutable, 0, nullptr);
-  return loader_->load_into(
-      header_.segment_base_offset + segment_offset,
-      tensor_layout.get().nbytes(),
-      info,
-      buffer);
+    ET_CHECK_OR_RETURN_ERROR(
+        size <= tensor_layout.get().nbytes(),
+        InvalidArgument,
+        "Buffer size %zu is smaller than tensor size %zu",
+        size,
+        tensor_layout.get().nbytes());
+
+    // Load mutable data.
+    DataLoader::SegmentInfo info = DataLoader::SegmentInfo(
+        DataLoader::SegmentInfo::Type::Mutable, 0, nullptr);
+    return loader_->load_into(
+        header_.segment_base_offset + segment_offset,
+        tensor_layout.get().nbytes(),
+        info,
+        buffer);
+  } else {
+    return merged_maps_[index]->load_data_into(key, buffer, size);
+  }
 }
 
 ET_NODISCARD Result<uint32_t> FlatTensorDataMap::get_num_keys() const {
-  return flat_tensor_->named_data()->size();
+  // Guaranteed safe, as the segment_index is a uint32_t, which means
+  // that there can't be more than uint32_t keys.
+  return static_cast<uint32_t>(key_to_map_index_.size());
 }
 
 ET_NODISCARD Result<const char*> FlatTensorDataMap::get_key(
@@ -190,7 +219,40 @@ ET_NODISCARD Result<const char*> FlatTensorDataMap::get_key(
       "Index %u out of range of size %u",
       index,
       num_keys);
-  return flat_tensor_->named_data()->Get(index)->key()->c_str();
+
+  uint32_t current_index = 0;
+  for (const auto& pair : key_to_map_index_) {
+    if (current_index == index) {
+      return pair.first.c_str();
+    }
+    current_index++;
+  }
+  return Error::NotFound;
+}
+
+ET_NODISCARD Error FlatTensorDataMap::merge(const NamedDataMap* other) {
+  ET_CHECK_OR_RETURN_ERROR(
+      other != nullptr, InvalidArgument, "Merge error: other is nullptr.");
+
+  // Check if any duplicate keys exist.
+  uint32_t num_keys = other->get_num_keys().get();
+
+  for (uint32_t i = 0; i < num_keys; i++) {
+    const char* key = other->get_key(i).get();
+    ET_CHECK_OR_RETURN_ERROR(
+        key_to_map_index_.find(key) == key_to_map_index_.end(),
+        InvalidArgument,
+        "Merge error: key %s already exists in the named_data_map.",
+        key);
+  }
+  // Place keys into the map.
+  for (uint32_t i = 0; i < num_keys; i++) {
+    const char* key = other->get_key(i).get();
+    key_to_map_index_[key] = static_cast<int64_t>(merged_maps_.size());
+  }
+
+  merged_maps_.push_back(other);
+  return Error::Ok;
 }
 
 /* static */ Result<FlatTensorDataMap> FlatTensorDataMap::load(
@@ -261,8 +323,18 @@ ET_NODISCARD Result<const char*> FlatTensorDataMap::get_key(
       InvalidExternalData,
       "FlatTensor segments is nullptr, malformed PTD file.");
 
+  // Add keys to the map.
+  std::unordered_map<std::string, int64_t> key_to_map_index;
+  for (int i = 0; i < flat_tensor->named_data()->size(); i++) {
+    const auto* named_data = flat_tensor->named_data()->Get(i);
+    key_to_map_index[named_data->key()->c_str()] = -1;
+  }
   return FlatTensorDataMap(
-      fh.get(), std::move(flat_tensor_data.get()), flat_tensor, loader);
+      fh.get(),
+      std::move(flat_tensor_data.get()),
+      flat_tensor,
+      loader,
+      std::move(key_to_map_index));
 }
 
 } // namespace extension

extension/flat_tensor/flat_tensor_data_map.h

@@ -8,17 +8,18 @@
 
 #pragma once
 
-#include <executorch/runtime/core/named_data_map.h>
-
 #include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>
 
+#include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/core/data_loader.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/result.h>
 #include <executorch/runtime/core/tensor_layout.h>
 #include <executorch/runtime/platform/compiler.h>
 
+#include <unordered_map>
 #include <utility>
+#include <vector>
 
 // Forward declare flatbuffer types. This is a public header and must not
 // include the generated flatbuffer header.
@@ -94,6 +95,17 @@ class FlatTensorDataMap final
   ET_NODISCARD executorch::runtime::Result<const char*> get_key(
       uint32_t index) const override;
 
+  /**
+   * Merge a named_data_map into the current one.
+   * @param[in] other The named_data_map to merge.
+   * @return Error indicating if the merge was successful or not.
+   *
+   * Note: The FlatTensorDataMap does not perform a deep copy; it holds a
+   * reference to other, so other must outlive the FlatTensorDataMap instance.
+   */
+  ET_NODISCARD executorch::runtime::Error merge(
+      const NamedDataMap* other) override;
+
   FlatTensorDataMap(FlatTensorDataMap&&) noexcept = default;
 
   ~FlatTensorDataMap() override = default;
@@ -103,11 +115,14 @@ class FlatTensorDataMap final
       const FlatTensorHeader& header,
       executorch::runtime::FreeableBuffer&& flat_tensor_data,
       const flat_tensor_flatbuffer::FlatTensor* flat_tensor,
-      executorch::runtime::DataLoader* loader)
+      executorch::runtime::DataLoader* loader,
+      std::unordered_map<std::string, int64_t> key_to_map_index)
       : header_(header),
         flat_tensor_data_(std::move(flat_tensor_data)),
         flat_tensor_(flat_tensor),
-        loader_(loader) {}
+        loader_(loader),
+        key_to_map_index_(std::move(key_to_map_index)),
+        merged_maps_({}) {}
 
   // Not copyable or assignable.
   FlatTensorDataMap(const FlatTensorDataMap& rhs) = delete;
@@ -125,6 +140,13 @@ class FlatTensorDataMap final
 
   // Data loader, used to load segment data.
   executorch::runtime::DataLoader* loader_;
+
+  // Cache of keys to data map index.
+  // index=-1 is used for the flat_tensor data map.
+  std::unordered_map<std::string, int64_t> key_to_map_index_;
+
+  // Other NamedDataMaps.
+  std::vector<const NamedDataMap*> merged_maps_;
 };
 
 } // namespace extension

extension/flat_tensor/test/CMakeLists.txt

@@ -21,21 +21,23 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake)
 add_custom_command(
   OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.pte"
          "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.ptd"
+         "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte"
+         "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.ptd"
   COMMAND
-    ${PYTHON_EXECUTABLE} -m test.models.export_program --modules "ModuleAddMul"
+    ${PYTHON_EXECUTABLE} -m test.models.export_program --modules "ModuleAddMul,ModuleLinear"
     --external-constants --outdir "${CMAKE_CURRENT_BINARY_DIR}" 2> /dev/null
   WORKING_DIRECTORY ${EXECUTORCH_ROOT}
 )
 
 add_custom_target(
   extension_flat_tensor_test_resources
-  DEPENDS "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.pte"
-          "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.ptd"
+  DEPENDS "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.ptd"
+          "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.ptd"
 )
 
 set(test_env
-    "ET_MODULE_ADD_MUL_PROGRAM_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.pte"
     "ET_MODULE_ADD_MUL_DATA_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.ptd"
+    "ET_MODULE_LINEAR_DATA_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.ptd"
 )
 
 set(_test_srcs flat_tensor_data_map_test.cpp flat_tensor_header_test.cpp)