Revert "[mlir] Fix DistinctAttributeUniquer deleting attribute storage when crash reproduction is enabled"

mlir/include/mlir/IR/MLIRContext.h

@@ -153,14 +153,6 @@ class MLIRContext {
     disableMultithreading(!enable);
   }
 
-  /// Set the flag specifying if thread-local storage should be used by storage
-  /// allocators in this context. Note that disabling mutlithreading implies
-  /// thread-local storage is also disabled.
-  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

@@ -24,7 +24,6 @@
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/Support/TrailingObjects.h"
-#include <mutex>
 
 namespace mlir {
 namespace detail {
@@ -402,8 +401,7 @@ class DistinctAttributeUniquer {
 /// is freed after the destruction of the distinct attribute allocator.
 class DistinctAttributeAllocator {
 public:
-  DistinctAttributeAllocator(bool threadingIsEnabled)
-      : threadingIsEnabled(threadingIsEnabled), useThreadLocalAllocator(true) {};
+  DistinctAttributeAllocator() = default;
 
   DistinctAttributeAllocator(DistinctAttributeAllocator &&) = delete;
   DistinctAttributeAllocator(const DistinctAttributeAllocator &) = delete;
@@ -413,49 +411,12 @@ class DistinctAttributeAllocator {
   /// Allocates a distinct attribute storage using a thread local bump pointer
   /// allocator to enable synchronization free parallel allocations.
   DistinctAttrStorage *allocate(Attribute referencedAttr) {
-    if (!useThreadLocalAllocator && threadingIsEnabled) {
-      std::scoped_lock<std::mutex> lock(allocatorMutex);
-      return allocateImpl(referencedAttr);
-    }
-    return allocateImpl(referencedAttr);
-  }
-
-  /// Sets a flag that stores if multithreading is enabled. The flag is used to
-  /// decide if locking is needed when using a non thread-safe allocator.
-  void disableMultiThreading(bool disable = true) {
-    threadingIsEnabled = !disable;
-  }
-
-  /// Sets a flag to disable using thread local bump pointer allocators and use
-  /// a single thread-safe allocator. Use this to persist allocated storage
-  /// beyond the lifetime of a child thread calling this function while ensuring
-  /// thread-safe allocation.
-  void disableThreadLocalStorage(bool disable = true) {
-    useThreadLocalAllocator = !disable;
-  }
-
-private:
-  DistinctAttrStorage *allocateImpl(Attribute referencedAttr) {
-    return new (getAllocatorInUse().Allocate<DistinctAttrStorage>())
+    return new (allocatorCache.get().Allocate<DistinctAttrStorage>())
         DistinctAttrStorage(referencedAttr);
   }
 
-  /// If threading is disabled on the owning MLIR context, a normal non
-  /// thread-local, non-thread safe bump pointer allocator is used instead to
-  /// prevent use-after-free errors whenever attribute storage created on a
-  /// crash recover thread is accessed after the thread joins.
-  llvm::BumpPtrAllocator &getAllocatorInUse() {
-    if (useThreadLocalAllocator)
-      return allocatorCache.get();
-    return allocator;
-  }
-
+private:
   ThreadLocalCache<llvm::BumpPtrAllocator> allocatorCache;
-  llvm::BumpPtrAllocator allocator;
-  std::mutex allocatorMutex;
-
-  bool threadingIsEnabled : 1;
-  bool useThreadLocalAllocator : 1;
 };
 } // namespace detail
 } // namespace mlir

mlir/lib/IR/MLIRContext.cpp

@@ -268,8 +268,7 @@ class MLIRContextImpl {
 
 public:
   MLIRContextImpl(bool threadingIsEnabled)
-      : threadingIsEnabled(threadingIsEnabled),
-        distinctAttributeAllocator(threadingIsEnabled) {
+      : threadingIsEnabled(threadingIsEnabled) {
     if (threadingIsEnabled) {
       ownedThreadPool = std::make_unique<llvm::DefaultThreadPool>();
       threadPool = ownedThreadPool.get();
@@ -597,7 +596,6 @@ void MLIRContext::disableMultithreading(bool disable) {
   // Update the threading mode for each of the uniquers.
   impl->affineUniquer.disableMultithreading(disable);
   impl->attributeUniquer.disableMultithreading(disable);
-  impl->distinctAttributeAllocator.disableMultiThreading(disable);
   impl->typeUniquer.disableMultithreading(disable);
 
   // Destroy thread pool (stop all threads) if it is no longer needed, or create
@@ -719,10 +717,6 @@ bool MLIRContext::isOperationRegistered(StringRef name) {
   return RegisteredOperationName::lookup(name, this).has_value();
 }
 
-void MLIRContext::disableThreadLocalStorage(bool disable) {
-  getImpl().distinctAttributeAllocator.disableThreadLocalStorage(disable);
-}
-
 void Dialect::addType(TypeID typeID, AbstractType &&typeInfo) {
   auto &impl = context->getImpl();
   assert(impl.multiThreadedExecutionContext == 0 &&

mlir/lib/Pass/PassCrashRecovery.cpp

@@ -414,15 +414,6 @@ struct FileReproducerStream : public mlir::ReproducerStream {
 
 LogicalResult PassManager::runWithCrashRecovery(Operation *op,
                                                 AnalysisManager am) {
-  // Notify the context to disable the use of thread-local storage while the
-  // pass manager is running in a crash recovery context thread. Re-enable the
-  // thread local storage upon function exit. This is required to persist any
-  // attribute storage allocated during passes beyond the lifetime of the
-  // recovery context thread.
-  MLIRContext *ctx = getContext();
-  ctx->disableThreadLocalStorage();
-  auto guard =
-      llvm::make_scope_exit([ctx]() { ctx->enableThreadLocalStorage(); });
   crashReproGenerator->initialize(getPasses(), op, verifyPasses);
 
   // Safely invoke the passes within a recovery context.