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

@@ -49,15 +49,16 @@ ET_NODISCARD Result<const TensorLayout> DataMap::get_metadata(
     const char* fqn) const {
   auto result = _name_to_tensor.find(fqn);
   if (result == _name_to_tensor.end()) {
-    return Error::NotFound;
+    return Error::InvalidArgument;
   }
+  // value is a tuple of (segment_index, offset, tensor_layout)
   return std::get<2>(result->second);
 }
 
 ET_NODISCARD Result<FreeableBuffer> DataMap::get_data(const char* fqn) const {
   auto result = _name_to_tensor.find(fqn);
   if (result == _name_to_tensor.end()) {
-    return Error::NotFound;
+    return Error::InvalidArgument;
   }
   int offset = std::get<1>(result->second);
   TensorLayout tensor = std::get<2>(result->second);
@@ -66,17 +67,17 @@ ET_NODISCARD Result<FreeableBuffer> DataMap::get_data(const char* fqn) const {
   return FreeableBuffer(data, tensor.nbytes(), nullptr);
 }
 
-ET_NODISCARD Error
+ET_NODISCARD Result<size_t>
 DataMap::load_data_into(const char* fqn, size_t size, void* buffer) const {
   return Error::NotImplemented;
 }
 
-ET_NODISCARD Result<int> DataMap::get_num_keys() const {
+ET_NODISCARD Result<size_t> DataMap::get_num_keys() const {
   return _name_to_tensor.size();
 }
 
-ET_NODISCARD Result<const char*> DataMap::get_key(int index) const {
-  if (index <= 0 || index >= _name_to_tensor.size()) {
+ET_NODISCARD Result<const char*> DataMap::get_key(size_t index) const {
+  if (index < 0 || index >= _name_to_tensor.size()) {
     return Error::InvalidArgument;
   }
 
@@ -158,7 +159,7 @@ ET_NODISCARD Result<const char*> DataMap::get_key(int index) const {
   assert(s_tensor_metadata != nullptr);
 
   std::unordered_map<std::string, std::tuple<int, int, TensorLayout>>
-      fqn_to_tensor_layout = {};
+      name_to_tensor = {};
   for (int i = 0; i < s_tensor_metadata->size(); i++) {
     // Create TensorLayouts.
     ScalarType scalar_type =
@@ -178,7 +179,7 @@ ET_NODISCARD Result<const char*> DataMap::get_key(int index) const {
     std::string fqn = s_tensor_metadata->Get(i)->fully_qualified_name()->str();
 
     auto val = std::make_tuple(segment_index, offset, tensor_layout);
-    fqn_to_tensor_layout.insert({fqn, std::move(val)});
+    name_to_tensor.insert({fqn, std::move(val)});
   }
 
   // Load constant data.
@@ -203,7 +204,7 @@ ET_NODISCARD Result<const char*> DataMap::get_key(int index) const {
 
   return DataMap(
       std::move(flat_tensor_data.get()),
-      std::move(fqn_to_tensor_layout),
+      std::move(name_to_tensor),
       std::move(_data_ro.get()));
 }
 

extension/flat_tensor/named_data_map/data_map.h

@@ -38,12 +38,13 @@ class DataMap final : public executorch::runtime::NamedDataMap {
   ET_NODISCARD
   executorch::runtime::Result<executorch::runtime::FreeableBuffer> get_data(
       const char* fqn) const override;
-  ET_NODISCARD runtime::Error
+  ET_NODISCARD executorch::runtime::Result<size_t>
   load_data_into(const char* fqn, size_t size, void* buffer) const override;
 
-  ET_NODISCARD executorch::runtime::Result<int> get_num_keys() const override;
+  ET_NODISCARD executorch::runtime::Result<size_t> get_num_keys()
+      const override;
   ET_NODISCARD executorch::runtime::Result<const char*> get_key(
-      int index) const override;
+      size_t index) const override;
 
   DataMap(DataMap&&) noexcept = default;
   ~DataMap() override;

extension/flat_tensor/test/data_map_test.cpp

@@ -32,29 +32,41 @@ class DataMapTest : public ::testing::Test {
     // Since these tests cause ET_LOG to be called, the PAL must be initialized
     // first.
     executorch::runtime::runtime_init();
-  }
-};
 
-TEST_F(DataMapTest, LoadDataMap) {
-  const char* path = std::getenv("ET_MODULE_LINEAR_DATA");
-  Result<FileDataLoader> loader = FileDataLoader::from(path);
-  ASSERT_EQ(loader.error(), Error::Ok);
+    // Load data map.
+    // The eager linear model is defined at:
+    // //executorch/test/models/linear_model.py
+    const char* path = std::getenv("ET_MODULE_LINEAR_DATA");
+    Result<FileDataLoader> loader = FileDataLoader::from(path);
+    ASSERT_EQ(loader.error(), Error::Ok);
 
-  Result<FreeableBuffer> header = loader->load(
-      /*offset=*/0,
-      FlatTensorHeader::kNumHeadBytes,
-      /*segment_info=*/
-      DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External));
+    Result<FreeableBuffer> header = loader->load(
+        /*offset=*/0,
+        FlatTensorHeader::kNumHeadBytes,
+        /*segment_info=*/
+        DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External));
 
-  ASSERT_EQ(header.error(), Error::Ok);
+    ASSERT_EQ(header.error(), Error::Ok);
+
+    data_map_loader_ =
+        std::make_unique<FileDataLoader>(std::move(loader.get()));
+  }
+  std::unique_ptr<FileDataLoader> data_map_loader_;
+};
 
-  auto data_map_loader_ =
-      std::make_unique<FileDataLoader>(std::move(loader.get()));
+TEST_F(DataMapTest, LoadDataMap) {
+  Result<DataMap> data_map = DataMap::load(data_map_loader_.get());
+  EXPECT_EQ(data_map.error(), Error::Ok);
+}
 
+TEST_F(DataMapTest, DataMap_GetMetadata) {
   Result<DataMap> data_map = DataMap::load(data_map_loader_.get());
   EXPECT_EQ(data_map.error(), Error::Ok);
 
-  // Check tensor metadata.
+  // Check tensor layouts are correct.
+  // From //executorch/test/models/linear_model.py, we have the tensors
+  // self.a = 3 * torch.ones(2, 2, dtype=torch.float)
+  // self.b = 2 * torch.ones(2, 2, dtype=torch.float)
   Result<const TensorLayout> const_a_res = data_map->get_metadata("a");
   assert(const_a_res.ok());
 
@@ -83,16 +95,50 @@ TEST_F(DataMapTest, LoadDataMap) {
   EXPECT_EQ(dim_order_b[0], 0);
   EXPECT_EQ(dim_order_b[1], 1);
 
-  // Check tensor data.
+  // Check get_metadata fails when key is not found.
+  Result<const TensorLayout> const_c_res = data_map->get_metadata("c");
+  EXPECT_EQ(const_c_res.error(), Error::InvalidArgument);
+}
+
+TEST_F(DataMapTest, DataMap_GetData) {
+  Result<DataMap> data_map = DataMap::load(data_map_loader_.get());
+  EXPECT_EQ(data_map.error(), Error::Ok);
+
+  // Check tensor data sizes are correct.
   Result<FreeableBuffer> data_a_res = data_map->get_data("a");
   assert(data_a_res.ok());
-  // Check we have the correct tensor data.
   FreeableBuffer data_a = std::move(data_a_res.get());
   EXPECT_EQ(data_a.size(), 16);
 
   Result<FreeableBuffer> data_b_res = data_map->get_data("b");
   assert(data_b_res.ok());
-  // Check we have the correct tensor data.
   FreeableBuffer data_b = std::move(data_b_res.get());
   EXPECT_EQ(data_b.size(), 16);
+
+  // Check get_data fails when key is not found.
+  Result<FreeableBuffer> data_c_res = data_map->get_data("c");
+  EXPECT_EQ(data_c_res.error(), Error::InvalidArgument);
+}
+
+TEST_F(DataMapTest, DataMap_Keys) {
+  Result<DataMap> data_map = DataMap::load(data_map_loader_.get());
+  EXPECT_EQ(data_map.error(), Error::Ok);
+
+  // Check num tensors is 2.
+  Result<size_t> num_tensors_res = data_map->get_num_keys();
+  assert(num_tensors_res.ok());
+  EXPECT_EQ(num_tensors_res.get(), 2);
+
+  // Check get_key returns the correct keys.
+  Result<const char*> key0_res = data_map->get_key(0);
+  assert(key0_res.ok());
+  EXPECT_EQ(strcmp(key0_res.get(), "b"), 0);
+
+  Result<const char*> key1_res = data_map->get_key(1);
+  assert(key1_res.ok());
+  EXPECT_EQ(strcmp(key1_res.get(), "a"), 0);
+
+  // Check get_key fails when out of bounds.
+  Result<const char*> key2_res = data_map->get_key(2);
+  EXPECT_EQ(key2_res.error(), Error::InvalidArgument);
 }

runtime/core/named_data_map.h

@@ -46,28 +46,30 @@ class ET_EXPERIMENTAL NamedDataMap {
    * Loads data corresponding to the fqn into the provided buffer.
    *
    * @param fqn Fully qualified name of the tensor.
-   * @param size The number of bytes to load.
+   * @param size The number of bytes to load. Use `get_metadata` to retrieve the
+   * size of the tensor for a given fqn.
    * @param buffer The buffer to load the data into. Must point to at least
    * `size` bytes of memory.
-   * @return An error code on if the load was successful.
+   * @return Result containing the number of bytes written on success.
    */
-  ET_NODISCARD virtual Error
+  ET_NODISCARD virtual Result<size_t>
   load_data_into(const char* fqn, size_t size, void* buffer) const = 0;
 
   /**
    * Get the number of keys in the NamedDataMap.
    *
    * @return Result containing the number of keys.
    */
-  ET_NODISCARD virtual Result<int> get_num_keys() const = 0;
+  ET_NODISCARD virtual Result<size_t> get_num_keys() const = 0;
 
   /**
    * Get the key at the given index.
    *
    * @param index The index of the key to retrieve.
-   * @return Result containing the key at the given index.
+   * @return Result containing the key at the given index. Note: the returned
+   * pointer is only valid for the lifetime of the DataMap.
    */
-  ET_NODISCARD virtual Result<const char*> get_key(int index) const = 0;
+  ET_NODISCARD virtual Result<const char*> get_key(size_t index) const = 0;
 };
 
 } // namespace runtime