[libc] Support 'assert.h' on the GPU

libc/config/gpu/api.td

@@ -3,6 +3,38 @@ include "config/public_api.td"
 include "spec/stdc.td"
 include "spec/posix.td"
 include "spec/gpu_ext.td"
+include "spec/gnu_ext.td"
+include "spec/llvm_libc_ext.td"
+
+def AssertMacro : MacroDef<"assert"> {
+  let Defn = [{
+    #undef assert
+
+    #ifdef NDEBUG
+    #define assert(e) (void)0
+    #else
+
+    #define assert(e)  \
+      ((e) ? (void)0 : __assert_fail(#e, __FILE__, __LINE__, __PRETTY_FUNCTION__))
+    #endif
+  }];
+}
+
+def StaticAssertMacro : MacroDef<"static_assert"> {
+  let Defn = [{
+    #ifndef __cplusplus
+    #undef static_assert
+    #define static_assert _Static_assert
+    #endif
+  }];
+}
+
+def AssertAPI : PublicAPI<"assert.h"> {
+  let Macros = [
+    AssertMacro,
+    StaticAssertMacro,
+  ];
+}
 
 def StringAPI : PublicAPI<"string.h"> {
   let Types = ["size_t"];

libc/config/gpu/entrypoints.txt

@@ -1,4 +1,7 @@
 set(TARGET_LIBC_ENTRYPOINTS
+    # assert.h entrypoints
+    libc.src.assert.__assert_fail
+
     # ctype.h entrypoints
     libc.src.ctype.isalnum
     libc.src.ctype.isalpha

libc/config/gpu/headers.txt

@@ -1,4 +1,5 @@
 set(TARGET_PUBLIC_HEADERS
+    libc.include.assert
     libc.include.ctype
     libc.include.string
     libc.include.inttypes

libc/docs/gpu/support.rst

@@ -130,7 +130,7 @@ fopen          |check|    |check|
 fread          |check|    |check|
 =============  =========  ============
 
-stdio.h
+time.h
 --------
 
 =============  =========  ============
@@ -139,3 +139,13 @@ Function Name  Available  RPC Required
 clock          |check|
 nanosleep      |check|
 =============  =========  ============
+
+assert.h
+--------
+
+=============  =========  ============
+Function Name  Available  RPC Required
+=============  =========  ============
+assert         |check|    |check|
+__assert_fail  |check|    |check|
+=============  =========  ============

libc/src/__support/GPU/utils.h

@@ -19,4 +19,19 @@
 #include "generic/utils.h"
 #endif
 
+namespace __llvm_libc {
+namespace gpu {
+/// Get the first active thread inside the lane.
+LIBC_INLINE uint64_t get_first_lane_id(uint64_t lane_mask) {
+  return __builtin_ffsl(lane_mask) - 1;
+}
+
+/// Conditional that is only true for a single thread in a lane.
+LIBC_INLINE bool is_first_lane(uint64_t lane_mask) {
+  return gpu::get_lane_id() == get_first_lane_id(lane_mask);
+}
+
+} // namespace gpu
+} // namespace __llvm_libc
+
 #endif // LLVM_LIBC_SRC_SUPPORT_OSUTIL_IO_H

libc/src/__support/RPC/rpc.h

@@ -214,7 +214,7 @@ template <bool Invert, typename Packet> struct Process {
     // restrict to a single thread to avoid one thread dropping the lock, then
     // an unrelated warp claiming the lock, then a second thread in this warp
     // dropping the lock again.
-    clear_nth(lock, index, rpc::is_first_lane(lane_mask));
+    clear_nth(lock, index, gpu::is_first_lane(lane_mask));
     gpu::sync_lane(lane_mask);
   }
 
@@ -546,7 +546,7 @@ template <uint16_t opcode> LIBC_INLINE Client::Port Client::open() {
       continue;
     }
 
-    if (is_first_lane(lane_mask)) {
+    if (gpu::is_first_lane(lane_mask)) {
       process.packet[index].header.opcode = opcode;
       process.packet[index].header.mask = lane_mask;
     }

libc/src/__support/RPC/rpc_util.h

@@ -30,16 +30,6 @@ LIBC_INLINE void sleep_briefly() {
 #endif
 }
 
-/// Get the first active thread inside the lane.
-LIBC_INLINE uint64_t get_first_lane_id(uint64_t lane_mask) {
-  return __builtin_ffsl(lane_mask) - 1;
-}
-
-/// Conditional that is only true for a single thread in a lane.
-LIBC_INLINE bool is_first_lane(uint64_t lane_mask) {
-  return gpu::get_lane_id() == get_first_lane_id(lane_mask);
-}
-
 /// Conditional to indicate if this process is running on the GPU.
 LIBC_INLINE constexpr bool is_process_gpu() {
 #if defined(LIBC_TARGET_ARCH_IS_GPU)

libc/src/assert/CMakeLists.txt

@@ -1,12 +1,18 @@
+if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${LIBC_TARGET_OS})
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/${LIBC_TARGET_OS})
+else()
+  add_subdirectory(generic)
+endif()
+
+if(TARGET libc.src.assert.${LIBC_TARGET_OS}.__assert_fail)
+  set(assert_fail_dep libc.src.assert.${LIBC_TARGET_OS}.__assert_fail)
+else()
+  set(assert_fail_dep libc.src.assert.generic.__assert_fail)
+endif()
+
 add_entrypoint_object(
   __assert_fail
-  SRCS
-    __assert_fail.cpp
-  HDRS
-    __assert_fail.h
-    assert.h
+  ALIAS
   DEPENDS
-    libc.include.assert
-    libc.src.__support.OSUtil.osutil
-    libc.src.stdlib.abort
+    ${assert_fail_dep}
 )

libc/src/assert/generic/CMakeLists.txt

@@ -0,0 +1,12 @@
+add_entrypoint_object(
+  __assert_fail
+  SRCS
+    __assert_fail.cpp
+  HDRS
+    ../__assert_fail.h
+    ../assert.h
+  DEPENDS
+    libc.include.assert
+    libc.src.__support.OSUtil.osutil
+    libc.src.stdlib.abort
+)

libc/src/assert/gpu/CMakeLists.txt

@@ -0,0 +1,14 @@
+add_entrypoint_object(
+  __assert_fail
+  SRCS
+    __assert_fail.cpp
+  HDRS
+    ../__assert_fail.h
+    ../assert.h
+  DEPENDS
+    libc.include.assert
+    libc.src.__support.OSUtil.osutil
+    libc.src.__support.GPU.utils
+    libc.src.__support.CPP.atomic
+    libc.src.stdlib.abort
+)

libc/src/assert/gpu/__assert_fail.cpp

@@ -0,0 +1,45 @@
+//===-- GPU definition of a libc internal assert macro ----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "src/assert/__assert_fail.h"
+
+#include "src/__support/CPP/atomic.h"
+#include "src/__support/GPU/utils.h"
+#include "src/__support/libc_assert.h"
+#include "src/stdlib/abort.h"
+
+namespace __llvm_libc {
+
+// A single-use lock to allow only a single thread to print the assertion.
+static cpp::Atomic<uint32_t> lock = 0;
+
+LLVM_LIBC_FUNCTION(void, __assert_fail,
+                   (const char *assertion, const char *file, unsigned line,
+                    const char *function)) {
+  uint64_t mask = gpu::get_lane_mask();
+  // We only want a single work group or warp to handle the assertion. Each
+  // group attempts to claim the lock, if it is already claimed we simply exit.
+  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();
+  }
+
+  // Only a single line should be printed if an assertion is hit.
+  if (gpu::is_first_lane(mask))
+    __llvm_libc::report_assertion_failure(assertion, file, line, function);
+  __llvm_libc::abort();
+}
+
+} // namespace __llvm_libc