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

[Fangrui Song via llvm-commits]

================
@@ -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);
----------------
MaskRay wrote:

Thanks for the review.

The remaining items, avoiding per-instance memory waste, and what a general non-BumpPtrAllocator-specific facility should look like, are follow-up work, and where it should live is best settled by whoever proposes that solution, against its actual users.

I remain unconvinced that's `Support` for this code, but nothing here depends on resolving it.

Landing this as the prerequisite for #209690 - which is more important and really shouldn't be blocked by the debate on whether PerThreadBumpPtrAllocator should be placed in Support.

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