[Mlir-commits] [mlir] [mlir] Expose loadAllAvailableDialects to mlir-opt (PR #86374)

[Matteo Franciolini]

https://github.com/mfrancio updated https://github.com/llvm/llvm-project/pull/86374

>From 9aa2020655b479797cf9769401510be3794eed89 Mon Sep 17 00:00:00 2001
From: Matteo Franciolini <mfranciolini at tesla.com>
Date: Fri, 22 Mar 2024 11:38:07 -0700
Subject: [PATCH] [mlir] Expose loadAllAvailableDialects to mlir-opt

The MLIRContext class supports preemptively loading all the available
dialects in the context. This option used to be exposed to the MLIR
entry point and it is useful for testing purposes. The patch exposes
back the functionality so that downstream implementations of mlir-opt
can access the option through MlirOptMainConfig.

The patch also exposes a new entry point for MlirOptMain so that client
dialects can interact with the main mlir driver by providing
inputFilename, outputFilename, dialect registry and configuration
while handling the command line options separately.
---
 .../include/mlir/Tools/mlir-opt/MlirOptMain.h | 24 +++++++++++++++
 mlir/lib/Tools/mlir-opt/MlirOptMain.cpp       | 29 ++++++++++++++-----
 2 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
index 4f7f83cdb47376..8a8f401efab776 100644
--- a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
+++ b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
@@ -62,6 +62,17 @@ class MlirOptMainConfig {
     return allowUnregisteredDialectsFlag;
   }
 
+  /// Load all available dialects in the context.
+  /// This option is for convenience during testing only and discouraged in
+  /// general.
+  MlirOptMainConfig &loadAllAvailableDialects(bool shouldLoad) {
+    loadAllAvailableDialectsFlag = shouldLoad;
+    return *this;
+  }
+  bool shouldLoadAllAvailableDialects() const {
+    return loadAllAvailableDialectsFlag;
+  }
+
   /// Set the debug configuration to use.
   MlirOptMainConfig &setDebugConfig(tracing::DebugConfig config) {
     debugConfig = std::move(config);
@@ -211,6 +222,9 @@ class MlirOptMainConfig {
   /// IRDL file to register before processing the input.
   std::string irdlFileFlag = "";
 
+  /// Load all available dialects in the context.
+  bool loadAllAvailableDialectsFlag = false;
+
   /// Location Breakpoints to filter the action logging.
   std::vector<tracing::BreakpointManager *> logActionLocationFilter;
 
@@ -274,6 +288,16 @@ LogicalResult MlirOptMain(llvm::raw_ostream &outputStream,
                           DialectRegistry &registry,
                           const MlirOptMainConfig &config);
 
+/// Implementation for tools like `mlir-opt`.
+/// - inputFilename is the name of the input mlir file.
+/// - outputFilename is the name of the output file.
+/// - registry should contain all the dialects that can be parsed in the source.
+/// - config contains the configuration options for the tool.
+LogicalResult MlirOptMain(llvm::StringRef inputFilename,
+                          llvm::StringRef outputFilename,
+                          DialectRegistry &registry,
+                          const MlirOptMainConfig &config);
+
 /// Implementation for tools like `mlir-opt`.
 /// - toolName is used for the header displayed by `--help`.
 /// - registry should contain all the dialects that can be parsed in the source.
diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
index 44c5e9826f3b79..a6a3adce0c361a 100644
--- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
+++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
@@ -119,6 +119,11 @@ struct MlirOptMainConfigCLOptions : public MlirOptMainConfig {
                  "parsing"),
         cl::location(useExplicitModuleFlag), cl::init(false));
 
+    static cl::opt<bool, /*ExternalStorage=*/true> loadAllAvailableDialects(
+        "load-all-available-dialects",
+        cl::desc("Load all available dialects in the context"),
+        cl::location(loadAllAvailableDialectsFlag), cl::init(false));
+
     static cl::opt<bool, /*ExternalStorage=*/true> runReproducer(
         "run-reproducer", cl::desc("Run the pipeline stored in the reproducer"),
         cl::location(runReproducerFlag), cl::init(false));
@@ -453,6 +458,10 @@ static LogicalResult processBuffer(raw_ostream &os,
   if (threadPool)
     context.setThreadPool(*threadPool);
 
+  // Load all the available dialects in context if requested.
+  if (config.shouldLoadAllAvailableDialects())
+    context.loadAllAvailableDialects();
+
   StringRef irdlFile = config.getIrdlFile();
   if (!irdlFile.empty() && failed(loadIRDLDialects(irdlFile, context)))
     return failure();
@@ -552,15 +561,10 @@ LogicalResult mlir::MlirOptMain(llvm::raw_ostream &outputStream,
                                config.outputSplitMarker());
 }
 
-LogicalResult mlir::MlirOptMain(int argc, char **argv,
-                                llvm::StringRef inputFilename,
+LogicalResult mlir::MlirOptMain(llvm::StringRef inputFilename,
                                 llvm::StringRef outputFilename,
-                                DialectRegistry &registry) {
-
-  InitLLVM y(argc, argv);
-
-  MlirOptMainConfig config = MlirOptMainConfig::createFromCLOptions();
-
+                                DialectRegistry &registry,
+                                const MlirOptMainConfig &config) {
   if (config.shouldShowDialects())
     return printRegisteredDialects(registry);
 
@@ -593,6 +597,15 @@ LogicalResult mlir::MlirOptMain(int argc, char **argv,
   return success();
 }
 
+LogicalResult mlir::MlirOptMain(int argc, char **argv,
+                                llvm::StringRef inputFilename,
+                                llvm::StringRef outputFilename,
+                                DialectRegistry &registry) {
+  InitLLVM y(argc, argv);
+  MlirOptMainConfig config = MlirOptMainConfig::createFromCLOptions();
+  return MlirOptMain(inputFilename, outputFilename, registry, config);
+}
+
 LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
                                 DialectRegistry &registry) {