Back out "Give DataLoader 64-bit source offsets" (pytorch#23360)

Summary:
Original commit changeset: 14d0a1991858

Original Phabricator Diff: D120705447

Differential Revision: D122670332

Co-authored-by: Anmol Gupta <agupta10@meta.com>

Author: cmt0

extension/data_loader/test/buffer_data_loader_test.cpp

@@ -8,10 +8,8 @@
 
 #include <executorch/extension/data_loader/buffer_data_loader.h>
 
-#include <array>
 #include <cstdint>
 #include <cstring>
-#include <limits>
 
 #include <gtest/gtest.h>
 
@@ -218,62 +216,3 @@ TEST_F(BufferDataLoaderTest, InBoundsLoadIntoSucceeds) {
     EXPECT_EQ(data[0], 1);
   }
 }
-
-TEST_F(BufferDataLoaderTest, WideInterfaceForwardsToLegacyMethods) {
-  std::array<uint8_t, 3> data{1, 2, 3};
-  BufferDataLoader loader(data.data(), data.size());
-  const DataLoader::SegmentInfo segment_info(
-      DataLoader::SegmentInfo::Type::Program);
-
-  Result<FreeableBuffer> loaded = loader.load_at_offset(1, 2, segment_info);
-  ASSERT_TRUE(loaded.ok());
-  const std::array<uint8_t, 2> expected{2, 3};
-  EXPECT_EQ(std::memcmp(loaded->data(), expected.data(), expected.size()), 0);
-
-  std::array<uint8_t, 2> destination{};
-  EXPECT_EQ(
-      loader.load_into_at_offset(
-          1, destination.size(), segment_info, destination.data()),
-      Error::Ok);
-  EXPECT_EQ(destination, expected);
-
-  Result<uint64_t> source_size = loader.source_size();
-  ASSERT_TRUE(source_size.ok());
-  EXPECT_EQ(source_size.get(), sizeof(data));
-}
-
-TEST_F(BufferDataLoaderTest, WideInterfaceRejectsUnrepresentableRange) {
-  std::array<uint8_t, 1> data{};
-  BufferDataLoader loader(data.data(), data.size());
-  const DataLoader::SegmentInfo segment_info(
-      DataLoader::SegmentInfo::Type::Program);
-
-  EXPECT_EQ(
-      loader
-          .load_at_offset(std::numeric_limits<uint64_t>::max(), 1, segment_info)
-          .error(),
-      Error::NotSupported);
-  EXPECT_EQ(
-      loader.load_into_at_offset(
-          std::numeric_limits<uint64_t>::max(), 1, segment_info, data.data()),
-      Error::NotSupported);
-}
-
-#if SIZE_MAX < UINT64_MAX
-TEST_F(BufferDataLoaderTest, WideInterfaceRejectsUnrepresentableOffsets) {
-  std::array<uint8_t, 1> data{};
-  BufferDataLoader loader(data.data(), data.size());
-  const DataLoader::SegmentInfo segment_info(
-      DataLoader::SegmentInfo::Type::Program);
-  constexpr uint64_t kUnrepresentableOffset =
-      static_cast<uint64_t>(SIZE_MAX) + 1;
-
-  EXPECT_EQ(
-      loader.load_at_offset(kUnrepresentableOffset, 0, segment_info).error(),
-      Error::NotSupported);
-  EXPECT_EQ(
-      loader.load_into_at_offset(
-          kUnrepresentableOffset, 0, segment_info, data.data()),
-      Error::NotSupported);
-}
-#endif

extension/flat_tensor/flat_tensor_data_map.cpp

@@ -22,7 +22,6 @@
 #include <executorch/runtime/platform/compiler.h>
 
 #include <cinttypes>
-#include <limits>
 
 using executorch::runtime::Error;
 using executorch::runtime::FreeableBuffer;
