MmapDataLoader can handle non-page-aligned offsets (#1975)

Summary:
Pull Request resolved: #1975

Instead of requiring that load offsets are aligned to the size of the OS page, map in the set of pages that overlap the requested offset + size.

This will let us load .pte files with 4096-aligned segments onto iOS devices, which use 16384-byte OS pages.

It also means that we could reduce the segment alignment to `max(delegate_alignment, constant_tensor_alignment)` to reduce the size of .pte files.

Reviewed By: shoumikhin

Differential Revision: D53772818

fbshipit-source-id: d08cef5345d1abaad060a6b40cd0db18a7f1b94b

Author: dbort

exir/capture/_config.py

@@ -62,8 +62,7 @@ class ExecutorchBackendConfig:
     extract_constant_segment: bool = True
 
     # When extracting segments, the starting offset of each segment will be
-    # aligned to this value (in bytes). When using mmap() to load segments, this
-    # should be a multiple of the OS page size.
+    # aligned to this value (in bytes). Must be a power of two.
     segment_alignment: int = 4096
 
     # If provided, the minimum alignment of tensor buffers in the program. Must

extension/data_loader/mmap_data_loader.cpp

@@ -26,6 +26,34 @@ namespace torch {
 namespace executor {
 namespace util {
 
+namespace {
+
+struct Range {
+  // Address or offset.
+  uintptr_t start;
+  // Size in bytes.
+  size_t size;
+};
+
+/**
+ * Given an address region, returns the start offset and byte size of the set of
+ * pages that completely covers the region.
+ */
+Range get_overlapping_pages(uintptr_t offset, size_t size, size_t page_size) {
+  size_t page_mask = ~(page_size - 1);
+  // The address of the page that starts at or before the beginning of the
+  // region.
+  uintptr_t start = offset & page_mask;
+  // The address of the page that starts after the end of the region.
+  uintptr_t end = (offset + size + ~page_mask) & page_mask;
+  return {
+      /*start=*/start,
+      /*size=*/static_cast<size_t>(end - start),
+  };
+}
+
+} // namespace
+
 MmapDataLoader::~MmapDataLoader() {
   // file_name_ can be nullptr if this instance was moved from, but freeing a
   // null pointer is safe.
@@ -41,7 +69,7 @@ Result<MmapDataLoader> MmapDataLoader::from(
   // Cache the page size.
   long page_size = sysconf(_SC_PAGESIZE);
   if (page_size < 0) {
-    ET_LOG(Error, "Could not get page size: %s (%d)", strerror(errno), errno);
+    ET_LOG(Error, "Could not get page size: %s (%d)", ::strerror(errno), errno);
     return Error::AccessFailed;
   }
   if ((page_size & ~(page_size - 1)) != page_size) {
@@ -53,7 +81,11 @@ Result<MmapDataLoader> MmapDataLoader::from(
   int fd = ::open(file_name, O_RDONLY);
   if (fd < 0) {
     ET_LOG(
-        Error, "Failed to open %s: %s (%d)", file_name, strerror(errno), errno);
+        Error,
+        "Failed to open %s: %s (%d)",
+        file_name,
+        ::strerror(errno),
+        errno);
     return Error::AccessFailed;
   }
 
@@ -89,9 +121,28 @@ Result<MmapDataLoader> MmapDataLoader::from(
 }
 
 namespace {
-/// FreeableBuffer::FreeFn-compatible callback.
-void MunmapSegment(__ET_UNUSED void* context, void* data, size_t size) {
-  ::munmap(data, size);
+/**
+ * FreeableBuffer::FreeFn-compatible callback.
+ *
+ * `context` is actually the OS page size as a uintptr_t.
+ */
+void MunmapSegment(void* context, void* data, size_t size) {
+  const uintptr_t page_size = reinterpret_cast<uintptr_t>(context);
+
+  Range range =
+      get_overlapping_pages(reinterpret_cast<uintptr_t>(data), size, page_size);
+  int ret = ::munmap(reinterpret_cast<void*>(range.start), range.size);
+  if (ret < 0) {
+    // Let the user know that something went wrong, but there's nothing we can
+    // do about it.
+    ET_LOG(
+        Error,
+        "munmap(0x%zx, %zu) failed: %s (ignored)",
+        (size_t)range.start,
+        range.size,
+        ::strerror(errno),
+        errno);
+  }
 }
 } // namespace
 
@@ -109,13 +160,6 @@ Result<FreeableBuffer> MmapDataLoader::Load(size_t offset, size_t size) {
       offset,
       size,
       file_size_);
-  ET_CHECK_OR_RETURN_ERROR(
-      (offset & ~(page_size_ - 1)) == offset,
-      InvalidArgument,
-      "File %s: offset 0x%zx not aligned to 0x%zx",
-      file_name_,
-      offset,
-      page_size_);
   ET_CHECK_OR_RETURN_ERROR(
       // Recommended by a lint warning.
       offset <= std::numeric_limits<off_t>::max(),
@@ -128,43 +172,64 @@ Result<FreeableBuffer> MmapDataLoader::Load(size_t offset, size_t size) {
     return FreeableBuffer(nullptr, 0, /*free_fn=*/nullptr);
   }
 
+  // Find the range of pages that covers the requested region.
+  Range range =
+      get_overlapping_pages(static_cast<uintptr_t>(offset), size, page_size_);
+
   // Map the pages read-only. MAP_PRIVATE vs. MAP_SHARED doesn't matter since
   // the data is read-only, but use PRIVATE just to further avoid accidentally
   // modifying the file.
-  void* pages = mmap(
-      nullptr, size, PROT_READ, MAP_PRIVATE, fd_, static_cast<off_t>(offset));
+  void* pages = ::mmap(
+      nullptr,
+      range.size,
+      PROT_READ,
+      MAP_PRIVATE,
+      fd_,
+      static_cast<off_t>(range.start));
   ET_CHECK_OR_RETURN_ERROR(
       pages != MAP_FAILED,
       AccessFailed,
       "Failed to map %s: mmap(..., size=%zd, ..., fd=%d, offset=0x%zx)",
       file_name_,
-      size,
+      range.size,
       fd_,
-      offset);
+      range.start);
 
   if (mlock_config_ == MlockConfig::UseMlock ||
       mlock_config_ == MlockConfig::UseMlockIgnoreErrors) {
-    int err = mlock(pages, size);
+    int err = ::mlock(pages, size);
     if (err < 0) {
       ET_LOG(
           Error,
           "File %s: mlock(%p, %zu) failed: %s (%d)",
           file_name_,
           pages,
           size,
-          strerror(errno),
+          ::strerror(errno),
           errno);
       if (mlock_config_ == MlockConfig::UseMlockIgnoreErrors) {
         ET_LOG(Info, "Ignoring mlock() error");
       } else {
-        munmap(pages, size);
+        ::munmap(pages, size);
         return Error::NotSupported;
       }
     }
     // No need to keep track of this. munmap() will unlock as a side effect.
   }
 
-  return FreeableBuffer(pages, size, MunmapSegment);
+  // The requested data is at an offset into the mapped pages.
+  const void* data = static_cast<const uint8_t*>(pages) + offset - range.start;
+
+  return FreeableBuffer(
+      // The callback knows to unmap the whole pages that encompass this region.
+      data,
+      size,
+      MunmapSegment,
+      /*free_fn_context=*/
+      reinterpret_cast<void*>(
+          // Pass the cached OS page size to the callback so it doesn't need to
+          // query it again.
+          static_cast<uintptr_t>(page_size_)));
 }
 
 Result<size_t> MmapDataLoader::size() const {

extension/data_loader/test/mmap_data_loader_test.cpp

@@ -153,6 +153,42 @@ void MmapDataLoaderTest::test_in_bounds_loads_succeed(
     EXPECT_EQ(fb->size(), size);
     EXPECT_EQ(0, std::memcmp(fb->data(), &contents[offset], fb->size()));
   }
+
+  //
+  // Unaligned offsets and sizes
+  //
+
+  // Load a single, partial page with an offset that is not a multiple of the
+  // page size.
+  {
+    const size_t offset = 128; // Small power of 2
+    EXPECT_LT(offset, page_size_);
+    const size_t size = page_size_ / 2;
+    Result<FreeableBuffer> fb = mdl->Load(offset, size);
+    ASSERT_EQ(fb.error(), Error::Ok);
+    EXPECT_EQ(fb->size(), size);
+    EXPECT_EQ(0, std::memcmp(fb->data(), &contents[offset], fb->size()));
+  }
+
+  // Load multiple pages from a non-page-aligned but power-of-two offset.
+  {
+    const size_t offset = page_size_ + 128; // Small power of 2
+    const size_t size = page_size_ * 3 + page_size_ / 2 + 1; // Odd size
+    Result<FreeableBuffer> fb = mdl->Load(offset, size);
+    ASSERT_EQ(fb.error(), Error::Ok);
+    EXPECT_EQ(fb->size(), size);
+    EXPECT_EQ(0, std::memcmp(fb->data(), &contents[offset], fb->size()));
+  }
+
+  // Load multiple pages from an offset that is not a power of 2.
+  {
+    const size_t offset = page_size_ * 2 + 3; // Not a power of 2
+    const size_t size = page_size_ * 3 + page_size_ / 2 + 1; // Odd size
+    Result<FreeableBuffer> fb = mdl->Load(offset, size);
+    ASSERT_EQ(fb.error(), Error::Ok);
+    EXPECT_EQ(fb->size(), size);
+    EXPECT_EQ(0, std::memcmp(fb->data(), &contents[offset], fb->size()));
+  }
 }
 
 TEST_F(MmapDataLoaderTest, InBoundsLoadsSucceedNoMlock) {
@@ -239,33 +275,6 @@ TEST_F(MmapDataLoaderTest, OutOfBoundsLoadFails) {
   }
 }
 
-TEST_F(MmapDataLoaderTest, UnalignedOffsetFails) {
-  // Create a multi-page file; contents don't matter.
-  const size_t contents_size = 8 * page_size_;
-  auto contents = std::make_unique<uint8_t[]>(contents_size);
-  memset(contents.get(), 0x55, contents_size);
-  TempFile tf(contents.get(), contents_size);
-
-  Result<MmapDataLoader> mdl = MmapDataLoader::from(tf.path().c_str());
-  ASSERT_EQ(mdl.error(), Error::Ok);
-
-  // Loading from an unaligned offset should fail, even if the offset is
-  // a power of 2 and size is a whole page.
-  {
-    Result<FreeableBuffer> fb =
-        mdl->Load(/*offset=*/page_size_ / 2, /*size=*/page_size_);
-    EXPECT_NE(fb.error(), Error::Ok);
-  }
-
-  // Loading from an unaligned offset should fail, even if the offset is
-  // a power of 2 and size ends at the next page.
-  {
-    Result<FreeableBuffer> fb =
-        mdl->Load(/*offset=*/page_size_ / 2, /*size=*/page_size_ / 2);
-    EXPECT_NE(fb.error(), Error::Ok);
-  }
-}
-
 TEST_F(MmapDataLoaderTest, FromMissingFileFails) {
   // Wrapping a file that doesn't exist should fail.
   Result<MmapDataLoader> mdl = MmapDataLoader::from(