[llvm] PerThreadBumpPtrAllocator: remove dependency on getThreadIndex (PR #209687)
Fangrui Song via llvm-commits
llvm-commits at lists.llvm.org
Sat Jul 18 14:54:02 PDT 2026
https://github.com/MaskRay updated https://github.com/llvm/llvm-project/pull/209687
>From 2ea7f103d931cc835d2cc495e6b5ecd666e56451 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i at maskray.me>
Date: Tue, 14 Jul 2026 23:32:25 -0700
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20initia?=
=?UTF-8?q?l=20version?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Created using spr 1.3.5-bogner
---
.../llvm/Support/PerThreadBumpPtrAllocator.h | 97 ++++++++++++-------
.../Parallel/DWARFLinkerTypeUnit.cpp | 1 +
llvm/lib/Support/Allocator.cpp | 10 ++
.../Support/PerThreadBumpPtrAllocatorTest.cpp | 27 +++++-
4 files changed, 101 insertions(+), 34 deletions(-)
diff --git a/llvm/include/llvm/Support/PerThreadBumpPtrAllocator.h b/llvm/include/llvm/Support/PerThreadBumpPtrAllocator.h
index 0448fb2a12a9b..846c92933c47e 100644
--- a/llvm/include/llvm/Support/PerThreadBumpPtrAllocator.h
+++ b/llvm/include/llvm/Support/PerThreadBumpPtrAllocator.h
@@ -10,26 +10,38 @@
#define LLVM_SUPPORT_PERTHREADBUMPPTRALLOCATOR_H
#include "llvm/Support/Allocator.h"
-#include "llvm/Support/Parallel.h"
+#include "llvm/Support/Compiler.h"
+#include "llvm/Support/raw_ostream.h"
+#include <memory>
+#include <mutex>
+#include <vector>
namespace llvm {
namespace parallel {
-/// PerThreadAllocator is used in conjunction with ThreadPoolExecutor to allow
-/// per-thread allocations. It wraps a possibly thread-unsafe allocator,
-/// e.g. BumpPtrAllocator. PerThreadAllocator must be used with only main thread
-/// or threads created by ThreadPoolExecutor, as it utilizes getThreadIndex,
-/// which is set by ThreadPoolExecutor. To work properly, ThreadPoolExecutor
-/// should be initialized before PerThreadAllocator is created.
-/// TODO: The same approach might be implemented for ThreadPool.
-
+namespace detail {
+/// Return a new process-unique PerThreadAllocator instance id. Ids are never
+/// reused.
+LLVM_ABI unsigned claimPerThreadAllocatorId();
+} // namespace detail
+
+/// PerThreadAllocator wraps a thread-unsafe allocator (e.g. BumpPtrAllocator)
+/// for lock-free concurrent allocation: each thread receives its own
+/// sub-allocator on first allocation, and the PerThreadAllocator owns all
+/// sub-allocators. Recommended to used with the thread pool in Parallel.h even
+/// if there is no dependency on it.
template <typename AllocatorTy>
class PerThreadAllocator
: public AllocatorBase<PerThreadAllocator<AllocatorTy>> {
+ // Heap-allocated so that the class stays movable while holding a mutex.
+ struct State {
+ std::mutex Mutex;
+ std::vector<std::unique_ptr<AllocatorTy>> Allocators;
+ };
+
public:
PerThreadAllocator()
- : NumOfAllocators(parallel::getThreadCount()),
- Allocators(std::make_unique<AllocatorTy[]>(NumOfAllocators)) {}
+ : S(std::make_unique<State>()), Id(detail::claimPerThreadAllocatorId()) {}
/// \defgroup Methods which could be called asynchronously:
///
@@ -41,25 +53,45 @@ class PerThreadAllocator
/// Allocate \a Size bytes of \a Alignment aligned memory.
void *Allocate(size_t Size, size_t Alignment) {
- assert(getThreadIndex() < NumOfAllocators);
- return Allocators[getThreadIndex()].Allocate(Size, Alignment);
+ return getThreadLocalAllocator().Allocate(Size, Alignment);
}
/// Deallocate \a Ptr to \a Size bytes of memory allocated by this
/// allocator.
void Deallocate(const void *Ptr, size_t Size, size_t Alignment) {
- assert(getThreadIndex() < NumOfAllocators);
- return Allocators[getThreadIndex()].Deallocate(Ptr, Size, Alignment);
+ return getThreadLocalAllocator().Deallocate(Ptr, Size, Alignment);
}
- /// Return allocator corresponding to the current thread.
+ /// Return the calling thread's sub-allocator, creating it on first use.
AllocatorTy &getThreadLocalAllocator() {
- assert(getThreadIndex() < NumOfAllocators);
- return Allocators[getThreadIndex()];
+ // The calling thread's sub-allocator of each instance, indexed by a
+ // process-unique instance id.
+ //
+ // mlir::ThreadLocalCache keys an analogous per-thread map on the instance
+ // pointer and reclaims a thread's slot once the instance dies, but pays a
+ // map lookup and shared_ptr bookkeeping per allocation. Instances here are
+ // few and short-lived, so we prefer the O(1) vector index and accept that a
+ // thread's Cache only grows with the number of instances created.
+ static thread_local std::vector<AllocatorTy *> Cache;
+ if (LLVM_UNLIKELY(Cache.size() <= Id))
+ Cache.resize(Id + 1);
+ AllocatorTy *&A = Cache[Id];
+ if (LLVM_UNLIKELY(!A)) {
+ // Heap-allocate sub-allocators so that their addresses are stable and
+ // different threads' bump pointers do not share a cache line.
+ auto New = std::make_unique<AllocatorTy>();
+ A = New.get();
+ std::lock_guard<std::mutex> Lock(S->Mutex);
+ S->Allocators.push_back(std::move(New));
+ }
+ return *A;
}
- // Return number of used allocators.
- size_t getNumberOfAllocators() const { return NumOfAllocators; }
+ /// Return the number of sub-allocators, i.e. threads that have allocated.
+ size_t getNumberOfAllocators() const {
+ std::lock_guard<std::mutex> Lock(S->Mutex);
+ return S->Allocators.size();
+ }
/// @}
/// \defgroup Methods which could not be called asynchronously:
@@ -68,38 +100,37 @@ class PerThreadAllocator
/// Reset state of allocators.
void Reset() {
- for (size_t Idx = 0; Idx < getNumberOfAllocators(); Idx++)
- Allocators[Idx].Reset();
+ for (const auto &A : S->Allocators)
+ A->Reset();
}
/// Return total memory size used by all allocators.
size_t getTotalMemory() const {
size_t TotalMemory = 0;
-
- for (size_t Idx = 0; Idx < getNumberOfAllocators(); Idx++)
- TotalMemory += Allocators[Idx].getTotalMemory();
-
+ for (const auto &A : S->Allocators)
+ TotalMemory += A->getTotalMemory();
return TotalMemory;
}
/// Set red zone for all allocators.
void setRedZoneSize(size_t NewSize) {
- for (size_t Idx = 0; Idx < getNumberOfAllocators(); Idx++)
- Allocators[Idx].setRedZoneSize(NewSize);
+ for (const auto &A : S->Allocators)
+ A->setRedZoneSize(NewSize);
}
/// Print statistic for each allocator.
void PrintStats() const {
- for (size_t Idx = 0; Idx < getNumberOfAllocators(); Idx++) {
- errs() << "\n Allocator " << Idx << "\n";
- Allocators[Idx].PrintStats();
+ size_t Idx = 0;
+ for (const auto &A : S->Allocators) {
+ errs() << "\n Allocator " << Idx++ << "\n";
+ A->PrintStats();
}
}
/// @}
protected:
- size_t NumOfAllocators;
- std::unique_ptr<AllocatorTy[]> Allocators;
+ std::unique_ptr<State> S;
+ unsigned Id;
};
using PerThreadBumpPtrAllocator = PerThreadAllocator<BumpPtrAllocator>;
diff --git a/llvm/lib/DWARFLinker/Parallel/DWARFLinkerTypeUnit.cpp b/llvm/lib/DWARFLinker/Parallel/DWARFLinkerTypeUnit.cpp
index 58de4e91570df..97c5e382b08d0 100644
--- a/llvm/lib/DWARFLinker/Parallel/DWARFLinkerTypeUnit.cpp
+++ b/llvm/lib/DWARFLinker/Parallel/DWARFLinkerTypeUnit.cpp
@@ -9,6 +9,7 @@
#include "DWARFLinkerTypeUnit.h"
#include "DIEGenerator.h"
#include "llvm/Support/LEB128.h"
+#include "llvm/Support/Parallel.h"
using namespace llvm;
using namespace dwarf_linker;
diff --git a/llvm/lib/Support/Allocator.cpp b/llvm/lib/Support/Allocator.cpp
index 6ff68f1be3d63..6c5dbef4f28b6 100644
--- a/llvm/lib/Support/Allocator.cpp
+++ b/llvm/lib/Support/Allocator.cpp
@@ -11,8 +11,11 @@
//===----------------------------------------------------------------------===//
#include "llvm/Support/Allocator.h"
+#include "llvm/Support/PerThreadBumpPtrAllocator.h"
#include "llvm/Support/raw_ostream.h"
+#include <atomic>
+
namespace llvm {
namespace detail {
@@ -33,4 +36,11 @@ void PrintRecyclerStats(size_t Size,
<< "Number of elements free for recycling: " << FreeListSize << '\n';
}
+namespace parallel::detail {
+unsigned claimPerThreadAllocatorId() {
+ static std::atomic<unsigned> Counter;
+ return Counter.fetch_add(1, std::memory_order_relaxed);
+}
+} // namespace parallel::detail
+
} // namespace llvm
diff --git a/llvm/unittests/Support/PerThreadBumpPtrAllocatorTest.cpp b/llvm/unittests/Support/PerThreadBumpPtrAllocatorTest.cpp
index 4ac5ac4e1ff3b..35d903bdb1924 100644
--- a/llvm/unittests/Support/PerThreadBumpPtrAllocatorTest.cpp
+++ b/llvm/unittests/Support/PerThreadBumpPtrAllocatorTest.cpp
@@ -10,6 +10,8 @@
#include "llvm/Support/Parallel.h"
#include "gtest/gtest.h"
#include <cstdlib>
+#include <thread>
+#include <vector>
using namespace llvm;
using namespace parallel;
@@ -48,7 +50,30 @@ TEST(PerThreadBumpPtrAllocatorTest, ParallelAllocation) {
});
EXPECT_LE(sizeof(uint64_t) * NumAllocations, Allocator.getTotalMemory());
- EXPECT_EQ(Allocator.getNumberOfAllocators(), parallel::getThreadCount());
+ // Sub-allocators are created lazily, one per thread that allocated.
+ EXPECT_GE(Allocator.getNumberOfAllocators(), 1u);
+ EXPECT_LE(Allocator.getNumberOfAllocators(), parallel::getThreadCount());
}
+#if LLVM_ENABLE_THREADS
+TEST(PerThreadBumpPtrAllocatorTest, ArbitraryThreads) {
+ PerThreadBumpPtrAllocator Allocator;
+
+ constexpr size_t NumThreads = 4;
+ std::vector<std::thread> Threads;
+ for (size_t I = 0; I != NumThreads; ++I)
+ Threads.emplace_back([&Allocator, I] {
+ uint64_t *Var =
+ (uint64_t *)Allocator.Allocate(sizeof(uint64_t), alignof(uint64_t));
+ *Var = I;
+ EXPECT_EQ(I, *Var);
+ });
+ for (std::thread &T : Threads)
+ T.join();
+
+ EXPECT_EQ(NumThreads, Allocator.getNumberOfAllocators());
+ EXPECT_LE(sizeof(uint64_t) * NumThreads, Allocator.getTotalMemory());
+}
+#endif
+
} // anonymous namespace
More information about the llvm-commits
mailing list