[mlir] Provide OpBuilder-based replacements for edsc::BlockBuilder

[ftynse]

Summary

The ScopedBuilder class in EDSC is being gradually phased out in favor of core
OpBuilder-based helpers with callbacks. Provide helper functions that are
compatible with edsc::ScopedContext and can be used to create and populate
blocks using callbacks that take block arguments as callback arguments. This
removes the need for edsc::BlockHandle, forward-declaration of Values used
for block arguments and the tag edsc::Append class, leading to noticable
reduction in the verbosity of the code using helper functions.

Remove "eager mode" construction tests that are only relevant to the
BlockBuilder-based approach.

edsc::BlockHandle and edsc::BlockBuilder are now deprecated and will be
removed soon.

[rriddle]

mlir/lib/Dialect/Linalg/EDSC/Builders.cpp

140

Can you just pass regionBuilder directly?

[ftynse]

mlir/lib/Dialect/Linalg/EDSC/Builders.cpp

140

Sure, thanks, I forgot I cleaned up those signatures!

[ftynse]

Address review and rebase

mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h

private:
   using BuilderType =
       typename std::conditional_t<std::is_same<LoopTy, scf::ParallelOp>::value,
                                   ParallelLoopNestBuilder,
                                   LoopOrAffineLoopBuilder>;
 
   std::unique_ptr<BuilderType> builder;
 };
 
-inline void defaultRegionBuilder(ArrayRef<BlockArgument> args) {}
+inline void defaultRegionBuilder(ValueRange args) {}
 
 /// Build a `linalg.generic` op with the specified `inputs`, `outputs` and
 /// `region`.
 ///
 /// `otherValues` and `otherAttributes` may be passed and will be appended as
 /// operands and attributes respectively.
 ///
 /// Prerequisites:
 /// =============
 ///
 /// 1. `inputs` may contain StructuredIndexed that capture either buffer or
 /// tensor values.
 /// 2. `outputs` may contain StructuredIndexed that capture either buffer values
 /// or tensor types. If both buffer values and tensor types are present, then
 /// all buffer values must appear before any tensor type. Without this
 /// restriction output tensor results would need to be reordered, which would
 /// result in surprising behavior when combined with region definition.
 Operation *makeGenericLinalgOp(
     ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
     ArrayRef<StructuredIndexed> outputs,
-    function_ref<void(ArrayRef<BlockArgument>)> regionBuilder =
-        defaultRegionBuilder,
+    function_ref<void(ValueRange)> regionBuilder = defaultRegionBuilder,
     ArrayRef<Value> otherValues = {}, ArrayRef<Attribute> otherAttributes = {});
 
 namespace ops {
 using edsc::StructuredIndexed;
 
 //===----------------------------------------------------------------------===//
 // EDSC builders for linalg generic operations.
 //===----------------------------------------------------------------------===//
 
 /// Build the body of a region to compute a scalar multiply, under the current
 /// ScopedContext, at the current insert point.
-void mulRegionBuilder(ArrayRef<BlockArgument> args);
+void mulRegionBuilder(ValueRange args);
 
 /// Build the body of a region to compute a scalar multiply-accumulate, under
 /// the current ScopedContext, at the current insert point.
-void macRegionBuilder(ArrayRef<BlockArgument> args);
+void macRegionBuilder(ValueRange args);
 
 /// TODO(ntv): In the future we should tie these implementations to something in
 /// Tablegen that generates the proper interfaces and the proper sugared named
 /// ops.
 
 /// Build a linalg.pointwise, under the current ScopedContext, at the current
 /// insert point, that computes:
 /// ```

 /// computes `O = max(I1, I2)`. The client is responsible for specifying the
 /// proper indexings when creating the StructuredIndexed.
 Operation *linalg_generic_pointwise_max(StructuredIndexed I1,
                                         StructuredIndexed I2,
                                         StructuredIndexed O);
 
 // TODO(ntv): Implement more useful pointwise operations on a per-need basis.
 
-using MatmulRegionBuilder = function_ref<void(ArrayRef<BlockArgument> args)>;
+using MatmulRegionBuilder = function_ref<void(ValueRange args)>;
 
 /// Build a linalg.generic, under the current ScopedContext, at the current
 /// insert point, that computes:
 /// ```
 ///    (m, n, k) = (par, par, seq)
 ///    |
 ///    |  C(m, n) += A(m, k) * B(k, n)
 /// ```

mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h

 using std_zero_extendi = ValueBuilder<ZeroExtendIOp>;
 using std_sign_extendi = ValueBuilder<SignExtendIOp>;
 
 /// Branches into the mlir::Block* captured by BlockHandle `b` with `operands`.
 ///
 /// Prerequisites:
 ///   All Handles have already captured previously constructed IR objects.
 BranchOp std_br(BlockHandle bh, ValueRange operands);
+BranchOp std_br(Block *block, ValueRange operands);
 
 /// Creates a new mlir::Block* and branches to it from the current block.
 /// Argument types are specified by `operands`.
 /// Captures the new block in `bh` and the actual `operands` in `captures`. To
 /// insert the new mlir::Block*, a local ScopedContext is constructed and
 /// released to the current block. The branch operation is then added to the
 /// new block.
 ///

 /// `trueOperands` if `cond` evaluates to `true` (resp. `falseBranch` and
 /// `falseOperand` if `cond` evaluates to `false`).
 ///
 /// Prerequisites:
 ///   All Handles have captured previously constructed IR objects.
 CondBranchOp std_cond_br(Value cond, BlockHandle trueBranch,
                          ValueRange trueOperands, BlockHandle falseBranch,
                          ValueRange falseOperands);
+CondBranchOp std_cond_br(Value cond, Block *trueBranch, ValueRange trueOperands,
+                         Block *falseBranch, ValueRange falseOperands);
 
 /// Eagerly creates new mlir::Block* with argument types specified by
 /// `trueOperands`/`falseOperands`.
 /// Captures the new blocks in `trueBranch`/`falseBranch` and the arguments in
 /// `trueCaptures/falseCaptures`.
 /// To insert the new mlir::Block*, a local ScopedContext is constructed and
 /// released. The branch operation is then added in the original location and
 /// targeting the eagerly constructed blocks.

mlir/include/mlir/EDSC/Builders.h

friend LoopBuilder makeLoopBuilder(Value *iv, Value lb, Value ub, Value step,
                                      ValueRange iterArgsInitValues);
   Operation *op;
 };
 
 // This class exists solely to handle the C++ vexing parse case when
 // trying to enter a Block that has already been constructed.
 class Append {};
 
