diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h index fa8a4873572ce..edc8ab489e12d 100644 --- a/mlir/include/mlir/IR/Operation.h +++ b/mlir/include/mlir/IR/Operation.h @@ -1102,6 +1102,29 @@ inline raw_ostream &operator<<(raw_ostream &os, const Operation &op) { return os; } +/// A wrapper class that allows for printing an operation with a set of flags, +/// useful to act as a "stream modifier" to customize printing an operation +/// with a stream using the operator<< overload, e.g.: +/// llvm::dbgs() << OpWithFlags(op, OpPrintingFlags().skipRegions()); +class OpWithFlags { +public: + OpWithFlags(Operation *op, OpPrintingFlags flags = {}) + : op(op), theFlags(flags) {} + OpPrintingFlags &flags() { return theFlags; } + const OpPrintingFlags &flags() const { return theFlags; } + +private: + Operation *op; + OpPrintingFlags theFlags; + friend raw_ostream &operator<<(raw_ostream &os, const OpWithFlags &op); +}; + +inline raw_ostream &operator<<(raw_ostream &os, + const OpWithFlags &opWithFlags) { + opWithFlags.op->print(os, opWithFlags.flags()); + return os; +} + } // namespace mlir namespace llvm { diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp index 94d52215b8f14..e622365bf7bae 100644 --- a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp +++ b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp @@ -80,9 +80,8 @@ ChangeResult Liveness::meet(const AbstractSparseLattice &other) { LogicalResult LivenessAnalysis::visitOperation(Operation *op, ArrayRef operands, ArrayRef results) { - LLVM_DEBUG(DBGS() << "[visitOperation] Enter: "; - op->print(llvm::dbgs(), OpPrintingFlags().skipRegions()); - llvm::dbgs() << "\n"); + LDBG() << "[visitOperation] Enter: " + << OpWithFlags(op, OpPrintingFlags().skipRegions()); // This marks values of type (1.a) and (4) liveness as "live". if (!isMemoryEffectFree(op) || op->hasTrait()) { LDBG() << "[visitOperation] Operation has memory effects or is " diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp index cfd7b5b6c22ff..5650de2862545 100644 --- a/mlir/lib/Transforms/RemoveDeadValues.cpp +++ b/mlir/lib/Transforms/RemoveDeadValues.cpp @@ -36,6 +36,7 @@ #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/Dialect.h" +#include "mlir/IR/Operation.h" #include "mlir/IR/OperationSupport.h" #include "mlir/IR/SymbolTable.h" #include "mlir/IR/Value.h" @@ -411,9 +412,8 @@ static void processRegionBranchOp(RegionBranchOpInterface regionBranchOp, RunLivenessAnalysis &la, DenseSet &nonLiveSet, RDVFinalCleanupList &cl) { - LLVM_DEBUG(DBGS() << "Processing region branch op: "; regionBranchOp->print( - llvm::dbgs(), OpPrintingFlags().skipRegions()); - llvm::dbgs() << "\n"); + LDBG() << "Processing region branch op: " + << OpWithFlags(regionBranchOp, OpPrintingFlags().skipRegions()); // Mark live results of `regionBranchOp` in `liveResults`. auto markLiveResults = [&](BitVector &liveResults) { liveResults = markLives(regionBranchOp->getResults(), nonLiveSet, la); diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp index d224f732a198b..466718b1be29d 100644 --- a/mlir/lib/Transforms/Utils/DialectConversion.cpp +++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp @@ -14,6 +14,7 @@ #include "mlir/IR/Dominance.h" #include "mlir/IR/IRMapping.h" #include "mlir/IR/Iterators.h" +#include "mlir/IR/Operation.h" #include "mlir/Interfaces/FunctionInterfaces.h" #include "mlir/Rewrite/PatternApplicator.h" #include "llvm/ADT/SmallPtrSet.h" @@ -2092,8 +2093,9 @@ OperationLegalizer::legalize(Operation *op, // If the operation has no regions, just print it here. if (!isIgnored && op->getNumRegions() == 0) { - op->print(logger.startLine(), OpPrintingFlags().printGenericOpForm()); - logger.getOStream() << "\n\n"; + logger.startLine() << OpWithFlags(op, + OpPrintingFlags().printGenericOpForm()) + << "\n"; } });