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 @@ -420,8 +420,12 @@ // Consider the operations within this block, ignoring the terminator if // requested. + bool hasTerminator = + !block->empty() && block->back().hasTrait(); auto range = llvm::make_range( - block->begin(), std::prev(block->end(), printBlockTerminator ? 0 : 1)); + block->begin(), + std::prev(block->end(), + (!hasTerminator || printBlockTerminator) ? 0 : 1)); for (Operation &op : range) print(&op); } @@ -2604,8 +2608,12 @@ } currentIndent += indentWidth; + bool hasTerminator = + !block->empty() && block->back().hasTrait(); auto range = llvm::make_range( - block->begin(), std::prev(block->end(), printBlockTerminator ? 0 : 1)); + block->begin(), + std::prev(block->end(), + (!hasTerminator || printBlockTerminator) ? 0 : 1)); for (auto &op : range) { print(&op); os << newLine; diff --git a/mlir/test/IR/region.mlir b/mlir/test/IR/region.mlir --- a/mlir/test/IR/region.mlir +++ b/mlir/test/IR/region.mlir @@ -101,3 +101,9 @@ func @foo2() { return } } ) : () -> () + +// CHECK: test.single_no_terminator_custom_asm_op +// CHECK-NEXT: important_dont_drop +test.single_no_terminator_custom_asm_op { + "important_dont_drop"() : () -> () +} diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp --- a/mlir/test/lib/Dialect/Test/TestDialect.cpp +++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp @@ -1006,6 +1006,27 @@ regions.push_back(RegionSuccessor(&elseRegion(), getElseArgs())); } +//===----------------------------------------------------------------------===// +// SingleNoTerminatorCustomAsmOp +//===----------------------------------------------------------------------===// + +static ParseResult parseSingleNoTerminatorCustomAsmOp(OpAsmParser &parser, + OperationState &state) { + Region *body = state.addRegion(); + if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{})) + return failure(); + return success(); +} + +static void print(SingleNoTerminatorCustomAsmOp op, OpAsmPrinter &printer) { + printer << op.getOperationName(); + printer.printRegion( + op.getRegion(), /*printEntryBlockArgs=*/false, + // This op has a single block without terminators. But explicitly mark + // as not printing block terminators for testing. + /*printBlockTerminators=*/false); +} + #include "TestOpEnums.cpp.inc" #include "TestOpInterfaces.cpp.inc" #include "TestOpStructs.cpp.inc" diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td --- a/mlir/test/lib/Dialect/Test/TestOps.td +++ b/mlir/test/lib/Dialect/Test/TestOps.td @@ -338,13 +338,22 @@ // NoTerminator Operation //===----------------------------------------------------------------------===// -def SingleNoTerminatorOp : TEST_Op<"single_no_terminator_op", GraphRegionNoTerminator.traits> { +def SingleNoTerminatorOp : TEST_Op<"single_no_terminator_op", + GraphRegionNoTerminator.traits> { let regions = (region SizedRegion<1>:$my_region); let assemblyFormat = "attr-dict `:` $my_region"; } -def VariadicNoTerminatorOp : TEST_Op<"variadic_no_terminator_op", GraphRegionNoTerminator.traits> { +def SingleNoTerminatorCustomAsmOp : TEST_Op<"single_no_terminator_custom_asm_op", + [SingleBlock, NoTerminator]> { + let regions = (region SizedRegion<1>); + let parser = [{ return ::parseSingleNoTerminatorCustomAsmOp(parser, result); }]; + let printer = [{ return ::print(*this, p); }]; +} + +def VariadicNoTerminatorOp : TEST_Op<"variadic_no_terminator_op", + GraphRegionNoTerminator.traits> { let regions = (region VariadicRegion>:$my_regions); let assemblyFormat = "attr-dict `:` $my_regions";