688 lines · cpp
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-exception6//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 for11// maximum parallel progress. Slab handling is done by a wait-free reference12// counted guard. The first use of a slab will create it from system memory for13// 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/algorithm.h"20#include "src/__support/CPP/atomic.h"21#include "src/__support/CPP/bit.h"22#include "src/__support/CPP/new.h"23#include "src/__support/GPU/fixedstack.h"24#include "src/__support/GPU/utils.h"25#include "src/__support/RPC/rpc_client.h"26#include "src/__support/threads/sleep.h"27#include "src/string/memory_utils/inline_memcpy.h"28 29namespace LIBC_NAMESPACE_DECL {30 31constexpr static uint64_t MAX_SIZE = /* 64 GiB */ 64ull * 1024 * 1024 * 1024;32constexpr static uint64_t SLAB_SIZE = /* 2 MiB */ 2ull * 1024 * 1024;33constexpr static uint64_t ARRAY_SIZE = MAX_SIZE / SLAB_SIZE;34constexpr static uint64_t SLAB_ALIGNMENT = SLAB_SIZE - 1;35constexpr static uint32_t BITS_IN_WORD = sizeof(uint32_t) * 8;36constexpr static uint32_t BITS_IN_DWORD = sizeof(uint64_t) * 8;37constexpr static uint32_t MIN_SIZE = 16;38constexpr static uint32_t MIN_ALIGNMENT = MIN_SIZE - 1;39 40// The number of times to attempt claiming an in-progress slab allocation.41constexpr static uint32_t MAX_TRIES = 1024;42 43// The number of previously allocated slabs we will keep in memory.44constexpr static uint32_t CACHED_SLABS = 8;45 46// Configuration for whether or not we will return unused slabs to memory.47constexpr static bool RECLAIM = true;48 49static_assert(!(ARRAY_SIZE & (ARRAY_SIZE - 1)), "Must be a power of two");50 51namespace impl {52// Allocates more memory from the system through the RPC interface. All53// allocations from the system MUST be aligned on a 2MiB barrier. The default54// HSA allocator has this behavior for any allocation >= 2MiB and the CUDA55// driver provides an alignment field for virtual memory allocations.56static void *rpc_allocate(uint64_t size) {57 void *ptr = nullptr;58 rpc::Client::Port port = rpc::client.open<LIBC_MALLOC>();59 port.send_and_recv(60 [=](rpc::Buffer *buffer, uint32_t) { buffer->data[0] = size; },61 [&](rpc::Buffer *buffer, uint32_t) {62 ptr = reinterpret_cast<void *>(buffer->data[0]);63 });64 port.close();65 return ptr;66}67 68// Deallocates the associated system memory.69static void rpc_free(void *ptr) {70 rpc::Client::Port port = rpc::client.open<LIBC_FREE>();71 port.send([=](rpc::Buffer *buffer, uint32_t) {72 buffer->data[0] = reinterpret_cast<uintptr_t>(ptr);73 });74 port.close();75}76 77// Convert a potentially disjoint bitmask into an increasing integer per-lane78// for use with indexing between gpu lanes.79static inline uint32_t lane_count(uint64_t lane_mask, uint32_t id) {80 return cpp::popcount(lane_mask & ((uint64_t(1) << id) - 1));81}82 83// Obtain an initial value to seed a random number generator. We use the rounded84// multiples of the golden ratio from xorshift* as additional spreading.85static inline uint32_t entropy() {86 return (static_cast<uint32_t>(gpu::processor_clock()) ^87 (gpu::get_thread_id_x() * 0x632be59b) ^88 (gpu::get_block_id_x() * 0x85157af5)) *89 0x9e3779bb;90}91 92// Generate a random number and update the state using the xorshift32* PRNG.93static inline uint32_t xorshift32(uint32_t &state) {94 state ^= state << 13;95 state ^= state >> 17;96 state ^= state << 5;97 return state * 0x9e3779bb;98}99 100// Rounds the input value to the closest permitted chunk size. Here we accept101// the sum of the closest three powers of two. For a 2MiB slab size this is 48102// different chunk sizes. This gives us average internal fragmentation of 87.5%.103static inline constexpr uint32_t get_chunk_size(uint32_t x) {104 uint32_t y = x < MIN_SIZE ? MIN_SIZE : x;105 uint32_t pow2 = BITS_IN_WORD - cpp::countl_zero(y - 1);106 107 uint32_t s0 = 0b0100 << (pow2 - 3);108 uint32_t s1 = 0b0110 << (pow2 - 3);109 uint32_t s2 = 0b0111 << (pow2 - 3);110 uint32_t s3 = 0b1000 << (pow2 - 3);111 112 if (s0 > y)113 return (s0 + MIN_ALIGNMENT) & ~MIN_ALIGNMENT;114 if (s1 > y)115 return (s1 + MIN_ALIGNMENT) & ~MIN_ALIGNMENT;116 if (s2 > y)117 return (s2 + MIN_ALIGNMENT) & ~MIN_ALIGNMENT;118 return (s3 + MIN_ALIGNMENT) & ~MIN_ALIGNMENT;119}120 121// Converts a chunk size into an index suitable for a statically sized array.122static inline constexpr uint32_t get_chunk_id(uint32_t x) {123 if (x <= MIN_SIZE)124 return 0;125 uint32_t y = x >> 4;126 if (x < MIN_SIZE << 2)127 return cpp::popcount(y);128 return cpp::popcount(y) + 3 * (BITS_IN_WORD - cpp::countl_zero(y)) - 7;129}130 131// Rounds to the nearest power of two.132template <uint32_t N, typename T>133static inline constexpr T round_up(const T x) {134 static_assert(((N - 1) & N) == 0, "N must be a power of two");135 return (x + N) & ~(N - 1);136}137 138// Perform a lane parallel memset on a uint32_t pointer.139void uniform_memset(uint32_t *s, uint32_t c, uint32_t n, uint64_t uniform) {140 uint64_t mask = gpu::get_lane_mask();141 uint32_t workers = cpp::popcount(uniform);142 for (uint32_t i = impl::lane_count(mask & uniform, gpu::get_lane_id()); i < n;143 i += workers)144 s[i] = c;145}146 147// Indicates that the provided value is a power of two.148static inline constexpr bool is_pow2(uint64_t x) {149 return x && (x & (x - 1)) == 0;150}151 152// Where this chunk size should start looking in the global array. Small153// allocations are much more likely than large ones, so we give them the most154// space. We use a cubic easing function normalized on the possible chunks.155static inline constexpr uint32_t get_start_index(uint32_t chunk_size) {156 constexpr uint32_t max_chunk = impl::get_chunk_id(SLAB_SIZE / 2);157 uint64_t norm =158 (1 << 16) - (impl::get_chunk_id(chunk_size) << 16) / max_chunk;159 uint64_t bias = (norm * norm * norm) >> 32;160 uint64_t inv = (1 << 16) - bias;161 return static_cast<uint32_t>(((ARRAY_SIZE - 1) * inv) >> 16);162}163 164// Returns the id of the lane below this one that acts as its leader.165static inline uint32_t get_leader_id(uint64_t ballot, uint32_t id) {166 uint64_t mask = id < BITS_IN_DWORD - 1 ? ~0ull << (id + 1) : 0;167 return BITS_IN_DWORD - cpp::countl_zero(ballot & ~mask) - 1;168}169 170// We use a sentinal value to indicate a failed or in-progress allocation.171template <typename T> bool is_sentinel(const T &x) {172 if constexpr (cpp::is_pointer_v<T>)173 return reinterpret_cast<uintptr_t>(x) ==174 cpp::numeric_limits<uintptr_t>::max();175 else176 return x == cpp::numeric_limits<T>::max();177}178 179} // namespace impl180 181/// A slab allocator used to hand out identically sized slabs of memory.182/// Allocation is done through random walks of a bitfield until a free bit is183/// encountered. This reduces contention and is highly parallel on a GPU.184///185/// 0 4 8 16 ... 2 MiB186/// ┌────────┬──────────┬────────┬──────────────────┬──────────────────────────┐187/// │ chunk │ index │ pad │ bitfield[] │ memory[] │188/// └────────┴──────────┴────────┴──────────────────┴──────────────────────────┘189///190/// The size of the bitfield is the slab size divided by the chunk size divided191/// by the number of bits per word. We pad the interface to ensure 16 byte192/// alignment and to indicate that if the pointer is not aligned by 2MiB it193/// belongs to a slab rather than the global allocator.194struct Slab {195 // Header metadata for the slab, aligned to the minimum alignment.196 struct alignas(MIN_SIZE) Header {197 uint32_t chunk_size;198 uint32_t global_index;199 uint32_t cached_chunk_size;200 };201 202 // Initialize the slab with its chunk size and index in the global table for203 // use when freeing.204 Slab(uint32_t chunk_size, uint32_t global_index) {205 Header *header = reinterpret_cast<Header *>(memory);206 header->cached_chunk_size = cpp::numeric_limits<uint32_t>::max();207 header->chunk_size = chunk_size;208 header->global_index = global_index;209 }210 211 // Reset the memory with a new index and chunk size, not thread safe.212 Slab *reset(uint32_t chunk_size, uint32_t global_index) {213 Header *header = reinterpret_cast<Header *>(memory);214 header->cached_chunk_size = header->chunk_size;215 header->chunk_size = chunk_size;216 header->global_index = global_index;217 return this;218 }219 220 // Set the necessary bitfield bytes to zero in parallel using many lanes. This221 // must be called before the bitfield can be accessed safely, memory is not222 // guaranteed to be zero initialized in the current implementation.223 void initialize(uint64_t uniform) {224 // If this is a re-used slab the memory is already set to zero.225 if (get_cached_chunk_size() <= get_chunk_size())226 return;227 228 uint32_t size = (bitfield_bytes(get_chunk_size()) + sizeof(uint32_t) - 1) /229 sizeof(uint32_t);230 impl::uniform_memset(get_bitfield(), 0, size, uniform);231 }232 233 // Get the number of chunks that can theoretically fit inside this slab.234 constexpr static uint32_t num_chunks(uint32_t chunk_size) {235 return SLAB_SIZE / chunk_size;236 }237 238 // Get the number of bytes needed to contain the bitfield bits.239 constexpr static uint32_t bitfield_bytes(uint32_t chunk_size) {240 return __builtin_align_up(241 ((num_chunks(chunk_size) + BITS_IN_WORD - 1) / BITS_IN_WORD) * 8,242 MIN_ALIGNMENT + 1);243 }244 245 // The actual amount of memory available excluding the bitfield and metadata.246 constexpr static uint32_t available_bytes(uint32_t chunk_size) {247 return SLAB_SIZE - bitfield_bytes(chunk_size) - sizeof(Header);248 }249 250 // The number of chunks that can be stored in this slab.251 constexpr static uint32_t available_chunks(uint32_t chunk_size) {252 return available_bytes(chunk_size) / chunk_size;253 }254 255 // The length in bits of the bitfield.256 constexpr static uint32_t usable_bits(uint32_t chunk_size) {257 return available_bytes(chunk_size) / chunk_size;258 }259 260 // Get the location in the memory where we will store the chunk size.261 uint32_t get_chunk_size() const {262 return reinterpret_cast<const Header *>(memory)->chunk_size;263 }264 265 // Get the chunk size that was previously used.266 uint32_t get_cached_chunk_size() const {267 return reinterpret_cast<const Header *>(memory)->cached_chunk_size;268 }269 270 // Get the location in the memory where we will store the global index.271 uint32_t get_global_index() const {272 return reinterpret_cast<const Header *>(memory)->global_index;273 }274 275 // Get a pointer to where the bitfield is located in the memory.276 uint32_t *get_bitfield() {277 return reinterpret_cast<uint32_t *>(memory + sizeof(Header));278 }279 280 // Get a pointer to where the actual memory to be allocated lives.281 uint8_t *get_memory(uint32_t chunk_size) {282 return reinterpret_cast<uint8_t *>(get_bitfield()) +283 bitfield_bytes(chunk_size);284 }285 286 // Get a pointer to the actual memory given an index into the bitfield.287 void *ptr_from_index(uint32_t index, uint32_t chunk_size) {288 return get_memory(chunk_size) + index * chunk_size;289 }290 291 // Convert a pointer back into its bitfield index using its offset.292 uint32_t index_from_ptr(void *ptr, uint32_t chunk_size) {293 return static_cast<uint32_t>(reinterpret_cast<uint8_t *>(ptr) -294 get_memory(chunk_size)) /295 chunk_size;296 }297 298 // Randomly walks the bitfield until it finds a free bit. Allocations attempt299 // to put lanes right next to each other for better caching and convergence.300 void *allocate(uint64_t uniform, uint32_t reserved) {301 uint32_t chunk_size = get_chunk_size();302 uint32_t state = impl::entropy();303 304 // Try to find the empty bit in the bitfield to finish the allocation. We305 // start at the number of allocations as this is guaranteed to be available306 // until the user starts freeing memory.307 uint64_t lane_mask = gpu::get_lane_mask();308 uint32_t start = gpu::shuffle(309 lane_mask, cpp::countr_zero(uniform & lane_mask), reserved);310 for (;;) {311 uint64_t lane_mask = gpu::get_lane_mask();312 313 // Each lane tries to claim one bit in a single contiguous mask.314 uint32_t id = impl::lane_count(uniform & lane_mask, gpu::get_lane_id());315 uint32_t index = (start + id) % usable_bits(chunk_size);316 uint32_t slot = index / BITS_IN_WORD;317 uint32_t bit = index % BITS_IN_WORD;318 319 // Get the mask of bits destined for the same slot and coalesce it.320 uint32_t leader = impl::get_leader_id(321 uniform & gpu::ballot(lane_mask, !id || index % BITS_IN_WORD == 0),322 gpu::get_lane_id());323 uint32_t length = cpp::popcount(uniform & lane_mask) -324 impl::lane_count(uniform & lane_mask, leader);325 uint32_t bitmask =326 static_cast<uint32_t>(327 (uint64_t(1) << cpp::min(length, BITS_IN_WORD)) - 1)328 << bit;329 330 uint32_t before = 0;331 if (gpu::get_lane_id() == leader)332 before = cpp::AtomicRef(get_bitfield()[slot])333 .fetch_or(bitmask, cpp::MemoryOrder::RELAXED);334 before = gpu::shuffle(lane_mask, leader, before);335 if (~before & (1 << bit)) {336 cpp::atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);337 return ptr_from_index(index, chunk_size);338 }339 340 // If the previous operation found an empty bit we move there, otherwise341 // we generate new random index to start at.342 uint32_t after = before | bitmask;343 start = gpu::shuffle(344 gpu::get_lane_mask(),345 cpp::countr_zero(uniform & gpu::get_lane_mask()),346 ~after ? __builtin_align_down(index, BITS_IN_WORD) +347 cpp::countr_zero(~after)348 : __builtin_align_down(impl::xorshift32(state), BITS_IN_WORD));349 sleep_briefly();350 }351 }352 353 // Deallocates memory by resetting its corresponding bit in the bitfield.354 void deallocate(void *ptr) {355 uint32_t chunk_size = get_chunk_size();356 uint32_t index = index_from_ptr(ptr, chunk_size);357 uint32_t slot = index / BITS_IN_WORD;358 uint32_t bit = index % BITS_IN_WORD;359 360 cpp::atomic_thread_fence(cpp::MemoryOrder::RELEASE);361 cpp::AtomicRef(get_bitfield()[slot])362 .fetch_and(~(1u << bit), cpp::MemoryOrder::RELAXED);363 }364 365 // The actual memory the slab will manage. All offsets are calculated at366 // runtime with the chunk size to keep the interface convergent when a warp or367 // wavefront is handling multiple sizes at once.368 uint8_t memory[SLAB_SIZE];369};370 371// A global cache of previously allocated slabs for efficient reuse.372static FixedStack<Slab *, CACHED_SLABS> slab_cache;373 374/// A wait-free guard around a pointer resource to be created dynamically if375/// space is available and freed once there are no more users.376struct GuardPtr {377private:378 struct RefCounter {379 // Indicates that the object is in its deallocation phase and thus invalid.380 static constexpr uint32_t INVALID = uint32_t(1) << 31;381 382 // If a read preempts an unlock call we indicate this so the following383 // unlock call can swap out the helped bit and maintain exclusive ownership.384 static constexpr uint32_t HELPED = uint32_t(1) << 30;385 386 // Resets the reference counter, cannot be reset to zero safely.387 void reset(uint32_t n, uint32_t &count) {388 counter.store(n, cpp::MemoryOrder::RELAXED);389 count = n;390 }391 392 // Acquire a slot in the reference counter if it is not invalid.393 bool acquire(uint32_t n, uint32_t &count) {394 count = counter.fetch_add(n, cpp::MemoryOrder::RELAXED) + n;395 return (count & INVALID) == 0;396 }397 398 // Release a slot in the reference counter. This function should only be399 // called following a valid acquire call.400 bool release(uint32_t n) {401 // If this thread caused the counter to reach zero we try to invalidate it402 // and obtain exclusive rights to deconstruct it. If the CAS failed either403 // another thread resurrected the counter and we quit, or a parallel read404 // helped us invalidating it. For the latter, claim that flag and return.405 if (counter.fetch_sub(n, cpp::MemoryOrder::RELAXED) == n && RECLAIM) {406 uint32_t expected = 0;407 if (counter.compare_exchange_strong(expected, INVALID,408 cpp::MemoryOrder::RELAXED,409 cpp::MemoryOrder::RELAXED))410 return true;411 else if ((expected & HELPED) &&412 (counter.exchange(INVALID, cpp::MemoryOrder::RELAXED) &413 HELPED))414 return true;415 }416 return false;417 }418 419 // Returns the current reference count, potentially helping a releasing420 // thread.421 uint64_t read() {422 auto val = counter.load(cpp::MemoryOrder::RELAXED);423 if (val == 0 && RECLAIM &&424 counter.compare_exchange_strong(val, INVALID | HELPED,425 cpp::MemoryOrder::RELAXED))426 return 0;427 return (val & INVALID) ? 0 : val;428 }429 430 cpp::Atomic<uint32_t> counter{0};431 };432 433 cpp::Atomic<Slab *> ptr;434 RefCounter ref;435 436 // Should be called be a single lane for each different pointer.437 template <typename... Args>438 Slab *try_lock_impl(uint32_t n, uint32_t &count, Args &&...args) {439 Slab *expected = ptr.load(cpp::MemoryOrder::RELAXED);440 if (!expected &&441 ptr.compare_exchange_strong(442 expected,443 reinterpret_cast<Slab *>(cpp::numeric_limits<uintptr_t>::max()),444 cpp::MemoryOrder::RELAXED, cpp::MemoryOrder::RELAXED)) {445 count = cpp::numeric_limits<uint32_t>::max();446 447 Slab *cached = nullptr;448 if (slab_cache.pop(cached))449 return cached->reset(cpp::forward<Args>(args)...);450 451 void *raw = impl::rpc_allocate(sizeof(Slab));452 if (!raw)453 return nullptr;454 return new (raw) Slab(cpp::forward<Args>(args)...);455 }456 457 // If there is a slab allocation in progress we retry a few times.458 for (uint32_t t = 0; impl::is_sentinel(expected) && t < MAX_TRIES; ++t) {459 sleep_briefly();460 expected = ptr.load(cpp::MemoryOrder::RELAXED);461 }462 463 if (!expected || impl::is_sentinel(expected))464 return nullptr;465 466 if (!ref.acquire(n, count))467 return nullptr;468 469 cpp::atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);470 return RECLAIM ? ptr.load(cpp::MemoryOrder::RELAXED) : expected;471 }472 473 // Finalize the associated memory and signal that it is ready to use by474 // resetting the counter.475 void finalize(Slab *mem, uint32_t n, uint32_t &count) {476 cpp::atomic_thread_fence(cpp::MemoryOrder::RELEASE);477 ptr.store(mem, cpp::MemoryOrder::RELAXED);478 cpp::atomic_thread_fence(cpp::MemoryOrder::ACQUIRE);479 if (!ref.acquire(n, count))480 ref.reset(n, count);481 }482 483public:484 // Attempt to lock access to the pointer, potentially creating it if empty.485 // The uniform mask represents which lanes share the same pointer. For each486 // uniform value we elect a leader to handle it on behalf of the other lanes.487 template <typename... Args>488 Slab *try_lock(uint64_t lane_mask, uint64_t uniform, uint32_t &count,489 Args &&...args) {490 count = 0;491 Slab *result = nullptr;492 if (gpu::get_lane_id() == uint32_t(cpp::countr_zero(uniform)))493 result = try_lock_impl(cpp::popcount(uniform), count,494 cpp::forward<Args>(args)...);495 result = gpu::shuffle(lane_mask, cpp::countr_zero(uniform), result);496 count = gpu::shuffle(lane_mask, cpp::countr_zero(uniform), count);497 498 if (!result)499 return nullptr;500 501 // We defer storing the newly allocated slab until now so that we can use502 // multiple lanes to initialize it and release it for use.503 if (impl::is_sentinel(count)) {504 result->initialize(uniform);505 if (gpu::get_lane_id() == uint32_t(cpp::countr_zero(uniform)))506 finalize(result, cpp::popcount(uniform), count);507 count =508 gpu::shuffle(gpu::get_lane_mask(), cpp::countr_zero(uniform), count);509 }510 511 if (!impl::is_sentinel(count))512 count = count - cpp::popcount(uniform) +513 impl::lane_count(uniform, gpu::get_lane_id());514 515 return result;516 }517 518 // Release the associated lock on the pointer, potentially destroying it.519 void unlock(uint64_t lane_mask, uint64_t mask) {520 cpp::atomic_thread_fence(cpp::MemoryOrder::RELEASE);521 if (gpu::get_lane_id() == uint32_t(cpp::countr_zero(mask)) &&522 ref.release(cpp::popcount(mask))) {523 Slab *p = ptr.load(cpp::MemoryOrder::RELAXED);524 if (!slab_cache.push(p)) {525 p->~Slab();526 impl::rpc_free(p);527 }528 cpp::atomic_thread_fence(cpp::MemoryOrder::RELEASE);529 ptr.store(nullptr, cpp::MemoryOrder::RELAXED);530 }531 gpu::sync_lane(lane_mask);532 }533 534 // Get the current value of the reference counter.535 uint64_t use_count() { return ref.read(); }536};537 538// The global array used to search for a valid slab to allocate from.539static GuardPtr slots[ARRAY_SIZE] = {};540 541// Keep a cache of the last successful slot for each chunk size. Initialize it542// to an even spread of the total size. Must be updated if the chunking scheme543// changes.544#define S(X) (impl::get_start_index(X))545static cpp::Atomic<uint32_t> indices[] = {546 S(16), S(32), S(48), S(64), S(96), S(112), S(128),547 S(192), S(224), S(256), S(384), S(448), S(512), S(768),548 S(896), S(1024), S(1536), S(1792), S(2048), S(3072), S(3584),549 S(4096), S(6144), S(7168), S(8192), S(12288), S(14336), S(16384),550 S(24576), S(28672), S(32768), S(49152), S(57344), S(65536), S(98304),551 S(114688), S(131072), S(196608), S(229376), S(262144), S(393216), S(458752),552 S(524288), S(786432), S(917504), S(1048576)};553#undef S554 555// Tries to find a slab in the table that can support the given chunk size.556static Slab *find_slab(uint32_t chunk_size, uint64_t &uniform,557 uint32_t &reserved) {558 // We start at the index of the last successful allocation for this kind.559 uint32_t chunk_id = impl::get_chunk_id(chunk_size);560 uint32_t start = indices[chunk_id].load(cpp::MemoryOrder::RELAXED);561 562 for (uint32_t offset = 0; offset <= ARRAY_SIZE; ++offset) {563 uint32_t index =564 !offset ? start565 : (impl::get_start_index(chunk_size) + offset - 1) % ARRAY_SIZE;566 567 if (!offset ||568 slots[index].use_count() < Slab::available_chunks(chunk_size)) {569 uint64_t lane_mask = gpu::get_lane_mask();570 571 Slab *slab = slots[index].try_lock(lane_mask, uniform & lane_mask,572 reserved, chunk_size, index);573 574 // If we find a slab with a matching chunk size then we store the result.575 // Otherwise, we need to free the claimed lock and continue. In the case576 // of out-of-memory we receive a sentinel value and return a failure.577 if (slab && reserved < Slab::available_chunks(chunk_size) &&578 slab->get_chunk_size() == chunk_size) {579 if (index != start)580 indices[chunk_id].store(index, cpp::MemoryOrder::RELAXED);581 uniform = uniform & gpu::get_lane_mask();582 return slab;583 } else if (slab && (reserved >= Slab::available_chunks(chunk_size) ||584 slab->get_chunk_size() != chunk_size)) {585 slots[index].unlock(gpu::get_lane_mask(),586 gpu::get_lane_mask() & uniform);587 } else if (!slab && impl::is_sentinel(reserved)) {588 uniform = uniform & gpu::get_lane_mask();589 return nullptr;590 } else {591 sleep_briefly();592 }593 }594 }595 return nullptr;596}597 598// Release the lock associated with a given slab.599static void release_slab(Slab *slab) {600 uint32_t index = slab->get_global_index();601 uint64_t lane_mask = gpu::get_lane_mask();602 uint64_t uniform = gpu::match_any(lane_mask, index);603 slots[index].unlock(lane_mask, uniform);604}605 606namespace gpu {607 608void *allocate(uint64_t size) {609 if (!size)610 return nullptr;611 612 // Allocations requiring a full slab or more go directly to memory.613 if (size >= SLAB_SIZE / 2)614 return impl::rpc_allocate(impl::round_up<SLAB_SIZE>(size));615 616 // Try to find a slab for the rounded up chunk size and allocate from it.617 uint32_t chunk_size = impl::get_chunk_size(static_cast<uint32_t>(size));618 uint64_t uniform = gpu::match_any(gpu::get_lane_mask(), chunk_size);619 uint32_t reserved = 0;620 Slab *slab = find_slab(chunk_size, uniform, reserved);621 if (!slab)622 return nullptr;623 624 void *ptr = slab->allocate(uniform, reserved);625 return ptr;626}627 628void deallocate(void *ptr) {629 if (!ptr)630 return;631 632 // All non-slab allocations will be aligned on a 2MiB boundary.633 if (__builtin_is_aligned(ptr, SLAB_ALIGNMENT + 1))634 return impl::rpc_free(ptr);635 636 // The original slab pointer is the 2MiB boundary using the given pointer.637 Slab *slab = cpp::launder(reinterpret_cast<Slab *>(638 (reinterpret_cast<uintptr_t>(ptr) & ~SLAB_ALIGNMENT)));639 slab->deallocate(ptr);640 release_slab(slab);641}642 643void *reallocate(void *ptr, uint64_t size) {644 if (ptr == nullptr)645 return gpu::allocate(size);646 647 // Non-slab allocations are considered foreign pointers so we fail.648 if (__builtin_is_aligned(ptr, SLAB_ALIGNMENT + 1))649 return nullptr;650 651 // The original slab pointer is the 2MiB boundary using the given pointer.652 Slab *slab = cpp::launder(reinterpret_cast<Slab *>(653 (reinterpret_cast<uintptr_t>(ptr) & ~SLAB_ALIGNMENT)));654 if (slab->get_chunk_size() >= size)655 return ptr;656 657 // If we need a new chunk we reallocate and copy it over.658 void *new_ptr = gpu::allocate(size);659 inline_memcpy(new_ptr, ptr, slab->get_chunk_size());660 gpu::deallocate(ptr);661 return new_ptr;662}663 664void *aligned_allocate(uint32_t alignment, uint64_t size) {665 // All alignment values must be a non-zero power of two.666 if (!impl::is_pow2(alignment))667 return nullptr;668 669 // If the requested alignment is less than what we already provide this is670 // just a normal allocation.671 if (alignment <= MIN_ALIGNMENT + 1)672 return gpu::allocate(size);673 674 // We can't handle alignments greater than 2MiB so we simply fail.675 if (alignment > SLAB_ALIGNMENT + 1)676 return nullptr;677 678 // Trying to handle allocation internally would break the assumption that each679 // chunk is identical to eachother. Allocate enough memory with worst-case680 // alignment and then round up. The index logic will round down properly.681 uint64_t rounded = size + alignment - MIN_ALIGNMENT;682 void *ptr = gpu::allocate(rounded);683 return __builtin_align_up(ptr, alignment);684}685 686} // namespace gpu687} // namespace LIBC_NAMESPACE_DECL688