[Mlir-commits] [mlir] [mlir] DistinctAttributeAllocator: fix race (PR #132935)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Emilio Cota (cota)

<details>
<summary>Changes</summary>

The bitfields added in #<!-- -->128566 (`threadingIsEnabled : 1` and `useThreadLocalAllocator : 1`) are accessed without synchronization. Example tsan report:

```
WARNING: ThreadSanitizer: data race (pid=337)
  Write of size 1 at 0x7260001d0ff0 by thread T224:
    #<!-- -->0 disableThreadLocalStorage third_party/llvm/llvm-project/mlir/lib/IR/AttributeDetail.h:434:29
    #<!-- -->1 mlir::MLIRContext::disableThreadLocalStorage(bool) third_party/llvm/llvm-project/mlir/lib/IR/MLIRContext.cpp:723:40
    #<!-- -->2 mlir::PassManager::runWithCrashRecovery(mlir::Operation*, mlir::AnalysisManager) third_party/llvm/llvm-project/mlir/lib/Pass/PassCrashRecovery.cpp:423:8
[...]
Previous write of size 1 at 0x7260001d0ff0 by thread T222:
    #<!-- -->0 disableThreadLocalStorage third_party/llvm/llvm-project/mlir/lib/IR/AttributeDetail.h:434:29
    #<!-- -->1 mlir::MLIRContext::disableThreadLocalStorage(bool) third_party/llvm/llvm-project/mlir/lib/IR/MLIRContext.cpp:723:40
    #<!-- -->2 mlir::PassManager::runWithCrashRecovery(mlir::Operation*, mlir::AnalysisManager) third_party/llvm/llvm-project/mlir/lib/Pass/PassCrashRecovery.cpp:423:8
```

Fix the race by serializing accesses to these fields with the existing `allocatorMutex`.

---
Full diff: https://github.com/llvm/llvm-project/pull/132935.diff


1 Files Affected:

- (modified) mlir/lib/IR/AttributeDetail.h (+12-6) 


``````````diff
diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 8fed18140433c..08ab3c0114265 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -413,16 +413,17 @@ class DistinctAttributeAllocator {
   /// Allocates a distinct attribute storage using a thread local bump pointer
   /// allocator to enable synchronization free parallel allocations.
   DistinctAttrStorage *allocate(Attribute referencedAttr) {
+    std::unique_lock<std::mutex> lock(allocatorMutex);
     if (!useThreadLocalAllocator && threadingIsEnabled) {
-      std::scoped_lock<std::mutex> lock(allocatorMutex);
-      return allocateImpl(referencedAttr);
+      return allocateImpl(referencedAttr, lock);
     }
-    return allocateImpl(referencedAttr);
+    return allocateImpl(referencedAttr, lock);
   }
 
   /// 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) {
+    std::scoped_lock<std::mutex> lock(allocatorMutex);
     threadingIsEnabled = !disable;
   }
 
@@ -431,12 +432,15 @@ class DistinctAttributeAllocator {
   /// beyond the lifetime of a child thread calling this function while ensuring
   /// thread-safe allocation.
   void disableThreadLocalStorage(bool disable = true) {
+    std::scoped_lock<std::mutex> lock(allocatorMutex);
     useThreadLocalAllocator = !disable;
   }
 
 private:
-  DistinctAttrStorage *allocateImpl(Attribute referencedAttr) {
-    return new (getAllocatorInUse().Allocate<DistinctAttrStorage>())
+  DistinctAttrStorage *allocateImpl(Attribute referencedAttr,
+                                    const std::unique_lock<std::mutex>& lock) {
+    assert(lock.owns_lock());
+    return new (getAllocatorInUse(lock).Allocate<DistinctAttrStorage>())
         DistinctAttrStorage(referencedAttr);
   }
 
@@ -444,7 +448,9 @@ class DistinctAttributeAllocator {
   /// 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() {
+  llvm::BumpPtrAllocator &getAllocatorInUse(
+      const std::unique_lock<std::mutex>& lock) {
+    assert(lock.owns_lock());
     if (useThreadLocalAllocator)
       return allocatorCache.get();
     return allocator;

``````````

</details>


https://github.com/llvm/llvm-project/pull/132935