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

[jhuber6]

Summary

This patch adds the necessary support to provide assert functionality
through the GPU libc implementation. This implementation creates a
special-case GPU implementation rather than relying on the common
version. This is because the GPU has special considerings for printing.
The assertion is printed out in chunks with write_to_stderr, however
when combined with the GPU execution model this causes 32+ threads to
all execute in-lock step. Meaning that we'll get a horribly fragmented
message. Furthermore, potentially thousands of threads could hit the
assertion at once and try to print even if we had it all in one
printf.

This is solved by having a one-time lock that each thread group / wave /
warp will attempt to claim. We only let one thread group pass through
while the others simply stop executing. Finally only the first thread in
that group will do the printing until we finally abort execution.

[JonChesterfield]

I don't think the multiple writes to stderr in the normal implementation is an inherently good feature, how about we change the assert implementation on all platforms to compose a single write to stderr?

[sivachandra]

libc/src/assert/CMakeLists.txt

2

For the sake of consistency with other directories, you should move the generic implementation to a generic directory and list only an alias here:

if(TARGET libc.src.asser.${LIBC_TARGET_OS}.__assert_fail)
  set(assert_fail_dep libc.src.asser.${LIBC_TARGET_OS}.__assert_fail)
else()
  set(assert_fail_dep libc.src.assert.generic.__assert_fail)
endif()

add_entrypoint_object(
  __assert_fail
      ALIAS
      DEPENDS
        $[assert_fail_dep}
)

[jhuber6]

Address comments.

[jhuber6]

In D159296#4631913, @JonChesterfield wrote:

I don't think the multiple writes to stderr in the normal implementation is an inherently good feature, how about we change the assert implementation on all platforms to compose a single write to stderr?

It'd probably still be sufficiently different because of the warp level operations. I'll leave it to @sivachandra if he wants single printer semantics for this.

[sivachandra]

Did you move the generic __assert_fail.cpp file?

In D159296#4632014, @jhuber6 wrote:

In D159296#4631913, @JonChesterfield wrote:

I don't think the multiple writes to stderr in the normal implementation is an inherently good feature, how about we change the assert implementation on all platforms to compose a single write to stderr?

It'd probably still be sufficiently different because of the warp level operations. I'll leave it to @sivachandra if he wants single printer semantics for this.

I will change it separate from this patch.

[jhuber6]

I knew I was forgetting something.

[jhuber6]

And headers

[jhuber6]

More explicitly require that only the first lane does the atomic loading so that broadcast_value works.

[sivachandra]

OK for structuring changes.

[JonChesterfield]

Instead of this bespoke locking stuff, how about

if (gpu::is_first_lane(mask)) {
  report_assertion_failure();
  abort();
} else {
  endpgm();
}

?

[jhuber6]

In D159296#4632353, @JonChesterfield wrote:

Instead of this bespoke locking stuff, how about

if (gpu::is_first_lane(mask)) {
  report_assertion_failure();
  abort();
} else {
  endpgm();
}

?

endpgm on AMDGPU kills the whole wavefront. So you'd need to make sure that the printing and everything is done strictly first.

libc/config/gpu/api.td

 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"];
 }
 
 def StdlibAPI : PublicAPI<"stdlib.h"> {
   let Types = [
     "div_t",

libc/config/gpu/entrypoints.txt

 set(TARGET_LIBC_ENTRYPOINTS
+    # assert.h entrypoints
+    libc.src.assert.__assert_fail
+
     # ctype.h entrypoints
     libc.src.ctype.isalnum
     libc.src.ctype.isalpha
     libc.src.ctype.isascii
     libc.src.ctype.isblank
     libc.src.ctype.iscntrl
     libc.src.ctype.isdigit
     libc.src.ctype.isgraph

libc/config/gpu/headers.txt

 set(TARGET_PUBLIC_HEADERS
+    libc.include.assert
     libc.include.ctype
     libc.include.string
     libc.include.inttypes
     libc.include.math
     libc.include.fenv
     libc.include.time
     libc.include.errno
     libc.include.stdlib
     libc.include.stdio
 
     # Header for RPC extensions
     libc.include.gpu_rpc
 )

libc/docs/gpu/support.rst

 =============  =========  ============
 puts           |check|    |check|
 fputs          |check|    |check|
 fclose         |check|    |check|
 fopen          |check|    |check|
 fread          |check|    |check|
 =============  =========  ============
 
-stdio.h
+time.h
 --------
 
 =============  =========  ============
 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

 #if defined(LIBC_TARGET_ARCH_IS_AMDGPU)
 #include "amdgpu/utils.h"
 #elif defined(LIBC_TARGET_ARCH_IS_NVPTX)
 #include "nvptx/utils.h"
 #else
 #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

[[clang::convergent]] LIBC_INLINE void unlock(uint64_t lane_mask,
 
     // Wait for other threads in the warp to finish using the lock
     gpu::sync_lane(lane_mask);
 
     // Use exactly one thread to clear the nth bit in the lock array Must
     // 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);
   }
 
   /// Number of bytes to allocate for an inbox or outbox.
   LIBC_INLINE static constexpr uint64_t mailbox_bytes(uint32_t port_count) {
     return port_count * sizeof(cpp::Atomic<uint32_t>);
   }
 
for (uint32_t index = 0;; ++index) {
 
     // Once we acquire the index we need to check if we are in a valid sending
     // state.
     if (process.buffer_unavailable(in, out)) {
       process.unlock(lane_mask, index);
       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;
     }
     gpu::sync_lane(lane_mask);
     return Port(process, lane_mask, index, out);
   }
 }
 

libc/src/__support/RPC/rpc_util.h

#elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)
   __builtin_amdgcn_s_sleep(2);
 #elif defined(LIBC_TARGET_ARCH_IS_X86)
   __builtin_ia32_pause();
 #else
   // Simply do nothing if sleeping isn't supported on this platform.
 #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)
   return true;
 #else
   return false;
 #endif
 }

libc/src/assert/CMakeLists.txt

+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

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

libc/src/assert/gpu/CMakeLists.txt

 add_entrypoint_object(
   __assert_fail
   SRCS
     __assert_fail.cpp
   HDRS
-    __assert_fail.h
-    assert.h
+    ../__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

+//===-- 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