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

[via libc-commits]

Author: Joseph Huber
Date: 2023-09-12T14:59:02-05:00
New Revision: 688019851e696f13ce4701fc9a8ec3e116221288

URL: https://github.com/llvm/llvm-project/commit/688019851e696f13ce4701fc9a8ec3e116221288
DIFF: https://github.com/llvm/llvm-project/commit/688019851e696f13ce4701fc9a8ec3e116221288.diff

LOG: [libc][NFC] Factor GPU exiting into a common function (#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.

Added: 
    

Modified: 
    libc/src/__support/GPU/amdgpu/utils.h
    libc/src/__support/GPU/generic/utils.h
    libc/src/__support/GPU/nvptx/utils.h
    libc/src/__support/OSUtil/gpu/quick_exit.cpp
    libc/src/assert/gpu/__assert_fail.cpp
    libc/src/stdlib/gpu/abort.cpp

Removed: 
    


################################################################################
diff  --git a/libc/src/__support/GPU/amdgpu/utils.h b/libc/src/__support/GPU/amdgpu/utils.h
index 532a47b4a99058e..f2fa84bc4259fb6 100644
--- a/libc/src/__support/GPU/amdgpu/utils.h
+++ b/libc/src/__support/GPU/amdgpu/utils.h
@@ -176,6 +176,9 @@ LIBC_INLINE uint64_t fixed_frequency_clock() {
     return 0;
 }
 
+/// Terminates execution of the associated wavefront.
+[[noreturn]] LIBC_INLINE void end_program() { __builtin_amdgcn_endpgm(); }
+
 } // 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..c409f99cc4a8ba6 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; }
 
+[[noreturn]] LIBC_INLINE void end_program() { __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..dd51f78ffc473d1 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.
+[[noreturn]] LIBC_INLINE void end_program() {
+  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..3c43f6034b191e1 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::end_program();
 }
 
 } // namespace __llvm_libc

diff  --git a/libc/src/assert/gpu/__assert_fail.cpp b/libc/src/assert/gpu/__assert_fail.cpp
index b8ee168b069d164..30c5f7fff602c1b 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::end_program();
 
   // 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..441e776efd2349e 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::end_program();
 }
 
 } // namespace __llvm_libc