[mlir][EDSC] Fix off-by-one BlockBuilder insertion point.

Summary:
In the particular case of an insertion in a block without a terminator, the BlockBuilder insertion point should be block->end().

Adding a unit test to exercise this.

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

diff --git a/mlir/include/mlir/EDSC/Builders.h b/mlir/include/mlir/EDSC/Builders.h
index 7e0fbb7..c143789 100644 (file)
--- a/mlir/include/mlir/EDSC/Builders.h
+++ b/mlir/include/mlir/EDSC/Builders.h
@@ -118,10 +118,9 @@ protected:
   /// As a consequence we must allocate a new OpBuilder + ScopedContext and
   /// let the escape.
   void enter(mlir::Block *block) {
-    bodyScope = new ScopedContext(
-        ScopedContext::getBuilderRef(),
-        OpBuilder::InsertPoint(block, std::prev(block->end())),
-        ScopedContext::getLocation());
+    bodyScope = new ScopedContext(ScopedContext::getBuilderRef(),
+                                  OpBuilder::InsertPoint(block, block->end()),
+                                  ScopedContext::getLocation());
     if (!block->empty()) {
       auto &termOp = block->back();
       if (termOp.isKnownTerminator())

diff --git a/mlir/test/EDSC/builder-api-test.cpp b/mlir/test/EDSC/builder-api-test.cpp
index 6499451..cca6c20 100644 (file)
--- a/mlir/test/EDSC/builder-api-test.cpp
+++ b/mlir/test/EDSC/builder-api-test.cpp
@@ -175,6 +175,33 @@ TEST_FUNC(builder_max_min_for) {
   f.erase();
 }
 
+TEST_FUNC(builder_block_append) {
+  using namespace edsc::op;
+  auto f = makeFunction("builder_blocks");
+
+  OpBuilder builder(f.getBody());
+  ScopedContext scope(builder, f.getLoc());
+
+  BlockHandle b1, functionBlock(&f.front());
+  BlockBuilder(&b1, {}, {})([&] { std_constant_index(0); });
+  BlockBuilder(b1, Append())([&] { std_constant_index(1); });
+  BlockBuilder(b1, Append())([&] { std_ret(); });
+  // Get back to entry block and add a branch into b1
+  BlockBuilder(functionBlock, Append())([&] { std_br(b1, {}); });
+
+  // clang-format off
+  // CHECK-LABEL: @builder_blocks
+  // CHECK-NEXT:   br ^bb1
+  // CHECK-NEXT: ^bb1: // pred: ^bb0
+  // CHECK-NEXT:   constant 0 : index
+  // CHECK-NEXT:   constant 1 : index
+  // CHECK-NEXT:   return
+  // CHECK-NEXT: }
+  // clang-format on
+  f.print(llvm::outs());
+  f.erase();
+}
+
 TEST_FUNC(builder_blocks) {
   using namespace edsc::op;
   auto f = makeFunction("builder_blocks");
@@ -1071,8 +1098,8 @@ TEST_FUNC(builder_loop_for_yield) {
   // CHECK:     [[sum:%[0-9]+]] = addf [[arg0]], [[arg1]] : f32
   // CHECK:     loop.yield [[arg1]], [[sum]] : f32, f32
   // CHECK:     addf [[res]]#0, [[res]]#1 : f32
-
   // clang-format on
+
   f.print(llvm::outs());
   f.erase();
 }