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

[Joseph Huber via libc-commits]

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

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.


>From 7c4109a7c0b1b3cd9eda47f8e360de2d4ce29556 Mon Sep 17 00:00:00 2001
From: Joseph Huber <jhuber6 at vols.utk.edu>
Date: Tue, 12 Sep 2023 09:43:08 -0500
Subject: [PATCH] [libc][NFC] Factor GPU exiting into a common function

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.
---
 libc/src/__support/GPU/amdgpu/utils.h        |  6 ++++++
 libc/src/__support/GPU/generic/utils.h       |  2 ++
 libc/src/__support/GPU/nvptx/utils.h         |  6 ++++++
 libc/src/__support/OSUtil/gpu/quick_exit.cpp |  9 +--------
 libc/src/assert/gpu/__assert_fail.cpp        | 10 ++--------
 libc/src/stdlib/gpu/abort.cpp                |  7 +------
 6 files changed, 18 insertions(+), 22 deletions(-)

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