[libc-commits] [libc] [libc][NFC] Factor GPU exiting into a common function (PR #66093)

[via libc-commits]

llvmbot wrote:

@llvm/pr-subscribers-libc

<details>
<summary>Changes</summary>

Summary:
We currently call the GPU routine to terminate the current thread in
three separate locations .This should be wrapped into a helper function
to simplify the implementation.

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

6 Files Affected:

- (modified) libc/src/__support/GPU/amdgpu/utils.h (+6) 
- (modified) libc/src/__support/GPU/generic/utils.h (+2) 
- (modified) libc/src/__support/GPU/nvptx/utils.h (+6) 
- (modified) libc/src/__support/OSUtil/gpu/quick_exit.cpp (+1-8) 
- (modified) libc/src/assert/gpu/__assert_fail.cpp (+2-8) 
- (modified) libc/src/stdlib/gpu/abort.cpp (+1-6) 


<pre>
diff --git a/libc/src/__support/GPU/amdgpu/utils.h b/libc/src/__support/GPU/amdgpu/utils.h
index 532a47b4a99058e..45a4192c0c04f20 100644
--- a/libc/src/__support/GPU/amdgpu/utils.h
+++ b/libc/src/__support/GPU/amdgpu/utils.h
@@ -176,6 +176,12 @@ LIBC_INLINE uint64_t fixed_frequency_clock() {
     return 0;
 }
 
+/// Terminates execution of the associated wavefront.
+LIBC_INLINE void exit() {
+  __builtin_amdgcn_endpgm();
+  __builtin_unreachable();
+}
+
 } // namespace gpu
 } // namespace __llvm_libc
 
diff --git a/libc/src/__support/GPU/generic/utils.h b/libc/src/__support/GPU/generic/utils.h
index 5c7913176b7ce7c..3316e759f50ce2c 100644
--- a/libc/src/__support/GPU/generic/utils.h
+++ b/libc/src/__support/GPU/generic/utils.h
@@ -73,6 +73,8 @@ LIBC_INLINE uint64_t processor_clock() { return 0; }
 
 LIBC_INLINE uint64_t fixed_frequency_clock() { return 0; }
 
+LIBC_INLINE void exit() { __builtin_unreachable(); }
+
 } // namespace gpu
 } // namespace __llvm_libc
 
diff --git a/libc/src/__support/GPU/nvptx/utils.h b/libc/src/__support/GPU/nvptx/utils.h
index 5c6cf3a3a8e12cf..94912184557e643 100644
--- a/libc/src/__support/GPU/nvptx/utils.h
+++ b/libc/src/__support/GPU/nvptx/utils.h
@@ -153,6 +153,12 @@ LIBC_INLINE uint64_t fixed_frequency_clock() {
   return nsecs;
 }
 
+/// Terminates execution of the calling thread.
+LIBC_INLINE void exit() {
+  LIBC_INLINE_ASM("exit;" ::: "memory");
+  __builtin_unreachable();
+}
+
 } // namespace gpu
 } // namespace __llvm_libc
 
diff --git a/libc/src/__support/OSUtil/gpu/quick_exit.cpp b/libc/src/__support/OSUtil/gpu/quick_exit.cpp
index 470b2bdbdaa70c7..18f7b4cc6b78e7f 100644
--- a/libc/src/__support/OSUtil/gpu/quick_exit.cpp
+++ b/libc/src/__support/OSUtil/gpu/quick_exit.cpp
@@ -25,14 +25,7 @@ void quick_exit(int status) {
   });
   port.close();
 
-#if defined(LIBC_TARGET_ARCH_IS_NVPTX)
-  LIBC_INLINE_ASM("exit;" ::: "memory");
-#elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)
-  // This will terminate the entire wavefront, may not be valid with divergent
-  // work items.
-  __builtin_amdgcn_endpgm();
-#endif
-  __builtin_unreachable();
+  gpu::exit();
 }
 
 } // namespace __llvm_libc
diff --git a/libc/src/assert/gpu/__assert_fail.cpp b/libc/src/assert/gpu/__assert_fail.cpp
index b8ee168b069d164..3421848d4ad873f 100644
--- a/libc/src/assert/gpu/__assert_fail.cpp
+++ b/libc/src/assert/gpu/__assert_fail.cpp
@@ -27,14 +27,8 @@ LLVM_LIBC_FUNCTION(void, __assert_fail,
   uint32_t claimed = gpu::is_first_lane(mask)
                          ? !lock.fetch_or(1, cpp::MemoryOrder::ACQUIRE)
                          : 0;
-  if (!gpu::broadcast_value(mask, claimed)) {
-#if defined(LIBC_TARGET_ARCH_IS_NVPTX)
-    LIBC_INLINE_ASM("exit;" ::: "memory");
-#elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)
-    __builtin_amdgcn_endpgm();
-#endif
-    __builtin_unreachable();
-  }
+  if (!gpu::broadcast_value(mask, claimed))
+    gpu::exit();
 
   // Only a single line should be printed if an assertion is hit.
   if (gpu::is_first_lane(mask))
diff --git a/libc/src/stdlib/gpu/abort.cpp b/libc/src/stdlib/gpu/abort.cpp
index 5e208dba88169db..4c029f18d864f60 100644
--- a/libc/src/stdlib/gpu/abort.cpp
+++ b/libc/src/stdlib/gpu/abort.cpp
@@ -20,12 +20,7 @@ LLVM_LIBC_FUNCTION(void, abort, ()) {
   port.send([&](rpc::Buffer *) {});
   port.close();
 
-#if defined(LIBC_TARGET_ARCH_IS_NVPTX)
-  LIBC_INLINE_ASM("exit;" ::: "memory");
-#elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)
-  __builtin_amdgcn_endpgm();
-#endif
-  __builtin_unreachable();
+  gpu::exit();
 }
 
 } // namespace __llvm_libc
</pre>

</details>

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