[Mlir-commits] [mlir] 4694097 - [mlir] Don't elide the last op if there is no terminator

[Lei Zhang]

Author: Lei Zhang
Date: 2021-05-28T07:25:49-04:00
New Revision: 4694097dab72cc8caed8b93b82c87de686e81df6

URL: https://github.com/llvm/llvm-project/commit/4694097dab72cc8caed8b93b82c87de686e81df6
DIFF: https://github.com/llvm/llvm-project/commit/4694097dab72cc8caed8b93b82c87de686e81df6.diff

LOG: [mlir] Don't elide the last op if there is no terminator

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D103264

Added: 
    

Modified: 
    mlir/lib/IR/AsmPrinter.cpp
    mlir/test/IR/region.mlir
    mlir/test/lib/Dialect/Test/TestDialect.cpp
    mlir/test/lib/Dialect/Test/TestOps.td

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 6d3bd0d768c8..dace47a9e886 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -420,8 +420,12 @@ class DummyAliasOperationPrinter : private OpAsmPrinter {
 
     // Consider the operations within this block, ignoring the terminator if
     // requested.
+    bool hasTerminator =
+        !block->empty() && block->back().hasTrait<OpTrait::IsTerminator>();
     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 @@ void OperationPrinter::print(Block *block, bool printBlockArgs,
   }
 
   currentIndent += indentWidth;
+  bool hasTerminator =
+      !block->empty() && block->back().hasTrait<OpTrait::IsTerminator>();
   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
index fc0f2da36af4..a1370ac8e244 100644
--- a/mlir/test/IR/region.mlir
+++ b/mlir/test/IR/region.mlir
@@ -101,3 +101,9 @@ func @named_region_has_wrong_number_of_blocks() {
     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
index 8e9e0f264522..396f71eac802 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -1006,6 +1006,27 @@ void RegionIfOp::getSuccessorRegions(
   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
index 54ea3bff9bfc..8c3e45600b9f 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -338,13 +338,22 @@ def SizedRegionOp : TEST_Op<"sized_region_op", []> {
 // 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<SizedRegion<1>>:$my_regions);
 
   let assemblyFormat = "attr-dict `:` $my_regions";