xref: /freebsd/contrib/llvm-project/libc/src/__support/GPU/allocator.cpp (revision bb722a7d0f1642bff6487f943ad0427799a6e5bf)
1 //===-- GPU memory allocator implementation ---------------------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements a parallel allocator intended for use on a GPU device.
10 // The core algorithm is slab allocator using a random walk over a bitfield for
11 // maximum parallel progress. Slab handling is done by a wait-free reference
12 // counted guard. The first use of a slab will create it from system memory for
13 // re-use. The last use will invalidate it and free the memory.
14 //
15 //===----------------------------------------------------------------------===//
16 
17 #include "allocator.h"
18 
19 #include "src/__support/CPP/atomic.h"
20 #include "src/__support/CPP/bit.h"
21 #include "src/__support/CPP/new.h"
22 #include "src/__support/GPU/utils.h"
23 #include "src/__support/RPC/rpc_client.h"
24 #include "src/__support/threads/sleep.h"
25 #include "src/string/memory_utils/inline_memcpy.h"
26 
27 namespace LIBC_NAMESPACE_DECL {
28 
29 constexpr static uint64_t MAX_SIZE = /* 64 GiB */ 64ull * 1024 * 1024 * 1024;
30 constexpr static uint64_t SLAB_SIZE = /* 2 MiB */ 2ull * 1024 * 1024;
31 constexpr static uint64_t ARRAY_SIZE = MAX_SIZE / SLAB_SIZE;
32 constexpr static uint64_t SLAB_ALIGNMENT = SLAB_SIZE - 1;
33 constexpr static uint32_t BITS_IN_WORD = sizeof(uint32_t) * 8;
34 constexpr static uint32_t MIN_SIZE = 16;
35 constexpr static uint32_t MIN_ALIGNMENT = MIN_SIZE - 1;
36 
37 // A sentinel used to indicate an invalid but non-null pointer value.
38 constexpr static uint64_t SENTINEL = cpp::numeric_limits<uint64_t>::max();
39 
40 // The number of times we will try starting on a single index before skipping
41 // past it.
42 constexpr static uint32_t MAX_TRIES = 512;
43 
44 static_assert(!(ARRAY_SIZE & (ARRAY_SIZE - 1)), "Must be a power of two");
45 
46 namespace impl {
47 // Allocates more memory from the system through the RPC interface. All
48 // allocations from the system MUST be aligned on a 2MiB barrier. The default
49 // HSA allocator has this behavior for any allocation >= 2MiB and the CUDA
50 // driver provides an alignment field for virtual memory allocations.
rpc_allocate(uint64_t size)51 static void *rpc_allocate(uint64_t size) {
52   void *ptr = nullptr;
53   rpc::Client::Port port = rpc::client.open<LIBC_MALLOC>();
54   port.send_and_recv(
55       [=](rpc::Buffer *buffer, uint32_t) { buffer->data[0] = size; },
56       [&](rpc::Buffer *buffer, uint32_t) {
57         ptr = reinterpret_cast<void *>(buffer->data[0]);
58       });
59   port.close();
60   return ptr;
61 }
62 
63 // Deallocates the associated system memory.
rpc_free(void * ptr)64 static void rpc_free(void *ptr) {
65   rpc::Client::Port port = rpc::client.open<LIBC_FREE>();
66   port.send([=](rpc::Buffer *buffer, uint32_t) {
67     buffer->data[0] = reinterpret_cast<uintptr_t>(ptr);
68   });
69   port.close();
70 }
71 
72 // Convert a potentially disjoint bitmask into an increasing integer per-lane
73 // for use with indexing between gpu lanes.
lane_count(uint64_t lane_mask)74 static inline uint32_t lane_count(uint64_t lane_mask) {
75   return cpp::popcount(lane_mask & ((uint64_t(1) << gpu::get_lane_id()) - 1));
76 }
77 
78 // Obtain an initial value to seed a random number generator. We use the rounded
79 // multiples of the golden ratio from xorshift* as additional spreading.
entropy()80 static inline uint32_t entropy() {
81   return (static_cast<uint32_t>(gpu::processor_clock()) ^
82           (gpu::get_thread_id_x() * 0x632be59b) ^
83           (gpu::get_block_id_x() * 0x85157af5)) *
84          0x9e3779bb;
85 }
86 
87 // Generate a random number and update the state using the xorshift32* PRNG.
xorshift32(uint32_t & state)88 static inline uint32_t xorshift32(uint32_t &state) {
89   state ^= state << 13;
90   state ^= state >> 17;
91   state ^= state << 5;
92   return state * 0x9e3779bb;
93 }
94 
95 // Final stage of murmurhash used to get a unique index for the global array
hash(uint32_t x)96 static inline uint32_t hash(uint32_t x) {
97   x ^= x >> 16;
98   x *= 0x85ebca6b;
99   x ^= x >> 13;
100   x *= 0xc2b2ae35;
101   x ^= x >> 16;
102   return x;
103 }
104 
105 // Rounds the input value to the closest permitted chunk size. Here we accept
106 // the sum of the closest three powers of two. For a 2MiB slab size this is 48
107 // different chunk sizes. This gives us average internal fragmentation of 87.5%.
get_chunk_size(uint32_t x)108 static inline uint32_t get_chunk_size(uint32_t x) {
109   uint32_t y = x < MIN_SIZE ? MIN_SIZE : x;
110   uint32_t pow2 = BITS_IN_WORD - cpp::countl_zero(y - 1);
111 
112   uint32_t s0 = 0b0100 << (pow2 - 3);
113   uint32_t s1 = 0b0110 << (pow2 - 3);
114   uint32_t s2 = 0b0111 << (pow2 - 3);
115   uint32_t s3 = 0b1000 << (pow2 - 3);
116 
117   if (s0 > y)
118     return (s0 + MIN_ALIGNMENT) & ~MIN_ALIGNMENT;
119   if (s1 > y)
120     return (s1 + MIN_ALIGNMENT) & ~MIN_ALIGNMENT;
121   if (s2 > y)
122     return (s2 + MIN_ALIGNMENT) & ~MIN_ALIGNMENT;
123   return (s3 + MIN_ALIGNMENT) & ~MIN_ALIGNMENT;
124 }
125 
126 // Rounds to the nearest power of two.
127 template <uint32_t N, typename T>
round_up(const T x)128 static inline constexpr T round_up(const T x) {
129   static_assert(((N - 1) & N) == 0, "N must be a power of two");
130   return (x + N) & ~(N - 1);
131 }
132 
133 // Perform a lane parallel memset on a uint32_t pointer.
uniform_memset(uint32_t * s,uint32_t c,uint32_t n,uint64_t uniform)134 void uniform_memset(uint32_t *s, uint32_t c, uint32_t n, uint64_t uniform) {
135   uint64_t mask = gpu::get_lane_mask();
136   uint32_t workers = cpp::popcount(uniform);
137   for (uint32_t i = impl::lane_count(mask & uniform); i < n; i += workers)
138     s[i] = c;
139 }
140 
141 // Indicates that the provided value is a power of two.
is_pow2(uint64_t x)142 static inline constexpr bool is_pow2(uint64_t x) {
143   return x && (x & (x - 1)) == 0;
144 }
145 
146 } // namespace impl
147 
148 /// A slab allocator used to hand out identically sized slabs of memory.
149 /// Allocation is done through random walks of a bitfield until a free bit is
150 /// encountered. This reduces contention and is highly parallel on a GPU.
151 ///
152 /// 0       4           8       16                 ...                     2 MiB
153 /// ┌────────┬──────────┬────────┬──────────────────┬──────────────────────────┐
154 /// │ chunk  │  index   │  pad   │    bitfield[]    │         memory[]         │
155 /// └────────┴──────────┴────────┴──────────────────┴──────────────────────────┘
156 ///
157 /// The size of the bitfield is the slab size divided by the chunk size divided
158 /// by the number of bits per word. We pad the interface to ensure 16 byte
159 /// alignment and to indicate that if the pointer is not aligned by 2MiB it
160 /// belongs to a slab rather than the global allocator.
161 struct Slab {
162   // Header metadata for the slab, aligned to the minimum alignment.
163   struct alignas(MIN_SIZE) Header {
164     uint32_t chunk_size;
165     uint32_t global_index;
166   };
167 
168   // Initialize the slab with its chunk size and index in the global table for
169   // use when freeing.
SlabLIBC_NAMESPACE_DECL::Slab170   Slab(uint32_t chunk_size, uint32_t global_index) {
171     Header *header = reinterpret_cast<Header *>(memory);
172     header->chunk_size = chunk_size;
173     header->global_index = global_index;
174   }
175 
176   // Set the necessary bitfield bytes to zero in parallel using many lanes. This
177   // must be called before the bitfield can be accessed safely, memory is not
178   // guaranteed to be zero initialized in the current implementation.
initializeLIBC_NAMESPACE_DECL::Slab179   void initialize(uint64_t uniform) {
180     uint32_t size = (bitfield_bytes(get_chunk_size()) + sizeof(uint32_t) - 1) /
181                     sizeof(uint32_t);
182     impl::uniform_memset(get_bitfield(), 0, size, uniform);
183   }
184 
185   // Get the number of chunks that can theoretically fit inside this slab.
num_chunksLIBC_NAMESPACE_DECL::Slab186   constexpr static uint32_t num_chunks(uint32_t chunk_size) {
187     return SLAB_SIZE / chunk_size;
188   }
189 
190   // Get the number of bytes needed to contain the bitfield bits.
bitfield_bytesLIBC_NAMESPACE_DECL::Slab191   constexpr static uint32_t bitfield_bytes(uint32_t chunk_size) {
192     return __builtin_align_up(
193         ((num_chunks(chunk_size) + BITS_IN_WORD - 1) / BITS_IN_WORD) * 8,
194         MIN_ALIGNMENT + 1);
195   }
196 
197   // The actual amount of memory available excluding the bitfield and metadata.
available_bytesLIBC_NAMESPACE_DECL::Slab198   constexpr static uint32_t available_bytes(uint32_t chunk_size) {
199     return SLAB_SIZE - bitfield_bytes(chunk_size) - sizeof(Header);
200   }
201 
202   // The number of chunks that can be stored in this slab.
available_chunksLIBC_NAMESPACE_DECL::Slab203   constexpr static uint32_t available_chunks(uint32_t chunk_size) {
204     return available_bytes(chunk_size) / chunk_size;
205   }
206 
207   // The length in bits of the bitfield.
usable_bitsLIBC_NAMESPACE_DECL::Slab208   constexpr static uint32_t usable_bits(uint32_t chunk_size) {
209     return available_bytes(chunk_size) / chunk_size;
210   }
211 
212   // Get the location in the memory where we will store the chunk size.
get_chunk_sizeLIBC_NAMESPACE_DECL::Slab213   uint32_t get_chunk_size() const {
214     return reinterpret_cast<const Header *>(memory)->chunk_size;
215   }
216 
217   // Get the location in the memory where we will store the global index.
get_global_indexLIBC_NAMESPACE_DECL::Slab218   uint32_t get_global_index() const {
219     return reinterpret_cast<const Header *>(memory)->global_index;
220   }
221 
222   // Get a pointer to where the bitfield is located in the memory.
get_bitfieldLIBC_NAMESPACE_DECL::Slab223   uint32_t *get_bitfield() {
224     return reinterpret_cast<uint32_t *>(memory + sizeof(Header));
225   }
226 
227   // Get a pointer to where the actual memory to be allocated lives.
get_memoryLIBC_NAMESPACE_DECL::Slab228   uint8_t *get_memory(uint32_t chunk_size) {
229     return reinterpret_cast<uint8_t *>(get_bitfield()) +
230            bitfield_bytes(chunk_size);
231   }
232 
233   // Get a pointer to the actual memory given an index into the bitfield.
ptr_from_indexLIBC_NAMESPACE_DECL::Slab234   void *ptr_from_index(uint32_t index, uint32_t chunk_size) {
235     return get_memory(chunk_size) + index * chunk_size;
236   }
237 
238   // Convert a pointer back into its bitfield index using its offset.
index_from_ptrLIBC_NAMESPACE_DECL::Slab239   uint32_t index_from_ptr(void *ptr, uint32_t chunk_size) {
240     return static_cast<uint32_t>(reinterpret_cast<uint8_t *>(ptr) -
241                                  get_memory(chunk_size)) /
242            chunk_size;
243   }
244 
245   // Randomly walks the bitfield until it finds a free bit. Allocations attempt
246   // to put lanes right next to each other for better caching and convergence.
allocateLIBC_NAMESPACE_DECL::Slab247   void *allocate(uint64_t lane_mask, uint64_t uniform) {
248     uint32_t chunk_size = get_chunk_size();
249     uint32_t state = impl::entropy();
250 
251     // The uniform mask represents which lanes contain a uniform target pointer.
252     // We attempt to place these next to each other.
253     void *result = nullptr;
254     for (uint64_t mask = lane_mask; mask;
255          mask = gpu::ballot(lane_mask, !result)) {
256       if (result)
257         continue;
258 
259       uint32_t start = gpu::broadcast_value(lane_mask, impl::xorshift32(state));
260 
261       uint32_t id = impl::lane_count(uniform & mask);
262       uint32_t index = (start + id) % usable_bits(chunk_size);
263       uint32_t slot = index / BITS_IN_WORD;
264       uint32_t bit = index % BITS_IN_WORD;
265 
266       // Get the mask of bits destined for the same slot and coalesce it.
267       uint64_t match = uniform & gpu::match_any(mask, slot);
268       uint32_t length = cpp::popcount(match);
269       uint32_t bitmask = static_cast<uint32_t>((uint64_t(1) << length) - 1)
270                          << bit;
271 
272       uint32_t before = 0;
273       if (gpu::get_lane_id() == static_cast<uint32_t>(cpp::countr_zero(match)))
274         before = cpp::AtomicRef(get_bitfield()[slot])
275                      .fetch_or(bitmask, cpp::MemoryOrder::RELAXED);
276       before = gpu::shuffle(mask, cpp::countr_zero(match), before);
277       if (~before & (1 << bit))
278         result = ptr_from_index(index, chunk_size);
279       else
280         sleep_briefly();
281     }
282 
283     cpp::atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
284     return result;
285   }
286 
287   // Deallocates memory by resetting its corresponding bit in the bitfield.
deallocateLIBC_NAMESPACE_DECL::Slab288   void deallocate(void *ptr) {
289     uint32_t chunk_size = get_chunk_size();
290     uint32_t index = index_from_ptr(ptr, chunk_size);
291     uint32_t slot = index / BITS_IN_WORD;
292     uint32_t bit = index % BITS_IN_WORD;
293 
294     cpp::atomic_thread_fence(cpp::MemoryOrder::RELEASE);
295     cpp::AtomicRef(get_bitfield()[slot])
296         .fetch_and(~(1u << bit), cpp::MemoryOrder::RELAXED);
297   }
298 
299   // The actual memory the slab will manage. All offsets are calculated at
300   // runtime with the chunk size to keep the interface convergent when a warp or
301   // wavefront is handling multiple sizes at once.
302   uint8_t memory[SLAB_SIZE];
303 };
304 
305 /// A wait-free guard around a pointer resource to be created dynamically if
306 /// space is available and freed once there are no more users.
307 struct GuardPtr {
308 private:
309   struct RefCounter {
310     // Indicates that the object is in its deallocation phase and thus invalid.
311     static constexpr uint64_t INVALID = uint64_t(1) << 63;
312 
313     // If a read preempts an unlock call we indicate this so the following
314     // unlock call can swap out the helped bit and maintain exclusive ownership.
315     static constexpr uint64_t HELPED = uint64_t(1) << 62;
316 
317     // Resets the reference counter, cannot be reset to zero safely.
resetLIBC_NAMESPACE_DECL::GuardPtr::RefCounter318     void reset(uint32_t n, uint64_t &count) {
319       counter.store(n, cpp::MemoryOrder::RELAXED);
320       count = n;
321     }
322 
323     // Acquire a slot in the reference counter if it is not invalid.
acquireLIBC_NAMESPACE_DECL::GuardPtr::RefCounter324     bool acquire(uint32_t n, uint64_t &count) {
325       count = counter.fetch_add(n, cpp::MemoryOrder::RELAXED) + n;
326       return (count & INVALID) == 0;
327     }
328 
329     // Release a slot in the reference counter. This function should only be
330     // called following a valid acquire call.
releaseLIBC_NAMESPACE_DECL::GuardPtr::RefCounter331     bool release(uint32_t n) {
332       // If this thread caused the counter to reach zero we try to invalidate it
333       // and obtain exclusive rights to deconstruct it. If the CAS failed either
334       // another thread resurrected the counter and we quit, or a parallel read
335       // helped us invalidating it. For the latter, claim that flag and return.
336       if (counter.fetch_sub(n, cpp::MemoryOrder::RELAXED) == n) {
337         uint64_t expected = 0;
338         if (counter.compare_exchange_strong(expected, INVALID,
339                                             cpp::MemoryOrder::RELAXED,
340                                             cpp::MemoryOrder::RELAXED))
341           return true;
342         else if ((expected & HELPED) &&
343                  (counter.exchange(INVALID, cpp::MemoryOrder::RELAXED) &
344                   HELPED))
345           return true;
346       }
347       return false;
348     }
349 
350     // Returns the current reference count, potentially helping a releasing
351     // thread.
readLIBC_NAMESPACE_DECL::GuardPtr::RefCounter352     uint64_t read() {
353       auto val = counter.load(cpp::MemoryOrder::RELAXED);
354       if (val == 0 && counter.compare_exchange_strong(
355                           val, INVALID | HELPED, cpp::MemoryOrder::RELAXED))
356         return 0;
357       return (val & INVALID) ? 0 : val;
358     }
359 
360     cpp::Atomic<uint64_t> counter{0};
361   };
362 
363   cpp::Atomic<Slab *> ptr{nullptr};
364   RefCounter ref{};
365 
366   // Should be called be a single lane for each different pointer.
367   template <typename... Args>
try_lock_implLIBC_NAMESPACE_DECL::GuardPtr368   Slab *try_lock_impl(uint32_t n, uint64_t &count, Args &&...args) {
369     Slab *expected = ptr.load(cpp::MemoryOrder::RELAXED);
370     if (!expected &&
371         ptr.compare_exchange_strong(
372             expected, reinterpret_cast<Slab *>(SENTINEL),
373             cpp::MemoryOrder::RELAXED, cpp::MemoryOrder::RELAXED)) {
374       count = cpp::numeric_limits<uint64_t>::max();
375       void *raw = impl::rpc_allocate(sizeof(Slab));
376       if (!raw)
377         return nullptr;
378       return new (raw) Slab(cpp::forward<Args>(args)...);
379     }
380 
381     if (!expected || expected == reinterpret_cast<Slab *>(SENTINEL))
382       return nullptr;
383 
384     if (!ref.acquire(n, count))
385       return nullptr;
386 
387     cpp::atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
388     return ptr.load(cpp::MemoryOrder::RELAXED);
389   }
390 
391   // Finalize the associated memory and signal that it is ready to use by
392   // resetting the counter.
finalizeLIBC_NAMESPACE_DECL::GuardPtr393   void finalize(Slab *mem, uint32_t n, uint64_t &count) {
394     cpp::atomic_thread_fence(cpp::MemoryOrder::RELEASE);
395     ptr.store(mem, cpp::MemoryOrder::RELAXED);
396     cpp::atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);
397     if (!ref.acquire(n, count))
398       ref.reset(n, count);
399   }
400 
401 public:
402   // Attempt to lock access to the pointer, potentially creating it if empty.
403   // The uniform mask represents which lanes share the same pointer. For each
404   // uniform value we elect a leader to handle it on behalf of the other lanes.
405   template <typename... Args>
try_lockLIBC_NAMESPACE_DECL::GuardPtr406   Slab *try_lock(uint64_t lane_mask, uint64_t uniform, uint64_t &count,
407                  Args &&...args) {
408     count = 0;
409     Slab *result = nullptr;
410     if (gpu::get_lane_id() == uint32_t(cpp::countr_zero(uniform)))
411       result = try_lock_impl(cpp::popcount(uniform), count,
412                              cpp::forward<Args>(args)...);
413     result = gpu::shuffle(lane_mask, cpp::countr_zero(uniform), result);
414     count = gpu::shuffle(lane_mask, cpp::countr_zero(uniform), count);
415 
416     if (!result)
417       return nullptr;
418 
419     // We defer storing the newly allocated slab until now so that we can use
420     // multiple lanes to initialize it and release it for use.
421     if (count == cpp::numeric_limits<uint64_t>::max()) {
422       result->initialize(uniform);
423       if (gpu::get_lane_id() == uint32_t(cpp::countr_zero(uniform)))
424         finalize(result, cpp::popcount(uniform), count);
425     }
426 
427     if (count != cpp::numeric_limits<uint64_t>::max())
428       count = count - cpp::popcount(uniform) + impl::lane_count(uniform) + 1;
429 
430     return result;
431   }
432 
433   // Release the associated lock on the pointer, potentially destroying it.
unlockLIBC_NAMESPACE_DECL::GuardPtr434   void unlock(uint64_t lane_mask, uint64_t mask) {
435     cpp::atomic_thread_fence(cpp::MemoryOrder::RELEASE);
436     if (gpu::get_lane_id() == uint32_t(cpp::countr_zero(mask)) &&
437         ref.release(cpp::popcount(mask))) {
438       Slab *p = ptr.load(cpp::MemoryOrder::RELAXED);
439       p->~Slab();
440       impl::rpc_free(p);
441       cpp::atomic_thread_fence(cpp::MemoryOrder::RELEASE);
442       ptr.store(nullptr, cpp::MemoryOrder::RELAXED);
443     }
444     gpu::sync_lane(lane_mask);
445   }
446 
447   // Get the current value of the reference counter.
use_countLIBC_NAMESPACE_DECL::GuardPtr448   uint64_t use_count() { return ref.read(); }
449 };
450 
451 // The global array used to search for a valid slab to allocate from.
452 static GuardPtr slots[ARRAY_SIZE] = {};
453 
454 // Tries to find a slab in the table that can support the given chunk size.
find_slab(uint32_t chunk_size)455 static Slab *find_slab(uint32_t chunk_size) {
456   // We start at a hashed value to spread out different chunk sizes.
457   uint32_t start = impl::hash(chunk_size);
458   uint64_t lane_mask = gpu::get_lane_mask();
459   uint64_t uniform = gpu::match_any(lane_mask, chunk_size);
460 
461   Slab *result = nullptr;
462   uint32_t nudge = 0;
463   for (uint64_t mask = lane_mask; mask;
464        mask = gpu::ballot(lane_mask, !result), ++nudge) {
465     uint32_t index = cpp::numeric_limits<uint32_t>::max();
466     for (uint32_t offset = nudge / MAX_TRIES;
467          gpu::ballot(lane_mask, index == cpp::numeric_limits<uint32_t>::max());
468          offset += cpp::popcount(uniform & lane_mask)) {
469       uint32_t candidate =
470           (start + offset + impl::lane_count(uniform & lane_mask)) % ARRAY_SIZE;
471       uint64_t available =
472           gpu::ballot(lane_mask, slots[candidate].use_count() <
473                                      Slab::available_chunks(chunk_size));
474       uint32_t new_index = gpu::shuffle(
475           lane_mask, cpp::countr_zero(available & uniform), candidate);
476 
477       // Each uniform group will use the first empty slot they find.
478       if ((index == cpp::numeric_limits<uint32_t>::max() &&
479            (available & uniform)))
480         index = new_index;
481 
482       // Guaruntees that this loop will eventuall exit if there is no space.
483       if (offset >= ARRAY_SIZE) {
484         result = reinterpret_cast<Slab *>(SENTINEL);
485         index = 0;
486       }
487     }
488 
489     // Try to claim a slot for the found slot.
490     if (!result) {
491       uint64_t reserved = 0;
492       Slab *slab = slots[index].try_lock(lane_mask & mask, uniform & mask,
493                                          reserved, chunk_size, index);
494       // If we find a slab with a matching chunk size then we store the result.
495       // Otherwise, we need to free the claimed lock and continue. In the case
496       // of out-of-memory we return a sentinel value.
497       if (slab && reserved <= Slab::available_chunks(chunk_size) &&
498           slab->get_chunk_size() == chunk_size) {
499         result = slab;
500       } else if (slab && (reserved > Slab::available_chunks(chunk_size) ||
501                           slab->get_chunk_size() != chunk_size)) {
502         if (slab->get_chunk_size() != chunk_size)
503           start = index + 1;
504         slots[index].unlock(gpu::get_lane_mask(),
505                             gpu::get_lane_mask() & uniform);
506       } else if (!slab && reserved == cpp::numeric_limits<uint64_t>::max()) {
507         result = reinterpret_cast<Slab *>(SENTINEL);
508       } else {
509         sleep_briefly();
510       }
511     }
512   }
513   return result;
514 }
515 
516 // Release the lock associated with a given slab.
release_slab(Slab * slab)517 static void release_slab(Slab *slab) {
518   uint32_t index = slab->get_global_index();
519   uint64_t lane_mask = gpu::get_lane_mask();
520   uint64_t uniform = gpu::match_any(lane_mask, index);
521   slots[index].unlock(lane_mask, uniform);
522 }
523 
524 namespace gpu {
525 
allocate(uint64_t size)526 void *allocate(uint64_t size) {
527   if (!size)
528     return nullptr;
529 
530   // Allocations requiring a full slab or more go directly to memory.
531   if (size >= SLAB_SIZE / 2)
532     return impl::rpc_allocate(impl::round_up<SLAB_SIZE>(size));
533 
534   // Try to find a slab for the rounded up chunk size and allocate from it.
535   uint32_t chunk_size = impl::get_chunk_size(static_cast<uint32_t>(size));
536   Slab *slab = find_slab(chunk_size);
537   if (!slab || slab == reinterpret_cast<Slab *>(SENTINEL))
538     return nullptr;
539 
540   uint64_t lane_mask = gpu::get_lane_mask();
541   uint64_t uniform = gpu::match_any(lane_mask, slab->get_global_index());
542   void *ptr = slab->allocate(lane_mask, uniform);
543   return ptr;
544 }
545 
deallocate(void * ptr)546 void deallocate(void *ptr) {
547   if (!ptr)
548     return;
549 
550   // All non-slab allocations will be aligned on a 2MiB boundary.
551   if (__builtin_is_aligned(ptr, SLAB_ALIGNMENT + 1))
552     return impl::rpc_free(ptr);
553 
554   // The original slab pointer is the 2MiB boundary using the given pointer.
555   Slab *slab = cpp::launder(reinterpret_cast<Slab *>(
556       (reinterpret_cast<uintptr_t>(ptr) & ~SLAB_ALIGNMENT)));
557   slab->deallocate(ptr);
558   release_slab(slab);
559 }
560 
reallocate(void * ptr,uint64_t size)561 void *reallocate(void *ptr, uint64_t size) {
562   if (ptr == nullptr)
563     return gpu::allocate(size);
564 
565   // Non-slab allocations are considered foreign pointers so we fail.
566   if (__builtin_is_aligned(ptr, SLAB_ALIGNMENT + 1))
567     return nullptr;
568 
569   // The original slab pointer is the 2MiB boundary using the given pointer.
570   Slab *slab = cpp::launder(reinterpret_cast<Slab *>(
571       (reinterpret_cast<uintptr_t>(ptr) & ~SLAB_ALIGNMENT)));
572   if (slab->get_chunk_size() >= size)
573     return ptr;
574 
575   // If we need a new chunk we reallocate and copy it over.
576   void *new_ptr = gpu::allocate(size);
577   inline_memcpy(new_ptr, ptr, slab->get_chunk_size());
578   gpu::deallocate(ptr);
579   return new_ptr;
580 }
581 
aligned_allocate(uint32_t alignment,uint64_t size)582 void *aligned_allocate(uint32_t alignment, uint64_t size) {
583   // All alignment values must be a non-zero power of two.
584   if (!impl::is_pow2(alignment))
585     return nullptr;
586 
587   // If the requested alignment is less than what we already provide this is
588   // just a normal allocation.
589   if (alignment <= MIN_ALIGNMENT + 1)
590     return gpu::allocate(size);
591 
592   // We can't handle alignments greater than 2MiB so we simply fail.
593   if (alignment > SLAB_ALIGNMENT + 1)
594     return nullptr;
595 
596   // Trying to handle allocation internally would break the assumption that each
597   // chunk is identical to eachother. Allocate enough memory with worst-case
598   // alignment and then round up. The index logic will round down properly.
599   uint64_t rounded = size + alignment - MIN_ALIGNMENT;
600   void *ptr = gpu::allocate(rounded);
601   return __builtin_align_up(ptr, alignment);
602 }
603 
604 } // namespace gpu
605 } // namespace LIBC_NAMESPACE_DECL
606