[mlir] Introduce callback-based builders for AffineForOp

[ftynse]

Summary

Similarly to scf::ForOp, introduce additional function_ref arguments to
AffineForOp::build that can be used to populate the body of the loop during
its construction. Provide compatibility functions for constructing affine loop
nests using edsc::ScopedContext.

edsc::AffineLoopNestBuilder and reletad functionality is now deprecated and
will be removed soon, users are expected to switch to affineLoopNestBuilder
that provides similar functionality with a simpler OpBuilder-based
implementation.

Depends On D81753

[rriddle]

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

106

You could use m_Constant(IntegerAttr*) or m_ConstantInt(APInt*) instead here.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

1209–1215

nit: Can we rename this bodyBuilderFn? It was slightly confusing seeing bodyBuilder and builder next to each other.

[ftynse]

Address review

[ftynse]

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

106

I thought about that, but it turns out to be more verbose than what I currently have.

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

 namespace edsc {
 
 /// Constructs a new AffineForOp and captures the associated induction
 /// variable. A Value pointer is passed as the first argument and is the
 /// *only* way to capture the loop induction variable.
 LoopBuilder makeAffineLoopBuilder(Value *iv, ArrayRef<Value> lbs,
                                   ArrayRef<Value> ubs, int64_t step);
 
+/// Deprecated. Use affineLoopNestBuilder instead.
+///
 /// Explicit nested LoopBuilder. Offers a compressed multi-loop builder to avoid
 /// explicitly writing all the loops in a nest. This simple functionality is
 /// also useful to write rank-agnostic custom ops.
 ///
 /// Usage:
 ///
 /// ```c++
 ///    AffineLoopNestBuilder({&i, &j, &k}, {lb, lb, lb}, {ub, ub, ub}, {1, 1,
AffineLoopNestBuilder(MutableArrayRef<Value> ivs, ArrayRef<Value> lbs,
                         ArrayRef<Value> ubs, ArrayRef<int64_t> steps);
 
   void operator()(function_ref<void(void)> fun = nullptr);
 
 private:
   SmallVector<LoopBuilder, 4> loops;
 };
 
