[mlir] Add a new `print-ir-after-failure` IR pass printing flag

[rriddle]

Summary

This flag will print the IR after a pass only in the case where the pass failed. This can be useful to more easily view the invalid IR, without needing to print after every pass in the pipeline.

[mehdi_amini]

mlir/include/mlir/Pass/PassManager.h

245

I would re-document the "This option should *not* be used with the other print-ir-after flags above" here.

[jpienaar]

Nice, LGTM with Mehdi's comments addressed.

mlir/lib/Pass/PassManagerOptions.cpp

53

I see this follows the style of the rest, but a more active formulation would read easier to me.

123

Update comment?

[rriddle]

Update

mlir/docs/PassManagement.md

 
 *** IR Dump After CSE ***
 func @simple_constant() -> (i32, i32) {
   %c1_i32 = constant 1 : i32
   return %c1_i32, %c1_i32 : i32, i32
 }
 ```
 
+*   `print-ir-after-failure`
+    *   Only print IR after a pass failure.
+    *   This option should *not* be used with the other `print-ir-after` flags
+        above.
+
+```shell
+$ mlir-opt foo.mlir -pass-pipeline='func(cse,bad-pass)' -print-ir-failure
+
+*** IR Dump After BadPass Failed ***
+func @simple_constant() -> (i32, i32) {
+  %c1_i32 = constant 1 : i32
+  return %c1_i32, %c1_i32 : i32, i32
+}
+```
+
 *   `print-ir-module-scope`
     *   Always print the top-level module operation, regardless of pass type or
         operation nesting level.
     *   Note: Printing at module scope should only be used when multi-threading
         is disabled(`-mlir-disable-threading`)
 
 ```shell
 $ mlir-opt foo.mlir -mlir-disable-threading -pass-pipeline='func(cse)' -print-ir-after=cse -print-ir-module-scope

mlir/include/mlir/Pass/PassManager.h

public:
     /// * 'printModuleScope' signals if the top-level module IR should always be
     ///   printed. This should only be set to true when multi-threading is
     ///   disabled, otherwise we may try to print IR that is being modified
     ///   asynchronously.
     /// * 'printAfterOnlyOnChange' signals that when printing the IR after a
     ///   pass, in the case of a non-failure, we should first check if any
     ///   potential mutations were made. This allows for reducing the number of
     ///   logs that don't contain meaningful changes.
+    /// * 'printAfterOnlyOnFailure' signals that when printing the IR after a
+    ///   pass, we only print in the case of a failure.
     /// * 'opPrintingFlags' sets up the printing flags to use when printing the
     ///   IR.
     explicit IRPrinterConfig(
         bool printModuleScope = false, bool printAfterOnlyOnChange = false,
+        bool printAfterOnlyOnFailure = false,
         OpPrintingFlags opPrintingFlags = OpPrintingFlags());
     virtual ~IRPrinterConfig();
 
     /// A hook that may be overridden by a derived config that checks if the IR
     /// of 'operation' should be dumped *before* the pass 'pass' has been
     /// executed. If the IR should be dumped, 'printCallback' should be invoked
     /// with the stream to dump into.
     virtual void printBeforeIfEnabled(Pass *pass, Operation *operation,
                                       PrintCallbackFn printCallback);
 
     /// A hook that may be overridden by a derived config that checks if the IR
     /// of 'operation' should be dumped *after* the pass 'pass' has been
     /// executed. If the IR should be dumped, 'printCallback' should be invoked
     /// with the stream to dump into.
     virtual void printAfterIfEnabled(Pass *pass, Operation *operation,
                                      PrintCallbackFn printCallback);
 
     /// Returns true if the IR should always be printed at the top-level scope.
     bool shouldPrintAtModuleScope() const { return printModuleScope; }
 
     /// Returns true if the IR should only printed after a pass if the IR
     /// "changed".
     bool shouldPrintAfterOnlyOnChange() const { return printAfterOnlyOnChange; }
 
