[libc-commits] [libc] [libc] Fix malloc Block alignment issue on riscv32 (PR #117815)

via libc-commits libc-commits at lists.llvm.org
Tue Nov 26 16:32:01 PST 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-libc

Author: Daniel Thornburgh (mysterymath)

<details>
<summary>Changes</summary>

Aligning blocks to max_align_t is neither necessary nor sufficient to ensure that the usable_space() is so aligned. Instead, we make this an invariant of Block and maintain it in init() and split().

This allows targets like riscv32 with small pointers and large max_align_t to maintain the property that all available blocks are aligned for malloc(). This change adjusts the tests to match and also updates them closer to llvm-libc style.

---

Patch is 53.15 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/117815.diff


7 Files Affected:

- (modified) libc/fuzzing/__support/freelist_heap_fuzz.cpp (+3-3) 
- (modified) libc/src/__support/CMakeLists.txt (+1) 
- (modified) libc/src/__support/block.h (+127-124) 
- (modified) libc/src/__support/freelist_heap.h (+6-17) 
- (modified) libc/test/src/__support/block_test.cpp (+128-193) 
- (modified) libc/test/src/__support/freelist_heap_test.cpp (+40-66) 
- (modified) libc/test/src/__support/freelist_malloc_test.cpp (+8-13) 


``````````diff
diff --git a/libc/fuzzing/__support/freelist_heap_fuzz.cpp b/libc/fuzzing/__support/freelist_heap_fuzz.cpp
index d152d0f35499f8..d5764adcbb2642 100644
--- a/libc/fuzzing/__support/freelist_heap_fuzz.cpp
+++ b/libc/fuzzing/__support/freelist_heap_fuzz.cpp
@@ -100,13 +100,13 @@ optional<AllocType> choose<AllocType>(const uint8_t *&data, size_t &remainder) {
       *raw % static_cast<uint8_t>(AllocType::NUM_ALLOC_TYPES));
 }
 
