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

[Joseph Huber via libc-commits]

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

>From da842a765a9815915d4de3ac6d1a9a0d39839218 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        |  3 +++
 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, 15 insertions(+), 22 deletions(-)

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