fixup: revert to synchronized allocator

Author: abulavin

mlir/CMakeLists.txt

@@ -69,7 +69,6 @@ list(INSERT CMAKE_MODULE_PATH 0
 
 include(AddMLIR)
 include(IRDLToCpp)
-include(CheckAtomic)
 
 # -BSymbolic is incompatible with TypeID
 if("${CMAKE_SHARED_LINKER_FLAGS}" MATCHES "-Bsymbolic[^-]")

mlir/include/mlir/IR/MLIRContext.h

@@ -153,13 +153,6 @@ class MLIRContext {
     disableMultithreading(!enable);
   }
 
-  /// Set the flag specifying if thread-local storage should be used by
-  /// attribute storage allocators in this context.
-  void disableThreadLocalStorage(bool disable = true);
-  void enableThreadLocalStorage(bool enable = true) {
-    disableThreadLocalStorage(!enable);
-  }
-
   /// Set a new thread pool to be used in this context. This method requires
   /// that multithreading is disabled for this context prior to the call. This
   /// allows to share a thread pool across multiple contexts, as well as

mlir/lib/IR/AttributeDetail.h

@@ -19,12 +19,8 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
-#include "mlir/Support/StorageUniquer.h"
-#include "mlir/Support/ThreadLocalCache.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/Support/Allocator.h"
-#include <atomic>
-#include <cstdint>
 #include <mutex>
 
 namespace mlir {
@@ -398,112 +394,30 @@ class DistinctAttributeUniquer {
                                               Attribute referencedAttr);
 };
 
-/// An allocator for distinct attribute storage instances. It uses thread-local
-/// bump pointer allocators stored in a thread-local cache to ensure the storage
-/// is freed after the destruction of the distinct attribute allocator. The way
-/// in which storage is allocated may be changed from a thread-local allocator
-/// to a shared, locked allocator. This can be used to prevent use-after-free
-/// errors if storage is allocated on a thread that has a lifetime less than the
-/// lifetime of the storage.
+/// An allocator for distinct attribute storage instances. Uses a synchronized
+/// BumpPtrAllocator to ensure allocated storage is deleted when the
+/// DistinctAttributeAllocator is destroyed
 class DistinctAttributeAllocator final {
 public:
-  /// The allocation strategy for DistinctAttribute storage.
-  enum class AllocationMode : uint8_t {
-    /// Use a thread-local allocator. Lock-free, but storage lifetime is tied to
-    /// a thread, which may have shorter lifetime than the storage allocated
-    /// here.
-    ThreadLocal,
-    /// Use a single, shared allocator protected by a mutex. Slower due to
-    /// locking, but storage persists for the lifetime of the MLIR context.
-    SharedLocked,
-  };
-
-  DistinctAttributeAllocator() {
-    setAllocationMode(AllocationMode::ThreadLocal);
-  };
-
+  DistinctAttributeAllocator() = default;
   DistinctAttributeAllocator(DistinctAttributeAllocator &&) = delete;
   DistinctAttributeAllocator(const DistinctAttributeAllocator &) = delete;
   DistinctAttributeAllocator &
   operator=(const DistinctAttributeAllocator &) = delete;
 
-  /// Allocates a distinct attribute storage. In most cases this will be
-  /// using a thread-local allocator for synchronization free parallel accesses,
-  /// however if the PassManager's crash recovery is enabled this will incur
-  /// synchronization cost in exchange for persisting the storage.
   DistinctAttrStorage *allocate(Attribute referencedAttr) {
-    return std::invoke(allocatorFn.load(), *this, referencedAttr);
-  }
-
-  /// Sets the allocation mode. This is used by the MLIRContext to switch
-  /// allocation strategies, for example, during crash recovery.
-  void setAllocationMode(AllocationMode mode) {
-    switch (mode) {
-    case AllocationMode::ThreadLocal:
-      allocatorFn.store(&DistinctAttributeAllocator::allocateThreadLocalWrapper,
-                        std::memory_order_release);
-      break;
-    case AllocationMode::SharedLocked:
-      allocatorFn.store(
-          &DistinctAttributeAllocator::allocateLockedSharedWrapper,
-          std::memory_order_release);
-      break;
-    }
-  }
-
-private:
-  // Define some static wrappers for allocating using either the thread-local
-  // allocator or the shared locked allocator. We use these static wrappers to
-  // ensure that the width of the function pointer to these functions is the
-  // size of regular pointer on the platform (which is not always the case for
-  // pointer-to-member-function), and so should be handled by atomic
-  // instructions.
-  static DistinctAttrStorage *
-  allocateThreadLocalWrapper(DistinctAttributeAllocator &self,
-                             Attribute referencedAttr) {
-    return self.allocateUsingThreadLocal(referencedAttr);
-  }
-
-  static DistinctAttrStorage *
-  allocateLockedSharedWrapper(DistinctAttributeAllocator &self,
-                              Attribute referencedAttr) {
-    return self.allocateUsingLockedShared(referencedAttr);
-  }
-
-  DistinctAttrStorage *allocateUsingLockedShared(Attribute referencedAttr) {
-    std::scoped_lock<std::mutex> guard(sharedAllocatorMutex);
-    return doAllocate(sharedAllocator, referencedAttr);
-  }
-
-  DistinctAttrStorage *allocateUsingThreadLocal(Attribute referencedAttr) {
-    return doAllocate(allocatorCache.get(), referencedAttr);
-  };
-
-  DistinctAttrStorage *doAllocate(llvm::BumpPtrAllocator &allocator,
-                                  Attribute referencedAttr) {
+    std::scoped_lock<std::mutex> guard(allocatorMutex);
     return new (allocator.Allocate<DistinctAttrStorage>())
         DistinctAttrStorage(referencedAttr);
   };
 
-  /// Used to allocate Distinct Attribute storage without synchronization
-  ThreadLocalCache<llvm::BumpPtrAllocator> allocatorCache;
-
-  /// Used to allocate Distict Attribute storage with synchronization,
-  /// but without bounding the lifetime of the allocated memory to the
-  /// lifetime of a thread.
-  llvm::BumpPtrAllocator sharedAllocator;
-
-  /// Used to lock access to the shared allocator
-  std::mutex sharedAllocatorMutex;
-
-  using AllocatorFuncPtr =
-      DistinctAttrStorage *(*)(DistinctAttributeAllocator &, Attribute);
+private:
+  /// Used to allocate Distict Attribute storage. When this allocator destroyed
+  /// in till also dealllocate any storage instances
+  llvm::BumpPtrAllocator allocator;
 
-  /// Atomic function pointer for the allocation strategy. Using a pointer to a
-  /// static member function as opposed to a non-static one should guarantee
-  /// that the function pointer fits into a standard pointer size, and so will
-  /// atomic instruction support for the corresponding width.
-  std::atomic<AllocatorFuncPtr> allocatorFn;
+  /// Used to lock access to the allocator
+  std::mutex allocatorMutex;
 };
 } // namespace detail
 } // namespace mlir

