[libc] Implement simple lock-free stack data structure

Summary:
This patch implements a `FixedStack` which represents a lock-free stack
implemented using a fixed-size memory buffer. The utility for this class
is to implement a data structure that the GPU implementation can use
incases where a mutex is normally used. We cannot implement a
general-purpose mutex on the GPU due to the lack of a guaranteed fair
thread scheduler.

This lock free stack is implemented as a pair of 'pointers' to a used
and free stack. The pointers here are simple 32-bit indexes into an
underlying fixed-size memory buffer. The free stack is initialized to
point to the entire buffer, i.e. `1, 2, 3, ..., N` with N being used as
a sentinel value. To perform a push operation, we pop a node off of the
free stack and then push it into the used stack and vice-versa for a pop
operation.

The underlying pop implementation relies on continually trying to update
the head pointer to the next node using atomic CAS. The CAS loop will
repeat until it reads the next pointer successfully and the head pointer
has not changed.

The underlying push ipmlementation relies on contiually trying to update
the head pointer to the new node. The CAS loop will reap until we write
to the head pointer and it has not changed.

Both of these implementations rely on detecting whether or not the head
pointer has changed. Simply using indexes we are succeptiable to the ABA
problem as the stack could have been pushed and popped until we are back
to the same index and we have not noticed. For this reason, the 'head'
pointer is augmented with a 32-bit ABA that increments each time it is
updated. This allows the head update to be done in a single 64-bit
atomic update which is supported by hardware.

Author: jhuber6

libc/src/__support/CMakeLists.txt

@@ -177,6 +177,15 @@ add_header_library(
     libc.src.__support.CPP.array
 )
 
