[llvm] PerThreadBumpPtrAllocator: remove dependency on getThreadIndex (PR #209687)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-support

Author: Fangrui Song (MaskRay)

<details>
<summary>Changes</summary>

PerThreadAllocator, only used by DWARFLinker, indexes a fixed
getThreadCount()-sized array by `getThreadIndex()`, so it only works on
ThreadPoolExecutor threads (and asserts elsewhere). It false-shares
adjacent threads' bump pointers.

Instead, create each thread's sub-allocator lazily on first use, keyed
by a process-unique instance id in a thread-local cache, with the
instance owning the heap-allocated sub-allocators. This drops the
getThreadIndex()/getThreadCount() dependence, unblocking the caller
participation change for `parallelFor`.

mlir::ThreadLocalCache solves the same problem by keying a per-thread map on
the instance pointer and reclaiming a thread's slot when an instance dies, but
costs a map lookup and shared_ptr bookkeeping per allocation; instances here
are few and short-lived, so an id-indexed vector is cheaper. The counter
behind claimPerThreadAllocatorId is defined out of line in Allocator.cpp so
Windows DLLs cannot duplicate it and alias two instances' sub-allocators.

Aided by Claude Fable 5


---
Full diff: https://github.com/llvm/llvm-project/pull/209687.diff


4 Files Affected:

- (modified) llvm/include/llvm/Support/PerThreadBumpPtrAllocator.h (+64-33) 
- (modified) llvm/lib/DWARFLinker/Parallel/DWARFLinkerTypeUnit.cpp (+1) 
- (modified) llvm/lib/Support/Allocator.cpp (+10) 
- (modified) llvm/unittests/Support/PerThreadBumpPtrAllocatorTest.cpp (+26-1) 


``````````diff
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

``````````

</details>


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