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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Matteo Franciolini (mfrancio)

<details>
<summary>Changes</summary>

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.


---
Full diff: https://github.com/llvm/llvm-project/pull/86374.diff


2 Files Affected:

- (modified) mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h (+24) 
- (modified) mlir/lib/Tools/mlir-opt/MlirOptMain.cpp (+21-8) 


``````````diff
diff --git a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
index 4f7f83cdb47376..edd7aca3e70f02 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 dialects available in the registry.
+  /// This option is for convenience during testing only and discouraged in
+  /// general.
+  MlirOptMainConfig &loadAllAvailableDialects(bool allow) {
+    loadAllAvailableDialectsFlag = allow;
+    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 `MLIRContext`.
+  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..9e03325fdc1a0f 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 in the context all available dialects"),
+        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) {
 

``````````

</details>


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