[Mlir-commits] [mlir] [mlir] Use non thread-local allocator for DistinctAttr when threading is disabled (PR #128566)

[Artemiy Bulavin]

https://github.com/abulavin created https://github.com/llvm/llvm-project/pull/128566

Currently, `DistinctAttr` uses a specific 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 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-disable-threading --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 just `BumpPtrAllocator` when threading is disabled in the MLIRContext. The absence of threading caused by `--mlir-disable-threading` means ThreadLocalCache is unnecessary. This change also avoids segmentation faults when generating local reproducers with crash reproduction.

I have added tests specifically for the `-test-distinct-attrs` pass and the `enable-debug-info-on-llvm-scope` passes that cover the distinct attr case speficifcally and also add a test case for a pass that creates these distinct attributes (the latter is how I discovered this issue.)

>From b4a909e27ecd6970611c55ebf7208ba62c2e24b1 Mon Sep 17 00:00:00 2001
From: Artemiy Bulavin <artemiyb at graphcore.ai>
Date: Thu, 30 Jan 2025 14:32:03 +0000
Subject: [PATCH] Use non thread-local allocator for DistinctAttr when
 threading is disabled

---
 mlir/lib/IR/AttributeDetail.h                 | 22 +++++++++++++++++--
 mlir/lib/IR/MLIRContext.cpp                   |  1 +
 ...fo-func-scope-with-crash-reproduction.mlir | 19 ++++++++++++++++
 ...istinct-attrs-with-crash-reproduction.mlir | 17 ++++++++++++++
 4 files changed, 57 insertions(+), 2 deletions(-)
 create mode 100644 mlir/test/Dialect/LLVMIR/add-debuginfo-func-scope-with-crash-reproduction.mlir
 create mode 100644 mlir/test/IR/test-builtin-distinct-attrs-with-crash-reproduction.mlir

diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 26d40ac3a38f6..6787dc6d3e698 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/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
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 87782e84dd6e4..04f2cf8f7b1ec 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/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
diff --git a/mlir/test/Dialect/LLVMIR/add-debuginfo-func-scope-with-crash-reproduction.mlir b/mlir/test/Dialect/LLVMIR/add-debuginfo-func-scope-with-crash-reproduction.mlir
new file mode 100644
index 0000000000000..5f49df419bccb
--- /dev/null
+++ b/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>
diff --git a/mlir/test/IR/test-builtin-distinct-attrs-with-crash-reproduction.mlir b/mlir/test/IR/test-builtin-distinct-attrs-with-crash-reproduction.mlir
new file mode 100644
index 0000000000000..5c3b78bc6e154
--- /dev/null
+++ b/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} : () -> ()
+}