Fix a data race reported by TSan on allowsUnregisteredDialects

[Benoit]

Summary

To fix this without introducing a big lock or breaking compatibility,
this variable is turned into an atomic, and a restriction is introduced
on its mutability: it may only be mutated from false to true, not
from true to false. This allows to ensure that threads can't set this
back to false under the feet of other threads relying on this being
true.

[Benoit]

FYI @benvanik @GMNGeoffrey @mravishankar

[rriddle]

This is going to turn into an assert soon, whenever I get around to landing: https://reviews.llvm.org/D112021

[Benoit]

clarify comment

[Benoit]

Thanks for the quick reply! Great to hear this is going to be resolved in the principled way. I'll switch to fixing the faulty call sites in IREE, then.

mlir/include/mlir/IR/MLIRContext.h

public:
 
   /// Get (or create) a dialect for the given derived dialect name.
   /// The dialect will be loaded from the registry if no dialect is found.
   /// If no dialect is loaded for this name and none is available in the
   /// registry, returns nullptr.
   Dialect *getOrLoadDialect(StringRef name);
 
   /// Return true if we allow to create operation for unregistered dialects.
+  /// This is a sequentially-consistent atomic operation.
   bool allowsUnregisteredDialects();
 
   /// Enables creating operations in unregistered dialects.
+  /// If the current value (returned by `allowsUnregisteredDialects`) is `true`
+  /// then it is illegal to pass `allow = false` here.
+  /// This is a sequentially-consistent atomic operation.
   void allowUnregisteredDialects(bool allow = true);
 
   /// Return true if multi-threading is enabled by the context.
   bool isMultithreadingEnabled();
 
   /// Set the flag specifying if multi-threading is disabled by the context.
   /// The command line debugging flag `--mlir-disable-threading` is overriding
   /// this call and making it a no-op!

mlir/lib/IR/MLIRContext.cpp

public:
   DiagnosticEngine diagEngine;
 
   //===--------------------------------------------------------------------===//
   // Options
   //===--------------------------------------------------------------------===//
 
   /// In most cases, creating operation in unregistered dialect is not desired
   /// and indicate a misconfiguration of the compiler. This option enables to
-  /// detect such use cases
-  bool allowUnregisteredDialects = false;
+  /// detect such use cases.
+  ///
+  /// This hasn't always been an atomic, but before it was, there was a data
+  /// race on this member as some users were setting this inside of passes,
+  /// and that wasn't convenient to fix by adding a lock.
+  ///
+  /// So we made it an atomic, which technically removed the data race on
+  /// `allowUnregisteredDialects` itself, but didn't by itself prevent threads
+  /// from changing its value under each others' feet. To avoid that without
+  /// needing to introduce a lock, we introduced this restriction in
+  /// `setAllowUnregisteredDialects`:
+  ///
+  ///     The only possible mutation is from `false` to `true`.
+  std::atomic<bool> allowUnregisteredDialects{false};
 
   /// Enable support for multi-threading within MLIR.
   bool threadingIsEnabled = true;
 
   /// Track if we are currently executing in a threaded execution environment
   /// (like the pass-manager): this is only a debugging feature to help reducing
   /// the chances of data races one some context APIs.
 #ifndef NDEBUG
llvm::hash_code MLIRContext::getRegistryHash() {
   hash = llvm::hash_combine(hash, impl->loadedDialects.size());
   hash = llvm::hash_combine(hash, impl->registeredAttributes.size());
   hash = llvm::hash_combine(hash, impl->registeredOperations.size());
   hash = llvm::hash_combine(hash, impl->registeredTypes.size());
   return hash;
 }
 
 bool MLIRContext::allowsUnregisteredDialects() {
-  return impl->allowUnregisteredDialects;
+  // Note: using the default sequentially-consistent memory order here.
+  // See the comment in `allowUnregisteredDialects(bool)`.
+  return impl->allowUnregisteredDialects.load();
 }
 
 void MLIRContext::allowUnregisteredDialects(bool allowing) {
-  impl->allowUnregisteredDialects = allowing;
+  // It is illegal to mutate this from `true` back to `false`.
+  // See the `allowUnregisteredDialects` class member comment.
+  assert(allowing || !impl->allowUnregisteredDialects.load());
+  // Note: using the default sequentially-consistent memory order here.
+  // Weakly-ordered atomics are error-prone, so they need to be motivated by a
+  // specific performance reason. This method is not frequently called.
+  impl->allowUnregisteredDialects.store(allowing);
 }
 
 /// Return true if multi-threading is enabled by the context.
 bool MLIRContext::isMultithreadingEnabled() {
   return impl->threadingIsEnabled && llvm::llvm_is_multithreaded();
 }
 
 /// Set the flag specifying if multi-threading is disabled by the context.

mlir/test/CAPI/ir.c

static void buildWithInsertionsAndPrint(MlirContext ctx) {
   // Append operations to the rest of blocks to make them non-empty and thus
   // printable.
   mlirBlockAppendOwnedOperation(block2, op5);
   mlirBlockAppendOwnedOperation(block3, op6);
   mlirBlockAppendOwnedOperation(block4, op7);
 
   mlirOperationDump(op);
   mlirOperationDestroy(op);
-  mlirContextSetAllowUnregisteredDialects(ctx, false);
   // clang-format off
   // CHECK-LABEL:  "insertion.order.test"
   // CHECK:      ^{{.*}}(%{{.*}}: i1
   // CHECK:        "dummy.op1"
   // CHECK-NEXT:   "dummy.op2"
   // CHECK-NEXT:   "dummy.op3"
   // CHECK-NEXT:   "dummy.op4"
   // CHECK:      ^{{.*}}(%{{.*}}: i2