Fix issues with named data map load_into

Summary: Fix a couple bugs in load_into and add tests

Differential Revision: D70186266

Author: JacobSzwejbka

extension/flat_tensor/flat_tensor_data_map.cpp

@@ -156,7 +156,7 @@ ET_NODISCARD Result<size_t> FlatTensorDataMap::load_data_into(
     return tensor_layout.error();
   }
   ET_CHECK_OR_RETURN_ERROR(
-      size < tensor_layout.get().nbytes(),
+      size <= tensor_layout.get().nbytes(),
       InvalidArgument,
       "Buffer size %zu is smaller than tensor size %zu",
       size,
@@ -170,12 +170,16 @@ ET_NODISCARD Result<size_t> FlatTensorDataMap::load_data_into(
   // Load mutable data.
   DataLoader::SegmentInfo info = DataLoader::SegmentInfo(
       DataLoader::SegmentInfo::Type::Mutable, 0, nullptr);
-  return loader_->load_into(
+  auto e = loader_->load_into(
       header_.segment_base_offset + segment_offset.get() +
           metadata.get()->offset(),
       tensor_layout.get().nbytes(),
       info,
       buffer);
+  if (e != Error::Ok) {
+    return e;
+  }
+  return tensor_layout.get().nbytes();
 }
 
 ET_NODISCARD Result<size_t> FlatTensorDataMap::get_num_keys() const {
@@ -187,6 +191,7 @@ ET_NODISCARD Result<const char*> FlatTensorDataMap::get_key(
   if (index < 0 || index >= flat_tensor_->tensors()->size()) {
     return Error::InvalidArgument;
   }
+
   return flat_tensor_->tensors()->Get(index)->fully_qualified_name()->c_str();
 }
 

extension/flat_tensor/test/flat_tensor_data_map_test.cpp

@@ -137,3 +137,26 @@ TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_Keys) {
   Result<const char*> key2_res = data_map->get_key(2);
   EXPECT_EQ(key2_res.error(), Error::InvalidArgument);
 }
+
+TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_LoadInto) {
+  Result<FlatTensorDataMap> data_map =
+      FlatTensorDataMap::load(data_map_loader_.get());
+  EXPECT_EQ(data_map.error(), Error::Ok);
+
+    // get the metadata
+    auto meta_data_res = data_map->get_metadata("a");
+    ASSERT_EQ(meta_data_res.error(), Error::Ok);
+
+    // get data blob
+    void* data = malloc(meta_data_res->nbytes());
+    auto load_into_res =
+        data_map->load_data_into("a", data, meta_data_res->nbytes());
+    ASSERT_EQ(load_into_res.error(), Error::Ok);
+
+    // Check tensor data is correct.
+    float* data_a = static_cast<float*>(data);
+    for (int i = 0; i < 4; i++) {
+      EXPECT_EQ(data_a[i], 3.0);
+    }
+    free(data);
+}