[libc] Implement 'atexit' on the GPU correctly

libc/docs/gpu/support.rst

@@ -102,6 +102,7 @@ atol           |check|
 atoll          |check|
 exit           |check|    |check|
 abort          |check|    |check|
+atexit         |check|
 labs           |check|
 llabs          |check|
 div            |check|

libc/src/stdlib/CMakeLists.txt

@@ -363,20 +363,22 @@ add_entrypoint_object(
     libc.src.__support.OSUtil.osutil
 )
 
-add_entrypoint_object(
-  atexit
-  SRCS
-    atexit.cpp
-  HDRS
-    atexit.h
-  CXX_STANDARD
-    20 # For constinit of the atexit callback list.
-  DEPENDS
-    libc.src.__support.fixedvector
-    libc.src.__support.blockstore
-    libc.src.__support.threads.mutex
-    libc.src.__support.CPP.new
-)
+if(NOT LIBC_TARGET_OS_IS_GPU)
+  add_entrypoint_object(
+    atexit
+    SRCS
+      atexit.cpp
+    HDRS
+      atexit.h
+    CXX_STANDARD
+      20 # For constinit of the atexit callback list.
+    DEPENDS
+      libc.src.__support.fixedvector
+      libc.src.__support.blockstore
+      libc.src.__support.threads.mutex
+      libc.src.__support.CPP.new
+  )
+endif()
 
 add_entrypoint_object(
   exit
@@ -398,6 +400,13 @@ add_entrypoint_object(
 )
 
 if(LIBC_TARGET_OS_IS_GPU)
+  add_entrypoint_object(
+    atexit
+    ALIAS
+    DEPENDS
+      .${LIBC_TARGET_OS}.atexit
+  )
+
   add_entrypoint_object(
     malloc
     ALIAS

libc/src/stdlib/atexit.cpp

@@ -28,14 +28,7 @@ struct AtExitUnit {
   constexpr AtExitUnit(AtExitCallback *c, void *p) : callback(c), payload(p) {}
 };
 
-#if defined(LIBC_TARGET_ARCH_IS_GPU)
-// The GPU build cannot handle the potentially recursive definitions required by
-// the BlockStore class. Additionally, the liklihood that someone exceeds this
-// while executing on the GPU is extremely small.
-// FIXME: It is not generally safe to use 'atexit' on the GPU because the
-//        mutexes simply passthrough. We will need a lock free stack.
-using ExitCallbackList = FixedVector<AtExitUnit, 64>;
-#elif defined(LIBC_COPT_PUBLIC_PACKAGING)
+#if defined(LIBC_COPT_PUBLIC_PACKAGING)
 using ExitCallbackList = cpp::ReverseOrderBlockStore<AtExitUnit, 32>;
 #else
 // BlockStore uses dynamic memory allocation. To avoid dynamic memory

libc/src/stdlib/gpu/CMakeLists.txt

@@ -30,3 +30,16 @@ add_entrypoint_object(
     libc.include.stdlib
     libc.src.__support.RPC.rpc_client
 )
+
+add_entrypoint_object(
+  atexit
+  SRCS
+    atexit.cpp
+  HDRS
+    ../atexit.h
+  CXX_STANDARD
+    20 # For constinit of the atexit callback list.
+  DEPENDS
+    libc.include.stdlib
+    libc.src.__support.fixedstack
+)

libc/src/stdlib/gpu/atexit.cpp

@@ -0,0 +1,63 @@
+//===-- GPU implementation of atexit --------------------------------------===//
+//
+// 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/stdlib/atexit.h"
+#include "src/__support/common.h"
+#include "src/__support/fixedstack.h"
+
+namespace LIBC_NAMESPACE {
+
+namespace {
+
+using AtExitCallback = void(void *);
+using StdCAtExitCallback = void(void);
+
+struct AtExitUnit {
+  AtExitCallback *callback = nullptr;
+  void *payload = nullptr;
+  constexpr AtExitUnit() = default;
+  constexpr AtExitUnit(AtExitCallback *c, void *p) : callback(c), payload(p) {}
+};
+
+// The GPU interface cannot use the standard implementation because it does not
+// support the Mutex type. Instead we use a lock free stack with a sufficiently
+// large size.
+constinit FixedStack<AtExitUnit, CALLBACK_LIST_SIZE_FOR_TESTS> exit_callbacks;
+
+void stdc_at_exit_func(void *payload) {
+  reinterpret_cast<StdCAtExitCallback *>(payload)();
+}
+
+} // namespace
+
+namespace internal {
+
+void call_exit_callbacks() {
+  AtExitUnit unit;
+  while (exit_callbacks.pop(unit))
+    unit.callback(unit.payload);
+}
+
+} // namespace internal
+
+static int add_atexit_unit(const AtExitUnit &unit) {
+  if (!exit_callbacks.push(unit))
+    return -1;
+  return 0;
+}
+
+extern "C" int __cxa_atexit(AtExitCallback *callback, void *payload, void *) {
+  return add_atexit_unit({callback, payload});
+}
+
+LLVM_LIBC_FUNCTION(int, atexit, (StdCAtExitCallback * callback)) {
+  return add_atexit_unit(
+      {&stdc_at_exit_func, reinterpret_cast<void *>(callback)});
+}
+
+} // namespace LIBC_NAMESPACE

libc/startup/gpu/amdgpu/CMakeLists.txt

@@ -5,6 +5,7 @@ add_startup_object(
   DEPENDS
     libc.src.__support.RPC.rpc_client
     libc.src.__support.GPU.utils
+    libc.src.__support.CPP.new
     libc.src.stdlib.exit
     libc.src.stdlib.atexit
   COMPILE_OPTIONS

libc/startup/gpu/nvptx/CMakeLists.txt

@@ -5,6 +5,7 @@ add_startup_object(
   DEPENDS
     libc.src.__support.RPC.rpc_client
     libc.src.__support.GPU.utils
+    libc.src.__support.CPP.new
     libc.src.stdlib.exit
     libc.src.stdlib.atexit
   COMPILE_OPTIONS