[libc-commits] [libc] 78c460b - [libc] Fix incorrect count when initializing slab

[Joseph Huber via libc-commits]

Author: Joseph Huber
Date: 2025-07-29T18:50:02-05:00
New Revision: 78c460bbe8f1fc17e2e66b37edf419ccecbfecba

URL: https://github.com/llvm/llvm-project/commit/78c460bbe8f1fc17e2e66b37edf419ccecbfecba
DIFF: https://github.com/llvm/llvm-project/commit/78c460bbe8f1fc17e2e66b37edf419ccecbfecba.diff

LOG: [libc] Fix incorrect count when initializing slab

Summary:
The initialization code should share the result with all of its
neighbors. Right now it sets them to the sentinel value and doesn't
shuffle them correctly. Shuffle them after initialization so we
correctly report that we succeeded in the allocation.

Added: 
    

Modified: 
    libc/src/__support/GPU/allocator.cpp

Removed: 
    


################################################################################
diff  --git a/libc/src/__support/GPU/allocator.cpp b/libc/src/__support/GPU/allocator.cpp
index 2b78c4dd1c33c..8fff4cca6607a 100644
--- a/libc/src/__support/GPU/allocator.cpp
+++ b/libc/src/__support/GPU/allocator.cpp
@@ -460,11 +460,13 @@ struct GuardPtr {
       result->initialize(uniform);
       if (gpu::get_lane_id() == uint32_t(cpp::countr_zero(uniform)))
         finalize(result, cpp::popcount(uniform), count);
+      count =
+          gpu::shuffle(gpu::get_lane_mask(), cpp::countr_zero(uniform), count);
     }
 
     if (!impl::is_sentinel(count))
       count = count - cpp::popcount(uniform) +
-              impl::lane_count(uniform, gpu::get_lane_id()) + 1;
+              impl::lane_count(uniform, gpu::get_lane_id());
 
     return result;
   }
@@ -536,13 +538,13 @@ static Slab *find_slab(uint32_t chunk_size, uint64_t &uniform) {
       // If we find a slab with a matching chunk size then we store the result.
       // Otherwise, we need to free the claimed lock and continue. In the case
       // of out-of-memory we receive a sentinel value and return a failure.
-      if (slab && reserved <= Slab::available_chunks(chunk_size) &&
+      if (slab && reserved < Slab::available_chunks(chunk_size) &&
           slab->get_chunk_size() == chunk_size) {
         if (index != start)
           indices[chunk_id].store(index, cpp::MemoryOrder::RELAXED);
         uniform = uniform & gpu::get_lane_mask();
         return slab;
-      } else if (slab && (reserved > Slab::available_chunks(chunk_size) ||
+      } else if (slab && (reserved >= Slab::available_chunks(chunk_size) ||
                           slab->get_chunk_size() != chunk_size)) {
         slots[index].unlock(gpu::get_lane_mask(),
                             gpu::get_lane_mask() & uniform);