[libc-commits] [libc] [libc] Sleep to maximize slab re-use in GPU malloc (PR #143609)

[via libc-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-libc

Author: Joseph Huber (jhuber6)

<details>
<summary>Changes</summary>

Summary:
This patch simply adds a long sleep after we set the reference counter
to zero. Hopefully this suspends the thread long enough that another
will come along and revive the counter instead of having it get
deallocated and then needing to make a new slab.


---
Full diff: https://github.com/llvm/llvm-project/pull/143609.diff


1 Files Affected:

- (modified) libc/src/__support/GPU/allocator.cpp (+14) 


``````````diff
diff --git a/libc/src/__support/GPU/allocator.cpp b/libc/src/__support/GPU/allocator.cpp
index 135ced3df704c..f664ffe180069 100644
--- a/libc/src/__support/GPU/allocator.cpp
+++ b/libc/src/__support/GPU/allocator.cpp
@@ -129,6 +129,16 @@ static inline constexpr T round_up(const T x) {
   return (x + N) & ~(N - 1);
 }
 
+// Sleep for an extended period of time to allow other threads to progress.
+static inline void sleep_extensively() {
+#if defined(LIBC_TARGET_ARCH_IS_NVPTX)
+  if (__nvvm_reflect("__CUDA_ARCH") >= 700)
+    LIBC_INLINE_ASM("nanosleep.u32 8192;" :: : "memory");
+#elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)
+  __builtin_amdgcn_s_sleep(128);
+#endif
+}
+
 } // namespace impl
 
 /// A slab allocator used to hand out identically sized slabs of memory.
@@ -313,6 +323,10 @@ template <typename T> struct GuardPtr {
       // another thread resurrected the counter and we quit, or a parallel read
       // helped us invalidating it. For the latter, claim that flag and return.
       if (counter.fetch_sub(n, cpp::MemoryOrder::RELAXED) == n) {
+        // Yield this thread here to maximize the chance that this CAS fails and
+        // we re-use the slab instead of deallocating it.
+        impl::sleep_extensively();
+
         uint64_t expected = 0;
         if (counter.compare_exchange_strong(expected, INVALID,
                                             cpp::MemoryOrder::RELAXED,

``````````

</details>


https://github.com/llvm/llvm-project/pull/143609