diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h index edc8ab489e12d..4f89f8b85861e 100644 --- a/mlir/include/mlir/IR/Operation.h +++ b/mlir/include/mlir/IR/Operation.h @@ -1125,6 +1125,26 @@ inline raw_ostream &operator<<(raw_ostream &os, return os; } +/// A wrapper class that allows for printing an operation with a custom +/// AsmState, useful to act as a "stream modifier" to customize printing an +/// operation with a stream using the operator<< overload, e.g.: +/// llvm::dbgs() << OpWithState(op, OpPrintingFlags().skipRegions()); +class OpWithState { +public: + OpWithState(Operation *op, AsmState &state) : op(op), theState(state) {} + +private: + Operation *op; + AsmState &theState; + friend raw_ostream &operator<<(raw_ostream &os, const OpWithState &op); +}; + +inline raw_ostream &operator<<(raw_ostream &os, + const OpWithState &opWithState) { + opWithState.op->print(os, const_cast(opWithState).theState); + return os; +} + } // namespace mlir namespace llvm { diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp index bdcdaa407e616..de714d8b740af 100644 --- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp +++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp @@ -501,8 +501,7 @@ performActions(raw_ostream &os, << "bytecode version while not emitting bytecode"; AsmState asmState(op.get(), OpPrintingFlags(), /*locationMap=*/nullptr, &fallbackResourceMap); - op.get()->print(os, asmState); - os << '\n'; + os << OpWithState(op.get(), asmState) << '\n'; return success(); }