mlir/lib/IR/CMakeLists.txt

@@ -68,7 +68,5 @@ add_mlir_library(MLIRIR
 
   LINK_LIBS PUBLIC
   MLIRSupport
-  PRIVATE
-  ${LLVM_ATOMIC_LIB}
   )
 

mlir/lib/IR/MLIRContext.cpp

@@ -31,6 +31,7 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/Mutex.h"
 #include "llvm/Support/RWMutex.h"
 #include "llvm/Support/ThreadPool.h"
@@ -711,13 +712,6 @@ bool MLIRContext::isOperationRegistered(StringRef name) {
   return RegisteredOperationName::lookup(name, this).has_value();
 }
 
-void MLIRContext::disableThreadLocalStorage(bool disable) {
-  const auto mode =
-      disable ? DistinctAttributeAllocator::AllocationMode::SharedLocked
-              : DistinctAttributeAllocator::AllocationMode::ThreadLocal;
-  getImpl().distinctAttributeAllocator.setAllocationMode(mode);
-}
-
 void Dialect::addType(TypeID typeID, AbstractType &&typeInfo) {
   auto &impl = context->getImpl();
   assert(impl.multiThreadedExecutionContext == 0 &&

mlir/lib/Pass/PassCrashRecovery.cpp

@@ -418,15 +418,6 @@ LogicalResult PassManager::runWithCrashRecovery(Operation *op,
   // thread-local storage during passes beyond the lifetime of the recovery
   // context thread.
   const bool threadingEnabled = getContext()->isMultithreadingEnabled();
-  if (threadingEnabled) {
-    crashRecoveryLock.lock();
-    getContext()->disableThreadLocalStorage();
-  }
-  auto guard = llvm::make_scope_exit([this]() {
-    getContext()->enableThreadLocalStorage();
-    crashRecoveryLock.unlock();
-  });
-
   crashReproGenerator->initialize(getPasses(), op, verifyPasses);
 
   // Safely invoke the passes within a recovery context.
@@ -439,8 +430,6 @@ LogicalResult PassManager::runWithCrashRecovery(Operation *op,
     recoveryContext.RunSafely(runPassesFn);
   crashReproGenerator->finalize(op, passManagerResult);
 
-  if (!threadingEnabled)
-    guard.release();
   return passManagerResult;
 }
 

mlir/unittests/IR/DistinctAttributeAllocatorTest.cpp

@@ -12,7 +12,6 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/MLIRContext.h"
-#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/CrashRecoveryContext.h"
 #include <thread>
 
@@ -28,17 +27,11 @@ TEST(DistinctAttributeAllocatorTest, TestAttributeWellFormedAfterThreadJoin) {
   OpBuilder builder(&ctx);
   DistinctAttr attr;
 
-  {
-    ctx.disableThreadLocalStorage();
-    auto guard =
-        llvm::make_scope_exit([&ctx]() { ctx.enableThreadLocalStorage(); });
-
-    std::thread t([&ctx, &attr]() {
-      attr = DistinctAttr::create(UnitAttr::get(&ctx));
-      ASSERT_TRUE(attr);
-    });
-    t.join();
-  }
+  std::thread t([&ctx, &attr]() {
+    attr = DistinctAttr::create(UnitAttr::get(&ctx));
+    ASSERT_TRUE(attr);
+  });
+  t.join();
 
   // If the attribute storage got deleted after the thread joins (which we don't
   // want) then trying to access it triggers an assert in Debug mode, and a