[mlir] Refactor affine loop nest builders

[ftynse]

Summary

Existing implementation of affine loop nest builders relies on EDSC
ScopedContext, which is not used pervasively. Provide a common OpBuilder-based
helper function to construct a perfect nest of affine loops with the body of
the innermost loop populated by a callback. Use this function to implement the
EDSC version.

Affine "for" loops differ from SCF "for" loops by (1) not allowing to yield
values and (2) supporting short-hand form for constant bounds, which justifies
a separate implementation of the loop nest builder for the same of simplicity.

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

 /// not an induction variable, then return nullptr.
 AffineForOp getForInductionVarOwner(Value val);
 
 /// Extracts the induction variables from a list of AffineForOps and places them
 /// in the output argument `ivs`.
 void extractForInductionVars(ArrayRef<AffineForOp> forInsts,
                              SmallVectorImpl<Value> *ivs);
 
+/// Builds a perfect nest of affine "for" loops, i.e. each loop except the
+/// innermost only contains another loop and a terminator. The loops iterate
+/// from "lbs" to "ubs" with "steps". The body of the innermost loop is
+/// populated by calling "bodyBuilderFn" and providing it with an OpBuilder, a
+/// Location and a list of loop induction variables.
+void buildAffineLoopNest(OpBuilder &builder, Location loc,
+                         ArrayRef<int64_t> lbs, ArrayRef<int64_t> ubs,
+                         ArrayRef<int64_t> steps,
+                         function_ref<void(OpBuilder &, Location, ValueRange)>
+                             bodyBuilderFn = nullptr);
+void buildAffineLoopNest(OpBuilder &builder, Location loc, ValueRange lbs,
+                         ValueRange ubs, ArrayRef<int64_t> steps,
+                         function_ref<void(OpBuilder &, Location, ValueRange)>
+                             bodyBuilderFn = nullptr);
+
 /// AffineBound represents a lower or upper bound in the for operation.
 /// This class does not own the underlying operands. Instead, it refers
 /// to the operands stored in the AffineForOp. Its life span should not exceed
 /// that of the for operation it refers to.
 class AffineBound {
 public:
   AffineForOp getAffineForOp() { return op; }
   AffineMap getMap() { return map; }

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

mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(Value *iv,
                                                          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;
-  }
+  // Wrap the body builder function into an interface compatible with the main
+  // builder.
+  auto wrappedBuilderFn = [&](OpBuilder &nestedBuilder, Location nestedLoc,
+                              ValueRange ivs) {
+    ScopedContext context(nestedBuilder, nestedLoc);
+    bodyBuilderFn(ivs);
+  };
+  function_ref<void(OpBuilder &, Location, ValueRange)> wrapper;
+    wrapper = wrappedBuilderFn;
 
-  // Fetch the builder and location.
+  // Extract the builder, location and construct the loop nest.
   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());
-  }
+  buildAffineLoopNest(builder, loc, lbs, ubs, steps, wrapper);
 }
 
 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();

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

 /// them.
 void mlir::extractForInductionVars(ArrayRef<AffineForOp> forInsts,
                                    SmallVectorImpl<Value> *ivs) {
   ivs->reserve(forInsts.size());
   for (auto forInst : forInsts)
     ivs->push_back(forInst.getInductionVar());
 }
 
+/// Builds an affine loop nest, using "loopCreatorFn" to create individual loop
+/// operations.
+template <typename BoundListTy, typename LoopCreatorTy>
+void buildAffineLoopNestImpl(
+    OpBuilder &builder, Location loc, BoundListTy lbs, BoundListTy ubs,
+    ArrayRef<int64_t> steps,
+    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn,
+    LoopCreatorTy &&loopCreatorFn) {
+  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.
+  OpBuilder::InsertionGuard guard(builder);
+  if (lbs.empty()) {
+    if (bodyBuilderFn)
+      bodyBuilderFn(builder, loc, ValueRange());
+    return;
+  }
+
+  // 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) {
+        OpBuilder::InsertionGuard nestedGuard(nestedBuilder);
+        bodyBuilderFn(nestedBuilder, nestedLoc, ivs);
+      }
+      nestedBuilder.create<AffineTerminatorOp>(nestedLoc);
+    };
+
+    // Delegate actual loop creation to the callback in order to dispatch
+    // between constant- and variable-bound loops.
+    auto loop = loopCreatorFn(builder, loc, lbs[i], ubs[i], steps[i], loopBody);
+    builder.setInsertionPointToStart(loop.getBody());
+  }
+}
+
+/// Creates an affine loop from the bounds known to be constants.
+static AffineForOp buildAffineLoopFromConstants(
+    OpBuilder &builder, Location loc, int64_t lb, int64_t ub, int64_t step,
+    function_ref<void(OpBuilder &, Location, Value)> bodyBuilderFn) {
+  return builder.create<AffineForOp>(loc, lb, ub, step, bodyBuilderFn);
+}
+
+/// Creates an affine loop from the bounds that may or may not be constants.
+static AffineForOp buildAffineLoopFromValues(
+    OpBuilder &builder, Location loc, Value lb, Value ub, int64_t step,
+    function_ref<void(OpBuilder &, Location, Value)> bodyBuilderFn) {
+  auto lbConst = lb.getDefiningOp<ConstantIndexOp>();
+  auto ubConst = ub.getDefiningOp<ConstantIndexOp>();
+  if (lbConst && ubConst)
+    return buildAffineLoopFromConstants(builder, loc, lbConst.getValue(),
+                                        ubConst.getValue(), step,
+                                        bodyBuilderFn);
+  return builder.create<AffineForOp>(loc, lb, builder.getDimIdentityMap(), ub,
+                                     builder.getDimIdentityMap(), step,
+                                     bodyBuilderFn);
+}
+
+void mlir::buildAffineLoopNest(
+    OpBuilder &builder, Location loc, ArrayRef<int64_t> lbs,
+    ArrayRef<int64_t> ubs, ArrayRef<int64_t> steps,
+    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
+  buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn,
+                          buildAffineLoopFromConstants);
+}
+
+void mlir::buildAffineLoopNest(
+    OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
+    ArrayRef<int64_t> steps,
+    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
+  buildAffineLoopNestImpl(builder, loc, lbs, ubs, steps, bodyBuilderFn,
+                          buildAffineLoopFromValues);
+}
+
 //===----------------------------------------------------------------------===//
 // AffineIfOp
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// Remove else blocks that have nothing other than the terminator.
 struct SimplifyDeadElse : public OpRewritePattern<AffineIfOp> {
   using OpRewritePattern<AffineIfOp>::OpRewritePattern;