[mlir] loop::ForOp: provide builders with callbacks for loop body

[ftynse]

Summary

Thanks to a recent change that made ::build functions take an instance of
OpBuilder, it is now possible to build operations within a region attached to
the operaiton about to be created. Exercise this on loop::ForOp by taking a
callback that populates the loop body while the loop is being created.

Additionally, provide helper functions to build perfect nests of ForOps,
with support for iteration arguments. These functions provide the same
functionality as EDSC LoopNestBuilder with simpler implementation, without
relying on edsc::ScopedContext, and using OpBuilder in an unambiguous way.
Compatibility functions for EDSC are provided, but may be removed in the
future.

[mehdi_amini]

Typo in the commit description: operaiton -> operation

mlir/lib/Dialect/LoopOps/LoopOps.cpp

56

Can you document the condition?

I assume it is something like "if there are iterArgs provided, we don't create a default terminator as it must take the same argument types."

[bondhugula]

mlir/include/mlir/Dialect/LoopOps/LoopOps.h

72

Did you mean "without iteration arguments (like for reductions)" instead of "without reductions"?

mlir/lib/Conversion/VectorToLoops/ConvertVectorToLoops.cpp

496

/*bodyBuilder=*/

570

/*bodyBuilder=*/

mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp

82

Likewise.

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

152

Likewise.

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

145

/*lbs=*/..., /*ubs=*/..., /*step=*/...)

[bondhugula]

Mostly minor comments. On another note, I've never seen ValueVector before - I don't find it upstream either. Is this based on another pending revision?

[ftynse]

On another note, I've never seen ValueVector before - I don't find it upstream either.

It's defined in line 47 of LoopOps.h / SCF.h in this commit.

mlir/include/mlir/Dialect/LoopOps/LoopOps.h

72

I think one implies the other, but okay to change.

mlir/lib/Conversion/VectorToLoops/ConvertVectorToLoops.cpp

496

We use argument-naming comments for literals when they can be interpreted ambiguously https://llvm.org/docs/CodingStandards.html#comment-formatting. (e.g. we don't add comment to getRegion(0).) This argument is neither a literal nor really ambiguous.

570

Same as above, this defies the purpose of loop nest constructor.

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

145

Same as above, we don't systematically do this and I don't see a strong reason for it, especially in this commit.

[ftynse]

Address reviews

[ftynse]

Drop debug output

[bondhugula]

mlir/lib/Conversion/VectorToLoops/ConvertVectorToLoops.cpp

496

That URL is about comment formatting rules, and not about the necessity of such argument documentation. The LLVM style is actually silent on recommending its usage except a passing reference there, and so is the MLIR style. I think we should discuss this separately on discord --- using such documenting names where it's not obvious what the argument is (to a new user/browser) greatly improves call site readability, and in this instance as well. Unfortunately, a lot of the new code is missing out on this due to the lack of a guideline. Fine for this commit, but I'll bring this up for discussion on the list.

[ftynse]

mlir/lib/Conversion/VectorToLoops/ConvertVectorToLoops.cpp

496

I linked the only relevant piece I found.

Please prefer discourse (forum) to discord (chat), I find it extremely annoying to use the chat because it prompts short immediate answers that must arrive during night time at my location, or are otherwise ignored due to the lack of context. I will spare my arguments for that discussion. I don't think this change needs to be blocked on it.

[bondhugula]

Looks good for the parts I saw. Please go ahead and commit if the primary reviewer (@nicolasvasilache) is fine.

[ftynse]

Rebase

[rriddle]

mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp

88

Was this ever removed?

mlir/lib/Dialect/SCF/SCF.cpp

286 ↗

(On Diff #264904)

nit: typo beacause

[bondhugula]

This comment was removed by bondhugula.

[ftynse]

195d8571b9e296e7894d51cf59b7aeed42eaa2d2

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

public:
   LoopNestBuilder(MutableArrayRef<Value> ivs, ArrayRef<Value> lbs,
                   ArrayRef<Value> ubs, ArrayRef<Value> steps);
   Operation::result_range operator()(std::function<void(void)> fun = nullptr);
 
 private:
   SmallVector<LoopBuilder, 4> loops;
 };
 
