[Mlir-commits] [mlir] MLIR: add flag to conditionally disable automatic dialect loading (PR #120100)

[William Moses]

https://github.com/wsmoses updated https://github.com/llvm/llvm-project/pull/120100

>From 5dfbfc6a976b554c62483a201b1263dd4c8d17fe Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh at wsmoses.com>
Date: Mon, 16 Dec 2024 11:00:39 -0500
Subject: [PATCH] MLIR: add flag to conditionally disable automatic dialect
 loading

---
 mlir/include/mlir/Pass/PassManager.h |  6 ++++++
 mlir/lib/Pass/Pass.cpp               | 17 ++++++++++++-----
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Pass/PassManager.h b/mlir/include/mlir/Pass/PassManager.h
index d9bab431e2e0cc..0c8063bc948bca 100644
--- a/mlir/include/mlir/Pass/PassManager.h
+++ b/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
diff --git a/mlir/lib/Pass/Pass.cpp b/mlir/lib/Pass/Pass.cpp
index 6fd51c1e3cb538..fd3798652ec31c 100644
--- a/mlir/lib/Pass/Pass.cpp
+++ b/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)