-constexpr size_t heap_size = 64 * 1024;
+constexpr size_t HEAP_SIZE = 64 * 1024;
 
 optional<size_t> choose_size(const uint8_t *&data, size_t &remainder) {
   auto raw = choose<size_t>(data, remainder);
   if (!raw)
     return nullopt;
-  return *raw % heap_size;
+  return *raw % HEAP_SIZE;
 }
 
 optional<size_t> choose_alloc_idx(const AllocVec &allocs, const uint8_t *&data,
@@ -126,7 +126,7 @@ optional<size_t> choose_alloc_idx(const AllocVec &allocs, const uint8_t *&data,
   TYPE NAME = *maybe_##NAME
 
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t remainder) {
-  LIBC_NAMESPACE::FreeListHeapBuffer<heap_size> heap;
+  LIBC_NAMESPACE::FreeListHeapBuffer<HEAP_SIZE> heap;
   AllocVec allocs(heap);
 
   uint8_t canary = 0;
diff --git a/libc/src/__support/CMakeLists.txt b/libc/src/__support/CMakeLists.txt
index 8f85740f70a06e..f11ea3e84af559 100644
--- a/libc/src/__support/CMakeLists.txt
+++ b/libc/src/__support/CMakeLists.txt
@@ -12,6 +12,7 @@ add_header_library(
     libc.src.__support.CPP.optional
     libc.src.__support.CPP.span
     libc.src.__support.CPP.type_traits
+    libc.src.__support.math_extras
 )
 
 add_object_library(
diff --git a/libc/src/__support/block.h b/libc/src/__support/block.h
index 9ca3f11530c4ba..4c51f26cf4776a 100644
--- a/libc/src/__support/block.h
+++ b/libc/src/__support/block.h
@@ -18,6 +18,7 @@
 #include "src/__support/CPP/type_traits.h"
 #include "src/__support/libc_assert.h"
 #include "src/__support/macros/config.h"
+#include "src/__support/math_extras.h"
 
 #include <stdint.h>
 
@@ -49,8 +50,7 @@ LIBC_INLINE constexpr T *align_down(T *value, size_t alignment) {
 
 /// Returns the value rounded up to the nearest multiple of alignment.
 LIBC_INLINE constexpr size_t align_up(size_t value, size_t alignment) {
-  __builtin_add_overflow(value, alignment - 1, &value);
-  return align_down(value, alignment);
+  return align_down(value + alignment - 1, alignment);
 }
 
 /// Returns the value rounded up to the nearest multiple of alignment.
@@ -69,7 +69,9 @@ using cpp::optional;
 /// is also the block's size.
 ///
 /// Blocks will always be aligned to a `ALIGNMENT` boundary. Block sizes will
-/// always be rounded up to a multiple of `ALIGNMENT`.
+/// always be rounded up to a multiple of `ALIGNMENT`. All blocks have their
+/// usable space aligned to some multiple of max_align_t. This also implies that
+/// block outer sizes are aligned to max_align_t.
 ///
 /// As an example, the diagram below represents two contiguous `Block`s. The
 /// indices indicate byte offsets:
@@ -122,15 +124,16 @@ class Block {
   static constexpr size_t SIZE_MASK = ~(PREV_FREE_MASK | LAST_MASK);
 
 public:
-  static constexpr size_t ALIGNMENT = cpp::max(alignof(max_align_t), size_t{4});
+  static constexpr size_t ALIGNMENT = alignof(size_t);
   static const size_t BLOCK_OVERHEAD;
 
   // No copy or move.
   Block(const Block &other) = delete;
   Block &operator=(const Block &other) = delete;
 
-  /// Creates the first block for a given memory region, followed by a sentinel
-  /// last block. Returns the first block.
+  /// Initializes a given memory region into a first block and a sentinel last
+  /// block. Returns the first block, which has its usable space aligned to
+  /// max_align_t.
   static optional<Block *> init(ByteSpan region);
 
   /// @returns  A pointer to a `Block`, given a pointer to the start of the
@@ -186,6 +189,9 @@ class Block {
   }
 
   /// @returns A pointer to the usable space inside this block.
+  ///
+  /// Unless specifically requested otherwise, this will be aligned to
+  /// max_align_t.
   LIBC_INLINE cpp::byte *usable_space() {
     return reinterpret_cast<cpp::byte *>(this) + BLOCK_OVERHEAD;
   }
@@ -201,11 +207,13 @@ class Block {
   /// Attempts to split this block.
   ///
   /// If successful, the block will have an inner size of at least
-  /// `new_inner_size`, rounded to ensure that the split point is on an
-  /// ALIGNMENT boundary. The remaining space will be returned as a new block.
-  /// Note that the prev_ field of the next block counts as part of the inner
-  /// size of the returnd block.
-  optional<Block *> split(size_t new_inner_size);
+  /// `new_inner_size`. The remaining space will be returned as a new block,
+  /// with usable space aligned to `usable_space_alignment` and with the block
+  /// boundary aligned to ALIGNMENT. Note that the prev_ field of the next block
+  /// counts as part of the inner size of the block. `usable_space_alignment`
+  /// must be a multiple of max_align_t.
+  optional<Block *> split(size_t new_inner_size,
+                          size_t usable_space_alignment = alignof(max_align_t));
 
   /// Merges this block with the one that comes after it.
   bool merge_next();
@@ -256,38 +264,39 @@ class Block {
     return reinterpret_cast<uintptr_t>(usable_space()) % alignment == 0;
   }
 
-  /// @returns The new inner size of this block that would give the usable
-  /// space of the next block the given alignment.
-  LIBC_INLINE size_t padding_for_alignment(size_t alignment) const {
-    if (is_usable_space_aligned(alignment))
-      return 0;
-
-    // We need to ensure we can always split this block into a "padding" block
-    // and the aligned block. To do this, we need enough extra space for at
-    // least one block.
-    //
-    // |block   |usable_space                          |
-    // |........|......................................|
-    //                            ^
-    //                            Alignment requirement
-    //
+  // Returns the minimum inner size necessary for a block of that size to
+  // always be able to allocate at the given size and alignment.
+  //
+  // Returns 0 if there is no such size.
+  LIBC_INLINE static size_t min_size_for_allocation(size_t alignment,
+                                                    size_t size) {
+    LIBC_ASSERT(alignment >= alignof(max_align_t) &&
+                alignment % alignof(max_align_t) == 0 &&
+                "alignment must be multiple of max_align_t");
+
+    if (alignment == alignof(max_align_t))
+      return size;
+
+    // We must create a padding block to align the usable space of the next
+    // block. The next block's usable space can be found by advancing by
+    // BLOCK_OVERHEAD then aligning up. The amount advanced is the amount of
+    // additional inner size required.
     //
-    // |block   |space   |block   |usable_space        |
-    // |........|........|........|....................|
-    //                            ^
-    //                            Alignment requirement
-    //
-    alignment = cpp::max(alignment, ALIGNMENT);
-    uintptr_t start = reinterpret_cast<uintptr_t>(usable_space());
-    uintptr_t next_usable_space = align_up(start + BLOCK_OVERHEAD, alignment);
-    uintptr_t next_block = next_usable_space - BLOCK_OVERHEAD;
-    return next_block - start + sizeof(prev_);
+    // The minimum advancment is BLOCK_OVERHEAD, since the resulting position
+    // may happen to be aligned. What is the maximum? Advancing by
+    // BLOCK_OVERHEAD leaves the result still aligned to max_align_t. So the
+    // most additional advancement required would be if the point is exactly
+    // alignof(max_align_t) past alignment. The remaining size to the next
+    // alignment would then be alignment - alignof(max_align_t). So the total
+    // maximum advancement required is BLOCK_OVERHEAD + alignment -
+    // alignof(max_align_t).
+    if (add_overflow(size, BLOCK_OVERHEAD, size))
+      return 0;
+    if (add_overflow(size, alignment - alignof(max_align_t), size))
+      return 0;
+    return size;
   }
 
-  // Check that we can `allocate` a block with a given alignment and size from
-  // this existing block.
-  bool can_allocate(size_t alignment, size_t size) const;
-
   // This is the return type for `allocate` which can split one block into up to
   // three blocks.
   struct BlockInfo {
@@ -309,25 +318,23 @@ class Block {
     Block *next;
   };
 
-  // Divide a block into up to 3 blocks according to `BlockInfo`. This should
-  // only be called if `can_allocate` returns true.
+  // Divide a block into up to 3 blocks according to `BlockInfo`. Behavior is
+  // undefined if allocation is not possible for the given size and alignment.
   static BlockInfo allocate(Block *block, size_t alignment, size_t size);
 
 private:
   /// Construct a block to represent a span of bytes. Overwrites only enough
   /// memory for the block header; the rest of the span is left alone.
   LIBC_INLINE static Block *as_block(ByteSpan bytes) {
+    LIBC_ASSERT(reinterpret_cast<uintptr_t>(bytes.data()) % ALIGNMENT == 0 &&
+                "block start must be suitably aligned");
     return ::new (bytes.data()) Block(bytes.size());
   }
 
-  /// Like `split`, but assumes the caller has already checked to parameters to
-  /// ensure the split will succeed.
-  Block *split_impl(size_t new_inner_size);
-
   /// Offset from this block to the previous block. 0 if this is the first
   /// block. This field is only alive when the previous block is free;
   /// otherwise, its memory is reused as part of the previous block's usable
-  /// space.
+  /// space
   size_t prev_ = 0;
 
   /// Offset from this block to the next block. Valid even if this is the last
@@ -343,93 +350,78 @@ class Block {
   ///   previous block is free.
   /// * If the `last` flag is set, the block is the sentinel last block. It is
   ///   summarily considered used and has no next block.
-} __attribute__((packed, aligned(cpp::max(alignof(max_align_t), size_t{4}))));
-
-inline constexpr size_t Block::BLOCK_OVERHEAD =
-    align_up(sizeof(Block), ALIGNMENT);
-
-LIBC_INLINE ByteSpan get_aligned_subspan(ByteSpan bytes, size_t alignment) {
-  if (bytes.data() == nullptr)
-    return ByteSpan();
+};
 
-  auto unaligned_start = reinterpret_cast<uintptr_t>(bytes.data());
-  auto aligned_start = align_up(unaligned_start, alignment);
-  auto unaligned_end = unaligned_start + bytes.size();
-  auto aligned_end = align_down(unaligned_end, alignment);
+static_assert(alignof(max_align_t) > 4,
+              "the low two bits of block sizes must be free for use as flags");
 
-  if (aligned_end <= aligned_start)
-    return ByteSpan();
-
-  return bytes.subspan(aligned_start - unaligned_start,
-                       aligned_end - aligned_start);
-}
+inline constexpr size_t Block::BLOCK_OVERHEAD = sizeof(Block);
 
 LIBC_INLINE
 optional<Block *> Block::init(ByteSpan region) {
-  optional<ByteSpan> result = get_aligned_subspan(region, ALIGNMENT);
-  if (!result)
+  if (!region.data())
     return {};
 
-  region = result.value();
-  // Two blocks are allocated: a free block and a sentinel last block.
-  if (region.size() < 2 * BLOCK_OVERHEAD)
-    return {};
+  uintptr_t start = reinterpret_cast<uintptr_t>(region.data());
+  uintptr_t end = start + region.size();
+
+  // Find the earliest point after start that is aligned to max_align_t and has
+  // at least BLOCK_OVERHEAD before it.
+  uintptr_t usable_space =
+      align_up(start + BLOCK_OVERHEAD, alignof(max_align_t));
+
+  uintptr_t block_start = usable_space - BLOCK_OVERHEAD;
+  LIBC_ASSERT(block_start % ALIGNMENT == 0 && "block start must be aligned");
 
-  if (cpp::numeric_limits<size_t>::max() < region.size())
+  // A last block must be able to follow the first block.
+  if (usable_space + align_up(BLOCK_OVERHEAD, alignof(max_align_t)) > end)
     return {};
 
-  Block *block = as_block(region.first(region.size() - BLOCK_OVERHEAD));
-  Block *last = as_block(region.last(BLOCK_OVERHEAD));
+  // Given that there is room, the last block starts BLOCK_OVERHEAD
+  // before the last aligned point in the region.
+  uintptr_t last_start = align_down(end, alignof(max_align_t)) - BLOCK_OVERHEAD;
+  LIBC_ASSERT(last_start % ALIGNMENT == 0 &&
+              "last block start must be aligned");
+
+  auto *block_start_ptr = reinterpret_cast<cpp::byte *>(block_start);
+  auto *last_start_ptr = reinterpret_cast<cpp::byte *>(last_start);
+
+  Block *block = as_block({block_start_ptr, last_start_ptr});
+  LIBC_ASSERT(block->is_usable_space_aligned(alignof(max_align_t)) &&
+              "newly initialized block should be aligned to max_align_t");
+  Block *last = as_block({last_start_ptr, BLOCK_OVERHEAD});
+  LIBC_ASSERT(last->is_usable_space_aligned(alignof(max_align_t)) &&
+              "last block should be aligned to max_align_t");
   block->mark_free();
   last->mark_last();
+  LIBC_ASSERT(block->outer_size() % alignof(max_align_t) == 0 &&
+              "outer sizes should be aligned to max_align_t");
   return block;
 }
 
-LIBC_INLINE
-bool Block::can_allocate(size_t alignment, size_t size) const {
-  if (inner_size() < size)
-    return false;
-  if (is_usable_space_aligned(alignment))
-    return true;
-
-  // Alignment isn't met, so a padding block is needed. Determine amount of
-  // inner_size() consumed by the padding block.
-  size_t padding_size = padding_for_alignment(alignment) - sizeof(prev_);
-
-  // Check that there is room for the allocation in the following aligned block.
-  size_t aligned_inner_size = inner_size() - padding_size - BLOCK_OVERHEAD;
-  return size <= aligned_inner_size;
-}
-
 LIBC_INLINE
 Block::BlockInfo Block::allocate(Block *block, size_t alignment, size_t size) {
-  LIBC_ASSERT(
-      block->can_allocate(alignment, size) &&
-      "Calls to this function for a given alignment and size should only be "
-      "done if `can_allocate` for these parameters returns true.");
+  LIBC_ASSERT(alignment % alignof(max_align_t) == 0 &&
+              "alignment must be a multiple of max_align_t");
 
   BlockInfo info{block, /*prev=*/nullptr, /*next=*/nullptr};
 
   if (!info.block->is_usable_space_aligned(alignment)) {
     Block *original = info.block;
-    optional<Block *> maybe_aligned_block =
-        original->split(info.block->padding_for_alignment(alignment));
+    // The padding block has no minimum size requirement.
+    optional<Block *> maybe_aligned_block = original->split(0, alignment);
     LIBC_ASSERT(maybe_aligned_block.has_value() &&
-                "This split should always result in a new block. The check in "
-                "`can_allocate` ensures that we have enough space here to make "
-                "two blocks.");
+                "it should always be possible to split for alignment");
 
     if (Block *prev = original->prev_free()) {
-      // If there is a free block before this, we can merge the current one with
-      // the newly created one.
+      // If there is a free block before this, we can merge the current one
+      // with the newly created one.
       prev->merge_next();
     } else {
       info.prev = original;
     }
 
     Block *aligned_block = *maybe_aligned_block;
-    LIBC_ASSERT(aligned_block->is_usable_space_aligned(alignment) &&
-                "The aligned block isn't aligned somehow.");
     info.block = aligned_block;
   }
 
@@ -441,33 +433,44 @@ Block::BlockInfo Block::allocate(Block *block, size_t alignment, size_t size) {
 }
 
 LIBC_INLINE
-optional<Block *> Block::split(size_t new_inner_size) {
+optional<Block *> Block::split(size_t new_inner_size,
+                               size_t usable_space_alignment) {
   if (used())
     return {};
-  // The prev_ field of the next block is always available, so there is a
-  // minimum size to a block created through splitting.
-  if (new_inner_size < sizeof(prev_))
-    new_inner_size = sizeof(prev_);
-
-  size_t old_inner_size = inner_size();
-  new_inner_size =
-      align_up(new_inner_size - sizeof(prev_), ALIGNMENT) + sizeof(prev_);
-  if (old_inner_size < new_inner_size)
-    return {};
 
-  if (old_inner_size - new_inner_size < BLOCK_OVERHEAD)
-    return {};
+  LIBC_ASSERT(usable_space_alignment % alignof(max_align_t) == 0 &&
+              "alignment must be a multiple of max_align_t");
 
-  return split_impl(new_inner_size);
-}
+  size_t old_outer_size = outer_size();
 
-LIBC_INLINE
-Block *Block::split_impl(size_t new_inner_size) {
-  size_t outer_size1 = outer_size(new_inner_size);
-  LIBC_ASSERT(outer_size1 % ALIGNMENT == 0 && "new size must be aligned");
-  ByteSpan new_region = region().subspan(outer_size1);
+  // Compute the minimum outer size that produces a block of at least
+  // new_inner_size.
+  new_inner_size = cpp::max(new_inner_size, sizeof(prev_));
+  size_t new_outer_size = outer_size(new_inner_size);
+
+  // Compute the minimum start of the next block's usable space.
+  uintptr_t start = reinterpret_cast<uintptr_t>(this);
+  uintptr_t next_usable_space = start + new_outer_size + BLOCK_OVERHEAD;
+
+  // Align up the next block's usable space to the requested alignment.
+  next_usable_space = align_up(next_usable_space, usable_space_alignment);
+
+  // Find the starting point of the next block.
+  uintptr_t next_block_start = next_usable_space - BLOCK_OVERHEAD;
+  LIBC_ASSERT(next_block_start % ALIGNMENT == 0 &&
+              "next block must be suitably aligned");
+
+  new_outer_size = next_block_start - start;
+  LIBC_ASSERT(new_outer_size % alignof(max_align_t) == 0 &&
+              "new size must be aligned to max_align_t");
+
+  if (old_outer_size < new_outer_size ||
+      old_outer_size - new_outer_size < BLOCK_OVERHEAD)
+    return {};
+
+  ByteSpan new_region = region().subspan(new_outer_size);
   next_ &= ~SIZE_MASK;
-  next_ |= outer_size1;
+  next_ |= new_outer_size;
 
   Block *new_block = as_block(new_region);
   mark_free(); // Free status for this block is now stored in new_block.
diff --git a/libc/src/__support/freelist_heap.h b/libc/src/__support/freelist_heap.h
index 8fa36257cb91ae..d58685194aeb81 100644
--- a/libc/src/__support/freelist_heap.h
+++ b/libc/src/__support/freelist_heap.h
@@ -89,28 +89,14 @@ LIBC_INLINE void *FreeListHeap::allocate_impl(size_t alignment, size_t size) {
   if (!is_initialized)
     init();
 
-  size_t request_size = size;
-
-  // TODO: usable_space should always be aligned to max_align_t.
-  if (alignment > alignof(max_align_t) ||
-      (Block::BLOCK_OVERHEAD % alignof(max_align_t) != 0)) {
-    // TODO: This bound isn't precisely calculated yet. It assumes one extra
-    // Block::ALIGNMENT to accomodate the possibility for padding block
-    // overhead. (alignment - 1) ensures that there is an aligned point
-    // somewhere in usable_space, but this isn't tight either, since
-    // usable_space is also already somewhat aligned.
-    if (add_overflow(size, (alignment - 1) + Block::ALIGNMENT, request_size))
-      return nullptr;
-  }
+  size_t request_size = Block::min_size_for_allocation(alignment, size);
+  if (!request_size)
+    return nullptr;
 
   Block *block = free_store.remove_best_fit(request_size);
   if (!block)
     return nullptr;
 
-  LIBC_ASSERT(block->can_allocate(alignment, size) &&
-              "block should always be large enough to allocate at the correct "
-              "alignment");
-
   auto block_info = Block::allocate(block, alignment, size);
   if (block_info.next)
     free_store.insert(block_info.next);
@@ -135,6 +121,9 @@ LIBC_INLINE void *FreeListHeap::aligned_allocate(size_t alignment,
   if (size % alignment != 0)
     return nullptr;
 
+  // The minimum alignment supported by Block is max_align_t.
+  alignment = cpp::max(alignment, alignof(max_align_t));
+
   return allocate_impl(alignment, size);
 }
 
diff --git a/libc/test/src/__support/block_test.cpp b/libc/test/src/__support/block_test.cpp
index 5e437db51b6092..15a33c4fe3c967 100644
--- a/libc/test/src/__support/block_test.cpp
+++ b/libc/test/src/__support/block_test.cpp
@@ -21,37 +21,47 @@ using LIBC_NAMESPACE::cpp::byte;
 using LIBC_NAMESPACE::cpp::span;
 
 TEST(LlvmLibcBlockTest, CanCreateSingleAlignedBlock) {
-  constexpr size_t kN = 1024;
-  alignas(Block::ALIGNMENT) array<byte, kN> bytes;
+  alignas(Block::ALIGNMENT) array<byte, 1024> bytes;
 
   auto result = Block::init(bytes);
   ASSERT_TRUE(result.has_value());
   Block *block = *result;
 
+  EXPECT_EQ(reinterpret_cas...
[truncated]

``````````

</details>


https://github.com/llvm/llvm-project/pull/117815


More information about the libc-commits mailing list