diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h --- a/mlir/include/mlir/IR/Operation.h +++ b/mlir/include/mlir/IR/Operation.h @@ -605,7 +605,7 @@ }; inline raw_ostream &operator<<(raw_ostream &os, Operation &op) { - op.print(os); + op.print(os, OpPrintingFlags().useLocalScope()); return os; } diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp --- a/mlir/lib/IR/AsmPrinter.cpp +++ b/mlir/lib/IR/AsmPrinter.cpp @@ -2344,23 +2344,23 @@ } void Operation::print(raw_ostream &os, OpPrintingFlags flags) { - // Handle top-level operations or local printing. - if (!getParent() || flags.shouldUseLocalScope()) { - AsmState state(this); - OperationPrinter(os, flags, state.getImpl()).print(this); - return; - } + // Find the operation to number from based upon the provided flags. + Operation *printedOp = this; + bool shouldUseLocalScope = flags.shouldUseLocalScope(); + do { + // If we are printing local scope, stop at the first operation that is + // isolated from above. + if (shouldUseLocalScope && printedOp->isKnownIsolatedFromAbove()) + break; - Operation *parentOp = getParentOp(); - if (!parentOp) { - os << "<>\n"; - return; - } - // Get the top-level op. - while (auto *nextOp = parentOp->getParentOp()) - parentOp = nextOp; + // Otherwise, traverse up to the next parent. + Operation *parentOp = printedOp->getParentOp(); + if (!parentOp) + break; + printedOp = parentOp; + } while (true); - AsmState state(parentOp); + AsmState state(printedOp); print(os, state, flags); } void Operation::print(raw_ostream &os, AsmState &state, OpPrintingFlags flags) {