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 @@ -9,6 +9,7 @@ #ifndef MLIR_PASS_PASSMANAGER_H #define MLIR_PASS_PASSMANAGER_H +#include "mlir/IR/OperationSupport.h" #include "mlir/Support/LogicalResult.h" #include "llvm/ADT/Optional.h" #include "llvm/ADT/SmallVector.h" @@ -172,8 +173,11 @@ /// 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. - explicit IRPrinterConfig(bool printModuleScope = false, - bool printAfterOnlyOnChange = false); + /// * 'opPrintingFlags' sets up the printing flags to use when printing the + /// IR. + explicit IRPrinterConfig( + bool printModuleScope = false, bool printAfterOnlyOnChange = false, + OpPrintingFlags opPrintingFlags = OpPrintingFlags()); virtual ~IRPrinterConfig(); /// A hook that may be overridden by a derived config that checks if the IR @@ -197,6 +201,9 @@ /// "changed". bool shouldPrintAfterOnlyOnChange() const { return printAfterOnlyOnChange; } + /// 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; @@ -204,6 +211,9 @@ /// A flag that indicates that the IR after a pass should only be printed if /// a change is detected. bool printAfterOnlyOnChange; + + /// Flags to control printing behavior. + OpPrintingFlags opPrintingFlags; }; /// Add an instrumentation to print the IR before and after pass execution, @@ -220,6 +230,8 @@ /// * '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. + /// * 'opPrintingFlags' sets up the printing flags to use when printing the + /// IR. /// * 'out' corresponds to the stream to output the printed IR to. void enableIRPrinting( std::function shouldPrintBeforePass = @@ -227,7 +239,8 @@ std::function shouldPrintAfterPass = [](Pass *, Operation *) { return true; }, bool printModuleScope = true, bool printAfterOnlyOnChange = true, - raw_ostream &out = llvm::errs()); + raw_ostream &out = llvm::errs(), + OpPrintingFlags opPrintingFlags = OpPrintingFlags()); //===--------------------------------------------------------------------===// // Pass Timing 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 @@ -141,7 +141,8 @@ config->printBeforeIfEnabled(pass, op, [&](raw_ostream &out) { out << formatv("// *** IR Dump Before {0} ***", pass->getName()); - printIR(op, config->shouldPrintAtModuleScope(), out, OpPrintingFlags()); + printIR(op, config->shouldPrintAtModuleScope(), out, + config->getOpPrintingFlags()); out << "\n\n"; }); } @@ -165,7 +166,8 @@ config->printAfterIfEnabled(pass, op, [&](raw_ostream &out) { out << formatv("// *** IR Dump After {0} ***", pass->getName()); - printIR(op, config->shouldPrintAtModuleScope(), out, OpPrintingFlags()); + printIR(op, config->shouldPrintAtModuleScope(), out, + config->getOpPrintingFlags()); out << "\n\n"; }); } @@ -190,9 +192,11 @@ /// Initialize the configuration. PassManager::IRPrinterConfig::IRPrinterConfig(bool printModuleScope, - bool printAfterOnlyOnChange) + bool printAfterOnlyOnChange, + OpPrintingFlags opPrintingFlags) : printModuleScope(printModuleScope), - printAfterOnlyOnChange(printAfterOnlyOnChange) {} + printAfterOnlyOnChange(printAfterOnlyOnChange), + opPrintingFlags(opPrintingFlags) {} PassManager::IRPrinterConfig::~IRPrinterConfig() {} /// A hook that may be overridden by a derived config that checks if the IR @@ -223,8 +227,10 @@ BasicIRPrinterConfig( std::function shouldPrintBeforePass, std::function shouldPrintAfterPass, - bool printModuleScope, bool printAfterOnlyOnChange, raw_ostream &out) - : IRPrinterConfig(printModuleScope, printAfterOnlyOnChange), + bool printModuleScope, bool printAfterOnlyOnChange, + OpPrintingFlags opPrintingFlags, raw_ostream &out) + : IRPrinterConfig(printModuleScope, printAfterOnlyOnChange, + opPrintingFlags), shouldPrintBeforePass(shouldPrintBeforePass), shouldPrintAfterPass(shouldPrintAfterPass), out(out) { assert((shouldPrintBeforePass || shouldPrintAfterPass) && @@ -267,8 +273,9 @@ void PassManager::enableIRPrinting( std::function shouldPrintBeforePass, std::function shouldPrintAfterPass, - bool printModuleScope, bool printAfterOnlyOnChange, raw_ostream &out) { + bool printModuleScope, bool printAfterOnlyOnChange, raw_ostream &out, + OpPrintingFlags opPrintingFlags) { enableIRPrinting(std::make_unique( std::move(shouldPrintBeforePass), std::move(shouldPrintAfterPass), - printModuleScope, printAfterOnlyOnChange, out)); + printModuleScope, printAfterOnlyOnChange, opPrintingFlags, out)); }