[llvm] d07c3cf - [Support] Introduce ThreadSafeAllocator
Steven Wu via llvm-commits
llvm-commits at lists.llvm.org
Fri Oct 6 13:55:20 PDT 2023
Author: Steven Wu
Date: 2023-10-06T13:54:53-07:00
New Revision: d07c3cf667160cc55944bf2b5ed337a435649649
URL: https://github.com/llvm/llvm-project/commit/d07c3cf667160cc55944bf2b5ed337a435649649
DIFF: https://github.com/llvm/llvm-project/commit/d07c3cf667160cc55944bf2b5ed337a435649649.diff
LOG: [Support] Introduce ThreadSafeAllocator
Add support for ThreadSafeAllocator, which is needed for a CAS
implementation, which requires thread safe allocation for data storage.
Reviewed By: dblaikie
Differential Revision: https://reviews.llvm.org/D133713
Added:
llvm/include/llvm/Support/ThreadSafeAllocator.h
llvm/unittests/Support/ThreadSafeAllocatorTest.cpp
Modified:
llvm/unittests/Support/CMakeLists.txt
Removed:
################################################################################
diff --git a/llvm/include/llvm/Support/ThreadSafeAllocator.h b/llvm/include/llvm/Support/ThreadSafeAllocator.h
new file mode 100644
index 000000000000000..3092287e691f754
--- /dev/null
+++ b/llvm/include/llvm/Support/ThreadSafeAllocator.h
@@ -0,0 +1,63 @@
+//===- ThreadSafeAllocator.h ------------------------------------*- 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_SUPPORT_THREADSAFEALLOCATOR_H
+#define LLVM_SUPPORT_THREADSAFEALLOCATOR_H
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/Support/Allocator.h"
+#include <atomic>
+
+namespace llvm {
+
+/// Thread-safe allocator adaptor. Uses a spin lock on the assumption that
+/// contention here is extremely rare.
+///
+/// TODO: Using a spin lock on every allocation can be quite expensive when
+/// contention is high. Since this is mainly used for BumpPtrAllocator and
+/// SpecificBumpPtrAllocator, it'd be better to have a specific thread-safe
+/// BumpPtrAllocator implementation that only use a fair lock when allocating a
+/// new slab but otherwise using atomic and be lock-free.
+template <class AllocatorType> class ThreadSafeAllocator {
+ struct LockGuard {
+ LockGuard(std::atomic_flag &Flag) : Flag(Flag) {
+ if (LLVM_UNLIKELY(Flag.test_and_set(std::memory_order_acquire)))
+ while (Flag.test_and_set(std::memory_order_acquire)) {
+ }
+ }
+ ~LockGuard() { Flag.clear(std::memory_order_release); }
+ std::atomic_flag &Flag;
+ };
+
+public:
+ auto Allocate(size_t N) {
+ return applyLocked([N](AllocatorType &Alloc) { return Alloc.Allocate(N); });
+ }
+
+ auto Allocate(size_t Size, size_t Align) {
+ return applyLocked([Size, Align](AllocatorType &Alloc) {
+ return Alloc.Allocate(Size, Align);
+ });
+ }
+
+ template <typename FnT,
+ typename T = typename llvm::function_traits<FnT>::result_t>
+ T applyLocked(FnT Fn) {
+ LockGuard Lock(Flag);
+ return Fn(Alloc);
+ }
+
+private:
+ AllocatorType Alloc;
+ std::atomic_flag Flag = ATOMIC_FLAG_INIT;
+};
+
+} // namespace llvm
+
+#endif // LLVM_SUPPORT_THREADSAFEALLOCATOR_H
diff --git a/llvm/unittests/Support/CMakeLists.txt b/llvm/unittests/Support/CMakeLists.txt
index 12f2e9959326045..dfd55b228900d8e 100644
--- a/llvm/unittests/Support/CMakeLists.txt
+++ b/llvm/unittests/Support/CMakeLists.txt
@@ -81,6 +81,7 @@ add_llvm_unittest(SupportTests
SwapByteOrderTest.cpp
TarWriterTest.cpp
ThreadPool.cpp
+ ThreadSafeAllocatorTest.cpp
Threading.cpp
TimerTest.cpp
TimeProfilerTest.cpp
diff --git a/llvm/unittests/Support/ThreadSafeAllocatorTest.cpp b/llvm/unittests/Support/ThreadSafeAllocatorTest.cpp
new file mode 100644
index 000000000000000..d9a85b435ebdb7e
--- /dev/null
+++ b/llvm/unittests/Support/ThreadSafeAllocatorTest.cpp
@@ -0,0 +1,138 @@
+//===- llvm/unittest/Support/ThreadSafeAllocatorTest.cpp ------------------===//
+//
+// 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 "llvm/Support/ThreadSafeAllocator.h"
+#include "llvm/Config/llvm-config.h"
+#include "llvm/Support/ThreadPool.h"
+#include "gtest/gtest.h"
+#include <atomic>
+#include <thread>
+
+using namespace llvm;
+
+namespace {
+
+struct AllocCondition {
+ std::mutex BusyLock, EndLock;
+ std::condition_variable Busy, End;
+ bool IsBusy = false, IsEnd = false;
+ std::atomic<unsigned> BytesAllocated = 0;
+
+ void startAllocation() {
+ {
+ std::lock_guard<std::mutex> Lock(BusyLock);
+ IsBusy = true;
+ }
+ Busy.notify_all();
+ }
+ void waitAllocationStarted() {
+ std::unique_lock<std::mutex> LBusy(BusyLock);
+ Busy.wait(LBusy, [&]() { return IsBusy; });
+ IsBusy = false;
+ }
+ void finishAllocation() {
+ {
+ std::lock_guard<std::mutex> Lock(EndLock);
+ IsEnd = true;
+ }
+ End.notify_all();
+ }
+ void waitAllocationFinished() {
+ std::unique_lock<std::mutex> LEnd(EndLock);
+ // Wait for end state.
+ End.wait(LEnd, [&]() { return IsEnd; });
+ IsEnd = false;
+ }
+};
+
+class MockAllocator : public AllocatorBase<MockAllocator> {
+public:
+ MockAllocator() = default;
+
+ void *Allocate(size_t Size, size_t Alignment) {
+ C.startAllocation();
+ C.waitAllocationFinished();
+ C.BytesAllocated += Size;
+ return Reserved;
+ }
+
+ AllocCondition &getAllocCondition() { return C; }
+
+private:
+ char Reserved[16];
+ AllocCondition C;
+};
+
+} // namespace
+
+#if (LLVM_ENABLE_THREADS)
+TEST(ThreadSafeAllocatorTest, AllocWait) {
+ ThreadSafeAllocator<MockAllocator> Alloc;
+ AllocCondition *C;
+ // Get the allocation from the allocator first since this requires a lock.
+ Alloc.applyLocked(
+ [&](MockAllocator &Alloc) { C = &Alloc.getAllocCondition(); });
+ ThreadPool Threads;
+ // First allocation of 1 byte.
+ Threads.async([&Alloc]() {
+ char *P = (char *)Alloc.Allocate(1, alignof(char));
+ P[0] = 0;
+ });
+ // No allocation yet.
+ EXPECT_EQ(C->BytesAllocated, 0u);
+ C->waitAllocationStarted(); // wait till 1st alloocation starts.
+ // Second allocation of 2 bytes.
+ Threads.async([&Alloc]() {
+ char *P = (char *)Alloc.Allocate(2, alignof(char));
+ P[1] = 0;
+ });
+ C->finishAllocation(); // finish 1st allocation.
+
+ C->waitAllocationStarted(); // wait till 2nd allocation starts.
+ // still 1 byte allocated since 2nd allocation is not finished yet.
+ EXPECT_EQ(C->BytesAllocated, 1u);
+ C->finishAllocation(); // finish 2nd allocation.
+
+ Threads.wait(); // all allocations done.
+ EXPECT_EQ(C->BytesAllocated, 3u);
+}
+
+TEST(ThreadSafeAllocatorTest, AllocWithAlign) {
+ ThreadSafeAllocator<BumpPtrAllocator> Alloc;
+ ThreadPool Threads;
+
+ for (unsigned Index = 1; Index < 100; ++Index)
+ Threads.async(
+ [&Alloc](unsigned I) {
+ int *P = (int *)Alloc.Allocate(sizeof(int) * I, alignof(int));
+ P[I - 1] = I;
+ },
+ Index);
+
+ Threads.wait();
+
+ Alloc.applyLocked([](BumpPtrAllocator &Alloc) {
+ EXPECT_EQ(4950U * sizeof(int), Alloc.getBytesAllocated());
+ });
+}
+
+TEST(ThreadSafeAllocatorTest, SpecificBumpPtrAllocator) {
+ ThreadSafeAllocator<SpecificBumpPtrAllocator<int>> Alloc;
+ ThreadPool Threads;
+
+ for (unsigned Index = 1; Index < 100; ++Index)
+ Threads.async(
+ [&Alloc](unsigned I) {
+ int *P = Alloc.Allocate(I);
+ P[I - 1] = I;
+ },
+ Index);
+
+ Threads.wait();
+}
+#endif
More information about the llvm-commits
mailing list