+add_header_library(
+  fixedstack
+  HDRS
+    fixedstack.h
+  DEPENDS
+    libc.src.__support.CPP.array
+    libc.src.__support.CPP.atomic
+)
+
 add_header_library(
   char_vector
   HDRS

libc/src/__support/fixedstack.h

@@ -0,0 +1,129 @@
+//===-- A lock-free data structure for a fixed capacity stack ---*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_SRC___SUPPORT_FIXEDSTACK_H
+#define LLVM_LIBC_SRC___SUPPORT_FIXEDSTACK_H
+
+#include "src/__support/CPP/array.h"
+#include "src/__support/CPP/atomic.h"
+#include "src/__support/threads/sleep.h"
+
+#include <stdint.h>
+
+namespace LIBC_NAMESPACE {
+
+// A lock-free fixed size stack backed by an underlying cpp::array data
+// structure. It supports push and pop operations in a thread safe manner.
+template <typename T, uint32_t CAPACITY> class alignas(16) FixedStack {
+  static_assert(CAPACITY < 1024 * 1024, "Invalid buffer size");
+
+  // The head of the free and used stacks. Represents as a 20-bit index combined
+  // with a 44-bit ABA tag that is updated in a single atomic operation.
+  uint64_t free;
+  uint64_t used;
+
+  // The stack is a linked list of indices into the underlying data
+  cpp::array<uint32_t, CAPACITY> next;
+  cpp::array<T, CAPACITY> data;
+
+  // Get the 20-bit index into the underlying array from the head.
+  static constexpr uint32_t get_node(uint64_t head) {
+    return static_cast<uint32_t>(head & 0xffff);
+  }
+
+  // Increment the old ABA tag and merge it into the new index.
+  static constexpr uint64_t make_new_head(uint64_t orig, uint32_t node) {
+    return static_cast<uint64_t>(node) | (((orig >> 20ul) + 1ul) << 20ul);
+  }
+
+  // Helper macros for the atomic operations. We cannot use the standard
+  // cpp::atomic helpers because the initializer will no longer be constexpr and
+  // the NVPTX backend cannot currently support all of the atomics.
+#define atomic_load(val, mem_order) __atomic_load_n(val, (int)mem_order)
+#define atomic_cas(val, expected, desired, success_order, failure_order)       \
+  __atomic_compare_exchange_n(val, expected, desired, /*weak=*/true,           \
+                              (int)success_order, (int)failure_order)
+
+  // Attempts to pop data from the given stack by making it point to the next
+  // node. We repeatedly attempt to write to the head using compare-and-swap,
+  // expecting that it has not been changed by any other thread.
+  uint32_t pop_impl(uint64_t *head) {
+    uint64_t orig = atomic_load(head, cpp::MemoryOrder::RELAXED);
+
+    for (;;) {
+      if (get_node(orig) == CAPACITY)
+        return CAPACITY;
+
+      uint32_t node =
+          atomic_load(&next[get_node(orig)], cpp::MemoryOrder::RELAXED);
+      if (atomic_cas(head, &orig, make_new_head(orig, node),
+                     cpp::MemoryOrder::ACQUIRE, cpp::MemoryOrder::RELAXED))
+        break;
+      sleep_briefly();
+    }
+    return get_node(orig);
+  }
+
+  // Attempts to push data to the given stack by making it point to the new
+  // node. We repeatedly attempt to write to the head using compare-and-swap,
+  // expecting that it has not been changed by any other thread.
+  uint32_t push_impl(uint64_t *head, uint32_t node) {
+    uint64_t orig = atomic_load(head, cpp::MemoryOrder::RELAXED);
+    for (;;) {
+      next[node] = get_node(orig);
+      if (atomic_cas(head, &orig, make_new_head(orig, node),
+                     cpp::MemoryOrder::RELEASE, cpp::MemoryOrder::RELAXED))
+        break;
+      sleep_briefly();
+    }
+    return get_node(*head);
+  }
+
+public:
+  // Initialize the free stack to be full and the used stack to be empty. We use
+  // the capacity of the stack as a sentinel value.
+  constexpr FixedStack() : free(0), used(CAPACITY), data{} {
+    for (uint32_t i = 0; i < CAPACITY; ++i)
+      next[i] = i + 1;
+  }
+
+  bool push(const T &val) {
+    uint32_t node = pop_impl(&free);
+    if (node == CAPACITY)
+      return false;
+
+    data[node] = val;
+    push_impl(&used, node);
+    return true;
+  }
+
+  bool pop(T &val) {
+    uint32_t node = pop_impl(&used);
+    if (node == CAPACITY)
+      return false;
+
+    val = data[node];
+    push_impl(&free, node);
+    return true;
+  }
+
+  bool empty() const {
+    return get_node(atomic_load(&used, cpp::MemoryOrder::RELAXED)) == CAPACITY;
+  }
+
+  bool full() const {
+    return get_node(atomic_load(&free, cpp::MemoryOrder::RELAXED)) == CAPACITY;
+  }
+
+#undef atomic_load
+#undef atomic_cas
+};
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC___SUPPORT_FIXEDSTACK_H

libc/test/integration/src/__support/CMakeLists.txt

@@ -1 +1,5 @@
+if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${LIBC_TARGET_OS})
+  add_subdirectory(${LIBC_TARGET_OS})
+endif()
+
 add_subdirectory(threads)

libc/test/integration/src/__support/gpu/CMakeLists.txt

@@ -0,0 +1,19 @@
+add_custom_target(support-gpu-integration-tests)
+add_dependencies(libc-integration-tests support-gpu-integration-tests)
+
+add_integration_test(
+  support_fixed_stack_test
+  SUITE support-gpu-integration-tests
+  SRCS
+    fixed_stack_test.cpp
+  DEPENDS
+    libc.src.__support.GPU.utils
+    libc.src.__support.fixedstack
+  LOADER_ARGS
+    --blocks-x 2
+    --blocks-y 2
+    --blocks-z 2
+    --threads-x 4
+    --threads-y 4
+    --threads-z 4
+)

libc/test/integration/src/__support/gpu/fixed_stack_test.cpp

@@ -0,0 +1,75 @@
+//===-- Integration test for the lock-free stack --------------------------===//
+//
+// 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/__support/GPU/utils.h"
+#include "src/__support/fixedstack.h"
+#include "test/IntegrationTest/test.h"
+
+using namespace LIBC_NAMESPACE;
+
+void single_thread() {
+  // FIXME: The NVPTX backend cannot handle atomic CAS on a local address space.
+#if defined(LIBC_TARGET_ARCH_IS_AMDGPU)
+  FixedStack<int, 16> local_stack;
+
+  for (int i = 0; i < 16; ++i)
+    EXPECT_TRUE(local_stack.push(i));
+  ASSERT_TRUE(local_stack.full());
+
+  int val;
+  for (int i = 0; i < 16; ++i) {
+    EXPECT_TRUE(local_stack.pop(val));
+    EXPECT_EQ(val, 16 - 1 - i);
+  }
+  ASSERT_TRUE(local_stack.empty());
+#endif
+}
+
+static FixedStack<uint32_t, 2048> global_stack;
+void multiple_threads() {
+  // We need enough space in the stack as threads in flight can temporarily
+  // consume memory before they finish comitting it back to the stack.
+  ASSERT_EQ(gpu::get_num_blocks() * gpu::get_num_threads(), 512);
+
+  uint32_t val;
+  uint32_t num_threads = static_cast<uint32_t>(gpu::get_num_threads());
+  for (int i = 0; i < 256; ++i) {
+    EXPECT_TRUE(global_stack.push(UINT32_MAX))
+    EXPECT_TRUE(global_stack.pop(val))
+    ASSERT_TRUE(val < num_threads || val == UINT32_MAX);
+  }
+
+  EXPECT_TRUE(global_stack.push(static_cast<uint32_t>(gpu::get_thread_id())));
+  EXPECT_TRUE(global_stack.push(static_cast<uint32_t>(gpu::get_thread_id())));
+  EXPECT_TRUE(global_stack.pop(val));
+  ASSERT_TRUE(val < num_threads || val == UINT32_MAX);
+
+  // Fill the rest of the stack with the default value.
+  while (!global_stack.push(UINT32_MAX))
+    ;
+}
+
+// Once all the threads have finished executing check the final state of the
+// stack. Destructors are always run with a single thread on the GPU.
+[[gnu::destructor]] void check_stack() {
+  ASSERT_FALSE(global_stack.empty());
+
+  while (!global_stack.empty()) {
+    uint32_t val;
+    ASSERT_TRUE(global_stack.pop(val));
+    ASSERT_TRUE(val < 64 || val == UINT32_MAX);
+  }
+}
+
+TEST_MAIN(int argc, char **argv, char **envp) {
+  single_thread();
+
+  multiple_threads();
+
+  return 0;
+}

libc/test/src/__support/CMakeLists.txt

@@ -122,6 +122,16 @@ add_libc_test(
     libc.src.__support.fixedvector
 )
 
+add_libc_test(
+  fixedstack_test
+  SUITE
+    libc-support-tests
+  SRCS
+    fixedstack_test.cpp
+  DEPENDS
+    libc.src.__support.fixedstack
+)
+
 add_libc_test(
   char_vector_test
   SUITE

libc/test/src/__support/fixedstack_test.cpp

@@ -0,0 +1,26 @@
+//===-- Unittests for FixedStack ------------------------------------------===//
+//
+// 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/__support/fixedstack.h"
+#include "test/UnitTest/Test.h"
+
+TEST(LlvmLibcFixedVectorTest, PushAndPop) {
+  static LIBC_NAMESPACE::FixedStack<int, 20> fixed_stack;
+  ASSERT_TRUE(fixed_stack.empty());
+  for (int i = 0; i < 20; i++)
+    ASSERT_TRUE(fixed_stack.push(i));
+  ASSERT_FALSE(fixed_stack.empty());
+  ASSERT_FALSE(fixed_stack.push(123));
+  int val;
+  for (int i = 20; i > 0; --i) {
+    ASSERT_TRUE(fixed_stack.pop(val));
+    ASSERT_EQ(val, i - 1);
+  }
+  ASSERT_FALSE(fixed_stack.pop(val));
+  ASSERT_TRUE(fixed_stack.empty());
+}