[mlir] Add asserts when changing various MLIRContext configurations

[rriddle]

Summary

This helps to prevent tsan failures when users inadvertantly mutate the
context in a non-safe way.

[mehdi_amini]

Thanks!

[nicolasvasilache]

mlir/lib/IR/MLIRContext.cpp

473

I have a general question here on the registration requirements with nested passes.

Is there a good rule of thumb for dependent dialects, external models and dynamically injected ops mixed with nested passes?
It seems that a mechanism to declare dependent passes would be a good QoL improvement: it does not seem manageable for nested pass writers to have to manually add registration for everything in dependent passes.

Am I missing something?

[mehdi_amini]

mlir/lib/IR/MLIRContext.cpp

473

Dependent dialects is supposed to be handle recursively, which admittedly can be tricky sometimes. I think the IREE HAL handles it properly as an example: https://github.com/google/iree/blob/main/iree/compiler/Dialect/HAL/Transforms/TranslateExecutables.cpp#L46-L53

External models are supposed to be injected by IR producers, if this is part of a nested pass, the pass that nest this can handle this in getDependentDialects() as well.

mlir/lib/IR/MLIRContext.cpp

 /// Returns the diagnostic engine for this context.
 DiagnosticEngine &MLIRContext::getDiagEngine() { return getImpl().diagEngine; }
 
 //===----------------------------------------------------------------------===//
 // Dialect and Operation Registration
 //===----------------------------------------------------------------------===//
 
 void MLIRContext::appendDialectRegistry(const DialectRegistry &registry) {
+  assert(impl->multiThreadedExecutionContext == 0 &&
+         "appending to the MLIRContext dialect registry while in a "
+         "multi-threaded execution context");
   registry.appendTo(impl->dialectsRegistry);
 
   // For the already loaded dialects, register the interfaces immediately.
   for (const auto &kvp : impl->loadedDialects)
     registry.registerDelayedInterfaces(kvp.second.get());
 }
 
 const DialectRegistry &MLIRContext::getDialectRegistry() {
   return impl->dialectsRegistry;
llvm::hash_code MLIRContext::getRegistryHash() {
   return hash;
 }
 
 bool MLIRContext::allowsUnregisteredDialects() {
   return impl->allowUnregisteredDialects;
 }
 
 void MLIRContext::allowUnregisteredDialects(bool allowing) {
+  assert(impl->multiThreadedExecutionContext == 0 &&
+         "changing MLIRContext `allow-unregistered-dialects` configuration "
+         "while in a multi-threaded execution context");
   impl->allowUnregisteredDialects = 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.
 void MLIRContext::disableMultithreading(bool disable) {
   // This API can be overridden by the global debugging flag
   // --mlir-disable-threading
   if (isThreadingGloballyDisabled())
     return;
+  assert(impl->multiThreadedExecutionContext == 0 &&
+         "changing MLIRContext `disable-threading` configuration while "
+         "in a multi-threaded execution context");
 
   impl->threadingIsEnabled = !disable;
 
   // Update the threading mode for each of the uniquers.
   impl->affineUniquer.disableMultithreading(disable);
   impl->attributeUniquer.disableMultithreading(disable);
   impl->typeUniquer.disableMultithreading(disable);
 

 /// Operation::emit.
 bool MLIRContext::shouldPrintOpOnDiagnostic() {
   return impl->printOpOnDiagnostic;
 }
 
 /// Set the flag specifying if we should attach the operation to diagnostics
 /// emitted via Operation::emit.
 void MLIRContext::printOpOnDiagnostic(bool enable) {
+  assert(impl->multiThreadedExecutionContext == 0 &&
+         "changing MLIRContext `print-op-on-diagnostic` configuration while in "
+         "a multi-threaded execution context");
   impl->printOpOnDiagnostic = enable;
 }
 
 /// Return true if we should attach the current stacktrace to diagnostics when
 /// emitted.
 bool MLIRContext::shouldPrintStackTraceOnDiagnostic() {
   return impl->printStackTraceOnDiagnostic;
 }
 
 /// Set the flag specifying if we should attach the current stacktrace when
 /// emitting diagnostics.
 void MLIRContext::printStackTraceOnDiagnostic(bool enable) {
+  assert(impl->multiThreadedExecutionContext == 0 &&
+         "changing MLIRContext `print-stacktrace-on-diagnostic` configuration "
+         "while in a multi-threaded execution context");
   impl->printStackTraceOnDiagnostic = enable;
 }
 
 /// Return information about all registered operations.  This isn't very
 /// efficient, typically you should ask the operations about their properties
 /// directly.
 std::vector<AbstractOperation *> MLIRContext::getRegisteredOperations() {
   // We just have the operations in a non-deterministic hash table order. Dump

mlir/test/lib/Dialect/Linalg/TestConvVectorization.cpp

void TestConvVectorization::runOnOperation() {
   linalg::populateConvVectorizationPatterns(context, stage1Patterns, tileSizes);
   SmallVector<FrozenRewritePatternSet, 4> frozenStage1Patterns;
   llvm::move(stage1Patterns, std::back_inserter(frozenStage1Patterns));
 
   RewritePatternSet stage2Patterns =
       linalg::getLinalgTilingCanonicalizationPatterns(context);
   scf::populateSCFForLoopCanonicalizationPatterns(stage2Patterns);
 
-  auto stage3Transforms = [](Operation *op) {
-    PassManager pm(op->getContext());
+  auto stage3Transforms = [this](Operation *op) {
+    OpPassManager pm(ModuleOp::getOperationName());
     pm.addPass(createLoopInvariantCodeMotionPass());
-    if (failed(pm.run(cast<ModuleOp>(op))))
+    if (failed(runPipeline(pm, op)))
       llvm_unreachable("Unexpected failure in cleanup pass pipeline.");
     op->walk([](FuncOp func) {
       promoteSingleIterationLoops(func);
       linalg::hoistRedundantVectorTransfers(func);
     });
     return success();
   };
 
void TestConvVectorization::runOnOperation() {
   // Pattern is not applied because rank-reducing vector transfer is not yet
   // supported as can be seen in splitFullAndPartialTransferPrecondition,
   // VectorTransforms.cpp
   vectorTransferPatterns.add<VectorTransferFullPartialRewriter>(
       context, vectorTransformOptions);
   (void)applyPatternsAndFoldGreedily(module, std::move(vectorTransferPatterns));
 
   // Programmatic controlled lowering of linalg.copy and linalg.fill.
-  PassManager pm(context);
+  OpPassManager pm(ModuleOp::getOperationName());
   pm.addNestedPass<FuncOp>(createConvertLinalgToLoopsPass());
-  if (failed(pm.run(module)))
+  if (failed(runPipeline(pm, module)))
     llvm_unreachable("Unexpected failure in linalg to loops pass.");
 
   // Programmatic controlled lowering of vector.contract only.
   RewritePatternSet vectorContractLoweringPatterns(context);
   populateVectorBroadcastLoweringPatterns(vectorContractLoweringPatterns);
   populateVectorContractLoweringPatterns(vectorContractLoweringPatterns,
                                          vectorTransformOptions);
   populateVectorMaskOpLoweringPatterns(vectorContractLoweringPatterns);

mlir/test/lib/Dialect/Linalg/TestLinalgFusionTransforms.cpp

struct TestLinalgGreedyFusion
   }
   void runOnFunction() override {
     MLIRContext *context = &getContext();
     RewritePatternSet patterns =
         linalg::getLinalgTilingCanonicalizationPatterns(context);
     patterns.add<ExtractSliceOfPadTensorSwapPattern>(context);
     scf::populateSCFForLoopCanonicalizationPatterns(patterns);
     FrozenRewritePatternSet frozenPatterns(std::move(patterns));
-    do {
-      (void)applyPatternsAndFoldGreedily(getFunction(), frozenPatterns);
-      PassManager pm(context);
-      pm.addPass(createLoopInvariantCodeMotionPass());
-      pm.addPass(createCanonicalizerPass());
-      pm.addPass(createCSEPass());
-      LogicalResult res = pm.run(getFunction()->getParentOfType<ModuleOp>());
-      if (failed(res))
+    OpPassManager pm(FuncOp::getOperationName());
+    do {
+      (void)applyPatternsAndFoldGreedily(getFunction(), frozenPatterns);
+      if (failed(runPipeline(pm, getFunction())))
         this->signalPassFailure();
     } while (succeeded(fuseLinalgOpsGreedily(getFunction())));
   }
 };
 
 /// Pass to test tile and fuse of sequence of operations. Intended only for
 /// testing.
 struct TestLinalgTileAndFuseSequencePass