+/// Deprecated. Use buildInNewBlock or appendToBlock instead.
+///
 /// A BlockBuilder is a NestedBuilder for mlir::Block*.
 /// This exists by opposition to LoopBuilder which is not related to an
 /// mlir::Block* but to a mlir::Value.
 /// It is meant to be used as a temporary object for representing any nested
 /// MLIR construct that is "related to" an mlir::Block*.
 class BlockBuilder : public NestedBuilder {
 public:
   /// Enters the mlir::Block* previously captured by `bh` and sets the insertion
public:
   /// scoped within a BlockBuilder.
   void operator()(function_ref<void(void)> fun = nullptr);
 
 private:
   BlockBuilder(BlockBuilder &) = delete;
   BlockBuilder &operator=(BlockBuilder &other) = delete;
 };
 
+/// Deprecated. Use Block * instead.
+///
 /// A BlockHandle represents a (potentially "delayed") Block abstraction.
 /// This extra abstraction is necessary because an mlir::Block is not an
 /// mlir::Value.
 /// A BlockHandle should be captured by pointer but otherwise passed by Value
 /// everywhere.
 class BlockHandle {
 public:
   /// A BlockHandle constructed without an mlir::Block* represents a "delayed"
public:
   operator bool() { return block != nullptr; }
   operator mlir::Block *() { return block; }
   mlir::Block *getBlock() { return block; }
 
 private:
   mlir::Block *block;
 };
 
