Meta alloc improvements.

1 files changed, 97 insertions, 95 deletions

diff --git a/src/zisp/gc/meta_alloc.zig b/src/zisp/gc/meta_alloc.zig
index d3b6be1..eb29028 100644
--- a/src/zisp/gc/meta_alloc.zig
+++ b/src/zisp/gc/meta_alloc.zig
@@ -34,17 +34,15 @@
 // have the following format:
 //
 //     {
-//       next_head_idx: u32,
-//       idx_count:     u32,
-//       _padding:      u64,
-//       idx_1:         u32,
-//       idx_2:         u32,
-//       ...
+//       next_head_idx:   u32,
+//       extra_idx_count: u32,
+//       pad_to_64_bytes: [7]u64,
+//       extra_idx_array: [32]u32,
 //     }
 //
 // So, it's not a typical free-list where each node simply points to the next
-// entry; rather, each node is the first of a chunk of N entries, where the
-// indexes of the remaining N-1 entries are inlined into this head node.
+// entry; rather, each node is the first of a chunk of 1+N entries, where the
+// indexes of the N extra entries are inlined into the head node.
 //
 // Each thread also has a small static thread-local array of indexes of most
 // recently freed slots.  This makes alloc/free extremely efficient in common
@@ -69,9 +67,8 @@ const std = @import("std");
 /// Size classes.
 const SIZES: [16]comptime_int = .{
     // Important: Ensure the smallest size class is inherently immune to false
-    // sharing, and no smaller than: 16 + 4 * ( RESERVE_MAX / 2 ) many bytes.
-    // That's 80 bytes assuming RESERVE_MAX = 32 (see below).
-    128,
+    // sharing, and no smaller than: 64 + 4 * ( RESERVE_MAX / 2 ) many bytes.
+    // That's 192 bytes assuming RESERVE_MAX = 64 (see below).
     256,
     512,
     1024,
@@ -87,20 +84,21 @@ const SIZES: [16]comptime_int = .{
     1024 * 1024,
     1024 * 1024 * 2,
     1024 * 1024 * 4,
+    1024 * 1024 * 8,
 };
 
 /// Gets the index [0,15] of a size class, from the size value, returning a
 /// value greater than 15 if the size is beyond the largest size class.
 inline fn getSizeClassIndex(size: usize) u8 {
-    const min = @as(u8, SIZES[0]);
+    const min: u32 = SIZES[0];
     return @ctz(size >> @ctz(min));
 }
 
-/// Size of each slab, in bytes.
-const SLAB_SIZE: usize = std.math.maxInt(i32) * 8;
+/// The length of each slab in 8-byte units.
+const SLAB_LEN: u32 = 0x8000_0000;
 
 /// Pointer to start of all 16 slabs as a contiguous vmem block.
-var slabs: [*]u8 = undefined;
+var slabs: *[SIZES.len][SLAB_LEN]u64 = undefined;
 
 /// Free-list head pointer with ABA counter.
 const FlHead = packed struct(u64) {
@@ -109,16 +107,13 @@ const FlHead = packed struct(u64) {
     idx: u32,
 };
 
-/// First invalid value for our 31-bit indexes; used as a NULL, since index
-/// value 0 is actually valid.
-const INVAL_IDX: u32 = 0x8000_0000;
-
 /// Global shared metadata per slab / size class.
 const SlabInfo = struct {
-    /// Force cache-line size alignment to prevent false sharing.
-    _: void align(std.atomic.cache_line) = undefined,
-    /// Head of global free-list for the size class, with ABA counter.
-    free_list: FlHead = .{ .aba = 0, .idx = INVAL_IDX },
+    /// Force cache line alignment to prevent false sharing.
+    _: void align(std.atomic.cache_line) = {},
+    /// Head of shared free-list for the size class; index value SLAB_LEN is
+    /// used to mean NULL, since it's an invalid slab index.
+    free_list: FlHead = .{ .aba = 0, .idx = SLAB_LEN },
     /// Global slab watermark: Start address of unused vmem, as 8-byte index.
     watermark: u32 = 0,
 };
@@ -127,35 +122,36 @@ const SlabInfo = struct {
 var slab_infos: [SIZES.len]SlabInfo = @splat(.{});
 
 /// Maximum number of slots reserved for a thread.
-const RESERVE_MAX = 32;
-
-/// Maps indexes [0,15] to { 32, 16, 8, 4 } in steps of 4, because we want more
-/// thread-local reserved slots for smaller size classes.  So, for example, the
-/// size classes from 128 to 1K will use 32 thread-local reserved slots, while
-/// classes 512K to 4M will use only 4 reserved slots.
+const RESERVE_MAX = 64;
+
+/// Given a size class index 0 to 15, returns the number of slots that threads
+/// should keep in reserve for that size class.  This starts at RESERVE_MAX for
+/// the smallest class, and is halved every four size classes.  This number is
+/// used both for deciding how much to bump the global watermark of a slab to
+/// reserve memory for a thread, and to determine the maximum count of entries
+/// stored in the thread's free slot cache for that size class.
 inline fn reserveCountForSizeClassIndex(sci: u8) u32 {
     return @as(u8, RESERVE_MAX) >> @intCast(sci / 4);
 }
 
 /// Thread-local metadata per slab / size class.
 const TlSlabInfo = struct {
-    /// Force cache-line size alignment to prevent false sharing.
-    _: void align(std.atomic.cache_line) = undefined,
+    /// Free slot cache; align to cache line to prevent false sharing and for
+    /// possible SIMD related benefits.
+    fc: [RESERVE_MAX]u32 align(std.atomic.cache_line) = @splat(0),
+    /// Current count of entries in free slot cache.
+    fc_count: u32 = 0,
     /// Start point of memory reserved for this thread.
     wm_lo: u32 = 0,
     /// End point of memory reserved for this thread.
     wm_hi: u32 = 0,
-    /// Current count of entries in free slot cache.
-    fc_count: u32 = 0,
-    /// Free slot cache; align to 16 for potential SIMD benefits.
-    fc: [RESERVE_MAX]u32 align(16) = undefined,
 };
 
 /// Thread-local metadata per slab / size class.
 threadlocal var tl_slab_infos: [SIZES.len]TlSlabInfo = @splat(.{});
 
 /// Wrapper around std.posix.mmap().
-fn mmap(size: usize) []u8 {
+fn mmap(size: usize) []align(4096) u8 {
     return std.posix.mmap(
         null,
         size,
@@ -168,7 +164,7 @@ fn mmap(size: usize) []u8 {
 
 /// Must call this once to initialize the slabs.
 pub fn init() void {
-    slabs = mmap(SIZES.len * SLAB_SIZE).ptr;
+    slabs = @ptrCast(mmap(SIZES.len * SLAB_LEN * 8));
 }
 
 /// Allocate a slot of the given size, which must be a power of two and greater
@@ -184,102 +180,109 @@ pub fn alloc(size: usize) []u8 {
     const sci = getSizeClassIndex(size);
     if (sci < SIZES.len) {
         @branchHint(.likely);
-        return alloc_size_class(sci, @intCast(size));
+        const slot_ptr = alloc_size_class(sci, size);
+        const ptr: [*]u8 = @ptrCast(slot_ptr);
+        return ptr[0..size];
     } else {
         return mmap(size);
     }
 }
 
-fn alloc_size_class(sci: u8, size: u32) []u8 {
+fn alloc_size_class(sci: u8, size: usize) [*]u64 {
     std.debug.assert(sci < SIZES.len);
 
     const tl = &tl_slab_infos[sci];
-    const slab = slabs + sci * SLAB_SIZE;
+    const slab: [*]u64 = &slabs[sci];
 
     // Reuse from thread-local free cache.
     if (tl.fc_count > 0) {
         tl.fc_count -= 1;
-        const ptr = slab + tl.fc[tl.fc_count] * 8;
-        return ptr[0..size];
+        const idx = tl.fc[tl.fc_count];
+        return slab[idx..];
     }
 
+    // The size in 8-byte slab/watermark units.
+    const wm_units: u32 = @intCast(size / 8);
+
     // Use part of the reserved memory for this thread.
     if (tl.wm_lo < tl.wm_hi) {
-        const ptr = slab + tl.wm_lo * 8;
-        tl.wm_lo += size;
-        return ptr[0..size];
+        const idx = tl.wm_lo;
+        tl.wm_lo += wm_units;
+        return slab[idx..];
     }
 
     // Try to transfer a chunk from the global free-list.
     const info = &slab_infos[sci];
-    if (alloc_from_fl(slab, info, tl)) |p| return p[0..size];
+    if (alloc_from_fl(slab, info, tl)) |p| return p;
 
     const res_n = reserveCountForSizeClassIndex(sci);
-    return alloc_fresh(slab, size, info, tl, res_n);
+    return alloc_fresh(size, slab, info, tl, wm_units, res_n);
 }
 
 fn alloc_from_fl(
-    slab: [*]u8,
+    slab: [*]u64,
     info: *SlabInfo,
     tl: *TlSlabInfo,
-) ?[*]u8 {
+) ?[*]u64 {
+    var ptr: [*]u64 = undefined;
+    var arr: [*]u32 = undefined;
+
     var fl_head = @atomicLoad(FlHead, &info.free_list, .acquire);
-    // Checking for < INVAL_IDX is optimal; it'll just test the sign bit.
-    while (fl_head.idx < INVAL_IDX) {
-        const ptr = slab + fl_head.idx * 8;
-        const arr: [*]u32 = @ptrCast(@alignCast(ptr));
+    while (true) : (std.atomic.spinLoopHint()) {
+        // Checking for >= SLAB_LEN is optimal; it'll just test the sign bit.
+        if (fl_head.idx >= SLAB_LEN) return null;
 
-        // Load this atomically since it's still globally accessible memory.
+        ptr = slab[fl_head.idx..];
+        arr = @ptrCast(ptr);
+
+        // Load this atomically since it's still globally accessible memory;
+        // unordered is fine since we discard it if the following CAS fails.
         const next_head = @atomicLoad(u32, &arr[0], .unordered);
 
-        if (@cmpxchgWeak(
+        fl_head = @cmpxchgWeak(
             FlHead,
             &info.free_list,
             fl_head,
             .{ .aba = fl_head.aba +% 1, .idx = next_head },
             .release,
             .acquire,
-        )) |head| {
-            fl_head = head;
-            std.atomic.spinLoopHint();
-            continue;
-        }
+        ) orelse break;
+    }
 
-        // Now detached from global; no need for atomics anymore.
-        tl.fc_count = arr[1];
+    // Now detached from global; no need for atomics anymore.
+    tl.fc_count = arr[1];
 
-        // We can safely copy more than needed since we set fc_count; this way
-        // it should compile neatly into SIMD instructions.
-        for (0..RESERVE_MAX / 2) |i| tl.fc[i] = arr[4 + i];
+    // We can safely copy more than needed since we set fc_count; this way it
+    // should compile neatly into SIMD instructions.
+    for (0..RESERVE_MAX / 2) |i| tl.fc[i] = arr[8 + i];
 
-        return ptr;
-    }
-    return null;
+    return ptr;
 }
 
 fn alloc_fresh(
-    slab: [*]u8,
-    size: u32,
+    size: usize,
+    slab: [*]u64,
     info: *SlabInfo,
     tl: *TlSlabInfo,
+    wm_units: u32,
     res_n: u32,
-) []u8 {
+) [*]u64 {
     // Have to use fresh memory from the top of the slab; bump it by a chunk,
     // atomically, to reserve memory for this thread.
-    const bump = res_n * size / 8;
+    const bump = res_n * wm_units;
     const old_wm = @atomicRmw(u32, &info.watermark, .Add, bump, .monotonic);
     const new_wm = old_wm + bump;
 
-    // Check only if the old WM was already overflown; chunk size is such that
-    // it's always safe to bump by a chunk if the old WM was still valid.
-    if (old_wm >= INVAL_IDX) {
+    // Check only if the old WM was already overflown; chunk count is such that
+    // it's always safe to bump by a chunk if the old WM was still valid, since
+    // wm_units, res_n, and slab length are all powers of two.
+    if (old_wm >= SLAB_LEN) {
         std.debug.panic("Exhausted slab for size class: {}", .{size});
     }
 
-    const ptr = slab + old_wm * 8;
-    tl.wm_lo = old_wm + size / 8;
     tl.wm_hi = new_wm;
-    return ptr[0..size];
+    tl.wm_lo = old_wm + wm_units;
+    return slab[old_wm..];
 }
 
 /// Free memory that was returned by alloc().
@@ -300,8 +303,10 @@ fn free_size_class(sci: u8, slot: []u8) void {
     std.debug.assert(sci < SIZES.len);
 
     const tl = &tl_slab_infos[sci];
-    const slab = slabs + sci * SLAB_SIZE;
-    const idx: u32 = @intCast((slot.ptr - slab) / 8);
+    const slab: [*]u64 = &slabs[sci];
+
+    const ptr: [*]u64 = @ptrCast(@alignCast(slot.ptr));
+    const idx: u32 = @intCast(ptr - slab);
 
     const res_n = reserveCountForSizeClassIndex(sci);
 
@@ -314,45 +319,42 @@ fn free_size_class(sci: u8, slot: []u8) void {
 
     // Local cache full; need to flush some to global free-list.
     const info = &slab_infos[sci];
-    free_into_fl(info, tl, slot, idx, res_n / 2);
+    free_into_fl(info, tl, ptr, idx, res_n);
 }
 
 fn free_into_fl(
     info: *SlabInfo,
     tl: *TlSlabInfo,
-    slot: []u8,
+    ptr: [*]u64,
     idx: u32,
-    num: u32,
+    res_n: u32,
 ) void {
-    const arr: [*]u32 = @ptrCast(@alignCast(slot.ptr));
+    const arr: [*]u32 = @ptrCast(ptr);
+    const half = res_n / 2;
 
     // Intuitively, this belongs to the end, but it doesn't matter; the rest of
     // this function is not allowed to fail anyway.
-    tl.fc_count = num;
+    tl.fc_count = half;
 
     // No need for atomic stores; we own this memory exclusively right now,
     // unless another thread calls free() on the same slot which is a bug.
-    arr[1] = num;
+    arr[1] = half;
     // Safe to copy more than needed; should compile into SIMD.
-    for (0..RESERVE_MAX / 2) |i| arr[4 + i] = tl.fc[num + i];
+    for (0..RESERVE_MAX / 2) |i| arr[8 + i] = tl.fc[half + i];
 
+    // Now atomically announce the new free-list head to global visibility.
     var fl_head = @atomicLoad(FlHead, &info.free_list, .acquire);
-    while (true) {
+    while (true) : (std.atomic.spinLoopHint()) {
         // Still owned exclusively; no need for atomic store.
         arr[0] = fl_head.idx;
-        if (@cmpxchgWeak(
+        fl_head = @cmpxchgWeak(
             FlHead,
             &info.free_list,
             fl_head,
             .{ .aba = fl_head.aba +% 1, .idx = idx },
             .release,
             .acquire,
-        )) |head| {
-            fl_head = head;
-            std.atomic.spinLoopHint();
-            continue;
-        }
-        return;
+        ) orelse break;
     }
 }