[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.
----------------
jhuber6 wrote:

Might be a good utility function overall, since you generally want something like this for every atomic spin loop.

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