diff --git a/mlir/include/mlir/Pass/PassManager.h b/mlir/include/mlir/Pass/PassManager.h --- a/mlir/include/mlir/Pass/PassManager.h +++ b/mlir/include/mlir/Pass/PassManager.h @@ -253,11 +253,15 @@ /// pass, we only print in the case of a failure. /// - This option should *not* be used with the other `printAfter` flags /// above. + /// * 'printCustomFormOnFailure' signals that when printing the IR after a + /// pass failure, the custom form should be used (unsafe) instead of the + /// generic form. /// * 'opPrintingFlags' sets up the printing flags to use when printing the /// IR. explicit IRPrinterConfig( bool printModuleScope = false, bool printAfterOnlyOnChange = false, bool printAfterOnlyOnFailure = false, + bool printCustomFormOnFailure = false, OpPrintingFlags opPrintingFlags = OpPrintingFlags()); virtual ~IRPrinterConfig(); @@ -288,6 +292,13 @@ return printAfterOnlyOnFailure; } + /// Returns true if the IR should be printed in custom form even on failure. + /// This is unsafe and there is no guaranee that the custom form printer + /// will not crash or print valid IR. + bool shouldPrintCustomFormOnFailure() const { + return printCustomFormOnFailure; + } + /// Returns the printing flags to be used to print the IR. OpPrintingFlags getOpPrintingFlags() const { return opPrintingFlags; } @@ -303,6 +314,10 @@ /// the pass failed. bool printAfterOnlyOnFailure; + /// A flag that indicates that the IR should be printed (or attempted to be + /// printed) in custom form even after a pass failure. + bool printCustomFormOnFailure; + /// Flags to control printing behavior. OpPrintingFlags opPrintingFlags; }; @@ -325,6 +340,9 @@ /// pass, we only print in the case of a failure. /// - This option should *not* be used with the other `printAfter` flags /// above. + /// * 'printCustomFormOnFailure' signals that when printing the IR after a + /// pass failure, the custom form should be used (unsafe) instead of the + /// generic form. /// * 'out' corresponds to the stream to output the printed IR to. /// * 'opPrintingFlags' sets up the printing flags to use when printing the /// IR. @@ -334,7 +352,8 @@ std::function shouldPrintAfterPass = [](Pass *, Operation *) { return true; }, bool printModuleScope = true, bool printAfterOnlyOnChange = true, - bool printAfterOnlyOnFailure = false, raw_ostream &out = llvm::errs(), + bool printAfterOnlyOnFailure = false, + bool printCustomFormOnFailure = false, raw_ostream &out = llvm::errs(), OpPrintingFlags opPrintingFlags = OpPrintingFlags()); //===--------------------------------------------------------------------===// diff --git a/mlir/lib/Pass/IRPrinting.cpp b/mlir/lib/Pass/IRPrinting.cpp --- a/mlir/lib/Pass/IRPrinting.cpp +++ b/mlir/lib/Pass/IRPrinting.cpp @@ -171,7 +171,9 @@ config->printAfterIfEnabled(pass, op, [&](raw_ostream &out) { out << formatv("// -----// IR Dump After {0} Failed", pass->getName()); printIR(op, config->shouldPrintAtModuleScope(), out, - OpPrintingFlags().printGenericOpForm()); + config->shouldPrintCustomFormOnFailure() + ? OpPrintingFlags() + : OpPrintingFlags().printGenericOpForm()); out << "\n\n"; }); } @@ -184,10 +186,12 @@ PassManager::IRPrinterConfig::IRPrinterConfig(bool printModuleScope, bool printAfterOnlyOnChange, bool printAfterOnlyOnFailure, + bool printCustomFormOnFailure, OpPrintingFlags opPrintingFlags) : printModuleScope(printModuleScope), printAfterOnlyOnChange(printAfterOnlyOnChange), printAfterOnlyOnFailure(printAfterOnlyOnFailure), + printCustomFormOnFailure(printCustomFormOnFailure), opPrintingFlags(opPrintingFlags) {} PassManager::IRPrinterConfig::~IRPrinterConfig() = default; @@ -220,10 +224,11 @@ std::function shouldPrintBeforePass, std::function shouldPrintAfterPass, bool printModuleScope, bool printAfterOnlyOnChange, - bool printAfterOnlyOnFailure, OpPrintingFlags opPrintingFlags, - raw_ostream &out) + bool printAfterOnlyOnFailure, bool printCustomFormOnFailure, + OpPrintingFlags opPrintingFlags, raw_ostream &out) : IRPrinterConfig(printModuleScope, printAfterOnlyOnChange, - printAfterOnlyOnFailure, opPrintingFlags), + printAfterOnlyOnFailure, printCustomFormOnFailure, + opPrintingFlags), shouldPrintBeforePass(std::move(shouldPrintBeforePass)), shouldPrintAfterPass(std::move(shouldPrintAfterPass)), out(out) { assert((this->shouldPrintBeforePass || this->shouldPrintAfterPass) && @@ -267,10 +272,10 @@ std::function shouldPrintBeforePass, std::function shouldPrintAfterPass, bool printModuleScope, bool printAfterOnlyOnChange, - bool printAfterOnlyOnFailure, raw_ostream &out, - OpPrintingFlags opPrintingFlags) { + bool printAfterOnlyOnFailure, bool printCustomFormOnFailure, + raw_ostream &out, OpPrintingFlags opPrintingFlags) { enableIRPrinting(std::make_unique( std::move(shouldPrintBeforePass), std::move(shouldPrintAfterPass), printModuleScope, printAfterOnlyOnChange, printAfterOnlyOnFailure, - opPrintingFlags, out)); + printCustomFormOnFailure, opPrintingFlags, out)); } diff --git a/mlir/lib/Pass/PassManagerOptions.cpp b/mlir/lib/Pass/PassManagerOptions.cpp --- a/mlir/lib/Pass/PassManagerOptions.cpp +++ b/mlir/lib/Pass/PassManagerOptions.cpp @@ -53,6 +53,13 @@ llvm::cl::desc( "When printing the IR after a pass, only print if the pass failed"), llvm::cl::init(false)}; + llvm::cl::opt printCustomFormAfterFailure{ + "mlir-print-custom-form-after-failure", + llvm::cl::desc( + "When printing the IR after a pass failure, print in custom form " + "instead of generic (WARNING: this is unsafe and there is no " + "guarantee of a crash-free or valid print"), + llvm::cl::init(false)}; llvm::cl::opt printModuleScope{ "mlir-print-ir-module-scope", llvm::cl::desc("When printing IR for print-ir-[before|after]{-all} " @@ -122,7 +129,7 @@ // Otherwise, add the IR printing instrumentation. pm.enableIRPrinting(shouldPrintBeforePass, shouldPrintAfterPass, printModuleScope, printAfterChange, printAfterFailure, - llvm::errs()); + printCustomFormAfterFailure, llvm::errs()); } void mlir::registerPassManagerCLOptions() {