112 lines · c
1//===-- A lock-free data structure for a fixed capacity stack ---*- C++ -*-===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#ifndef LLVM_LIBC_SRC___SUPPORT_GPU_FIXEDSTACK_H10#define LLVM_LIBC_SRC___SUPPORT_GPU_FIXEDSTACK_H11 12#include "src/__support/CPP/atomic.h"13#include "src/__support/threads/sleep.h"14 15#include <stdint.h>16 17namespace LIBC_NAMESPACE_DECL {18 19// A lock-free fixed size stack backed by an underlying array of data. It20// supports push and pop operations in a completely lock-free manner.21template <typename T, uint32_t CAPACITY> struct alignas(16) FixedStack {22 // The index is stored as a 20-bit value and cannot index into any more.23 static_assert(CAPACITY < 1024 * 1024, "Invalid buffer size");24 25 // The head of the free and used stacks. Represents as a 20-bit index combined26 // with a 44-bit ABA tag that is updated in a single atomic operation.27 uint64_t free;28 uint64_t used;29 30 // The stack is a linked list of indices into the underlying data31 uint32_t next[CAPACITY];32 T data[CAPACITY];33 34 // Get the 20-bit index into the underlying array from the head.35 LIBC_INLINE static constexpr uint32_t get_node(uint64_t head) {36 return static_cast<uint32_t>(head & 0xfffff);37 }38 39 // Increment the old ABA tag and merge it into the new index.40 LIBC_INLINE static constexpr uint64_t make_head(uint64_t orig,41 uint32_t node) {42 return static_cast<uint64_t>(node) | (((orig >> 20ul) + 1ul) << 20ul);43 }44 45 // Attempts to pop data from the given stack by making it point to the next46 // node. We repeatedly attempt to write to the head using compare-and-swap,47 // expecting that it has not been changed by any other thread.48 LIBC_INLINE uint32_t pop_impl(cpp::AtomicRef<uint64_t> head) {49 uint64_t orig = head.load(cpp::MemoryOrder::RELAXED);50 51 for (;;) {52 if (get_node(orig) == CAPACITY)53 return CAPACITY;54 55 uint32_t node =56 cpp::AtomicRef(next[get_node(orig)]).load(cpp::MemoryOrder::RELAXED);57 if (head.compare_exchange_strong(orig, make_head(orig, node),58 cpp::MemoryOrder::ACQUIRE,59 cpp::MemoryOrder::RELAXED))60 break;61 }62 return get_node(orig);63 }64 65 // Attempts to push data to the given stack by making it point to the new66 // node. We repeatedly attempt to write to the head using compare-and-swap,67 // expecting that it has not been changed by any other thread.68 LIBC_INLINE uint32_t push_impl(cpp::AtomicRef<uint64_t> head, uint32_t node) {69 uint64_t orig = head.load(cpp::MemoryOrder::RELAXED);70 for (;;) {71 next[node] = get_node(orig);72 if (head.compare_exchange_strong(orig, make_head(orig, node),73 cpp::MemoryOrder::RELEASE,74 cpp::MemoryOrder::RELAXED))75 break;76 }77 return get_node(head.load(cpp::MemoryOrder::RELAXED));78 }79 80public:81 // Initialize the free stack to be full and the used stack to be empty. We use82 // the capacity of the stack as a sentinel value.83 LIBC_INLINE constexpr FixedStack() : free(0), used(CAPACITY), data{} {84 for (uint32_t i = 0; i < CAPACITY; ++i)85 next[i] = i + 1;86 }87 88 LIBC_INLINE bool push(const T &val) {89 uint32_t node = pop_impl(cpp::AtomicRef(free));90 if (node == CAPACITY)91 return false;92 93 data[node] = val;94 push_impl(cpp::AtomicRef(used), node);95 return true;96 }97 98 LIBC_INLINE bool pop(T &val) {99 uint32_t node = pop_impl(cpp::AtomicRef(used));100 if (node == CAPACITY)101 return false;102 103 val = data[node];104 push_impl(cpp::AtomicRef(free), node);105 return true;106 }107};108 109} // namespace LIBC_NAMESPACE_DECL110 111#endif // LLVM_LIBC_SRC___SUPPORT_GPU_FIXEDSTACK_H112