+/// Adapters for building loop nests using the builder and the location stored
+/// in ScopedContext. Actual builders are in loop::buildLoopNest.
+loop::ValueVector loopNestBuilder(ValueRange lbs, ValueRange ubs,
+                                  ValueRange steps,
+                                  function_ref<void(ValueRange)> fun = nullptr);
+loop::ValueVector loopNestBuilder(Value lb, Value ub, Value step,
+                                  function_ref<void(Value)> fun = nullptr);
+loop::ValueVector loopNestBuilder(
+    Value lb, Value ub, Value step, ValueRange iterArgInitValues,
+    function_ref<loop::ValueVector(Value, ValueRange)> fun = nullptr);
+
 } // namespace edsc
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LOOPOPS_EDSC_BUILDERS_H_

mlir/include/mlir/Dialect/LoopOps/LoopOps.h

 // block first.
 void ensureLoopTerminator(Region &region, Builder &builder, Location loc);
 
 /// Returns the loop parent of an induction variable. If the provided value is
 /// not an induction variable, then return nullptr.
 ForOp getForInductionVarOwner(Value val);
 
 /// Returns the parallel loop parent of an induction variable. If the provided
-// value is not an induction variable, then return nullptr.
+/// value is not an induction variable, then return nullptr.
 ParallelOp getParallelForInductionVarOwner(Value val);
 
+/// An owning vector of values, handy to return from functions.
+using ValueVector = std::vector<Value>;
+
+/// Creates a perfect nest of "for" loops, i.e. all loops but the innermost
+/// contain only another loop and a terminator. The lower, upper bounds and
+/// steps are provided as `lbs`, `ubs` and `steps`, which are expected to be of
+/// the same size. `iterArgs` points to the initial values of the loop iteration
+/// arguments, which will be forwarded through the nest to the innermost loop.
+/// The body of the loop is populated using `bodyBuilder`, which accepts an
+/// ordered list of induction variables of all loops, followed by a list of
+/// iteration arguments of the innermost loop, in the same order as provided to
+/// `iterArgs`. This function is expected to return as many values as
+/// `iterArgs`, of the same type and in the same order, that will be treated as
+/// yielded from the loop body and forwarded back through the loop nest. If the
+/// function is not provided, the loop nest is not expected to have iteration
+/// arguments, the body of the innermost loop will be left empty, containing
+/// only the zero-operand terminator. Returns the values yielded by the
+/// outermost loop. If bound arrays are empty, the body builder will be called
+/// once to construct the IR outside of the loop with an empty list of induction
+/// variables.
+ValueVector buildLoopNest(
+    OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
+    ValueRange steps, ValueRange iterArgs,
+    function_ref<ValueVector(OpBuilder &, Location, ValueRange, ValueRange)>
+        bodyBuilder = nullptr);
+
+/// A convenience version for building loop nests without reductions. Does not
+/// take the initial value of reductions or expect the body building functions
+/// to return their current value.
+ValueVector buildLoopNest(OpBuilder &builder, Location loc, ValueRange lbs,
+                          ValueRange ubs, ValueRange steps,
+                          function_ref<void(OpBuilder &, Location, ValueRange)>
+                              bodyBuilder = nullptr);
+
 } // end namespace loop
 } // end namespace mlir
 #endif // MLIR_LOOPOPS_OPS_H_

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

let arguments = (ins Index:$lowerBound,
                        Variadic<AnyType>:$initArgs);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
               "Value lowerBound, Value upperBound, Value step, "
