Fix arena allocation alignment at the block boundary (#24773)

This fixes an arena edge case where a request fits in a block before
alignment, but no longer fits after rounding. Such requests now use the
oversized allocation path, and oversized allocations keep their returned
pointers naturally aligned.

Added tests for both boundary cases.

Assisted by: Codex Sol

---------

Signed-off-by: ShiroKSH <kushidashiro@gmail.com>

runtime/src/iree/base/internal/arena.c

diff --git a/runtime/src/iree/base/internal/arena.c b/runtime/src/iree/base/internal/arena.c
index 98266ad..42dc8f5 100644
--- a/runtime/src/iree/base/internal/arena.c
+++ b/runtime/src/iree/base/internal/arena.c
@@ -204,14 +204,22 @@
 
   iree_arena_block_pool_t* block_pool = arena->block_pool;
 
-  if (byte_length > block_pool->usable_block_size) {
+  // Pad allocations so each subsequent allocation starts aligned.
+  iree_host_size_t aligned_length = 0;
+  if (!iree_host_size_checked_align(byte_length, iree_max_align_t,
+                                    &aligned_length)) {
+    return iree_make_status(IREE_STATUS_OUT_OF_RANGE, "alignment overflow");
+  }
+
+  if (aligned_length > block_pool->usable_block_size) {
     // Oversized allocation that can't be handled by the block pool. We'll
     // allocate directly from the system allocator and track it ourselves for
     // freeing during reset.
     IREE_TRACE_ZONE_BEGIN_NAMED(z0, "iree_arena_allocate_oversize");
     iree_host_size_t allocation_size = 0;
-    if (!iree_host_size_checked_add(sizeof(iree_arena_oversized_allocation_t),
-                                    byte_length, &allocation_size)) {
+    if (!iree_host_size_checked_add(
+            iree_sizeof_struct(iree_arena_oversized_allocation_t), byte_length,
+            &allocation_size)) {
       IREE_TRACE_ZONE_END(z0);
       return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
                               "oversized allocation size overflow");
@@ -225,19 +233,12 @@
     arena->allocation_head = allocation;
     arena->total_allocation_size += allocation_size;
     arena->used_allocation_size += byte_length;
-    *out_ptr = (uint8_t*)allocation + sizeof(iree_arena_oversized_allocation_t);
+    *out_ptr = (uint8_t*)allocation +
+               iree_sizeof_struct(iree_arena_oversized_allocation_t);
     IREE_TRACE_ZONE_END(z0);
     return iree_ok_status();
   }
 
-  // Pad length allocated so that each pointer bump is always ending at an
-  // aligned address and the next allocation will start aligned.
-  iree_host_size_t aligned_length = 0;
-  if (!iree_host_size_checked_align(byte_length, iree_max_align_t,
-                                    &aligned_length)) {
-    return iree_make_status(IREE_STATUS_OUT_OF_RANGE, "alignment overflow");
-  }
-
   // Check to see if the current block (if any) has space - if not, get another.
   if (arena->block_head == NULL ||
       arena->block_bytes_remaining < aligned_length) {

runtime/src/iree/base/internal/arena.h

diff --git a/runtime/src/iree/base/internal/arena.h b/runtime/src/iree/base/internal/arena.h
index 3dd6ee5..19c715b 100644
--- a/runtime/src/iree/base/internal/arena.h
+++ b/runtime/src/iree/base/internal/arena.h
@@ -153,7 +153,7 @@
 
 // Allocates |byte_length| contiguous bytes from the arena.
 // The returned bytes will have undefined contents and must be initialized by
-// the caller.
+// the caller. The returned pointer is aligned to iree_max_align_t.
 iree_status_t iree_arena_allocate(iree_arena_allocator_t* arena,
                                   iree_host_size_t byte_length, void** out_ptr);
 

runtime/src/iree/base/internal/arena_test.cc

diff --git a/runtime/src/iree/base/internal/arena_test.cc b/runtime/src/iree/base/internal/arena_test.cc
index 20218cb..a21d505 100644
--- a/runtime/src/iree/base/internal/arena_test.cc
+++ b/runtime/src/iree/base/internal/arena_test.cc
@@ -158,12 +158,38 @@
   void* ptr = NULL;
   IREE_ASSERT_OK(iree_arena_allocate(&arena, kBlockSize * 4, &ptr));
   ASSERT_NE(ptr, nullptr);
+  EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr) % iree_max_align_t, 0u);
   memset(ptr, 0xEF, kBlockSize * 4);
 
   iree_arena_deinitialize(&arena);
   iree_arena_block_pool_deinitialize(&pool);
 }
 
+TEST(Arena, RoundedSizeUsesOversizedAllocation) {
+  iree_arena_block_pool_t pool;
+  iree_arena_block_pool_initialize(kBlockSize + 1, iree_allocator_system(),
+                                   &pool);
+  const iree_host_size_t remainder = pool.usable_block_size % iree_max_align_t;
+  ASSERT_NE(remainder, 0);
+  iree_arena_allocator_t arena;
+  iree_arena_initialize(&pool, &arena);
+
+  // A request that fits before rounding may not fit after natural alignment.
+  const iree_host_size_t request_size = pool.usable_block_size - remainder + 1;
+  void* ptr = NULL;
+  IREE_ASSERT_OK(iree_arena_allocate(&arena, request_size, &ptr));
+  ASSERT_NE(ptr, nullptr);
+  EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr) % iree_max_align_t, 0u);
+  EXPECT_EQ(arena.block_head, nullptr);
+  EXPECT_EQ(arena.block_tail, nullptr);
+  EXPECT_EQ(arena.block_bytes_remaining, 0);
+  EXPECT_NE(arena.allocation_head, nullptr);
+  memset(ptr, 0xEF, request_size);
+
+  iree_arena_deinitialize(&arena);
+  iree_arena_block_pool_deinitialize(&pool);
+}
+
 TEST(Arena, Reset) {
   iree_arena_block_pool_t pool;
   iree_arena_block_pool_initialize(kBlockSize, iree_allocator_system(), &pool);