+/// Creates a block in the region that contains the insertion block of the
+/// OpBuilder currently at the top of ScopedContext stack (appends the block to
+/// the region). Be aware that this will NOT update the insertion point of the
+/// builder to insert into the newly constructed block.
+Block *createBlock(TypeRange argTypes = llvm::None);
+
+/// Creates a block in the specified region using OpBuilder at the top of
+/// ScopedContext stack (appends the block to the region). Be aware that this
+/// will NOT update the insertion point of the builder to insert into the newly
+/// constructed block.
+Block *createBlockInRegion(Region &region, TypeRange argTypes = llvm::None);
+
+/// Calls "builderFn" with ScopedContext reconfigured to insert into "block" and
+/// passes in the block arguments. If the block has a terminator, the operations
+/// are inserted before the terminator, otherwise appended to the block.
+void appendToBlock(Block *block, function_ref<void(ValueRange)> builderFn);
+
+/// Creates a block in the region that contains the insertion block of the
+/// OpBuilder currently at the top of ScopedContext stack, and calls "builderFn"
+/// to populate the body of the block while passing it the block arguments.
+Block *buildInNewBlock(TypeRange argTypes,
+                       function_ref<void(ValueRange)> builderFn);
+
+/// Creates a block in the specified region using OpBuilder at the top of
+/// ScopedContext stack, and calls "builderFn" to populate the body of the block
+/// while passing it the block arguments.
+Block *buildInNewBlock(Region &region, TypeRange argTypes,
+                       function_ref<void(ValueRange)> builderFn);
+
 /// A StructuredIndexed represents an indexable quantity that is either:
 /// 1. a captured value, which is suitable for buffer and tensor operands, or;
 /// 2. a captured type, which is suitable for tensor return values.
 ///
 /// A StructuredIndexed itself is indexed and passed to `makeGenericLinalgOp`.
 /// It enable an idiomatic syntax for index expressions such as:
 ///
 /// ```

mlir/lib/Dialect/Linalg/EDSC/Builders.cpp

 }
 
 } // namespace edsc
 } // namespace mlir
 
 Operation *mlir::edsc::makeGenericLinalgOp(
     ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
     ArrayRef<StructuredIndexed> outputs,
-    function_ref<void(ArrayRef<BlockArgument>)> regionBuilder,
-    ArrayRef<Value> otherValues, ArrayRef<Attribute> otherAttributes) {
+    function_ref<void(ValueRange)> regionBuilder, ArrayRef<Value> otherValues,
+    ArrayRef<Attribute> otherAttributes) {
   for (unsigned i = 0, e = outputs.size(); i + 1 < e; ++i)
     assert(!(outputs[i].getType().isa<RankedTensorType>() &&
              outputs[i + 1].getType().isa<MemRefType>()) &&
            "output tensors must be passed after output buffers");
   auto &builder = edsc::ScopedContext::getBuilderRef();
   auto *ctx = builder.getContext();
   unsigned nInputs = inputs.size();
   unsigned nOutputs = outputs.size();
for (auto it : llvm::enumerate(values))
     blockTypes.push_back((it.index() < nViews)
                              ? getElementTypeOrSelf(it.value())
                              : it.value().getType());
 
   assert(op->getNumRegions() == 1);
   assert(op->getRegion(0).empty());
   OpBuilder opBuilder(op);
   ScopedContext scope(opBuilder, op->getLoc());
-  BlockHandle b;
-  SmallVector<Value, 8> handles(blockTypes.size());
-  BlockBuilder(&b, op->getRegion(0), blockTypes,
-               handles)([&] { regionBuilder(b.getBlock()->getArguments()); });
+  buildInNewBlock(op->getRegion(0), blockTypes,
+                  [&](ValueRange args) { regionBuilder(args); });
   assert(op->getRegion(0).getBlocks().size() == 1);
   return op;
 }
 