-              "ValueRange iterArgs = llvm::None">
+              "ValueRange iterArgs = llvm::None, "
+              "function_ref<void(OpBuilder &, Location, Value, ValueRange)>"
+              "    = nullptr">
   ];
 
   let extraClassDeclaration = [{
+    using BodyBuilderFn =
+        function_ref<void(OpBuilder &, Location, Value, ValueRange)>;
+
     Value getInductionVar() { return getBody()->getArgument(0); }
     Block::BlockArgListType getRegionIterArgs() {
       return getBody()->getArguments().drop_front();
     }
     Operation::operand_range getIterOperands() {
       return getOperands().drop_front(getNumControlOperands());
     }
 

mlir/lib/Conversion/VectorToLoops/ConvertVectorToLoops.cpp

LogicalResult VectorTransferRewriter<TransferReadOp>::matchAndRewrite(
   steps.reserve(vectorBoundsCapture.getSteps().size());
   for (auto step : vectorBoundsCapture.getSteps())
     steps.push_back(std_constant_index(step));
 
   // 2. Emit alloc-copy-load-dealloc.
   Value tmp = std_alloc(tmpMemRefType(transfer));
   StdIndexedValue local(tmp);
   Value vec = vector_type_cast(tmp);
-  SmallVector<Value, 8> ivs(lbs.size());
-  LoopNestBuilder(ivs, lbs, ubs, steps)([&] {
+  loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) {
+    auto ivs = llvm::to_vector<8>(loopIvs);
     // Swap the ivs which will reorder memory accesses.
     if (coalescedIdx >= 0)
       std::swap(ivs.back(), ivs[coalescedIdx]);
     // Computes clippedScalarAccessExprs in the loop nest scope (ivs exist).
     local(ivs) = remote(clip(transfer, memRefBoundsCapture, ivs));
   });
   Value vectorValue = std_load(vec);
   (std_dealloc(tmp)); // vexing parse
LogicalResult VectorTransferRewriter<TransferWriteOp>::matchAndRewrite(
   for (auto step : vectorBoundsCapture.getSteps())
     steps.push_back(std_constant_index(step));
 
   // 2. Emit alloc-store-copy-dealloc.
   Value tmp = std_alloc(tmpMemRefType(transfer));
   StdIndexedValue local(tmp);
   Value vec = vector_type_cast(tmp);
   std_store(vectorValue, vec);
-  SmallVector<Value, 8> ivs(lbs.size());
-  LoopNestBuilder(ivs, lbs, ubs, steps)([&] {
+  loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) {
+    auto ivs = llvm::to_vector<8>(loopIvs);
     // Swap the ivs which will reorder memory accesses.
     if (coalescedIdx >= 0)
       std::swap(ivs.back(), ivs[coalescedIdx]);
     // Computes clippedScalarAccessExprs in the loop nest scope (ivs exist).
     remote(clip(transfer, memRefBoundsCapture, ivs)) = local(ivs);
   });
   (std_dealloc(tmp)); // vexing parse...
 

mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp

for (unsigned i = 0; i < 3; ++i) {
     threadIds.push_back(
         builder.create<gpu::ThreadIdOp>(loc, indexType, dimName));
     blockDims.push_back(
         builder.create<gpu::BlockDimOp>(loc, indexType, dimName));
   }
 
   // Produce the loop nest with copies.
   SmallVector<Value, 8> ivs(lbs.size());
-  LoopNestBuilder(ivs, lbs, ubs, steps)([&]() {
+  loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) {
+    ivs.assign(loopIvs.begin(), loopIvs.end());
     auto activeIvs = llvm::makeArrayRef(ivs).take_back(rank);
     StdIndexedValue fromHandle(from), toHandle(to);
     toHandle(activeIvs) = fromHandle(activeIvs);
   });
+  ivs[0].getParentBlock()->dump();
 
   // Map the innermost loops to threads in reverse order.
   for (auto en :
        llvm::enumerate(llvm::reverse(llvm::makeArrayRef(ivs).take_back(
            GPUDialect::getNumWorkgroupDimensions())))) {
     Value v = en.value();
     auto loop = cast<loop::ForOp>(v.getParentRegion()->getParentOp());
     mapLoopToProcessorIds(loop, {threadIds[en.index()]},

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

 //===- Builders.cpp - MLIR Declarative Builder Classes --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/LoopOps/EDSC/Builders.h"
+#include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
 
 mlir::edsc::ParallelLoopNestBuilder::ParallelLoopNestBuilder(
     MutableArrayRef<Value> ivs, ArrayRef<Value> lbs, ArrayRef<Value> ubs,
mlir::edsc::makeLoopBuilder(Value *iv, Value lb, Value ub, Value step,
   for (size_t i = 0, e = iterArgsHandles.size(); i < e; ++i) {
     // Skipping the induction variable.
     iterArgsHandles[i] = body->getArgument(i + 1);
   }
   result.setOp(forOp);
   result.enter(body);
   return result;
 }
+
+mlir::loop::ValueVector
+mlir::edsc::loopNestBuilder(ValueRange lbs, ValueRange ubs, ValueRange steps,
+                            function_ref<void(ValueRange)> fun) {
+  // Delegates actual construction to loop::buildLoopNest by wrapping `fun` into
+  // the expected function interface.
+  assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
+  return mlir::loop::buildLoopNest(
+      ScopedContext::getBuilderRef(), ScopedContext::getLocation(), lbs, ubs,
+      steps, [&](OpBuilder &builder, Location loc, ValueRange ivs) {
+        ScopedContext context(builder, loc);
+        if (fun)
+          fun(ivs);
+      });
+}
+
+mlir::loop::ValueVector
+mlir::edsc::loopNestBuilder(Value lb, Value ub, Value step,
+                            function_ref<void(Value)> fun) {
+  // Delegates to the ValueRange-based version by wrapping the lambda.
+  auto wrapper = [&](ValueRange ivs) {
+    assert(ivs.size() == 1);
+    if (fun)
+      fun(ivs[0]);
+  };
+  return loopNestBuilder(ValueRange(lb), ValueRange(ub), ValueRange(step),
+                         wrapper);
+}
+
+mlir::loop::ValueVector mlir::edsc::loopNestBuilder(
+    Value lb, Value ub, Value step, ValueRange iterArgInitValues,
+    function_ref<loop::ValueVector(Value, ValueRange)> fun) {
+  // Delegates actual construction to loop::buildLoopNest by wrapping `fun` into
+  // the expected function interface.
+  assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
+  return mlir::loop::buildLoopNest(
+      ScopedContext::getBuilderRef(), ScopedContext::getLocation(), lb, ub,
+      step, iterArgInitValues,
+      [&](OpBuilder &builder, Location loc, ValueRange ivs, ValueRange args) {
+        assert(ivs.size() == 1 && "expected one induction variable");
+        ScopedContext context(builder, loc);
+        if (fun)
+          return fun(ivs[0], args);
+        return loop::ValueVector(iterArgInitValues.begin(),
+                                 iterArgInitValues.end());
+      });
+}

mlir/lib/Dialect/LoopOps/LoopOps.cpp

#include "mlir/Dialect/LoopOps/LoopOps.cpp.inc"
       >();
 }
 
 //===----------------------------------------------------------------------===//
 // ForOp
 //===----------------------------------------------------------------------===//
 
 void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
-                  Value ub, Value step, ValueRange iterArgs) {
+                  Value ub, Value step, ValueRange iterArgs,
+                  BodyBuilderFn bodyBuilder) {
   result.addOperands({lb, ub, step});
   result.addOperands(iterArgs);
   for (Value v : iterArgs)
     result.addTypes(v.getType());
   Region *bodyRegion = result.addRegion();
-  bodyRegion->push_back(new Block());
-  if (iterArgs.empty())
-    ForOp::ensureTerminator(*bodyRegion, builder, result.location);
-  bodyRegion->front().addArgument(builder.getIndexType());
+  bodyRegion->push_back(new Block);
+  Block &bodyBlock = bodyRegion->front();
+  bodyBlock.addArgument(builder.getIndexType());
   for (Value v : iterArgs)
-    bodyRegion->front().addArgument(v.getType());
+    bodyBlock.addArgument(v.getType());
+
+  if (iterArgs.empty() && !bodyBuilder) {
+    ForOp::ensureTerminator(*bodyRegion, builder, result.location);
+  } else if (bodyBuilder) {
+    OpBuilder::InsertionGuard guard(builder);
+    builder.setInsertionPointToStart(&bodyBlock);
+    bodyBuilder(builder, result.location, bodyBlock.getArgument(0),
+                bodyBlock.getArguments().drop_front());
+  }
 }
 
 static LogicalResult verify(ForOp op) {
   if (auto cst = dyn_cast_or_null<ConstantIndexOp>(op.step().getDefiningOp()))
     if (cst.getValue() <= 0)
       return op.emitOpError("constant step operand must be positive");
 
   // Check that the body defines as single block argument for the induction
void ForOp::getSuccessorRegions(Optional<unsigned> index,
   }
 
   // Otherwise, the loop may branch back to itself or the parent operation.
   assert(index.getValue() == 0 && "expected loop region");
   regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
   regions.push_back(RegionSuccessor(getResults()));
 }
 
+ValueVector mlir::loop::buildLoopNest(
+    OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
+    ValueRange steps, ValueRange iterArgs,
+    function_ref<ValueVector(OpBuilder &, Location, ValueRange, ValueRange)>
+        bodyBuilder) {
+  assert(lbs.size() == ubs.size() &&
+         "expected the same number of lower and upper bounds");
+  assert(lbs.size() == steps.size() &&
+         "expected the same number of lower bounds and steps");
+
+  // If there are no bounds, call the body-building function and return early.
+  if (lbs.empty()) {
+    ValueVector results =
+        bodyBuilder ? bodyBuilder(builder, loc, ValueRange(), iterArgs)
+                    : ValueVector();
+    assert(results.size() == iterArgs.size() &&
+           "loop nest body must return as many values as loop has iteration "
+           "arguments");
+    return results;
+  }
+
+  // First, create the loop structure iteratively using the body-builder
+  // callback of `ForOp::build`. Do not create `YieldOp`s yet.
+  OpBuilder::InsertionGuard guard(builder);
+  SmallVector<loop::ForOp, 4> loops;
+  SmallVector<Value, 4> ivs;
+  loops.reserve(lbs.size());
+  ivs.reserve(lbs.size());
+  ValueRange currentIterArgs = iterArgs;
+  Location currentLoc = loc;
+  for (unsigned i = 0, e = lbs.size(); i < e; ++i) {
+    auto loop = builder.create<loop::ForOp>(
+        currentLoc, lbs[i], ubs[i], steps[i], currentIterArgs,
+        [&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv,
+            ValueRange args) {
+          ivs.push_back(iv);
+          // It is safe to store ValueRange args because it points to block
+          // arguments of a loop operation that we also own.
+          currentIterArgs = args;
+          currentLoc = nestedLoc;
+        });
+    // Set the builder to point to the body of the newly created loop. We don't
+    // do this in the callback beacause the builder is reset when the callback
+    // returns.
+    builder.setInsertionPointToStart(loop.getBody());
+    loops.push_back(loop);
+  }
+
+  // For all loops but the innermost, yield the results of the nested loop.
+  for (unsigned i = 0, e = loops.size() - 1; i < e; ++i) {
+    builder.setInsertionPointToEnd(loops[i].getBody());
+    builder.create<loop::YieldOp>(loc, loops[i + 1].getResults());
+  }
+
+  // In the body of the innermost loop, call the body building function if any
+  // and yield its results.
+  builder.setInsertionPointToStart(loops.back().getBody());
+  ValueVector results = bodyBuilder
+                            ? bodyBuilder(builder, currentLoc, ivs,
+                                          loops.back().getRegionIterArgs())
+                            : ValueVector();
+  assert(results.size() == iterArgs.size() &&
+         "loop nest body must return as many values as loop has iteration "
+         "arguments");
+  builder.setInsertionPointToEnd(loops.back().getBody());
+  builder.create<loop::YieldOp>(loc, results);
+
+  // Return the results of the outermost loop.
+  return ValueVector(loops.front().result_begin(), loops.front().result_end());
+}
+
+ValueVector mlir::loop::buildLoopNest(
+    OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
+    ValueRange steps,
+    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder) {
+  // Delegate to the main function by wrapping the body builder.
+  return buildLoopNest(builder, loc, lbs, ubs, steps, llvm::None,
+                       [&bodyBuilder](OpBuilder &nestedBuilder,
+                                      Location nestedLoc, ValueRange ivs,
+                                      ValueRange) -> ValueVector {
+                         if (bodyBuilder)
+                           bodyBuilder(nestedBuilder, nestedLoc, ivs);
+                         return {};
+                       });
+}
+
 //===----------------------------------------------------------------------===//
 // IfOp
 //===----------------------------------------------------------------------===//
 
 void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,
                  bool withElseRegion) {
   build(builder, result, /*resultTypes=*/llvm::None, cond, withElseRegion);
 }

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

 
 TEST_FUNC(builder_loop_for) {
   auto indexType = IndexType::get(&globalContext());
   auto f = makeFunction("builder_loop_for", {},
                         {indexType, indexType, indexType, indexType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
-  Value i, a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
+  Value a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
       d(f.getArgument(3));
   using namespace edsc::op;
-  LoopNestBuilder(&i, a - b, c + d, a)();
+  loopNestBuilder(a - b, c + d, a);
 
   // clang-format off
   // CHECK-LABEL: func @builder_loop_for(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
   // CHECK-DAG:    [[r0:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 - s1)>()[%{{.*}}, %{{.*}}]
   // CHECK-DAG:    [[r1:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%{{.*}}, %{{.*}}]
   // CHECK-NEXT:   loop.for %{{.*}} = [[r0]] to [[r1]] step {{.*}} {
   // clang-format on
   f.print(llvm::outs());
TEST_FUNC(builder_loop_for_yield) {
   auto f32Type = FloatType::getF32(&globalContext());
   auto f = makeFunction("builder_loop_for_yield", {},
                         {indexType, indexType, indexType, indexType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   Value init0 = std_constant_float(llvm::APFloat(1.0f), f32Type);
   Value init1 = std_constant_float(llvm::APFloat(2.0f), f32Type);
-  Value i, a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
+  Value a(f.getArgument(0)), b(f.getArgument(1)), c(f.getArgument(2)),
       d(f.getArgument(3));
-  Value args01[2];
-  Value &arg0 = args01[0], &arg1 = args01[1];
   using namespace edsc::op;
-  auto results =
-      LoopNestBuilder(&i, a - b, c + d, a, args01, {init0, init1})([&] {
-        auto sum = arg0 + arg1;
-        loop_yield(ArrayRef<Value>{arg1, sum});
+  auto results = loopNestBuilder(a - b, c + d, a, {init0, init1},
+                                 [&](Value iv, ValueRange args) {
+                                   Value sum = args[0] + args[1];
+                                   return loop::ValueVector{args[1], sum};
                                  });
   results[0] + results[1];
 
   // clang-format off
   // CHECK-LABEL: func @builder_loop_for_yield(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
   // CHECK:     [[init0:%.*]] = constant
   // CHECK:     [[init1:%.*]] = constant
   // CHECK-DAG:    [[r0:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 - s1)>()[%{{.*}}, %{{.*}}]
   // CHECK-DAG:    [[r1:%[0-9]+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%{{.*}}, %{{.*}}]