[mlir] Fix DistinctAttributeUniquer deleting attribute storage when crash reproduction is enabled (#128566)

Currently, `DistinctAttr` uses an allocator wrapped in a
`ThreadLocalCache` to manage attribute storage allocations. This ensures
all allocations are freed when the allocator is destroyed.

However, this setup can cause use-after-free errors when
`mlir::PassManager` runs its passes on a separate thread as a result of
crash reproduction being enabled. Distinct attribute storages are
created in the child thread's local storage and freed once the thread
joins. Attempting to access these attributes after this can result in
segmentation faults, such as during printing or alias analysis.

Example: This invocation of `mlir-opt` demonstrates the segfault issue
due to distinct attributes being created in a child thread and their
storage being freed once the thread joins:
```
mlir-opt --mlir-pass-pipeline-crash-reproducer=. --test-distinct-attrs mlir/test/IR/test-builtin-distinct-attrs.mlir
```

This pull request changes the distinct attribute allocator to use
different allocators depending on whether or not threading is enabled
and whether or not the pass manager is running its passes in a separate
thread. If multithreading is disabled, a non thread-local allocator is
used. If threading remains enabled and the pass manager invokes its pass
pipelines in a child thread, then a non-thread local but synchronised
allocator is used. This ensures that the lifetime of allocated storage
persists beyond the lifetime of the child thread.

I have added two tests for the `-test-distinct-attrs` pass and the
`-enable-debug-info-on-llvm-scope` passes that run them with crash
reproduction enabled.

Author: abulavin

mlir/include/mlir/IR/MLIRContext.h

@@ -153,6 +153,14 @@ 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,6 +24,7 @@
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/Support/TrailingObjects.h"
+#include <mutex>
 
 namespace mlir {
 namespace detail {
@@ -401,7 +402,8 @@ class DistinctAttributeUniquer {
 /// is freed after the destruction of the distinct attribute allocator.
 class DistinctAttributeAllocator {
 public:
-  DistinctAttributeAllocator() = default;
+  DistinctAttributeAllocator(bool threadingIsEnabled)
+      : threadingIsEnabled(threadingIsEnabled), useThreadLocalAllocator(true) {};
 
   DistinctAttributeAllocator(DistinctAttributeAllocator &&) = delete;
   DistinctAttributeAllocator(const DistinctAttributeAllocator &) = delete;
@@ -411,12 +413,49 @@ class DistinctAttributeAllocator {
   /// Allocates a distinct attribute storage using a thread local bump pointer
   /// allocator to enable synchronization free parallel allocations.
   DistinctAttrStorage *allocate(Attribute referencedAttr) {
-    return new (allocatorCache.get().Allocate<DistinctAttrStorage>())
-        DistinctAttrStorage(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>())
+        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;
+  }
+
   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,7 +268,8 @@ class MLIRContextImpl {
 
 public:
   MLIRContextImpl(bool threadingIsEnabled)
-      : threadingIsEnabled(threadingIsEnabled) {
+      : threadingIsEnabled(threadingIsEnabled),
+        distinctAttributeAllocator(threadingIsEnabled) {
     if (threadingIsEnabled) {
       ownedThreadPool = std::make_unique<llvm::DefaultThreadPool>();
       threadPool = ownedThreadPool.get();
@@ -596,6 +597,7 @@ 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
@@ -717,6 +719,10 @@ 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,6 +414,15 @@ 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.

mlir/test/Dialect/LLVMIR/add-debuginfo-func-scope-with-crash-reproduction.mlir

@@ -0,0 +1,22 @@
+// Test that the enable-debug-info-scope-on-llvm-func pass can create its
+// distinct attributes when running in the crash reproducer thread.
+
+// RUN: mlir-opt --mlir-disable-threading --mlir-pass-pipeline-crash-reproducer=. \
+// RUN:          --pass-pipeline="builtin.module(ensure-debug-info-scope-on-llvm-func)" \
+// RUN:          --mlir-print-debuginfo %s | FileCheck %s
+
+// RUN: mlir-opt --mlir-pass-pipeline-crash-reproducer=. \
+// RUN:          --pass-pipeline="builtin.module(ensure-debug-info-scope-on-llvm-func)" \
+// RUN:          --mlir-print-debuginfo %s | FileCheck %s
+
+module {
+  llvm.func @func_no_debug() {
+    llvm.return loc(unknown)
+  } loc(unknown)
+} loc(unknown)
+
+// CHECK-LABEL: llvm.func @func_no_debug()
+// CHECK: llvm.return loc(#loc
+// CHECK: loc(#loc[[LOC:[0-9]+]])
+// CHECK: #di_compile_unit = #llvm.di_compile_unit<id = distinct[{{.*}}]<>,
+// CHECK: #di_subprogram = #llvm.di_subprogram<id = distinct[{{.*}}]<>

mlir/test/IR/test-builtin-distinct-attrs-with-crash-reproduction.mlir

@@ -0,0 +1,18 @@
+// This test verifies that when running with crash reproduction enabled, distinct
+// attribute storage is not allocated in thread-local storage. Since crash
+// reproduction runs the pass manager in a separate thread, using thread-local
+// storage for distinct attributes causes use-after-free errors once the thread
+// that runs the pass manager joins.
+
+// RUN: mlir-opt --mlir-disable-threading --mlir-pass-pipeline-crash-reproducer=. %s -test-distinct-attrs | FileCheck %s
+// RUN: mlir-opt --mlir-pass-pipeline-crash-reproducer=. %s -test-distinct-attrs | FileCheck %s
+
+// CHECK: #[[DIST0:.*]] = distinct[0]<42 : i32>
+// CHECK: #[[DIST1:.*]] = distinct[1]<42 : i32>
+#distinct = distinct[0]<42 : i32>
+
+// CHECK: @foo_1
+func.func @foo_1() {
+  // CHECK: "test.op"() {distinct.input = #[[DIST0]], distinct.output = #[[DIST1]]}
+  "test.op"() {distinct.input = #distinct} : () -> ()
+}