[libc-commits] [libc] [libc] Implement simple lock-free stack data structure (PR #83026)
Joseph Huber via libc-commits
libc-commits at lists.llvm.org
Wed Feb 28 11:12:10 PST 2024
https://github.com/jhuber6 updated https://github.com/llvm/llvm-project/pull/83026
>From 5e1851b96d1eabfae3f22843fff361f400c4bc90 Mon Sep 17 00:00:00 2001
From: Joseph Huber <huberjn at outlook.com>
Date: Mon, 26 Feb 2024 10:02:42 -0600
Subject: [PATCH] [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.
---
libc/src/__support/CMakeLists.txt | 9 ++
libc/src/__support/fixedstack.h | 130 ++++++++++++++++++
.../integration/src/__support/CMakeLists.txt | 4 +
.../src/__support/gpu/CMakeLists.txt | 19 +++
.../src/__support/gpu/fixed_stack_test.cpp | 75 ++++++++++
libc/test/src/__support/CMakeLists.txt | 10 ++
libc/test/src/__support/fixedstack_test.cpp | 26 ++++
7 files changed, 273 insertions(+)
create mode 100644 libc/src/__support/fixedstack.h
create mode 100644 libc/test/integration/src/__support/gpu/CMakeLists.txt
create mode 100644 libc/test/integration/src/__support/gpu/fixed_stack_test.cpp
create mode 100644 libc/test/src/__support/fixedstack_test.cpp
diff --git a/libc/src/__support/CMakeLists.txt b/libc/src/__support/CMakeLists.txt
index 1a4b3e9a2145c0..5288c808ca6bbb 100644
--- a/libc/src/__support/CMakeLists.txt
+++ b/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
diff --git a/libc/src/__support/fixedstack.h b/libc/src/__support/fixedstack.h
new file mode 100644
index 00000000000000..ee79a3e2ceac44
--- /dev/null
+++ b/libc/src/__support/fixedstack.h
@@ -0,0 +1,130 @@
+//===-- 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 {
+ // The index is stored as a 20-bit value and cannot index into any more.
+ 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
diff --git a/libc/test/integration/src/__support/CMakeLists.txt b/libc/test/integration/src/__support/CMakeLists.txt
index 7c853ff10259f5..4eaa8f50269815 100644
--- a/libc/test/integration/src/__support/CMakeLists.txt
+++ b/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)
diff --git a/libc/test/integration/src/__support/gpu/CMakeLists.txt b/libc/test/integration/src/__support/gpu/CMakeLists.txt
new file mode 100644
index 00000000000000..d83a762cacbd8d
--- /dev/null
+++ b/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
+)
diff --git a/libc/test/integration/src/__support/gpu/fixed_stack_test.cpp b/libc/test/integration/src/__support/gpu/fixed_stack_test.cpp
new file mode 100644
index 00000000000000..52fad148023602
--- /dev/null
+++ b/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;
+}
diff --git a/libc/test/src/__support/CMakeLists.txt b/libc/test/src/__support/CMakeLists.txt
index 7200ac276fe502..4810aa5a31c5ab 100644
--- a/libc/test/src/__support/CMakeLists.txt
+++ b/libc/test/src/__support/CMakeLists.txt
@@ -117,6 +117,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
diff --git a/libc/test/src/__support/fixedstack_test.cpp b/libc/test/src/__support/fixedstack_test.cpp
new file mode 100644
index 00000000000000..cbbffb91311a29
--- /dev/null
+++ b/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());
+}
More information about the libc-commits
mailing list