[libc-commits] [libc] [libc] Implement simple lock-free stack data structure (PR #83026)

[Joseph Huber via libc-commits]

================
@@ -0,0 +1,141 @@
+//===-- 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 <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 < UINT32_MAX, "Invalid buffer size");
+
+  // The head of the free and used stacks. Represents as a 32-bit index combined
+  // with a 32-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 32-bit index into the underlying array from the head.
+  static constexpr uint32_t get_node(uint64_t head) {
+    return static_cast<uint32_t>(head & 0xffffffff);
+  }
+
+  // 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 >> 32ul) + 1ul) << 32ul);
+  }
+
+  void sleep_briefly() {
+#if defined(LIBC_TARGET_ARCH_IS_NVPTX)
+    if (__nvvm_reflect("__CUDA_ARCH") >= 700)
+      LIBC_INLINE_ASM("nanosleep.u32 32;" ::: "memory");
+#elif defined(LIBC_TARGET_ARCH_IS_AMDGPU)
+    __builtin_amdgcn_s_sleep(1);
+#elif defined(LIBC_TARGET_ARCH_IS_X86)
+    __builtin_ia32_pause();
+#else
+    // Simply do nothing if sleeping isn't supported on this platform.
+#endif
+  }
+
+  // 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,           \
----------------
jhuber6 wrote:

The header comment explains why. Here's why
1. When the type signature is atomic it cannot be intialized constexpr by the constructor
2. When used as a function it invokes a switch-case for all combinations of atomic, many of which are not implemented for the NVPTX backend
3. We don't actually need the atomicity all the time, when we pop a node off we effectively own it so we do not need it to be atomic.

https://github.com/llvm/llvm-project/pull/83026