commit 2199c1d769de214bbc0a7c0054357f82c432b3d6
author Ben Vanik <ben.vanik@gmail.com> Tue Jan 07 08:17:42 2025 -0800
committer Ben Vanik <ben.vanik@gmail.com> Tue Jan 07 09:10:41 2025 -0800
tree bc39eeb35473f0cb4339004619c9882c7489a2af
parent ea462c896f074b4eafe3deb45bc43c9df75a7d53

Adding iree_arena_block_pool_preallocate.
Arena users can preallocate to ensure they have blocks available
that were allocated (maybe) with some spatial locality.

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 ba8f9e8..17f361d 100644
--- a/runtime/src/iree/base/internal/arena.c
+++ b/runtime/src/iree/base/internal/arena.c
@@ -41,6 +41,27 @@
   IREE_TRACE_ZONE_END(z0);
 }
 
+iree_status_t iree_arena_block_pool_preallocate(
+    iree_arena_block_pool_t* block_pool, iree_host_size_t count) {
+  IREE_TRACE_ZONE_BEGIN(z0);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, count);
+
+  for (iree_host_size_t i = 0; i < count; ++i) {
+    uint8_t* block_base = NULL;
+    IREE_RETURN_AND_END_ZONE_IF_ERROR(
+        z0, iree_allocator_malloc_uninitialized(block_pool->block_allocator,
+                                                block_pool->total_block_size,
+                                                (void**)&block_base));
+    iree_arena_block_t* block =
+        iree_arena_block_trailer(block_pool, block_base);
+    iree_atomic_arena_block_slist_concat(&block_pool->available_slist, block,
+                                         block);
+  }
+
+  IREE_TRACE_ZONE_END(z0);
+  return iree_ok_status();
+}
+
 void iree_arena_block_pool_trim(iree_arena_block_pool_t* block_pool) {
   IREE_TRACE_ZONE_BEGIN(z0);
 
@@ -125,6 +146,7 @@
     } while (head);
     arena->allocation_head = NULL;
   }
+
   if (arena->block_head != NULL) {
 #if defined(IREE_SANITIZER_ADDRESS)
     iree_arena_block_t* block = arena->block_head;
@@ -141,6 +163,10 @@
     arena->block_tail = NULL;
   }
 
+  arena->total_allocation_size = 0;
+  arena->used_allocation_size = 0;
+  arena->block_bytes_remaining = 0;
+
   IREE_TRACE_ZONE_END(z0);
 }
 

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 199e3fc..ad252b5 100644
--- a/runtime/src/iree/base/internal/arena.h
+++ b/runtime/src/iree/base/internal/arena.h
@@ -77,6 +77,10 @@
 // back to it.
 void iree_arena_block_pool_deinitialize(iree_arena_block_pool_t* block_pool);
 
+// Preallocates |count| blocks and adds them to the pool free list.
+iree_status_t iree_arena_block_pool_preallocate(
+    iree_arena_block_pool_t* block_pool, iree_host_size_t count);
+
 // Trims the pool by freeing unused blocks back to the allocator.
 // Acquired blocks are not freed and remain valid.
 void iree_arena_block_pool_trim(iree_arena_block_pool_t* block_pool);