[Mlir-commits] [mlir] [mlir] Add --list-passes option to mlir-opt (PR #100420)

[llvmlistbot at llvm.org]

https://github.com/Natan-GabrielTiutiuIntel created https://github.com/llvm/llvm-project/pull/100420

Currently, the only way to see the passes that were registered is by calling “mlir-opt --help”. However, for compilers with 500+ passes, the help message becomes too long and sometimes hard to understand. In this PR I add a new "--list-passes" option to mlir-opt, which can be used for printing only the registered passes, a feature that would be extremely useful.

>From 2ffca36b3fd7c284c2629de55d53b76e1e822841 Mon Sep 17 00:00:00 2001
From: "Tiutiu, Natan-Gabriel" <natan-gabriel.tiutiu at intel.com>
Date: Wed, 24 Jul 2024 16:16:33 +0000
Subject: [PATCH] added --list-passes option to mlir-opt

---
 mlir/include/mlir/Pass/PassRegistry.h         |  3 +++
 .../include/mlir/Tools/mlir-opt/MlirOptMain.h | 10 +++++++
 mlir/lib/Pass/PassRegistry.cpp                | 26 +++++++++++++++++++
 mlir/lib/Tools/mlir-opt/MlirOptMain.cpp       | 16 ++++++++++++
 4 files changed, 55 insertions(+)

diff --git a/mlir/include/mlir/Pass/PassRegistry.h b/mlir/include/mlir/Pass/PassRegistry.h
index 08874f0121991..f9cdee683d3ff 100644
--- a/mlir/include/mlir/Pass/PassRegistry.h
+++ b/mlir/include/mlir/Pass/PassRegistry.h
@@ -44,6 +44,9 @@ using PassAllocatorFunction = std::function<std::unique_ptr<Pass>()>;
 // PassRegistry
 //===----------------------------------------------------------------------===//
 
+/// Prints the passes that were previously registered and stored in passRegistry
+void printRegisteredPasses();
+
 /// Structure to group information about a passes and pass pipelines (argument
 /// to invoke via mlir-opt, description, pass pipeline builder).
 class PassRegistryEntry {
diff --git a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
index d0ca188abcc27..ef976c1155795 100644
--- a/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
+++ b/mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h
@@ -119,6 +119,13 @@ class MlirOptMainConfig {
     return success();
   }
 
+  /// List the registered passes and return.
+  MlirOptMainConfig &listPasses(bool list) {
+    listPassesFlag = list;
+    return *this;
+  }
+  bool shouldListPasses() const { return listPassesFlag; }
+
   /// Enable running the reproducer information stored in resources (if
   /// present).
   MlirOptMainConfig &runReproducer(bool enableReproducer) {
@@ -219,6 +226,9 @@ class MlirOptMainConfig {
   /// The callback to populate the pass manager.
   std::function<LogicalResult(PassManager &)> passPipelineCallback;
 
+  /// List the registered passes and return.
+  bool listPassesFlag = false;
+
   /// Enable running the reproducer.
   bool runReproducerFlag = false;
 
diff --git a/mlir/lib/Pass/PassRegistry.cpp b/mlir/lib/Pass/PassRegistry.cpp
index 483cbe80faba6..cb7c15580ec54 100644
--- a/mlir/lib/Pass/PassRegistry.cpp
+++ b/mlir/lib/Pass/PassRegistry.cpp
@@ -68,6 +68,32 @@ static void printOptionHelp(StringRef arg, StringRef desc, size_t indent,
 // PassRegistry
 //===----------------------------------------------------------------------===//
 
+/// Prints the passes that were previously registered and stored in passRegistry
+void mlir::printRegisteredPasses() {
+  size_t maxWidth = 0;
+  for (auto &entry : *passRegistry)
+    maxWidth = std::max(maxWidth, entry.second.getOptionWidth() + 4);
+
+  // Functor used to print the ordered entries of a registration map.
+  auto printOrderedEntries = [&](StringRef header, auto &map) {
+    llvm::SmallVector<PassRegistryEntry *, 32> orderedEntries;
+    for (auto &kv : map)
+      orderedEntries.push_back(&kv.second);
+    llvm::array_pod_sort(
+        orderedEntries.begin(), orderedEntries.end(),
+        [](PassRegistryEntry *const *lhs, PassRegistryEntry *const *rhs) {
+          return (*lhs)->getPassArgument().compare((*rhs)->getPassArgument());
+        });
+
+    llvm::outs().indent(0) << header << ":\n";
+    for (PassRegistryEntry *entry : orderedEntries)
+      entry->printHelpStr(/*indent=*/2, maxWidth);
+  };
+
+  // Print the available passes.
+  printOrderedEntries("Passes", *passRegistry);
+}
+
 /// Print the help information for this pass. This includes the argument,
 /// description, and any pass options. `descIndent` is the indent that the
 /// descriptions should be aligned.
diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
index 831e1e675a44a..0b88d31f41a7c 100644
--- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
+++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
@@ -30,6 +30,7 @@
 #include "mlir/Parser/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
+#include "mlir/Pass/PassRegistry.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/Timing.h"
 #include "mlir/Support/ToolUtilities.h"
@@ -118,6 +119,10 @@ struct MlirOptMainConfigCLOptions : public MlirOptMainConfig {
                  "parsing"),
         cl::location(useExplicitModuleFlag), cl::init(false));
 
+    static cl::opt<bool, /*ExternalStorage=*/true> listPasses(
+        "list-passes", cl::desc("Print the list of registered passes and exit"),
+        cl::location(listPassesFlag), 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));
@@ -522,6 +527,11 @@ static LogicalResult printRegisteredDialects(DialectRegistry &registry) {
   return success();
 }
 
+static LogicalResult printRegisteredPassesAndReturn() {
+  mlir::printRegisteredPasses();
+  return success();
+}
+
 LogicalResult mlir::MlirOptMain(llvm::raw_ostream &outputStream,
                                 std::unique_ptr<llvm::MemoryBuffer> buffer,
                                 DialectRegistry &registry,
@@ -529,6 +539,9 @@ LogicalResult mlir::MlirOptMain(llvm::raw_ostream &outputStream,
   if (config.shouldShowDialects())
     return printRegisteredDialects(registry);
 
+  if (config.shouldListPasses())
+    return printRegisteredPassesAndReturn();
+
   // The split-input-file mode is a very specific mode that slices the file
   // up into small pieces and checks each independently.
   // We use an explicit threadpool to avoid creating and joining/destroying
@@ -565,6 +578,9 @@ LogicalResult mlir::MlirOptMain(int argc, char **argv,
   if (config.shouldShowDialects())
     return printRegisteredDialects(registry);
 
+  if (config.shouldListPasses())
+    return printRegisteredPassesAndReturn();
+
   // When reading from stdin and the input is a tty, it is often a user mistake
   // and the process "appears to be stuck". Print a message to let the user know
   // about it!