@@ -55,14 +54,11 @@ Result<const flat_tensor_flatbuffer::NamedData*> get_named_data(
         flatbuffers::Offset<flat_tensor_flatbuffer::NamedData>>* named_data,
     const flatbuffers::Vector<
         flatbuffers::Offset<flat_tensor_flatbuffer::DataSegment>>* segments,
-    uint64_t segment_data_size) {
+    uint64_t segment_end_offset) {
   // Linear search by name.
   if (named_data == nullptr) {
     return Error::NotFound;
   }
-  if (segments == nullptr) {
-    return Error::InvalidExternalData;
-  }
   for (flatbuffers::uoffset_t i = 0; i < named_data->size(); ++i) {
     if (key.size() == named_data->Get(i)->key()->size() &&
         std::strncmp(
@@ -71,12 +67,12 @@ Result<const flat_tensor_flatbuffer::NamedData*> get_named_data(
             named_data->Get(i)->key()->size()) == 0) {
       const auto* found = named_data->Get(i);
       // Validate the named_data.
-      const flatbuffers::uoffset_t segment_index = found->segment_index();
+      size_t segment_index = found->segment_index();
       ET_CHECK_OR_RETURN_ERROR(
-          segment_index < segments->size(),
+          segment_index >= 0 && segment_index < segments->size(),
           InvalidExternalData,
           "Segment index %zu for key %.*s is out of bounds for segment size %d. Malformed PTD file.",
-          static_cast<size_t>(segment_index),
+          segment_index,
           static_cast<int>(key.size()),
           key.data(),
           segments->size());
@@ -87,13 +83,13 @@ Result<const flat_tensor_flatbuffer::NamedData*> get_named_data(
               static_cast<uint64_t>(segments->Get(segment_index)->offset()),
               static_cast<uint64_t>(segments->Get(segment_index)->size()),
               &seg_end) &&
-              seg_end <= segment_data_size,
+              seg_end <= segment_end_offset,
           InvalidExternalData,
-          "Segment offset %" PRIu64 " + size %" PRIu64
-          " exceeds segment_data_size %" PRIu64 "; malformed PTD file.",
+          "Invalid segment offset %" PRIu64
+          " is larger than the segment_base_offset + segment_data_size %" PRIu64
+          "; malformed PTD file.",
           segments->Get(segment_index)->offset(),
-          segments->Get(segment_index)->size(),
-          segment_data_size);
+          segment_end_offset);
       return found;
     }
   }
@@ -115,21 +111,6 @@ Result<uint64_t> get_segment_end_offset(const FlatTensorHeader& header) {
   return segment_end_offset;
 }
 
-Result<uint64_t> get_absolute_segment_offset(
-    const FlatTensorHeader& header,
-    uint64_t segment_offset) {
-  uint64_t absolute_offset = 0;
-  ET_CHECK_OR_RETURN_ERROR(
-      !c10::add_overflows(
-          header.segment_base_offset, segment_offset, &absolute_offset),
-      InvalidExternalData,
-      "segment_base_offset %" PRIu64 " + segment offset %" PRIu64
-      " overflows uint64_t; malformed PTD file.",
-      header.segment_base_offset,
-      segment_offset);
-  return absolute_offset;
-}
-
 Result<const TensorLayout> create_tensor_layout(
     const flat_tensor_flatbuffer::TensorLayout* tensor_layout) {
   ScalarType scalar_type =
@@ -155,7 +136,7 @@ ET_NODISCARD Result<const TensorLayout> FlatTensorDataMap::get_tensor_layout(
       key,
       flat_tensor_->named_data(),
       flat_tensor_->segments(),
-      header_.segment_data_size);
+      segment_end_offset.get());
   if (!named_data.ok()) {
     return named_data.error();
   }
@@ -172,7 +153,7 @@ ET_NODISCARD Result<FreeableBuffer> FlatTensorDataMap::get_data(
       key,
       flat_tensor_->named_data(),
       flat_tensor_->segments(),
-      header_.segment_data_size);
+      segment_end_offset.get());
   if (!named_data.ok()) {
     return named_data.error();
   }
@@ -182,21 +163,9 @@ ET_NODISCARD Result<FreeableBuffer> FlatTensorDataMap::get_data(
       flat_tensor_->segments()->Get(segment_index)->offset();
   uint64_t segment_size = flat_tensor_->segments()->Get(segment_index)->size();
 
-  Result<uint64_t> absolute_offset =
-      get_absolute_segment_offset(header_, segment_offset);
-  if (!absolute_offset.ok()) {
-    return absolute_offset.error();
-  }
-  ET_CHECK_OR_RETURN_ERROR(
-      segment_size <= std::numeric_limits<size_t>::max(),
-      NotSupported,
-      "Segment size %" PRIu64 " exceeds the maximum load size %zu",
+  return loader_->load(
+      /*offset=*/header_.segment_base_offset + segment_offset,
       segment_size,
-      std::numeric_limits<size_t>::max());
-
-  return loader_->load_at_offset(
-      absolute_offset.get(),
-      static_cast<size_t>(segment_size),
       DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Constant));
 }
 
@@ -212,21 +181,14 @@ ET_NODISCARD Error FlatTensorDataMap::load_data_into(
       key,
       flat_tensor_->named_data(),
       flat_tensor_->segments(),
-      header_.segment_data_size);
+      segment_end_offset.get());
   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();
