Add an instrumentation for conditionally printing the IR before and after pass execution. This instrumentation can be added directly to the PassManager via 'enableIRPrinting'. mlir-opt exposes access to this instrumentation via the following flags:

* print-ir-before=(comma-separated-pass-list)
  - Print the IR before each of the passes provided within the pass list.
* print-ir-before-all
  - Print the IR before every pass in the pipeline.
* print-ir-after=(comma-separated-pass-list)
  - Print the IR after each of the passes provided within the pass list.
* print-ir-after-all
  - Print the IR after every pass in the pipeline.
* print-ir-module-scope
  - Always print the Module IR, even for non module passes.

PiperOrigin-RevId: 238523649

mlir/g3doc/WritingAPass.md

@@ -462,3 +462,72 @@ $ mlir-opt foo.mlir -cse -canonicalize -convert-to-llvmir -pass-timing
    0.0016 (  7.9%)   0.0006 (  8.9%)   0.0022 (  8.2%)   0.0022 (  8.2%)  ModuleVerifier
    0.0199 (100.0%)   0.0073 (100.0%)   0.0272 (100.0%)   0.0272 (100.0%)  Total
 ```
+
+#### IR Printing
+
+When debugging it is often useful to dump the IR at various stages of a pass
+pipeline. This is where the IR printing instrumentation comes into play. This
+instrumentation allows for conditionally printing the IR before and after pass
+execution by optionally filtering on the pass being executed. This
+instrumentation can be added directly to the PassManager via the
+`enableIRPrinting` method. `mlir-opt` provides a few useful flags for utilizing
+this instrumentation:
+
+*   `print-ir-before=(comma-separated-pass-list)`
+    *   Print the IR before each of the passes provided within the pass list.
+*   `print-ir-before-all`
+    *   Print the IR before every pass in the pipeline.
+
+```shell
+$ mlir-opt foo.mlir -cse -print-ir-before=cse
+
+*** IR Dump Before CSE ***
+func @simple_constant() -> (i32, i32) {
+  %c1_i32 = constant 1 : i32
+  %c1_i32_0 = constant 1 : i32
+  return %c1_i32, %c1_i32_0 : i32, i32
+}
+```
+
+*   `print-ir-after=(comma-separated-pass-list)`
+    *   Print the IR after each of the passes provided within the pass list.
+*   `print-ir-after-all`
+    *   Print the IR after every pass in the pipeline.
+
+```shell
+$ mlir-opt foo.mlir -cse -print-ir-after=cse
+
+*** IR Dump After CSE ***
+func @simple_constant() -> (i32, i32) {
+  %c1_i32 = constant 1 : i32
+  return %c1_i32, %c1_i32 : i32, i32
+}
+```
+
+*   `print-ir-module-scope`
+    *   Always print the Module IR, even for non module passes.
+
+```shell
+$ mlir-opt foo.mlir -cse -print-ir-after=cse -print-ir-module-scope
+
+*** IR Dump After CSE ***  (function: bar)
+func @bar(%arg0: f32, %arg1: f32) -> f32 {
+  ...
+}
+
+func @simple_constant() -> (i32, i32) {
+  %c1_i32 = constant 1 : i32
+  %c1_i32_0 = constant 1 : i32
+  return %c1_i32, %c1_i32_0 : i32, i32
+}
+
+*** IR Dump After CSE ***  (function: simple_constant)
+func @bar(%arg0: f32, %arg1: f32) -> f32 {
+  ...
+}
+
+func @simple_constant() -> (i32, i32) {
+  %c1_i32 = constant 1 : i32
+  return %c1_i32, %c1_i32 : i32, i32
+}
+```

mlir/include/mlir/Pass/PassManager.h

@@ -21,6 +21,10 @@
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/SmallVector.h"
 
+namespace llvm {
+class Any;
+} // end namespace llvm
+
 namespace mlir {
 class FunctionPassBase;
 class Module;
@@ -83,6 +87,17 @@ public:
   /// over the given pointer.
   void addInstrumentation(PassInstrumentation *pi);
 
+  /// Add an instrumentation to print the IR before and after pass execution.
+  /// * 'shouldPrintBeforePass' and 'shouldPrintAfterPass' correspond to filter
+  ///   functions that take a 'Pass *'. These function should return true if the
+  ///   IR should be printed or not.
+  /// * 'printModuleScope' signals if the module IR should be printed, even for
+  ///   non module passes.
+  /// * 'out' corresponds to the stream to output the printed IR to.
+  void enableIRPrinting(std::function<bool(Pass *)> shouldPrintBeforePass,
+                        std::function<bool(Pass *)> shouldPrintAfterPass,
+                        bool printModuleScope, raw_ostream &out);
+
   /// Add an instrumentation to time the execution of passes and the computation
   /// of analyses.
   /// Note: Timing should be enabled after all other instrumentations to avoid

mlir/include/mlir/Pass/PassRegistry.h

@@ -156,6 +156,8 @@ struct PassPipelineRegistration {
 struct PassNameParser : public llvm::cl::parser<const PassRegistryEntry *> {
   PassNameParser(llvm::cl::Option &opt);
 
+  void initialize();
+
   void printOptionInfo(const llvm::cl::Option &O,
                        size_t GlobalWidth) const override;
 };

mlir/lib/Pass/IRPrinting.cpp

@@ -0,0 +1,128 @@
+//===- IRPrinting.cpp -----------------------------------------------------===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+
+#include "PassDetail.h"
+#include "mlir/IR/Module.h"
+#include "mlir/Pass/PassManager.h"
+#include "llvm/Support/Format.h"
+#include "llvm/Support/FormatVariadic.h"
+
+using namespace mlir;
+using namespace mlir::detail;
+
+namespace {
+class IRPrinterInstrumentation : public PassInstrumentation {
+public:
+  /// A filter function to decide if the given ir should be printed. Returns
+  /// true if the ir should be printed, false otherwise.
+  using ShouldPrintFn = std::function<bool(Pass *)>;
+
+  IRPrinterInstrumentation(ShouldPrintFn &&shouldPrintBeforePass,
+                           ShouldPrintFn &&shouldPrintAfterPass,
+                           bool printModuleScope, raw_ostream &out)
+      : shouldPrintBeforePass(shouldPrintBeforePass),
+        shouldPrintAfterPass(shouldPrintAfterPass),
+        printModuleScope(printModuleScope), out(out) {
+    assert((shouldPrintBeforePass || shouldPrintAfterPass) &&
+           "expected atleast one valid filter function");
+  }
+
+private:
+  /// Instrumentation hooks.
+  void runBeforePass(Pass *pass, const llvm::Any &ir) override;
+  void runAfterPass(Pass *pass, const llvm::Any &ir) override;
+  void runAfterPassFailed(Pass *pass, const llvm::Any &ir) override;
+
+  /// Filter functions for before and after pass execution.
+  ShouldPrintFn shouldPrintBeforePass, shouldPrintAfterPass;
+
+  /// Flag to toggle if the printer should always print at module scope.
+  bool printModuleScope;
+
+  /// The stream to output to.
+  raw_ostream &out;
+};
+} // end anonymous namespace
+
+static void printIR(const llvm::Any &ir, bool printModuleScope,
+                    raw_ostream &out) {
+  // Check for printing at module scope.
+  if (printModuleScope && llvm::any_isa<Function *>(ir)) {
+    Function *function = llvm::any_cast<Function *>(ir);
+
+    // Print the function name and a newline before the Module.
+    out << " (function: " << function->getName() << ")\n";
+    function->getModule()->print(out);
+    return;
+  }
+
+  // Print a newline before the IR.
+  out << "\n";
+
+  // Print the given function.
+  if (llvm::any_isa<Function *>(ir)) {
+    llvm::any_cast<Function *>(ir)->print(out);
+    return;
+  }
+
+  // Print the given module.
+  assert(llvm::any_isa<Module *>(ir) && "unexpected IR unit");
+  llvm::any_cast<Module *>(ir)->print(out);
+}
+
+/// Instrumentation hooks.
+void IRPrinterInstrumentation::runBeforePass(Pass *pass, const llvm::Any &ir) {
+  // Skip adaptor passes and passes that the user filtered out.
+  if (!shouldPrintBeforePass || isAdaptorPass(pass) ||
+      !shouldPrintBeforePass(pass))
+    return;
+  out << formatv("*** IR Dump Before {0} ***", pass->getName());
+  printIR(ir, printModuleScope, out);
+}
+
+void IRPrinterInstrumentation::runAfterPass(Pass *pass, const llvm::Any &ir) {
+  // Skip adaptor passes and passes that the user filtered out.
+  if (!shouldPrintAfterPass || isAdaptorPass(pass) ||
+      !shouldPrintAfterPass(pass))
+    return;
+  out << formatv("*** IR Dump After {0} ***", pass->getName());
+  printIR(ir, printModuleScope, out);
+}
+
+void IRPrinterInstrumentation::runAfterPassFailed(Pass *pass,
+                                                  const llvm::Any &ir) {
+  // Skip adaptor passes and passes that the user filtered out.
+  if (!shouldPrintAfterPass || isAdaptorPass(pass) ||
+      !shouldPrintAfterPass(pass))
+    return;
+  out << formatv("*** IR Dump After {0} Failed ***", pass->getName());
+  printIR(ir, printModuleScope, out);
+}
+
+//===----------------------------------------------------------------------===//
+// PassManager
+//===----------------------------------------------------------------------===//
+
+/// Add an instrumentation to print the IR before and after pass execution.
+void PassManager::enableIRPrinting(
+    std::function<bool(Pass *)> shouldPrintBeforePass,
+    std::function<bool(Pass *)> shouldPrintAfterPass, bool printModuleScope,
+    raw_ostream &out) {
+  addInstrumentation(new IRPrinterInstrumentation(
+      std::move(shouldPrintBeforePass), std::move(shouldPrintAfterPass),
+      printModuleScope, out));
+}

mlir/lib/Pass/PassDetail.h

@@ -119,6 +119,13 @@ private:
   FunctionPassExecutor fpe;
 };
 
+/// Utility function to return if a pass refers to an adaptor pass. Adaptor
+/// passes are those that internally execute a pipeline, such as the
+/// ModuleToFunctionPassAdaptor.
+inline bool isAdaptorPass(Pass *pass) {
+  return isa<ModuleToFunctionPassAdaptor>(pass);
+}
+
 } // end namespace detail
 } // end namespace mlir
 #endif // MLIR_PASS_PASSDETAIL_H_

mlir/lib/Pass/PassRegistry.cpp

@@ -88,7 +88,11 @@ const PassInfo *mlir::Pass::lookupPassInfo(const PassID *passID) {
 //===----------------------------------------------------------------------===//
 
 PassNameParser::PassNameParser(llvm::cl::Option &opt)
-    : llvm::cl::parser<const PassRegistryEntry *>(opt) {
+    : llvm::cl::parser<const PassRegistryEntry *>(opt) {}
+
+void PassNameParser::initialize() {
+  llvm::cl::parser<const PassRegistryEntry *>::initialize();
+
   /// Add the pass entries.
   for (const auto &kv : *passRegistry) {
     addLiteralOption(kv.second.getPassArgument(), &kv.second,

mlir/lib/Pass/PassTiming.cpp

@@ -30,13 +30,6 @@ using namespace mlir::detail;
 constexpr llvm::StringLiteral kPassTimingDescription =
     "... Pass execution timing report ...";
 
-/// Utility function to return if a pass refers to an adaptor pass. Adaptor
-/// passes are those that internally execute a pipeline, such as the
-/// ModuleToFunctionPassAdaptor.
-static bool isAdaptorPass(Pass *pass) {
-  return isa<ModuleToFunctionPassAdaptor>(pass);
-}
-
 namespace {
 struct PassTiming : public PassInstrumentation {
   PassTiming(PassTimingDisplayMode displayMode) : displayMode(displayMode) {}

mlir/test/Pass/ir-printing.mlir

@@ -0,0 +1,54 @@
+// RUN: mlir-opt %s -cse -canonicalize -print-ir-before=cse  -o /dev/null 2>&1 | FileCheck -check-prefix=BEFORE %s
+// RUN: mlir-opt %s -cse -canonicalize -print-ir-before-all -o /dev/null 2>&1 | FileCheck -check-prefix=BEFORE_ALL %s
+// RUN: mlir-opt %s -cse -canonicalize -print-ir-after=cse -o /dev/null 2>&1 | FileCheck -check-prefix=AFTER %s
+// RUN: mlir-opt %s -cse -canonicalize -print-ir-after-all -o /dev/null 2>&1 | FileCheck -check-prefix=AFTER_ALL %s
+// RUN: mlir-opt %s -cse -canonicalize -print-ir-before=cse -print-ir-module-scope -o /dev/null 2>&1 | FileCheck -check-prefix=BEFORE_MODULE %s
+
+func @foo() {
+  return
+}
+
+func @bar() {
+  return
+}
+
+// BEFORE: *** IR Dump Before CSE ***
+// BEFORE-NEXT: func @foo()
+// BEFORE: *** IR Dump Before CSE ***
+// BEFORE-NEXT: func @bar()
+// BEFORE-NOT: *** IR Dump Before Canonicalizer ***
+// BEFORE-NOT: *** IR Dump After
+
+// BEFORE_ALL: *** IR Dump Before CSE ***
+// BEFORE_ALL-NEXT: func @foo()
+// BEFORE_ALL: *** IR Dump Before Canonicalizer ***
+// BEFORE_ALL-NEXT: func @foo()
+// BEFORE_ALL: *** IR Dump Before CSE ***
+// BEFORE_ALL-NEXT: func @bar()
+// BEFORE_ALL: *** IR Dump Before Canonicalizer ***
+// BEFORE_ALL-NEXT: func @bar()
+// BEFORE_ALL-NOT: *** IR Dump After
+
+// AFTER-NOT: *** IR Dump Before
+// AFTER: *** IR Dump After CSE ***
+// AFTER-NEXT: func @foo()
+// AFTER: *** IR Dump After CSE ***
+// AFTER-NEXT: func @bar()
+// AFTER-NOT: *** IR Dump After Canonicalizer ***
+
+// AFTER_ALL-NOT: *** IR Dump Before
+// AFTER_ALL: *** IR Dump After CSE ***
+// AFTER_ALL-NEXT: func @foo()
+// AFTER_ALL: *** IR Dump After Canonicalizer ***
+// AFTER_ALL-NEXT: func @foo()
+// AFTER_ALL: *** IR Dump After CSE ***
+// AFTER_ALL-NEXT: func @bar()
+// AFTER_ALL: *** IR Dump After Canonicalizer ***
+// AFTER_ALL-NEXT: func @bar()
+
+// BEFORE_MODULE: *** IR Dump Before CSE *** (function: foo)
+// BEFORE_MODULE: func @foo()
+// BEFORE_MODULE: func @bar()
+// BEFORE_MODULE: *** IR Dump Before CSE *** (function: bar)
+// BEFORE_MODULE: func @foo()
+// BEFORE_MODULE: func @bar()

mlir/tools/mlir-opt/mlir-opt.cpp

@@ -85,10 +85,76 @@ static cl::opt<PassTimingDisplayMode> passTimingDisplayMode(
                clEnumValN(PassTimingDisplayMode::Pipeline, "pipeline",
                           "display the results with a nested pipeline view")));
 
+namespace {
+typedef llvm::cl::list<const mlir::PassRegistryEntry *, bool, PassNameParser>
+    PassOptionList;
+}
+
+// Print IR out before/after specified passes.
+static PassOptionList
+    printBefore("print-ir-before",
+                llvm::cl::desc("Print IR before specified passes"));
+
+static PassOptionList
+    printAfter("print-ir-after",
+               llvm::cl::desc("Print IR after specified passes"));
+
+static cl::opt<bool> printBeforeAll("print-ir-before-all",
+                                    llvm::cl::desc("Print IR before each pass"),
+                                    cl::init(false));
+static cl::opt<bool> printAfterAll("print-ir-after-all",
+                                   llvm::cl::desc("Print IR after each pass"),
+                                   cl::init(false));
+
+static cl::opt<bool> printModuleScope(
+    "print-ir-module-scope",
+    cl::desc("When printing IR for print-ir-[before|after]{-all} always print "
+             "a module IR"),
+    cl::init(false));
+
 static std::vector<const mlir::PassRegistryEntry *> *passList;
 
 enum OptResult { OptSuccess, OptFailure };
 
+/// Add an IR printing instrumentation if enabled by any 'print-ir' flags.
+static void addPrinterInstrumentation(PassManager &pm) {
+  std::function<bool(Pass *)> shouldPrintBeforePass, shouldPrintAfterPass;
+
+  // Handle print-before.
+  if (printBeforeAll) {
+    // If we are printing before all, then just return true for the filter.
+    shouldPrintBeforePass = [](Pass *) { return true; };
+  } else if (printBefore.getNumOccurrences() != 0) {
+    // Otherwise if there are specific passes to print before, then check to see
+    // if the pass info for the current pass is included in the list.
+    shouldPrintBeforePass = [&](Pass *pass) {
+      auto *passInfo = pass->lookupPassInfo();
+      return passInfo && llvm::is_contained(printBefore, passInfo);
+    };
+  }
+
+  // Handle print-after.
+  if (printAfterAll) {
+    // If we are printing after all, then just return true for the filter.
+    shouldPrintAfterPass = [](Pass *) { return true; };
+  } else if (printAfter.getNumOccurrences() != 0) {
+    // Otherwise if there are specific passes to print after, then check to see
+    // if the pass info for the current pass is included in the list.
+    shouldPrintAfterPass = [&](Pass *pass) {
+      auto *passInfo = pass->lookupPassInfo();
+      return passInfo && llvm::is_contained(printAfter, passInfo);
+    };
+  }
+
+  // If there are no valid printing filters, then just return.
+  if (!shouldPrintBeforePass && !shouldPrintAfterPass)
+    return;
+
+  // Otherwise, add the IR printing instrumentation.
+  pm.enableIRPrinting(shouldPrintBeforePass, shouldPrintAfterPass,
+                      printModuleScope, llvm::dbgs());
+}
+
 /// Given a MemoryBuffer along with a line and column within it, return the
 /// location being referenced.
 static SMLoc getLocFromLineAndCol(MemoryBuffer &membuf, unsigned lineNo,
@@ -148,6 +214,10 @@ static OptResult performActions(SourceMgr &sourceMgr, MLIRContext *context) {
     passEntry->addToPipeline(pm);
 
   // Add any necessary instrumentations.
+
+  // Add the IR printing instrumentation.
+  addPrinterInstrumentation(pm);
+
   // Note: The pass timing instrumentation should be added last to avoid any
   // potential "ghost" timing from other instrumentations being unintentionally
   // included in the timing results.
@@ -388,8 +458,7 @@ int main(int argc, char **argv) {
   InitLLVM y(argc, argv);
 
   // Parse pass names in main to ensure static initialization completed.
-  llvm::cl::list<const mlir::PassRegistryEntry *, bool, mlir::PassNameParser>
-      passList("", llvm::cl::desc("Compiler passes to run"));
+  PassOptionList passList("", llvm::cl::desc("Compiler passes to run"));
   ::passList = &passList;
   cl::ParseCommandLineOptions(argc, argv, "MLIR modular optimizer driver\n");