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

[William Moses]

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

None

>From ffb9d5d97778baba0fd625bb284ce49a2466b7e2 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..235078bbe0b225 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
   //===--------------------------------------------------------------------===//
@@ -354,6 +357,9 @@ class PassManager : public OpPassManager {
 
     /// Flags to control printing behavior.
     OpPrintingFlags opPrintingFlags;
+
+    /// A flag to disable dependent dialect registration.
+    bool loadDialects;
   };
 
   /// Add an instrumentation to print the IR before and after pass execution,
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)