-  uint64_t segment_size = flat_tensor_->segments()->Get(segment_index)->size();
-
-  Result<uint64_t> absolute_offset =
-      get_absolute_segment_offset(header_, segment_offset);
-  if (!absolute_offset.ok()) {
-    return absolute_offset.error();
-  }
 
   Result<const TensorLayout> tensor_layout =
       create_tensor_layout(named_data.get()->tensor_layout());
@@ -238,21 +200,18 @@ ET_NODISCARD Error FlatTensorDataMap::load_data_into(
   ET_CHECK_OR_RETURN_ERROR(
       size <= tensor_layout.get().nbytes(),
       InvalidArgument,
-      "Requested size %zu exceeds tensor size %zu",
+      "Buffer size %zu is smaller than tensor size %zu",
       size,
       tensor_layout.get().nbytes());
-  ET_CHECK_OR_RETURN_ERROR(
-      static_cast<uint64_t>(size) <= segment_size,
-      InvalidExternalData,
-      "Requested size %zu exceeds segment size %" PRIu64,
-      size,
-      segment_size);
 
   // Load mutable data.
   DataLoader::SegmentInfo info = DataLoader::SegmentInfo(
       DataLoader::SegmentInfo::Type::Mutable, 0, nullptr);
-  return loader_->load_into_at_offset(
-      absolute_offset.get(), size, info, buffer);
+  return loader_->load_into(
+      header_.segment_base_offset + segment_offset,
+      tensor_layout.get().nbytes(),
+      info,
+      buffer);
 }
 
 ET_NODISCARD Result<uint32_t> FlatTensorDataMap::get_num_keys() const {
@@ -274,7 +233,7 @@ ET_NODISCARD Result<const char*> FlatTensorDataMap::get_key(
 /* static */ Result<FlatTensorDataMap> FlatTensorDataMap::load(
     DataLoader* loader) {
   // Check header.
-  Result<FreeableBuffer> header = loader->load_at_offset(
+  Result<FreeableBuffer> header = loader->load(
       /*offset=*/0,
       FlatTensorHeader::kNumHeadBytes,
       DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program));
@@ -291,41 +250,19 @@ ET_NODISCARD Result<const char*> FlatTensorDataMap::get_key(
       "Failed to parse FlatTensor header with error code %u. File may be corrupt.",
       static_cast<uint32_t>(fh.error()));
 
-  Result<uint64_t> expected_size = get_segment_end_offset(fh.get());
-  if (!expected_size.ok()) {
-    return expected_size.error();
-  }
-  Result<uint64_t> actual_size = loader->source_size();
-  if (!actual_size.ok()) {
-    return actual_size.error();
-  }
+  size_t expected_size = fh->segment_base_offset + fh->segment_data_size;
+  size_t actual_size = loader->size().get();
   ET_CHECK_OR_RETURN_ERROR(
-      expected_size.get() <= actual_size.get(),
+      expected_size <= actual_size,
       InvalidExternalData,
-      "File size is too small; file may be corrupted or truncated. Expected %" PRIu64
-      " from flat_tensor header, received %" PRIu64 " from data loader",
-      expected_size.get(),
-      actual_size.get());
-
-  uint64_t flat_tensor_data_size = 0;
-  ET_CHECK_OR_RETURN_ERROR(
-      !c10::add_overflows(
-          fh->flatbuffer_offset, fh->flatbuffer_size, &flat_tensor_data_size),
-      InvalidExternalData,
-      "flatbuffer_offset %" PRIu64 " + flatbuffer_size %" PRIu64
-      " overflows uint64_t; malformed PTD file.",
-      fh->flatbuffer_offset,
-      fh->flatbuffer_size);
-  ET_CHECK_OR_RETURN_ERROR(
-      flat_tensor_data_size <= std::numeric_limits<size_t>::max(),
-      NotSupported,
-      "FlatTensor metadata size exceeds the addressable buffer size %zu",
-      std::numeric_limits<size_t>::max());
+      "File size is too small; file may be corrupted or truncated. Expected %zu from flat_tensor header, received %zu from data loader",
+      expected_size,
+      actual_size);
 
   // Load flatbuffer data as a segment.
-  Result<FreeableBuffer> flat_tensor_data = loader->load_at_offset(
+  Result<FreeableBuffer> flat_tensor_data = loader->load(
       /*offset=*/0,
-      static_cast<size_t>(flat_tensor_data_size),
+      fh->flatbuffer_offset + fh->flatbuffer_size,
       DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program));
   if (!flat_tensor_data.ok()) {
     ET_LOG(Error, "Failed to load flat_tensor data.");