+/// Creates a perfect nest of affine "for" loops, given the list of lower
+/// bounds, upper bounds and steps. The three lists are expected to contain the
+/// same number of elements. Uses the OpBuilder and Location stored in
+/// ScopedContext and assumes they are non-null. The optional "bodyBuilderFn"
+/// callback is called to construct the body of the innermost loop and is passed
+/// the list of loop induction variables, in order from outermost to innermost.
+/// The function is expected to use the builder and location stored in
+/// ScopedContext at the moment of the call. The function should not create
+/// the affine terminator op, which will be added regardless of the
+/// "bodyBuilderFn" being present.
+void affineLoopNestBuilder(
+    ValueRange lbs, ValueRange ubs, ArrayRef<int64_t> steps,
+    function_ref<void(ValueRange)> bodyBuilderFn = nullptr);
+
+/// Creates a single affine "for" loop, iterating from max(lbs) to min(ubs) with
+/// the given step. Uses the OpBuilder and Location stored in ScopedContext and
+/// assumes they are non-null. The optional "bodyBuilderFn" callback is called
+/// to construct the body of the loop and is passed the induction variable. The
+/// function is expected to use the builder and location stored in ScopedContext
+/// at the moment of the call. The function should not create the affine
+/// terminator op, which will be added regardless of the "bodyBuilderFn" being
+/// present.
+void affineLoopBuilder(ValueRange lbs, ValueRange ubs, int64_t step,
+                       function_ref<void(Value)> bodyBuilderFn = nullptr);
+
 namespace op {
 
 Value operator+(Value lhs, Value rhs);
 Value operator-(Value lhs, Value rhs);
 Value operator*(Value lhs, Value rhs);
 Value operator/(Value lhs, Value rhs);
 Value operator%(Value lhs, Value rhs);
 Value floorDiv(Value lhs, Value rhs);

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

let description = [{
     ```
   }];
   let arguments = (ins Variadic<AnyType>);
   let regions = (region SizedRegion<1>:$region);
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
-              "int64_t lowerBound, int64_t upperBound, int64_t step = 1">,
+              "int64_t lowerBound, int64_t upperBound, int64_t step = 1, "
+              "function_ref<void(OpBuilder &, Location, Value)> bodyBuilder "
+              "    = nullptr">,
     OpBuilder<"OpBuilder &builder, OperationState &result, "
               "ValueRange lbOperands, AffineMap lbMap, "
               "ValueRange ubOperands, AffineMap ubMap, "
-              "int64_t step = 1">
+              "int64_t step = 1, "
+              "function_ref<void(OpBuilder &, Location, Value)> bodyBuilder "
+              "    = nullptr">
   ];
 
   let extraClassDeclaration = [{
     static StringRef getStepAttrName() { return "step"; }
     static StringRef getLowerBoundAttrName() { return "lower_bound"; }
     static StringRef getUpperBoundAttrName() { return "upper_bound"; }
 
     Value getInductionVar() { return getBody()->getArgument(0); }

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

for (unsigned idx = 0; idx < numUnrolledInstances; ++idx) {
                       indices.drop_front(leadingRank).take_front(majorRank),
                       indices.take_back(minorRank), memrefBoundsCapture);
     }
   } else {
     VectorBoundsCapture vectorBoundsCapture(majorVectorType);
     auto majorLbs = vectorBoundsCapture.getLbs();
     auto majorUbs = vectorBoundsCapture.getUbs();
     auto majorSteps = vectorBoundsCapture.getSteps();
-    SmallVector<Value, 8> majorIvs(vectorBoundsCapture.rank());
-    AffineLoopNestBuilder(majorIvs, majorLbs, majorUbs, majorSteps)([&] {
+    affineLoopNestBuilder(
+        majorLbs, majorUbs, majorSteps, [&](ValueRange majorIvs) {
           ValueRange indices(xferOp.indices());
           loopBodyBuilder(majorIvs, indices.take_front(leadingRank),
                           indices.drop_front(leadingRank).take_front(majorRank),
                           indices.take_back(minorRank), memrefBoundsCapture);
         });
   }
 }
 
 template <typename ConcreteOp>
 Value NDTransferOpHelper<ConcreteOp>::emitInBoundsCondition(
     ValueRange majorIvs, ValueRange majorOffsets,
     MemRefBoundsCapture &memrefBounds,
     SmallVectorImpl<Value> &majorIvsPlusOffsets) {

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

 
 mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(Value *iv,
                                                          ArrayRef<Value> lbs,
                                                          ArrayRef<Value> ubs,
                                                          int64_t step) {
   loops.emplace_back(makeAffineLoopBuilder(iv, lbs, ubs, step));
 }
 
+void mlir::edsc::affineLoopNestBuilder(
+    ValueRange lbs, ValueRange ubs, ArrayRef<int64_t> steps,
+    function_ref<void(ValueRange)> bodyBuilderFn) {
+  assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
+  assert(lbs.size() == ubs.size() && "Mismatch in number of arguments");
+  assert(lbs.size() == steps.size() && "Mismatch in number of arguments");
+
+  // If there are no loops to be constructed, construct the body anyway.
+  if (lbs.empty()) {
+    if (bodyBuilderFn)
+      bodyBuilderFn(ValueRange());
+    return;
+  }
+
+  // Fetch the builder and location.
+  OpBuilder &builder = ScopedContext::getBuilderRef();
+  OpBuilder::InsertionGuard guard(builder);
+  Location loc = ScopedContext::getLocation();
+  AffineMap identity = builder.getDimIdentityMap();
+
+  // Create the loops iteratively and store the induction variables.
+  SmallVector<Value, 4> ivs;
+  ivs.reserve(lbs.size());
+  for (unsigned i = 0, e = lbs.size(); i < e; ++i) {
+    // Callback for creating the loop body, always creates the terminator.
+    auto loopBody = [&](OpBuilder &nestedBuilder, Location nestedLoc,
+                        Value iv) {
+      ivs.push_back(iv);
+      // In the innermost loop, call the body builder.
+      if (i == e - 1 && bodyBuilderFn) {
+        ScopedContext nestedContext(nestedBuilder, loc);
+        OpBuilder::InsertionGuard nestedGuard(nestedBuilder);
+        bodyBuilderFn(ivs);
+      }
+      nestedBuilder.create<AffineTerminatorOp>(nestedLoc);
+    };
+
+    // Create the loop. If the bounds are known to be constants, use the
+    // constant form of the loop.
+    auto lbConst = lbs[i].getDefiningOp<ConstantIndexOp>();
+    auto ubConst = ubs[i].getDefiningOp<ConstantIndexOp>();
+    auto loop = lbConst && ubConst
+                    ? builder.create<AffineForOp>(loc, lbConst.getValue(),
+                                                  ubConst.getValue(), steps[i],
+                                                  loopBody)
+                    : builder.create<AffineForOp>(loc, lbs[i], identity, ubs[i],
+                                                  identity, steps[i], loopBody);
+    builder.setInsertionPointToStart(loop.getBody());
+  }
+}
+
+void mlir::edsc::affineLoopBuilder(ValueRange lbs, ValueRange ubs, int64_t step,
+                                   function_ref<void(Value)> bodyBuilderFn) {
+  // Fetch the builder and location.
+  assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
+  OpBuilder &builder = ScopedContext::getBuilderRef();
+  Location loc = ScopedContext::getLocation();
+
+  // Create the actual loop and call the body builder, if provided, after
+  // updating the scoped context.
+  builder.create<AffineForOp>(
+      loc, lbs, builder.getMultiDimIdentityMap(lbs.size()), ubs,
+      builder.getMultiDimIdentityMap(ubs.size()), step,
+      [&](OpBuilder &nestedBuilder, Location nestedLoc, Value iv) {
+        if (bodyBuilderFn) {
+          ScopedContext nestedContext(nestedBuilder, nestedLoc);
+          OpBuilder::InsertionGuard guard(nestedBuilder);
+          bodyBuilderFn(iv);
+        }
+        nestedBuilder.create<AffineTerminatorOp>(nestedLoc);
+      });
+}
+
 mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(
     MutableArrayRef<Value> ivs, ArrayRef<Value> lbs, ArrayRef<Value> ubs,
     ArrayRef<int64_t> steps) {
   assert(ivs.size() == lbs.size() && "Mismatch in number of arguments");
   assert(ivs.size() == ubs.size() && "Mismatch in number of arguments");
   assert(ivs.size() == steps.size() && "Mismatch in number of arguments");
   for (auto it : llvm::zip(ivs, lbs, ubs, steps))
     loops.emplace_back(makeAffineLoopBuilder(&std::get<0>(it), std::get<1>(it),

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

LogicalResult AffineDmaWaitOp::fold(ArrayRef<Attribute> cstOperands,
   /// dma_wait(memrefcast) -> dma_wait
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // AffineForOp
 //===----------------------------------------------------------------------===//
 
-void AffineForOp::build(OpBuilder &builder, OperationState &result,
-                        ValueRange lbOperands, AffineMap lbMap,
-                        ValueRange ubOperands, AffineMap ubMap, int64_t step) {
+void AffineForOp::build(
+    OpBuilder &builder, OperationState &result, ValueRange lbOperands,
+    AffineMap lbMap, ValueRange ubOperands, AffineMap ubMap, int64_t step,
+    function_ref<void(OpBuilder &, Location, Value)> bodyBuilder) {
   assert(((!lbMap && lbOperands.empty()) ||
           lbOperands.size() == lbMap.getNumInputs()) &&
          "lower bound operand count does not match the affine map");
   assert(((!ubMap && ubOperands.empty()) ||
           ubOperands.size() == ubMap.getNumInputs()) &&
          "upper bound operand count does not match the affine map");
   assert(step > 0 && "step has to be a positive integer constant");
 
   // Add an attribute for the step.
   result.addAttribute(getStepAttrName(),
                       builder.getIntegerAttr(builder.getIndexType(), step));
 
   // Add the lower bound.
   result.addAttribute(getLowerBoundAttrName(), AffineMapAttr::get(lbMap));
   result.addOperands(lbOperands);
 
   // Add the upper bound.
   result.addAttribute(getUpperBoundAttrName(), AffineMapAttr::get(ubMap));
   result.addOperands(ubOperands);
 
   // Create a region and a block for the body.  The argument of the region is
   // the loop induction variable.
   Region *bodyRegion = result.addRegion();
-  Block *body = new Block();
-  body->addArgument(IndexType::get(builder.getContext()));
+  Block *body = new Block;
+  Value inductionVar = body->addArgument(IndexType::get(builder.getContext()));
   bodyRegion->push_back(body);
+  if (bodyBuilder) {
+    OpBuilder::InsertionGuard guard(builder);
+    builder.setInsertionPointToStart(body);
+    bodyBuilder(builder, result.location, inductionVar);
+  } else {
     ensureTerminator(*bodyRegion, builder, result.location);
   }
+}
 
-void AffineForOp::build(OpBuilder &builder, OperationState &result, int64_t lb,
-                        int64_t ub, int64_t step) {
+void AffineForOp::build(
+    OpBuilder &builder, OperationState &result, int64_t lb, int64_t ub,
+    int64_t step,
+    function_ref<void(OpBuilder &, Location, Value)> bodyBuilder) {
   auto lbMap = AffineMap::getConstantMap(lb, builder.getContext());
   auto ubMap = AffineMap::getConstantMap(ub, builder.getContext());
-  return build(builder, result, {}, lbMap, {}, ubMap, step);
+  return build(builder, result, {}, lbMap, {}, ubMap, step, bodyBuilder);
 }
 
 static LogicalResult verify(AffineForOp op) {
   // Check that the body defines as single block argument for the induction
   // variable.
   auto *body = op.getBody();
   if (body->getNumArguments() != 1 || !body->getArgument(0).getType().isIndex())
     return op.emitOpError(

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

 TEST_FUNC(builder_dynamic_for_func_args) {
   auto indexType = IndexType::get(&globalContext());
   auto f32Type = FloatType::getF32(&globalContext());
   auto f =
       makeFunction("builder_dynamic_for_func_args", {}, {indexType, indexType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
-  Value i, j, lb(f.getArgument(0)), ub(f.getArgument(1));
+  Value lb(f.getArgument(0)), ub(f.getArgument(1));
   Value f7(std_constant_float(llvm::APFloat(7.0f), f32Type));
   Value f13(std_constant_float(llvm::APFloat(13.0f), f32Type));
   Value i7(std_constant_int(7, 32));
   Value i13(std_constant_int(13, 32));
-  AffineLoopNestBuilder(&i, lb, ub, 3)([&] {
+  affineLoopBuilder(lb, ub, 3, [&](Value i) {
     using namespace edsc::op;
     lb *std_constant_index(3) + ub;
     lb + std_constant_index(3);
-    AffineLoopNestBuilder(&j, lb, ub, 2)([&] {
+    affineLoopBuilder(lb, ub, 2, [&](Value j) {
       ceilDiv(std_constant_index(31) * floorDiv(i + j * std_constant_index(3),
                                                 std_constant_index(32)),
               std_constant_index(32));
       ((f7 + f13) / f7) % f13 - f7 *f13;
       ((i7 + i13) / i7) % i13 - i7 *i13;
     });
   });
 
TEST_FUNC(builder_dynamic_for) {
   auto f = makeFunction("builder_dynamic_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)),
       d(f.getArgument(3));
   using namespace edsc::op;
-  AffineLoopNestBuilder(&i, a - b, c + d, 2)();
+  affineLoopBuilder(a - b, c + d, 2);
 
   // clang-format off
   // CHECK-LABEL: func @builder_dynamic_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:   affine.for %{{.*}} = affine_map<(d0) -> (d0)>([[r0]]) to affine_map<(d0) -> (d0)>([[r1]]) step 2 {
   // clang-format on
   f.print(llvm::outs());

 
 TEST_FUNC(builder_max_min_for) {
   auto indexType = IndexType::get(&globalContext());
   auto f = makeFunction("builder_max_min_for", {},
                         {indexType, indexType, indexType, indexType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
-  Value i, lb1(f.getArgument(0)), lb2(f.getArgument(1)), ub1(f.getArgument(2)),
+  Value lb1(f.getArgument(0)), lb2(f.getArgument(1)), ub1(f.getArgument(2)),
       ub2(f.getArgument(3));
-  AffineLoopNestBuilder(&i, {lb1, lb2}, {ub1, ub2}, 1)();
+  affineLoopBuilder({lb1, lb2}, {ub1, ub2}, 1);
   std_ret();
 
   // clang-format off
   // CHECK-LABEL: func @builder_max_min_for(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
   // CHECK:  affine.for %{{.*}} = max affine_map<(d0, d1) -> (d0, d1)>(%{{.*}}, %{{.*}}) to min affine_map<(d0, d1) -> (d0, d1)>(%{{.*}}, %{{.*}}) {
   // CHECK:  return
   // clang-format on
   f.print(llvm::outs());
TEST_FUNC(builder_helpers) {
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   // clang-format off
   Value f7 = std_constant_float(llvm::APFloat(7.0f), f32Type);
   MemRefBoundsCapture vA(f.getArgument(0)), vB(f.getArgument(1)),
       vC(f.getArgument(2));
   AffineIndexedValue A(f.getArgument(0)), B(f.getArgument(1)), C(f.getArgument(2));
-  Value ivs[2];
-  Value &i = ivs[0], &j = ivs[1];
-  Value k1, k2, lb0, lb1, lb2, ub0, ub1, ub2;
+  Value lb0, lb1, lb2, ub0, ub1, ub2;
   int64_t step0, step1, step2;
   std::tie(lb0, ub0, step0) = vA.range(0);
   std::tie(lb1, ub1, step1) = vA.range(1);
   lb2 = vA.lb(2);
   ub2 = vA.ub(2);
   step2 = vA.step(2);
-  AffineLoopNestBuilder(ivs, {lb0, lb1}, {ub0, ub1}, {step0, step1})([&]{
-    AffineLoopNestBuilder(&k1, lb2, ub2, step2)([&]{
+  affineLoopNestBuilder({lb0, lb1}, {ub0, ub1}, {step0, step1}, [&](ValueRange ivs) {
+    Value i = ivs[0];
+    Value j = ivs[1];
+    affineLoopBuilder(lb2, ub2, step2, [&](Value k1){
       C(i, j, k1) = f7 + A(i, j, k1) + B(i, j, k1);
     });
-    AffineLoopNestBuilder(&k2, lb2, ub2, step2)([&]{
+    affineLoopBuilder(lb2, ub2, step2, [&](Value k2){
       C(i, j, k2) += A(i, j, k2) + B(i, j, k2);
     });
   });
 
   // CHECK-LABEL: @builder_helpers
   //      CHECK:   affine.for %{{.*}} = affine_map<(d0) -> (d0)>({{.*}}) to affine_map<(d0) -> (d0)>({{.*}}) {
   // CHECK-NEXT:     affine.for %{{.*}} = affine_map<(d0) -> (d0)>({{.*}}) to affine_map<(d0) -> (d0)>({{.*}}) {
   // CHECK-NEXT:       affine.for %{{.*}} = affine_map<(d0) -> (d0)>({{.*}}) to affine_map<(d0) -> (d0)>({{.*}}) {
TEST_FUNC(select_op_i32) {
   using namespace edsc::op;
   auto f32Type = FloatType::getF32(&globalContext());
   auto memrefType = MemRefType::get(
       {ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
   auto f = makeFunction("select_op", {}, {memrefType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
-  // clang-format off
   Value zero = std_constant_index(0), one = std_constant_index(1);
   MemRefBoundsCapture vA(f.getArgument(0));
   AffineIndexedValue A(f.getArgument(0));
-  Value ivs[2];
-  Value &i = ivs[0], &j = ivs[1];
-  AffineLoopNestBuilder(ivs, {zero, zero}, {one, one}, {1, 1})([&]{
-    std_select(eq(i, zero), A(zero, zero), A(i, j));
+  affineLoopNestBuilder({zero, zero}, {one, one}, {1, 1}, [&](ValueRange ivs) {
+    std_select(eq(ivs[0], zero), A(zero, zero), A(ivs[0], ivs[1]));
   });
 
+  // clang-format off
   // CHECK-LABEL: @select_op
   //      CHECK: affine.for %{{.*}} = 0 to 1 {
   // CHECK-NEXT:   affine.for %{{.*}} = 0 to 1 {
   //  CHECK-DAG:     {{.*}} = cmpi "eq"
   //  CHECK-DAG:     {{.*}} = affine.load
   //  CHECK-DAG:     {{.*}} = affine.load
   // CHECK-NEXT:     {{.*}} = select
   // clang-format on
   f.print(llvm::outs());
   f.erase();
 }
 
 TEST_FUNC(select_op_f32) {
   auto f32Type = FloatType::getF32(&globalContext());
   auto memrefType = MemRefType::get(
       {ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
   auto f = makeFunction("select_op", {}, {memrefType, memrefType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   // clang-format off
   Value zero = std_constant_index(0), one = std_constant_index(1);
   MemRefBoundsCapture vA(f.getArgument(0)), vB(f.getArgument(1));
   AffineIndexedValue A(f.getArgument(0)), B(f.getArgument(1));
-  Value ivs[2];
-  Value &i = ivs[0], &j = ivs[1];
-  AffineLoopNestBuilder(ivs, {zero, zero}, {one, one}, {1, 1})([&]{
+  affineLoopNestBuilder({zero, zero}, {one, one}, {1, 1}, [&](ValueRange ivs) {
     using namespace edsc::op;
+    Value i = ivs[0], j = ivs[1];
     std_select(eq(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
     std_select(ne(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
     std_select(B(i, j) >= B(i + one, j), A(zero, zero), A(i, j));
     std_select(B(i, j) <= B(i + one, j), A(zero, zero), A(i, j));
     std_select(B(i, j) < B(i + one, j), A(zero, zero), A(i, j));
     std_select(B(i, j) > B(i + one, j), A(zero, zero), A(i, j));
   });
 
TEST_FUNC(tile_2d) {
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   Value zero = std_constant_index(0);
   MemRefBoundsCapture vA(f.getArgument(0)), vB(f.getArgument(1)),
       vC(f.getArgument(2));
   AffineIndexedValue A(f.getArgument(0)), B(f.getArgument(1)),
       C(f.getArgument(2));
-  Value ivs[2];
-  Value &i = ivs[0], &j = ivs[1];
-  Value k1, k2;
+  Value i, j, k1, k2;
   Value M(vC.ub(0)), N(vC.ub(1)), O(vC.ub(2));
 
   // clang-format off
   using namespace edsc::op;
-  AffineLoopNestBuilder(ivs, {zero, zero}, {M, N}, {1, 1})([&]{
-    AffineLoopNestBuilder(&k1, zero, O, 1)([&]{
+  affineLoopNestBuilder({zero, zero}, {M, N}, {1, 1}, [&](ValueRange ivs) {
+    i = ivs[0];
+    j = ivs[1];
+    affineLoopBuilder(zero, O, 1, [&](Value k) {
+      k1 = k;
       C(i, j, k1) = A(i, j, k1) + B(i, j, k1);
     });
-    AffineLoopNestBuilder(&k2, zero, O, 1)([&]{
+    affineLoopBuilder(zero, O, 1, [&](Value k) {
+      k2 = k;
       C(i, j, k2) = A(i, j, k2) + B(i, j, k2);
     });
   });
   // clang-format on
 
   auto li = getForInductionVarOwner(i), lj = getForInductionVarOwner(j),
        lk1 = getForInductionVarOwner(k1), lk2 = getForInductionVarOwner(k2);
   auto indicesL1 = mlir::tile({li, lj}, {512, 1024}, {lk1, lk2});
auto f = makeFunction("indirect_access", {},
                         {memrefType, memrefType, memrefType, memrefType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   Value zero = std_constant_index(0);
   MemRefBoundsCapture vC(f.getArgument(2));
   AffineIndexedValue B(f.getArgument(1)), D(f.getArgument(3));
   StdIndexedValue A(f.getArgument(0)), C(f.getArgument(2));
-  Value i, N(vC.ub(0));
+  Value N(vC.ub(0));
 
   // clang-format off
-  AffineLoopNestBuilder(&i, zero, N, 1)([&]{
+  affineLoopBuilder(zero, N, 1, [&](Value i) {
       C((Value)D(i)) = A((Value)B(i));
   });
   // clang-format on
 
   // clang-format off
   // CHECK-LABEL: func @indirect_access
   // CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: memref<?xf32>, %[[ARG2:.*]]: memref<?xf32>, %[[ARG3:.*]]: memref<?xf32>)
   // CHECK-DAG:  [[B:%.*]] = affine.load %[[ARG1]]
TEST_FUNC(empty_map_load_store) {
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   Value zero = std_constant_index(0);
   Value one = std_constant_index(1);
   AffineIndexedValue input(f.getArgument(0)), res(f.getArgument(1));
 
   // clang-format off
-  Value iv;
-  AffineLoopNestBuilder(&iv, zero, one, 1)([&]{
+  affineLoopBuilder(zero, one, 1, [&](Value) {
       res() = input();
   });
   // clang-format on
 
   // clang-format off
   // CHECK-LABEL: func @empty_map_load_store(
   // CHECK:  [[A:%.*]] = affine.load %{{.*}}[]
   // CHECK:  affine.store [[A]], %{{.*}}[]