MLIR: add flag to conditionally disable automatic dialect loading

Author: wsmoses

mlir/include/mlir/Pass/PassManager.h

@@ -274,6 +274,9 @@ class PassManager : public OpPassManager {
   /// Runs the verifier after each individual pass.
   void enableVerifier(bool enabled = true);
 
+  /// Whether dependent dialects should be automatically loaded.
+  void setAutomaticDialectLoading(bool shouldLoad);
+
   //===--------------------------------------------------------------------===//
   // Instrumentations
   //===--------------------------------------------------------------------===//
@@ -496,6 +499,9 @@ class PassManager : public OpPassManager {
 
   /// A flag that indicates if the IR should be verified in between passes.
   bool verifyPasses : 1;
+
+  /// A flag to disable dependent dialect registration.
+  bool loadDialects : 1;
 };
 
 /// Register a set of useful command-line options that can be used to configure

mlir/lib/Pass/Pass.cpp

@@ -853,11 +853,13 @@ LogicalResult PassManager::run(Operation *op) {
            << op->getName() << "' op";
 
   // Register all dialects for the current pipeline.
-  DialectRegistry dependentDialects;
-  getDependentDialects(dependentDialects);
-  context->appendDialectRegistry(dependentDialects);
-  for (StringRef name : dependentDialects.getDialectNames())
-    context->getOrLoadDialect(name);
+  if (loadDialects) {
+    DialectRegistry dependentDialects;
+    getDependentDialects(dependentDialects);
+    context->appendDialectRegistry(dependentDialects);
+    for (StringRef name : dependentDialects.getDialectNames())
+      context->getOrLoadDialect(name);
+  }
 
   // Before running, make sure to finalize the pipeline pass list.
   if (failed(getImpl().finalizePassList(context)))
@@ -893,6 +895,11 @@ LogicalResult PassManager::run(Operation *op) {
   return result;
 }
 
+
+void PassManager::setAutomaticDialectLoading(bool shouldLoad) {
+  loadDialects = shouldLoad;
+}
+
 /// Add the provided instrumentation to the pass manager.
 void PassManager::addInstrumentation(std::unique_ptr<PassInstrumentation> pi) {
   if (!instrumentor)