Add an assertion to protect against missing Dialect registration in a pass pipeline (NFC)

[mehdi_amini]

Improve error message

[mehdi_amini]

Fix build with stdlibc++

[mehdi_amini]

Rename the API to be generic for any multi-threaded context and not mention the pass pipeline

[rriddle]

mlir/include/mlir/IR/MLIRContext.h

161

nit: on misuses of some APIs

mlir/lib/IR/MLIRContext.cpp

267

nit: executing in a

270

Can you wrap this within NDEBUG? It should be a complete NOOP in release.

[rriddle]

mlir/lib/IR/MLIRContext.cpp

497

Can you add asserts to the other non-thread safe APIs?

[rriddle]

mlir/include/mlir/IR/MLIRContext.h

26

nit: This shouldn't be necessary anymore.

[mehdi_amini]

Address comments

mlir/include/mlir/IR/MLIRContext.h

 namespace mlir {
 class AbstractOperation;
 class DiagnosticEngine;
 class Dialect;
 class DialectRegistry;
 class InFlightDiagnostic;
 class Location;
 class MLIRContextImpl;
 class StorageUniquer;
 DialectRegistry &getGlobalDialectRegistry();
 
 /// MLIRContext is the top-level object for a collection of MLIR modules.  It
 /// holds immortal uniqued objects like types, and the tables used to unique
 /// them.
 ///
 /// MLIRContext gets a redundant "MLIR" prefix because otherwise it ends up with
 /// a very generic name ("Context") and because it is uncommon for clients to
public:
   /// Returns the storage uniquer used for constructing type storage instances.
   /// This should not be used directly.
   StorageUniquer &getTypeUniquer();
 
   /// Returns the storage uniquer used for constructing attribute storage
   /// instances. This should not be used directly.
   StorageUniquer &getAttributeUniquer();
 
+  /// These APIs are tracking whether the context will be used in a
+  /// multithreading environment: this has no effect other than enabling
+  /// assertions on misuses of some APIs.
+  void enterMultiThreadedExecution();
+  void exitMultiThreadedExecution();
+
 private:
   const std::unique_ptr<MLIRContextImpl> impl;
 
   /// Get a dialect for the provided namespace and TypeID: abort the program if
   /// a dialect exist for this namespace with different TypeID. If a dialect has
   /// not been loaded for this namespace/TypeID yet, use the provided ctor to
   /// create one on the fly and load it. Returns a pointer to the dialect owned
   /// by the context.

mlir/lib/IR/MLIRContext.cpp

public:
   /// 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;
 
   /// 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
+  std::atomic<int> multiThreadedExecutionContext{0};
+#endif
+
   /// If the operation should be attached to diagnostics printed via the
   /// Operation::emit methods.
   bool printOpOnDiagnostic = true;
 
   /// If the current stack trace should be attached when emitting diagnostics.
   bool printStackTraceOnDiagnostic = false;
 
   //===--------------------------------------------------------------------===//
MLIRContext::getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID,
                               function_ref<std::unique_ptr<Dialect>()> ctor) {
   auto &impl = getImpl();
   // Get the correct insertion position sorted by namespace.
   std::unique_ptr<Dialect> &dialect = impl.loadedDialects[dialectNamespace];
 
   if (!dialect) {
     LLVM_DEBUG(llvm::dbgs()
                << "Load new dialect in Context" << dialectNamespace);
+#ifndef NDEBUG
+    if (impl.multiThreadedExecutionContext != 0) {
+      llvm::errs() << "Loading a dialect (" << dialectNamespace
+                   << ") while in a multi-threaded execution context (maybe "
+                      "the PassManager): this can indicate a "
+                      "missing `dependentDialects` in a pass for example.";
+      abort();
+    }
+#endif
     dialect = ctor();
     assert(dialect && "dialect ctor failed");
     return dialect.get();
   }
 
   // Abort if dialect with namespace has already been registered.
   if (dialect->getTypeID() != dialectID)
     llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +
void MLIRContext::disableMultithreading(bool disable) {
   impl->threadingIsEnabled = !disable;
 
   // Update the threading mode for each of the uniquers.
   impl->affineUniquer.disableMultithreading(disable);
   impl->attributeUniquer.disableMultithreading(disable);
   impl->typeUniquer.disableMultithreading(disable);
 }
 
+void MLIRContext::enterMultiThreadedExecution() {
+#ifndef NDEBUG
+  ++impl->multiThreadedExecutionContext;
+#endif
+}
+void MLIRContext::exitMultiThreadedExecution() {
+#ifndef NDEBUG
+  --impl->multiThreadedExecutionContext;
+#endif
+}
+
 /// Return true if we should attach the operation to diagnostics emitted via
 /// Operation::emit.
 bool MLIRContext::shouldPrintOpOnDiagnostic() {
   return impl->printOpOnDiagnostic;
 }
 
 /// Set the flag specifying if we should attach the operation to diagnostics
 /// emitted via Operation::emit.
bool MLIRContext::isOperationRegistered(StringRef name) {
   return impl->registeredOperations.count(name);
 }
 
 void Dialect::addOperation(AbstractOperation opInfo) {
   assert((getNamespace().empty() || opInfo.dialect.name == getNamespace()) &&
          "op name doesn't start with dialect namespace");
   assert(&opInfo.dialect == this && "Dialect object mismatch");
   auto &impl = context->getImpl();
+  assert(impl.multiThreadedExecutionContext == 0 &&
+         "Registering a new operation kind while in a multi-threaded execution "
+         "context");
   StringRef opName = opInfo.name;
   if (!impl.registeredOperations.insert({opName, std::move(opInfo)}).second) {
     llvm::errs() << "error: operation named '" << opInfo.name
                  << "' is already registered.\n";
     abort();
   }
 }
 
 void Dialect::addType(TypeID typeID, AbstractType &&typeInfo) {
   auto &impl = context->getImpl();
+  assert(impl.multiThreadedExecutionContext == 0 &&
+         "Registering a new type kind while in a multi-threaded execution "
+         "context");
   auto *newInfo =
       new (impl.abstractDialectSymbolAllocator.Allocate<AbstractType>())
           AbstractType(std::move(typeInfo));
   if (!impl.registeredTypes.insert({typeID, newInfo}).second)
     llvm::report_fatal_error("Dialect Type already registered.");
 }
 
 void Dialect::addAttribute(TypeID typeID, AbstractAttribute &&attrInfo) {
   auto &impl = context->getImpl();
+  assert(impl.multiThreadedExecutionContext == 0 &&
+         "Registering a new attribute kind while in a multi-threaded execution "
+         "context");
   auto *newInfo =
       new (impl.abstractDialectSymbolAllocator.Allocate<AbstractAttribute>())
           AbstractAttribute(std::move(attrInfo));
   if (!impl.registeredAttributes.insert({typeID, newInfo}).second)
     llvm::report_fatal_error("Dialect Attribute already registered.");
 }
 
 /// Get the dialect that registered the attribute with the provided typeid.

mlir/lib/Pass/Pass.cpp

LogicalResult PassManager::run(ModuleOp module) {
   // Register all dialects for the current pipeline.
   DialectRegistry dependentDialects;
   getDependentDialects(dependentDialects);
   dependentDialects.loadAll(module.getContext());
 
   // Construct an analysis manager for the pipeline.
   ModuleAnalysisManager am(module, instrumentor.get());
 
+  // Notify the context that we start running a pipeline for book keeping.
+  module.getContext()->enterMultiThreadedExecution();
+
   // If reproducer generation is enabled, run the pass manager with crash
   // handling enabled.
   LogicalResult result = crashReproducerFileName
                              ? runWithCrashRecovery(module, am)
                              : OpPassManager::run(module, am);
 
+  // Notify the context that the run is done.
+  module.getContext()->exitMultiThreadedExecution();
+
   // Dump all of the pass statistics if necessary.
   if (passStatisticsMode)
     dumpStatistics();
   return result;
 }
 
 /// Enable support for the pass manager to generate a reproducer on the event
 /// of a crash or a pass failure. `outputFile` is a .mlir filename used to write

mlir/test/lib/Transforms/TestConvertCallOp.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     rewriter.replaceOpWithNewOp<LLVM::NullOp>(op, getVoidPtrType());
     return success();
   }
 };
 
 class TestConvertCallOp
     : public PassWrapper<TestConvertCallOp, OperationPass<ModuleOp>> {
 public:
+  void getDependentDialects(DialectRegistry &registry) const final {
+    registry.insert<LLVM::LLVMDialect>();
+  }
+
   void runOnOperation() override {
     ModuleOp m = getOperation();
 
     // Populate type conversions.
     LLVMTypeConverter type_converter(m.getContext());
     type_converter.addConversion([&](TestType type) {
       return LLVM::LLVMType::getInt8PtrTy(m.getContext());
     });