lib/gpalloc/src/cache/local.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const builtin = @import("builtin");
  3 const tls = @import("gpalloc_thread_cache_slot");
  4 const sys = @import("sys");
  5 const page_mod = @import("../page/root.zig");
  6 const pointer = @import("pointer.zig");
  7 const size_class = @import("../root.zig").class;
  8 const thread = @import("thread.zig");
  9 
 10 const sys_thread = sys.thread;
 11 const FreeNode = page_mod.FreeNode;
 12 const invalid_extra_class_index = std.math.maxInt(u8);
 13 const ThreadCacheBin = pointer.PointerBin(thread.capacity);
 14 const ExtraThreadCacheBin = pointer.PointerBin(thread.extra_bin_capacity);
 15 const linked_class_count = size_class.count - thread.fixed_class_count;
 16 const preload_build = builtin.output_mode == .Lib and builtin.link_mode == .dynamic;
 17 const preload_extra_class_lookup = buildPreloadExtraClassLookup();
 18 
 19 pub const ThreadCache = struct {
 20     owner: ?*anyopaque = null,
 21     storage_owner: ?*anyopaque = null,
 22     next_retained: ?*ThreadCache = null,
 23     bins: [thread.fixed_class_count]ThreadCacheBin = @as([thread.fixed_class_count]ThreadCacheBin, @splat(.{})),
 24     extra_bins: [thread.extra_class_count]ExtraThreadCacheBin =
 25         @as([thread.extra_class_count]ExtraThreadCacheBin, @splat(.{})),
 26     linked_heads: [linked_class_count]?*FreeNode = @as([linked_class_count]?*FreeNode, @splat(null)),
 27     linked_counts: [linked_class_count]u16 = @as([linked_class_count]u16, @splat(0)),
 28 
 29     pub fn reset(cache: *ThreadCache, owner: *anyopaque) void {
 30         cache.owner = owner;
 31         cache.storage_owner = owner;
 32         cache.next_retained = null;
 33         for (&cache.bins) |*bin| bin.count = 0;
 34         for (&cache.extra_bins) |*bin| bin.count = 0;
 35         @memset(&cache.linked_heads, null);
 36         @memset(&cache.linked_counts, 0);
 37     }
 38 };
 39 
 40 pub inline fn loadThreadCache() ?*ThreadCache {
 41     return @ptrCast(@alignCast(tls.load()));
 42 }
 43 
 44 pub inline fn loadThreadCacheForSmallFastPath() ?*ThreadCache {
 45     return loadThreadCache();
 46 }
 47 
 48 pub inline fn loadThreadCacheForAllocFastPath() ?*ThreadCache {
 49     return loadThreadCache();
 50 }
 51 
 52 pub inline fn storeThreadCache(cache: ?*ThreadCache) void {
 53     tls.store(cache);
 54 }
 55 
 56 pub inline fn localThreadCacheBinCount(cache: *const ThreadCache, class_index: usize) usize {
 57     if (class_index < thread.fixed_class_count) {
 58         var count: usize = cache.bins[class_index].count;
 59         if (threadExtraClassIndex(class_index)) |extra_index| {
 60             count += cache.extra_bins[extra_index].count;
 61         }
 62         return count;
 63     }
 64     return cache.linked_counts[class_index - thread.fixed_class_count];
 65 }
 66 
 67 pub inline fn popLocalThreadCacheBin(cache: *ThreadCache, class_index: usize) ?[*]u8 {
 68     if (class_index < thread.fixed_class_count) {
 69         if (cache.bins[class_index].pop()) |ptr| return ptr;
 70         if (threadExtraClassIndex(class_index)) |extra_index| {
 71             if (cache.extra_bins[extra_index].pop()) |ptr| return ptr;
 72         }
 73         return null;
 74     }
 75 
 76     const linked_index = class_index - thread.fixed_class_count;
 77     const node = cache.linked_heads[linked_index] orelse return null;
 78     cache.linked_heads[linked_index] = node.next;
 79     cache.linked_counts[linked_index] -= 1;
 80     return @ptrCast(node);
 81 }
 82 
 83 pub inline fn pushLocalThreadCacheBin(
 84     cache: *ThreadCache,
 85     class_index: usize,
 86     ptr: [*]u8,
 87     cache_limit: u16,
 88 ) bool {
 89     if (class_index < thread.fixed_class_count) {
 90         const base_limit = @min(thread.capacity, cache_limit);
 91         if (cache.bins[class_index].push(ptr, base_limit)) return true;
 92 
 93         const extra_limit = cache_limit - base_limit;
 94         if (extra_limit == 0) return false;
 95         const extra_index = threadExtraClassIndex(class_index) orelse return false;
 96         return cache.extra_bins[extra_index].push(ptr, extra_limit);
 97     }
 98 
 99     const linked_index = class_index - thread.fixed_class_count;
100     const count = cache.linked_counts[linked_index];
101     if (count >= cache_limit) return false;
102     const node: *FreeNode = @ptrCast(@alignCast(ptr));
103     node.* = .{ .next = cache.linked_heads[linked_index] };
104     cache.linked_heads[linked_index] = node;
105     cache.linked_counts[linked_index] = count + 1;
106     return true;
107 }
108 
109 pub inline fn threadExtraClassIndex(class_index: usize) ?usize {
110     if (comptime preload_build) {
111         const extra_index = preload_extra_class_lookup[class_index];
112         return if (extra_index == invalid_extra_class_index) null else extra_index;
113     }
114     return thread.extraClassIndex(class_index);
115 }
116 
117 pub fn cpuCountThreadCacheActiveLimit() usize {
118     return sys_thread.cpuCount();
119 }
120 
121 pub fn defaultThreadCacheActiveLimit() usize {
122     return cpuCountThreadCacheActiveLimit();
123 }
124 
125 fn buildPreloadExtraClassLookup() [size_class.count]u8 {
126     var lookup = @as([size_class.count]u8, @splat(invalid_extra_class_index));
127     for (0..size_class.count) |class_index| {
128         if (thread.extraClassIndex(class_index)) |extra_index| {
129             lookup[class_index] = @intCast(extra_index);
130         }
131     }
132     return lookup;
133 }
134 
135 test "thread-cache bins size local overflow separately by class" {
136     var cache: ThreadCache = .{};
137     const tiny_class = size_class.indexFor(32, .@"1").?;
138     const small_class = size_class.indexFor(48, .@"1").?;
139     const linked_class = size_class.indexFor(2048, .@"1").?;
140 
141     try std.testing.expect(@sizeOf(ThreadCache) < (size_class.count * thread.capacity * @sizeOf([*]u8)));
142     try std.testing.expectEqual(thread.capacity, cache.bins[0].items.len);
143     try std.testing.expectEqual(thread.extra_bin_capacity, cache.extra_bins[0].items.len);
144     try std.testing.expectEqual(@as(usize, 1024), size_class.size(thread.fixed_class_count - 1));
145     try std.testing.expectEqual(@as(usize, 1280), size_class.size(thread.fixed_class_count));
146     try std.testing.expectEqual(
147         @as(u16, thread.capacity + thread.extraCapacity(tiny_class)),
148         thread.classLimit(tiny_class),
149     );
150     try std.testing.expectEqual(
151         @as(u16, @min(
152             thread.target_bytes / size_class.size(small_class),
153             thread.capacity + thread.extraCapacity(small_class),
154         )),
155         thread.classLimit(small_class),
156     );
157 
158     var storage: usize = 0;
159     const ptr: [*]u8 = @ptrCast(&storage);
160     const cache_limit = thread.classLimit(small_class);
161     var index: usize = 0;
162     while (index < cache_limit) : (index += 1) {
163         try std.testing.expect(pushLocalThreadCacheBin(&cache, small_class, ptr, cache_limit));
164     }
165     try std.testing.expectEqual(@as(usize, cache_limit), localThreadCacheBinCount(&cache, small_class));
166     try std.testing.expect(!pushLocalThreadCacheBin(&cache, small_class, ptr, cache_limit));
167 
168     var linked_storage: [thread.capacity]usize = undefined;
169     const linked_limit = thread.classLimit(linked_class);
170     index = 0;
171     while (index < linked_limit) : (index += 1) {
172         const linked_ptr: [*]u8 = @ptrCast(&linked_storage[index]);
173         try std.testing.expect(pushLocalThreadCacheBin(&cache, linked_class, linked_ptr, linked_limit));
174     }
175     try std.testing.expectEqual(@as(usize, linked_limit), localThreadCacheBinCount(&cache, linked_class));
176     try std.testing.expect(!pushLocalThreadCacheBin(&cache, linked_class, ptr, linked_limit));
177 }