[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:

I don't think ConcurrentHashTableByPtr fits here. It locks on every lookup (`insert()`), hit or miss. `Allocate()` would then acquire a mutex per allocation, which the point of a per-thread allocator.

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