Update on "[executorch][flat_tensor] DataMap implementation"

DataMap implementation that
* Loads a flat_tensor file
* Populates a map with {fqn: tensor} and {fqn: TensorLayout}.
* Makes tensor information available via the named_data_map.h interface.

For now, DataMap doesn't store the DataLoader.
- If/when tensors are in their own segments, DataMap should also store a DataLoader.

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

[ghstack-poisoned]

Author: lucylq

extension/flat_tensor/named_data_map/data_map.cpp

@@ -5,18 +5,17 @@
  * 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/extension/flat_tensor/named_data_map/data_map.h>
 #include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>
 #include <executorch/extension/flat_tensor/serialize/schema_generated.h>
-
 #include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 #include <executorch/runtime/core/freeable_buffer.h>
 #include <executorch/runtime/core/result.h>
 #include <executorch/runtime/core/span.h>
 #include <executorch/runtime/platform/compiler.h>
 
-#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
-
 #include <tuple>
 #include <unordered_map>
 
@@ -48,16 +47,16 @@ bool IsAligned(const void* data) {
 
 ET_NODISCARD Result<const TensorLayout> DataMap::get_metadata(
     const char* fqn) const {
-  auto result = _tensor_map.find(fqn);
-  if (result == _tensor_map.end()) {
+  auto result = _name_to_tensor.find(fqn);
+  if (result == _name_to_tensor.end()) {
     return Error::NotFound;
   }
   return std::get<2>(result->second);
 }
 
 ET_NODISCARD Result<FreeableBuffer> DataMap::get_data(const char* fqn) const {
-  auto result = _tensor_map.find(fqn);
-  if (result == _tensor_map.end()) {
+  auto result = _name_to_tensor.find(fqn);
+  if (result == _name_to_tensor.end()) {
     return Error::NotFound;
   }
   int offset = std::get<1>(result->second);
@@ -68,15 +67,15 @@ ET_NODISCARD Result<FreeableBuffer> DataMap::get_data(const char* fqn) const {
 }
 
 ET_NODISCARD Result<int> DataMap::get_num_keys() const {
-  return _tensor_map.size();
+  return _name_to_tensor.size();
 }
 
 ET_NODISCARD Result<const char*> DataMap::get_key(int index) const {
-  if (index <= 0 || index >= _tensor_map.size()) {
+  if (index <= 0 || index >= _name_to_tensor.size()) {
     return Error::InvalidArgument;
   }
 
-  auto iter = _tensor_map.begin();
+  auto iter = _name_to_tensor.begin();
   for (int i = 0; i < index; ++i) {
     ++iter;
   }

extension/flat_tensor/named_data_map/data_map.h

@@ -8,16 +8,13 @@
 
 #pragma once
 
-#include <cstddef>
-
 #include <executorch/runtime/core/data_loader.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/core/result.h>
+#include <executorch/runtime/core/tensor_layout.h>
 #include <executorch/runtime/platform/compiler.h>
 
-#include <executorch/runtime/core/exec_aten/exec_aten.h>
-
-#include <executorch/runtime/core/tensor_layout.h>
 #include <unordered_map>
 #include <utility>
 
@@ -54,10 +51,11 @@ class DataMap final : public executorch::runtime::NamedDataMap {
       executorch::runtime::FreeableBuffer&& flat_tensor_data,
       std::unordered_map<
           std::string,
-          std::tuple<int, int, executorch::runtime::TensorLayout>> tensor_map,
+          std::tuple<int, int, executorch::runtime::TensorLayout>>
+          name_to_tensor,
       executorch::runtime::FreeableBuffer&& data_ro)
       : _flat_tensor_data(std::move(flat_tensor_data)),
-        _tensor_map(std::move(tensor_map)),
+        _name_to_tensor(std::move(name_to_tensor)),
         _data_ro(std::move(data_ro)) {}
 
   // Not copyable or assignable.
@@ -66,14 +64,14 @@ class DataMap final : public executorch::runtime::NamedDataMap {
   DataMap& operator=(const DataMap& rhs) = delete;
 
   // FlatTensor flatbuffer data. Contains the data backing up
-  // TensorLayout information in the _tensor_map; must outlive it.
+  // TensorLayout information in the _name_to_tensor map; must outlive it.
   executorch::runtime::FreeableBuffer _flat_tensor_data;
 
-  // Map of fqn to {segment index, offset, TensorLayout}.
+  // Map of name to {segment index, offset, TensorLayout}.
   std::unordered_map<
       std::string,
       std::tuple<int, int, executorch::runtime::TensorLayout>>
-      _tensor_map;
+      _name_to_tensor;
 
   // Raw, read-only tensor data.
   executorch::runtime::FreeableBuffer _data_ro;

extension/flat_tensor/test/data_map_test.cpp

@@ -6,18 +6,16 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/extension/flat_tensor/named_data_map/data_map.h>
-
 #include <executorch/extension/data_loader/buffer_data_loader.h>
 #include <executorch/extension/data_loader/file_data_loader.h>
+#include <executorch/extension/flat_tensor/named_data_map/data_map.h>
+#include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>
+#include <executorch/extension/flat_tensor/serialize/schema_generated.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/result.h>
 #include <executorch/runtime/platform/runtime.h>
 #include <executorch/test/utils/DeathTest.h>
 
-#include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>
-#include <executorch/extension/flat_tensor/serialize/schema_generated.h>
-
 #include <gtest/gtest.h>
 
 using namespace ::testing;

extension/flat_tensor/test/targets.bzl

@@ -32,12 +32,12 @@ def define_common_targets(is_fbcode=False):
                 "data_map_test.cpp",
             ],
             deps = [
-                "//executorch/extension/flat_tensor/serialize:schema",
-                "//executorch/extension/flat_tensor/serialize:generated_headers",
-                "//executorch/extension/flat_tensor/named_data_map:data_map",
-                "//executorch/extension/flat_tensor/serialize:flat_tensor_header",
                 "//executorch/extension/data_loader:buffer_data_loader",
                 "//executorch/extension/data_loader:file_data_loader",
+                "//executorch/extension/flat_tensor/named_data_map:data_map",
+                "//executorch/extension/flat_tensor/serialize:flat_tensor_header",
+                "//executorch/extension/flat_tensor/serialize:generated_headers",
+                "//executorch/extension/flat_tensor/serialize:schema",
                 "//executorch/runtime/core/exec_aten:lib",
             ],
             env = modules_env,