+    /// Returns true if the IR should only printed after a pass if the pass
+    /// "failed".
+    bool shouldPrintAfterOnlyOnFailure() const {
+      return printAfterOnlyOnFailure;
+    }
+
     /// Returns the printing flags to be used to print the IR.
     OpPrintingFlags getOpPrintingFlags() const { return opPrintingFlags; }
 
   private:
     /// A flag that indicates if the IR should be printed at module scope.
     bool printModuleScope;
 
     /// A flag that indicates that the IR after a pass should only be printed if
     /// a change is detected.
     bool printAfterOnlyOnChange;
 
+    /// A flag that indicates that the IR after a pass should only be printed if
+    /// the pass failed.
+    bool printAfterOnlyOnFailure;
+
     /// Flags to control printing behavior.
     OpPrintingFlags opPrintingFlags;
   };
 
   /// Add an instrumentation to print the IR before and after pass execution,
   /// using the provided configuration.
   void enableIRPrinting(std::unique_ptr<IRPrinterConfig> config);
 
public:
   ///   IR.
   /// * 'out' corresponds to the stream to output the printed IR to.
   void enableIRPrinting(
       std::function<bool(Pass *, Operation *)> shouldPrintBeforePass =
           [](Pass *, Operation *) { return true; },
       std::function<bool(Pass *, Operation *)> shouldPrintAfterPass =
           [](Pass *, Operation *) { return true; },
       bool printModuleScope = true, bool printAfterOnlyOnChange = true,
-      raw_ostream &out = llvm::errs(),
+      bool printAfterOnlyOnFailure = false, raw_ostream &out = llvm::errs(),
       OpPrintingFlags opPrintingFlags = OpPrintingFlags());
 
   //===--------------------------------------------------------------------===//
   // Pass Timing
 
   /// A configuration struct provided to the pass timing feature.
   class PassTimingConfig {
   public:

mlir/lib/Pass/IRPrinting.cpp

printIR(op, config->shouldPrintAtModuleScope(), out,
             config->getOpPrintingFlags());
     out << "\n\n";
   });
 }
 
 void IRPrinterInstrumentation::runAfterPass(Pass *pass, Operation *op) {
   if (isa<OpToOpPassAdaptor>(pass))
     return;
+
+  // Check to see if we are only printing on failure.
+  if (config->shouldPrintAfterOnlyOnFailure())
+    return;
+
   // If the config asked to detect changes, compare the current fingerprint with
   // the previous.
   if (config->shouldPrintAfterOnlyOnChange()) {
     auto fingerPrintIt = beforePassFingerPrints.find(pass);
     assert(fingerPrintIt != beforePassFingerPrints.end() &&
            "expected valid fingerprint");
     // If the fingerprints are the same, we don't print the IR.
     if (fingerPrintIt->second == OperationFingerPrint(op)) {

 
 //===----------------------------------------------------------------------===//
 // IRPrinterConfig
 //===----------------------------------------------------------------------===//
 
 /// Initialize the configuration.
 PassManager::IRPrinterConfig::IRPrinterConfig(bool printModuleScope,
                                               bool printAfterOnlyOnChange,
+                                              bool printAfterOnlyOnFailure,
                                               OpPrintingFlags opPrintingFlags)
     : printModuleScope(printModuleScope),
       printAfterOnlyOnChange(printAfterOnlyOnChange),
+      printAfterOnlyOnFailure(printAfterOnlyOnFailure),
       opPrintingFlags(opPrintingFlags) {}
 PassManager::IRPrinterConfig::~IRPrinterConfig() {}
 
 /// A hook that may be overridden by a derived config that checks if the IR
 /// of 'operation' should be dumped *before* the pass 'pass' has been
 /// executed. If the IR should be dumped, 'printCallback' should be invoked
 /// with the stream to dump into.
 void PassManager::IRPrinterConfig::printBeforeIfEnabled(

 
 namespace {
 /// Simple wrapper config that allows for the simpler interface defined above.
 struct BasicIRPrinterConfig : public PassManager::IRPrinterConfig {
   BasicIRPrinterConfig(
       std::function<bool(Pass *, Operation *)> shouldPrintBeforePass,
       std::function<bool(Pass *, Operation *)> shouldPrintAfterPass,
       bool printModuleScope, bool printAfterOnlyOnChange,
-      OpPrintingFlags opPrintingFlags, raw_ostream &out)
+      bool printAfterOnlyOnFailure, OpPrintingFlags opPrintingFlags,
+      raw_ostream &out)
       : IRPrinterConfig(printModuleScope, printAfterOnlyOnChange,
-                        opPrintingFlags),
+                        printAfterOnlyOnFailure, opPrintingFlags),
         shouldPrintBeforePass(shouldPrintBeforePass),
         shouldPrintAfterPass(shouldPrintAfterPass), out(out) {
     assert((shouldPrintBeforePass || shouldPrintAfterPass) &&
            "expected at least one valid filter function");
   }
 
   void printBeforeIfEnabled(Pass *pass, Operation *operation,
                             PrintCallbackFn printCallback) final {
void PassManager::enableIRPrinting(std::unique_ptr<IRPrinterConfig> config) {
   addInstrumentation(
       std::make_unique<IRPrinterInstrumentation>(std::move(config)));
 }
 
 /// Add an instrumentation to print the IR before and after pass execution.
 void PassManager::enableIRPrinting(
     std::function<bool(Pass *, Operation *)> shouldPrintBeforePass,
     std::function<bool(Pass *, Operation *)> shouldPrintAfterPass,
-    bool printModuleScope, bool printAfterOnlyOnChange, raw_ostream &out,
+    bool printModuleScope, bool printAfterOnlyOnChange,
+    bool printAfterOnlyOnFailure, raw_ostream &out,
     OpPrintingFlags opPrintingFlags) {
   enableIRPrinting(std::make_unique<BasicIRPrinterConfig>(
       std::move(shouldPrintBeforePass), std::move(shouldPrintAfterPass),
-      printModuleScope, printAfterOnlyOnChange, opPrintingFlags, out));
+      printModuleScope, printAfterOnlyOnChange, printAfterOnlyOnFailure,
+      opPrintingFlags, out));
 }

mlir/lib/Pass/PassManagerOptions.cpp

struct PassManagerOptions {
   llvm::cl::opt<bool> printAfterAll{"print-ir-after-all",
                                     llvm::cl::desc("Print IR after each pass"),
                                     llvm::cl::init(false)};
   llvm::cl::opt<bool> printAfterChange{
       "print-ir-after-change",
       llvm::cl::desc(
           "When printing the IR after a pass, only print if the IR changed"),
       llvm::cl::init(false)};
+  llvm::cl::opt<bool> printAfterFailure{
+      "print-ir-after-failure",
+      llvm::cl::desc(
+          "When printing the IR after a pass, only print if the pass failed"),
+      llvm::cl::init(false)};
   llvm::cl::opt<bool> printModuleScope{
       "print-ir-module-scope",
       llvm::cl::desc("When printing IR for print-ir-[before|after]{-all} "
                      "always print the top-level operation"),
       llvm::cl::init(false)};
 
   /// Add an IR printing instrumentation if enabled by any 'print-ir' flags.
   void addPrinterInstrumentation(PassManager &pm);
if (printBeforeAll) {
     // if the pass info for the current pass is included in the list.
     shouldPrintBeforePass = [&](Pass *pass, Operation *) {
       auto *passInfo = pass->lookupPassInfo();
       return passInfo && printBefore.contains(passInfo);
     };
   }
 
   // Handle print-after.
-  if (printAfterAll) {
+  if (printAfterAll || printAfterFailure) {
     // If we are printing after all, then just return true for the filter.
     shouldPrintAfterPass = [](Pass *, Operation *) { return true; };
   } else if (printAfter.hasAnyOccurrences()) {
     // 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, Operation *) {
       auto *passInfo = pass->lookupPassInfo();
       return passInfo && printAfter.contains(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, printAfterChange, llvm::errs());
+                      printModuleScope, printAfterChange, printAfterFailure,
+                      llvm::errs());
 }
 
 /// Add a pass timing instrumentation if enabled by 'pass-timing' flags.
 void PassManagerOptions::addTimingInstrumentation(PassManager &pm) {
   if (passTiming)
     pm.enableTiming(
         std::make_unique<PassManager::PassTimingConfig>(passTimingDisplayMode));
 }

mlir/test/Pass/ir-printing.mlir

 // RUN: mlir-opt %s -mlir-disable-threading=true -pass-pipeline='func(cse,canonicalize)' -print-ir-before=cse  -o /dev/null 2>&1 | FileCheck -check-prefix=BEFORE %s
 // RUN: mlir-opt %s -mlir-disable-threading=true -pass-pipeline='func(cse,canonicalize)' -print-ir-before-all -o /dev/null 2>&1 | FileCheck -check-prefix=BEFORE_ALL %s
 // RUN: mlir-opt %s -mlir-disable-threading=true -pass-pipeline='func(cse,canonicalize)' -print-ir-after=cse -o /dev/null 2>&1 | FileCheck -check-prefix=AFTER %s
 // RUN: mlir-opt %s -mlir-disable-threading=true -pass-pipeline='func(cse,canonicalize)' -print-ir-after-all -o /dev/null 2>&1 | FileCheck -check-prefix=AFTER_ALL %s
 // RUN: mlir-opt %s -mlir-disable-threading=true -pass-pipeline='func(cse,canonicalize)' -print-ir-before=cse -print-ir-module-scope -o /dev/null 2>&1 | FileCheck -check-prefix=BEFORE_MODULE %s
 // RUN: mlir-opt %s -mlir-disable-threading=true -pass-pipeline='func(cse,cse)' -print-ir-after-all -print-ir-after-change -o /dev/null 2>&1 | FileCheck -check-prefix=AFTER_ALL_CHANGE %s
+// RUN: not mlir-opt %s -mlir-disable-threading=true -pass-pipeline='func(cse,test-pass-failure)' -print-ir-after-failure -o /dev/null 2>&1 | FileCheck -check-prefix=AFTER_FAILURE %s
 
 func @foo() {
   %0 = constant 0 : i32
   return
 }
 
 func @bar() {
   return

 // BEFORE_MODULE: func @foo()
 // BEFORE_MODULE: func @bar()
 
 // AFTER_ALL_CHANGE: *** IR Dump After{{.*}}CSE ***
 // AFTER_ALL_CHANGE-NEXT: func @foo()
 // AFTER_ALL_CHANGE-NOT: *** IR Dump After{{.*}}CSE ***
 // We expect that only 'foo' changed during CSE, and the second run of CSE did
 // nothing.
+
+// AFTER_FAILURE-NOT: *** IR Dump After{{.*}}CSE
+// AFTER_FAILURE: *** IR Dump After{{.*}}TestFailurePass Failed ***

mlir/test/lib/Pass/TestPassManager.cpp

 
 /// A test pass that always aborts to enable testing the crash recovery
 /// mechanism of the pass manager.
 class TestCrashRecoveryPass
     : public PassWrapper<TestCrashRecoveryPass, OperationPass<>> {
   void runOnOperation() final { abort(); }
 };
 
+/// A test pass that always fails to enable testing the failure recovery
+/// mechanisms of the pass manager.
+class TestFailurePass : public PassWrapper<TestFailurePass, OperationPass<>> {
+  void runOnOperation() final { signalPassFailure(); }
+};
+
 /// A test pass that contains a statistic.
 struct TestStatisticPass
     : public PassWrapper<TestStatisticPass, OperationPass<>> {
   TestStatisticPass() = default;
   TestStatisticPass(const TestStatisticPass &) {}
 
   Statistic opCount{this, "num-ops", "Number of operations counted"};
 
void registerPassManagerTestPass() {
   PassRegistration<TestModulePass>("test-module-pass",
                                    "Test a module pass in the pass manager");
 
   PassRegistration<TestFunctionPass>(
       "test-function-pass", "Test a function pass in the pass manager");
 
   PassRegistration<TestCrashRecoveryPass>(
       "test-pass-crash", "Test a pass in the pass manager that always crashes");
+  PassRegistration<TestFailurePass>(
+      "test-pass-failure", "Test a pass in the pass manager that always fails");
 
   PassRegistration<TestStatisticPass> unusedStatP("test-stats-pass",
                                                   "Test pass statistics");
 
   PassPipelineRegistration<>("test-pm-nested-pipeline",
                              "Test a nested pipeline in the pass manager",
                              testNestedPipeline);
   PassPipelineRegistration<>("test-textual-pm-nested-pipeline",