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

mlir/lib/IR/AttributeDetail.h

@@ -23,6 +23,7 @@
 #include "mlir/Support/ThreadLocalCache.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/PointerIntPair.h"
+#include "llvm/Support/Allocator.h"
 #include "llvm/Support/TrailingObjects.h"
 
 namespace mlir {
@@ -409,14 +410,31 @@ class DistinctAttributeAllocator {
   operator=(const DistinctAttributeAllocator &) = delete;
 
   /// Allocates a distinct attribute storage using a thread local bump pointer
-  /// allocator to enable synchronization free parallel allocations.
+  /// allocator to enable synchronization free parallel allocations. If
+  /// threading is disabled on the owning MLIR context, a normal non
+  /// thread-local 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.
   DistinctAttrStorage *allocate(Attribute referencedAttr) {
-    return new (allocatorCache.get().Allocate<DistinctAttrStorage>())
+    return new (getAllocatorInUse().Allocate<DistinctAttrStorage>())
         DistinctAttrStorage(referencedAttr);
   }
 
+  void disableMultithreading(bool disable = true) {
+    useThreadLocalCache = !disable;
+  };
+
 private:
+  llvm::BumpPtrAllocator &getAllocatorInUse() {
+    if (useThreadLocalCache) {
+      return allocatorCache.get();
+    }
+    return allocator;
+  }
+
   ThreadLocalCache<llvm::BumpPtrAllocator> allocatorCache;
+  llvm::BumpPtrAllocator allocator;
+  bool useThreadLocalCache : 1;
 };
 } // namespace detail
 } // namespace mlir

mlir/lib/IR/MLIRContext.cpp

@@ -596,6 +596,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

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

@@ -0,0 +1,19 @@
+// Test that the enable-debug-info-scope-on-llvm-func pass can create its
+// DI 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
+
+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_file = #llvm.di_file<"<unknown>" in "">
+// CHECK: #di_subprogram = #llvm.di_subprogram<id = distinct[{{.*}}]<>, compileUnit = #di_compile_unit, scope = #di_file, name = "func_no_debug", linkageName = "func_no_debug", file = #di_file, line = 1, scopeLine = 1, subprogramFlags = "Definition|Optimized", type = #di_subroutine_type>
+// CHECK: #loc[[LOC]] = loc(fused<#di_subprogram>

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

@@ -0,0 +1,17 @@
+// This is a regression test that 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
+
+// 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} : () -> ()
+}