-void mlir::edsc::ops::mulRegionBuilder(ArrayRef<BlockArgument> args) {
+void mlir::edsc::ops::mulRegionBuilder(ValueRange args) {
   using edsc::op::operator+;
   using edsc::op::operator*;
   assert(args.size() == 2 && "expected 2 block arguments");
   Value a(args[0]), b(args[1]);
   linalg_yield(a * b);
 }
 
-void mlir::edsc::ops::macRegionBuilder(ArrayRef<BlockArgument> args) {
+void mlir::edsc::ops::macRegionBuilder(ValueRange args) {
   using edsc::op::operator+;
   using edsc::op::operator*;
   assert(args.size() == 3 && "expected 3 block arguments");
   Value a(args[0]), b(args[1]), c(args[2]);
   linalg_yield(c + a * b);
 }
 
 Operation *mlir::edsc::ops::linalg_generic_pointwise(
     UnaryPointwiseOpBuilder unaryOp, StructuredIndexed I, StructuredIndexed O) {
   SmallVector<IteratorType, 4> iterTypes(O.getExprs().size(),
                                          IteratorType::Parallel);
   if (O.getType().isa<RankedTensorType>()) {
-    auto fun = [&unaryOp](ArrayRef<BlockArgument> args) {
+    auto fun = [&unaryOp](ValueRange args) {
       assert(args.size() == 1 && "expected 1 block arguments");
       Value a(args[0]);
       linalg_yield(unaryOp(a));
     };
     return makeGenericLinalgOp(iterTypes, {I}, {O}, fun);
   }
-  auto fun = [&unaryOp](ArrayRef<BlockArgument> args) {
+  auto fun = [&unaryOp](ValueRange args) {
     assert(args.size() == 2 && "expected 2 block arguments");
     Value a(args[0]);
     linalg_yield(unaryOp(a));
   };
   return makeGenericLinalgOp(iterTypes, {I}, {O}, fun);
 }
 
 Operation *mlir::edsc::ops::linalg_generic_pointwise_tanh(StructuredIndexed I,
                                                           StructuredIndexed O) {
   UnaryPointwiseOpBuilder unOp([](Value a) -> Value { return std_tanh(a); });
   return linalg_generic_pointwise(unOp, I, O);
 }
 
 /// Binary pointwise operation (with broadcast) entry point.
 Operation *mlir::edsc::ops::linalg_generic_pointwise(
     BinaryPointwiseOpBuilder binaryOp, StructuredIndexed I1,
     StructuredIndexed I2, StructuredIndexed O) {
   SmallVector<IteratorType, 4> iterTypes(O.getExprs().size(),
                                          IteratorType::Parallel);
   if (O.getType().isa<RankedTensorType>()) {
-    auto fun = [&binaryOp](ArrayRef<BlockArgument> args) {
+    auto fun = [&binaryOp](ValueRange args) {
       assert(args.size() == 2 && "expected 2 block arguments");
       Value a(args[0]), b(args[1]);
       linalg_yield(binaryOp(a, b));
     };
     return makeGenericLinalgOp(iterTypes, {I1, I2}, {O}, fun);
   }
-  auto fun = [&binaryOp](ArrayRef<BlockArgument> args) {
+  auto fun = [&binaryOp](ValueRange args) {
     assert(args.size() == 3 && "expected 3 block arguments");
     Value a(args[0]), b(args[1]);
     linalg_yield(binaryOp(a, b));
   };
   return makeGenericLinalgOp(iterTypes, {I1, I2}, {O}, fun);
 }
 
 Operation *mlir::edsc::ops::linalg_generic_pointwise_add(StructuredIndexed I1,

mlir/lib/Dialect/StandardOps/EDSC/Intrinsics.cpp

 using namespace mlir::edsc;
 
 BranchOp mlir::edsc::intrinsics::std_br(BlockHandle bh, ValueRange operands) {
   assert(bh && "Expected already captured BlockHandle");
   SmallVector<Value, 4> ops(operands.begin(), operands.end());
   return OperationBuilder<BranchOp>(bh.getBlock(), ops);
 }
 
+BranchOp mlir::edsc::intrinsics::std_br(Block *block, ValueRange operands) {
+  return OperationBuilder<BranchOp>(block, operands);
+}
+
 BranchOp mlir::edsc::intrinsics::std_br(BlockHandle *bh, ArrayRef<Type> types,
                                         MutableArrayRef<Value> captures,
                                         ValueRange operands) {
   assert(!*bh && "Unexpected already captured BlockHandle");
   BlockBuilder(bh, types, captures)(/* no body */);
   SmallVector<Value, 4> ops(operands.begin(), operands.end());
   return OperationBuilder<BranchOp>(bh->getBlock(), ops);
 }
 
+CondBranchOp mlir::edsc::intrinsics::std_cond_br(Value cond, Block *trueBranch,
+                                                 ValueRange trueOperands,
+                                                 Block *falseBranch,
+                                                 ValueRange falseOperands) {
+  return OperationBuilder<CondBranchOp>(cond, trueBranch, trueOperands,
+                                        falseBranch, falseOperands);
+}
+
 CondBranchOp mlir::edsc::intrinsics::std_cond_br(Value cond,
                                                  BlockHandle trueBranch,
                                                  ValueRange trueOperands,
                                                  BlockHandle falseBranch,
                                                  ValueRange falseOperands) {
   SmallVector<Value, 4> trueOps(trueOperands.begin(), trueOperands.end());
   SmallVector<Value, 4> falseOps(falseOperands.begin(), falseOperands.end());
   return OperationBuilder<CondBranchOp>(cond, trueBranch.getBlock(), trueOps,

mlir/lib/EDSC/Builders.cpp

BlockHandle mlir::edsc::BlockHandle::createInRegion(Region &region,
   // We do not want this behavior when using declarative builders with nesting.
   OpBuilder::InsertionGuard g(ScopedContext::getBuilderRef());
   ScopedContext::getBuilderRef().setInsertionPoint(res.block,
                                                    res.block->begin());
   res.block->addArguments(argTypes);
   return res;
 }
 
+Block *mlir::edsc::createBlock(TypeRange argTypes) {
+  assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
+  OpBuilder &builder = ScopedContext::getBuilderRef();
+  Block *block = builder.getInsertionBlock();
+  assert(block != nullptr &&
+         "insertion point not set up in the builder within ScopedContext");
+
+  return createBlockInRegion(*block->getParent(), argTypes);
+}
+
+Block *mlir::edsc::createBlockInRegion(Region &region, TypeRange argTypes) {
+  assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
+  OpBuilder &builder = ScopedContext::getBuilderRef();
+
+  OpBuilder::InsertionGuard guard(builder);
+  return builder.createBlock(&region, {}, argTypes);
+}
+
+void mlir::edsc::appendToBlock(Block *block,
+                               function_ref<void(ValueRange)> builderFn) {
+  assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
+  OpBuilder &builder = ScopedContext::getBuilderRef();
+
+  OpBuilder::InsertionGuard guard(builder);
+  if (block->empty() || block->back().isKnownNonTerminator())
+    builder.setInsertionPointToEnd(block);
+  else
+    builder.setInsertionPoint(&block->back());
+  builderFn(block->getArguments());
+}
+
+Block *mlir::edsc::buildInNewBlock(TypeRange argTypes,
+                                   function_ref<void(ValueRange)> builderFn) {
+  assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
+  OpBuilder &builder = ScopedContext::getBuilderRef();
+  Block *block = builder.getInsertionBlock();
+  assert(block != nullptr &&
+         "insertion point not set up in the builder within ScopedContext");
+  return buildInNewBlock(*block->getParent(), argTypes, builderFn);
+}
+
+Block *mlir::edsc::buildInNewBlock(Region &region, TypeRange argTypes,
+                                   function_ref<void(ValueRange)> builderFn) {
+  assert(ScopedContext::getContext() != nullptr && "ScopedContext not set up");
+  OpBuilder &builder = ScopedContext::getBuilderRef();
+
+  Block *block = createBlockInRegion(region, argTypes);
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(block);
+  builderFn(block->getArguments());
+  return block;
+}
+
 void mlir::edsc::LoopBuilder::operator()(function_ref<void(void)> fun) {
   // Call to `exit` must be explicit and asymmetric (cannot happen in the
   // destructor) because of ordering wrt comma operator.
   /// The particular use case concerns nested blocks:
   ///
   /// ```c++
   ///    For (&i, lb, ub, 1)({
   ///      /--- destructor for this `For` is not always called before ...

mlir/test/EDSC/builder-api-test.cpp

 
 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, {}); });
+  Block *b =
+      buildInNewBlock(TypeRange(), [&](ValueRange) { std_constant_index(0); });
+  appendToBlock(b, [&](ValueRange) { std_constant_index(1); });
+  appendToBlock(b, [&](ValueRange) { std_ret(); });
+  // Get back to entry block and add a branch into "b".
+  appendToBlock(&f.front(), [&](ValueRange) { std_br(b, {}); });
 
   // 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");
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   Value c1(std_constant_int(42, 32)), c2(std_constant_int(1234, 32));
   ReturnOp ret = std_ret();
 
-  Value r;
-  Value args12[2];
-  Value &arg1 = args12[0], &arg2 = args12[1];
-  Value args34[2];
-  Value &arg3 = args34[0], &arg4 = args34[1];
-  BlockHandle b1, b2, functionBlock(&f.front());
-  BlockBuilder(&b1, {c1.getType(), c1.getType()}, args12)(
-      // b2 has not yet been constructed, need to come back later.
-      // This is a byproduct of non-structured control-flow.
-  );
-  BlockBuilder(&b2, {c1.getType(), c1.getType()}, args34)([&] {
-    std_br(b1, {arg3, arg4});
-  });
+  Block *b1 = createBlock({c1.getType(), c1.getType()});
+  Block *b2 = buildInNewBlock({c1.getType(), c1.getType()},
+                              [&](ValueRange args) { std_br(b1, args); });
   // The insertion point within the toplevel function is now past b2, we will
   // need to get back the entry block.
-  // This is what happens with unstructured control-flow..
-  BlockBuilder(b1, Append())([&] {
-    r = arg1 + arg2;
-    std_br(b2, {arg1, r});
+  // This is what happens with unstructured control-flow.
+  appendToBlock(b1, [&](ValueRange args) {
+    Value r = args[0] + args[1];
+    std_br(b2, {args[0], r});
   });
-  // Get back to entry block and add a branch into b1
-  BlockBuilder(functionBlock, Append())([&] { std_br(b1, {c1, c2}); });
+  // Get back to entry block and add a branch into b1.
+  appendToBlock(&f.front(), [&](ValueRange) { std_br(b1, {c1, c2}); });
   ret.erase();
 
   // clang-format off
   // CHECK-LABEL: @builder_blocks
   // CHECK:        %{{.*}} = constant 42 : i32
   // CHECK-NEXT:   %{{.*}} = constant 1234 : i32
   // CHECK-NEXT:   br ^bb1(%{{.*}}, %{{.*}} : i32, i32)
   // CHECK-NEXT: ^bb1(%{{.*}}: i32, %{{.*}}: i32):   // 2 preds: ^bb0, ^bb2
   // CHECK-NEXT:   %{{.*}} = addi %{{.*}}, %{{.*}} : i32
   // CHECK-NEXT:   br ^bb2(%{{.*}}, %{{.*}} : i32, i32)
   // CHECK-NEXT: ^bb2(%{{.*}}: i32, %{{.*}}: i32):   // pred: ^bb1
   // CHECK-NEXT:   br ^bb1(%{{.*}}, %{{.*}} : i32, i32)
   // CHECK-NEXT: }
   // clang-format on
   f.print(llvm::outs());
   f.erase();
 }
 
-TEST_FUNC(builder_blocks_eager) {
-  using namespace edsc::op;
-  auto f = makeFunction("builder_blocks_eager");
-
-  OpBuilder builder(f.getBody());
-  ScopedContext scope(builder, f.getLoc());
-  Value c1(std_constant_int(42, 32)), c2(std_constant_int(1234, 32));
-  Value res;
-  Value args1And2[2], args3And4[2];
-  Value &arg1 = args1And2[0], &arg2 = args1And2[1], &arg3 = args3And4[0],
-        &arg4 = args3And4[1];
-
-  // clang-format off
-  BlockHandle b1, b2;
-  { // Toplevel function scope.
-    // Build a new block for b1 eagerly.
-    std_br(&b1, {c1.getType(), c1.getType()}, args1And2, {c1, c2});
-    // Construct a new block b2 explicitly with a branch into b1.
-    BlockBuilder(&b2, {c1.getType(), c1.getType()}, args3And4)([&]{
-        std_br(b1, {arg3, arg4});
-    });
-    /// And come back to append into b1 once b2 exists.
-    BlockBuilder(b1, Append())([&]{
-        res = arg1 + arg2;
-        std_br(b2, {arg1, res});
-    });
-  }
-
-  // CHECK-LABEL: @builder_blocks_eager
-  // CHECK:        %{{.*}} = constant 42 : i32
-  // CHECK-NEXT:   %{{.*}} = constant 1234 : i32
-  // CHECK-NEXT:   br ^bb1(%{{.*}}, %{{.*}} : i32, i32)
-  // CHECK-NEXT: ^bb1(%{{.*}}: i32, %{{.*}}: i32):   // 2 preds: ^bb0, ^bb2
-  // CHECK-NEXT:   %{{.*}} = addi %{{.*}}, %{{.*}} : i32
-  // CHECK-NEXT:   br ^bb2(%{{.*}}, %{{.*}} : i32, i32)
-  // CHECK-NEXT: ^bb2(%{{.*}}: i32, %{{.*}}: i32):   // pred: ^bb1
-  // CHECK-NEXT:   br ^bb1(%{{.*}}, %{{.*}} : i32, i32)
-  // CHECK-NEXT: }
-  // clang-format on
-  f.print(llvm::outs());
-  f.erase();
-}
-
 TEST_FUNC(builder_cond_branch) {
   auto f = makeFunction("builder_cond_branch", {},
                         {IntegerType::get(1, &globalContext())});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
-  Value funcArg(f.getArgument(0));
   Value c32(std_constant_int(32, 32)), c64(std_constant_int(64, 64)),
       c42(std_constant_int(42, 32));
   ReturnOp ret = std_ret();
 
-  Value arg1;
-  Value args23[2];
-  BlockHandle b1, b2, functionBlock(&f.front());
-  BlockBuilder(&b1, c32.getType(), arg1)([&] { std_ret(); });
-  BlockBuilder(&b2, {c64.getType(), c32.getType()}, args23)([&] { std_ret(); });
-  // Get back to entry block and add a conditional branch
-  BlockBuilder(functionBlock, Append())([&] {
-    std_cond_br(funcArg, b1, {c32}, b2, {c64, c42});
+  Block *b1 = buildInNewBlock(c32.getType(), [&](ValueRange) { std_ret(); });
+  Block *b2 = buildInNewBlock({c64.getType(), c32.getType()},
+                              [&](ValueRange) { std_ret(); });
+  // Get back to entry block and add a conditional branch.
+  appendToBlock(&f.front(), [&](ValueRange args) {
+    std_cond_br(args[0], b1, {c32}, b2, {c64, c42});
   });
   ret.erase();
 
   // clang-format off
   // CHECK-LABEL: @builder_cond_branch
   // CHECK:   %{{.*}} = constant 32 : i32
   // CHECK-NEXT:   %{{.*}} = constant 64 : i64
   // CHECK-NEXT:   %{{.*}} = constant 42 : i32
   // CHECK-NEXT:   cond_br %{{.*}}, ^bb1(%{{.*}} : i32), ^bb2(%{{.*}}, %{{.*}} : i64, i32)
   // CHECK-NEXT: ^bb1(%{{.*}}: i32):   // pred: ^bb0
   // CHECK-NEXT:   return
   // CHECK-NEXT: ^bb2(%{{.*}}: i64, %{{.*}}: i32):  // pred: ^bb0
   // CHECK-NEXT:   return
   // clang-format on
   f.print(llvm::outs());
   f.erase();
 }
 
-TEST_FUNC(builder_cond_branch_eager) {
-  using namespace edsc::op;
-  auto f = makeFunction("builder_cond_branch_eager", {},
-                        {IntegerType::get(1, &globalContext())});
-
-  OpBuilder builder(f.getBody());
-  ScopedContext scope(builder, f.getLoc());
-  Value arg0(f.getArgument(0));
-  Value c32(std_constant_int(32, 32)), c64(std_constant_int(64, 64)),
-      c42(std_constant_int(42, 32));
-
-  // clang-format off
-  BlockHandle b1, b2;
-  Value arg1[1], args2And3[2];
-  std_cond_br(arg0,
-              &b1, c32.getType(), arg1, c32,
-              &b2, {c64.getType(), c32.getType()}, args2And3, {c64, c42});
-  BlockBuilder(b1, Append())([]{
-      std_ret();
-  });
-  BlockBuilder(b2, Append())([]{
-      std_ret();
-  });
-
-  // CHECK-LABEL: @builder_cond_branch_eager
-  // CHECK:   %{{.*}} = constant 32 : i32
-  // CHECK-NEXT:   %{{.*}} = constant 64 : i64
-  // CHECK-NEXT:   %{{.*}} = constant 42 : i32
-  // CHECK-NEXT:   cond_br %{{.*}}, ^bb1(%{{.*}} : i32), ^bb2(%{{.*}}, %{{.*}} : i64, i32)
-  // CHECK-NEXT: ^bb1(%{{.*}}: i32):   // pred: ^bb0
-  // CHECK-NEXT:   return
-  // CHECK-NEXT: ^bb2(%{{.*}}: i64, %{{.*}}: i32):  // pred: ^bb0
-  // CHECK-NEXT:   return
-  // clang-format on
-  f.print(llvm::outs());
-  f.erase();
-}
-
 TEST_FUNC(builder_helpers) {
   using namespace edsc::op;
   auto f32Type = FloatType::getF32(&globalContext());
   auto memrefType =
       MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize,
                        ShapedType::kDynamicSize},
                       f32Type, {}, 0);
   auto f =

 
 TEST_FUNC(insertion_in_block) {
   using namespace edsc::op;
   auto indexType = IndexType::get(&globalContext());
   auto f = makeFunction("insertion_in_block", {}, {indexType, indexType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
-  BlockHandle b1;
-  // clang-format off
   std_constant_int(0, 32);
-  (BlockBuilder(&b1))([]{
-    std_constant_int(1, 32);
-  });
+  buildInNewBlock({}, [&](ValueRange) { std_constant_int(1, 32); });
   std_constant_int(2, 32);
+  // clang-format off
   // CHECK-LABEL: @insertion_in_block
   // CHECK: {{.*}} = constant 0 : i32
   // CHECK: {{.*}} = constant 2 : i32
   // CHECK: ^bb1:   // no predecessors
   // CHECK: {{.*}} = constant 1 : i32
   // clang-format on
   f.print(llvm::outs());
   f.erase();
TEST_FUNC(memref_vector_matmul_test) {
   auto typeC =
       MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize},
                       mnVectorType, {}, 0);
   auto f = makeFunction("memref_vector_matmul_test", {}, {typeA, typeB, typeC});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   Value A(f.getArgument(0)), B(f.getArgument(1)), C(f.getArgument(2));
-  auto contractionBuilder = [](ArrayRef<BlockArgument> args) {
+  auto contractionBuilder = [](ValueRange args) {
     assert(args.size() == 3 && "expected 3 block arguments");
     (linalg_yield(vector_contraction_matmul(args[0], args[1], args[2])));
   };
   linalg_generic_matmul(A, B, C, contractionBuilder);
 
   f.print(llvm::outs());
   f.erase();
 }