[mlir] NFC Async: always use 'b' for the current builder

[ezhulenev]

Summary

Currently some of the nested IR building inconsistently uses nb and b, it's very easy to call wrong builder outside of the current scope, so for simplicity all builders are always called b, and in nested IR building regions they just shadow the "parent" builder.

mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp

static ParallelComputeFunction createParallelComputeFunction(
   SmallVector<Value> isBlockFirstCoord(op.getNumLoops());
   SmallVector<Value> isBlockLastCoord(op.getNumLoops());
 
   // Builds inner loop nest inside async.execute operation that does all the
   // work concurrently.
   LoopNestBuilder workLoopBuilder = [&](size_t loopIdx) -> LoopBodyBuilder {
     return [&, loopIdx](OpBuilder &nestedBuilder, Location loc, Value iv,
                         ValueRange args) {
-      ImplicitLocOpBuilder nb(loc, nestedBuilder);
+      ImplicitLocOpBuilder b(loc, nestedBuilder);
 
       // Compute induction variable for `loopIdx`.
-      computeBlockInductionVars[loopIdx] = nb.create<arith::AddIOp>(
-          lowerBounds[loopIdx], nb.create<arith::MulIOp>(iv, steps[loopIdx]));
+      computeBlockInductionVars[loopIdx] = b.create<arith::AddIOp>(
+          lowerBounds[loopIdx], b.create<arith::MulIOp>(iv, steps[loopIdx]));
 
       // Check if we are inside first or last iteration of the loop.
-      isBlockFirstCoord[loopIdx] = nb.create<arith::CmpIOp>(
+      isBlockFirstCoord[loopIdx] = b.create<arith::CmpIOp>(
           arith::CmpIPredicate::eq, iv, blockFirstCoord[loopIdx]);
-      isBlockLastCoord[loopIdx] = nb.create<arith::CmpIOp>(
+      isBlockLastCoord[loopIdx] = b.create<arith::CmpIOp>(
           arith::CmpIPredicate::eq, iv, blockLastCoord[loopIdx]);
 
       // Check if the previous loop is in its first or last iteration.
       if (loopIdx > 0) {
-        isBlockFirstCoord[loopIdx] = nb.create<arith::AndIOp>(
+        isBlockFirstCoord[loopIdx] = b.create<arith::AndIOp>(
             isBlockFirstCoord[loopIdx], isBlockFirstCoord[loopIdx - 1]);
-        isBlockLastCoord[loopIdx] = nb.create<arith::AndIOp>(
+        isBlockLastCoord[loopIdx] = b.create<arith::AndIOp>(
             isBlockLastCoord[loopIdx], isBlockLastCoord[loopIdx - 1]);
       }
 
       // Keep building loop nest.
       if (loopIdx < op.getNumLoops() - 1) {
         if (loopIdx + 1 >= op.getNumLoops() - numBlockAlignedInnerLoops) {
           // For block aligned loops we always iterate starting from 0 up to
           // the loop trip counts.
-          nb.create<scf::ForOp>(c0, tripCounts[loopIdx + 1], c1, ValueRange(),
-                                workLoopBuilder(loopIdx + 1));
+          b.create<scf::ForOp>(c0, tripCounts[loopIdx + 1], c1, ValueRange(),
 
         } else {
           // Select nested loop lower/upper bounds depending on our position in
           // the multi-dimensional iteration space.
-          auto lb = nb.create<arith::SelectOp>(
-              isBlockFirstCoord[loopIdx], blockFirstCoord[loopIdx + 1], c0);
+          auto lb = b.create<arith::SelectOp>(isBlockFirstCoord[loopIdx],
+                                              blockFirstCoord[loopIdx + 1], c0);
 
-          auto ub = nb.create<arith::SelectOp>(isBlockLastCoord[loopIdx],
-                                               blockEndCoord[loopIdx + 1],
-                                               tripCounts[loopIdx + 1]);
+          auto ub = b.create<arith::SelectOp>(isBlockLastCoord[loopIdx],
 
-          nb.create<scf::ForOp>(lb, ub, c1, ValueRange(),
-                                workLoopBuilder(loopIdx + 1));
+          b.create<scf::ForOp>(lb, ub, c1, ValueRange(),
         }
 
-        nb.create<scf::YieldOp>(loc);
+        b.create<scf::YieldOp>(loc);
         return;
       }
 
       // Copy the body of the parallel op into the inner-most loop.
       BlockAndValueMapping mapping;
       mapping.map(op.getInductionVars(), computeBlockInductionVars);
       mapping.map(computeFuncType.captures, captures);
 
       for (auto &bodyOp : op.getLoopBody().getOps())
-        nb.clone(bodyOp, mapping);
+        b.clone(bodyOp, mapping);
     };
   };
 
   b.create<scf::ForOp>(blockFirstCoord[0], blockEndCoord[0], c1, ValueRange(),
                        workLoopBuilder(0));
   b.create<ReturnOp>(ValueRange());
 
   return {op.getNumLoops(), func, std::move(computeFuncType.captures)};
static void doAsyncDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
 
   // Check if the block size is one, in this case we can skip the async dispatch
   // completely. If this will be known statically, then canonicalization will
   // erase async group operations.
   Value isSingleBlock =
       b.create<arith::CmpIOp>(arith::CmpIPredicate::eq, blockCount, c1);
 
   auto syncDispatch = [&](OpBuilder &nestedBuilder, Location loc) {
-    ImplicitLocOpBuilder nb(loc, nestedBuilder);
+    ImplicitLocOpBuilder b(loc, nestedBuilder);
 
     // Call parallel compute function for the single block.
     SmallVector<Value> operands = {c0, blockSize};
     appendBlockComputeOperands(operands);
 
-    nb.create<CallOp>(parallelComputeFunction.func.sym_name(),
-                      parallelComputeFunction.func.getCallableResults(),
-                      operands);
-    nb.create<scf::YieldOp>();
+    b.create<CallOp>(parallelComputeFunction.func.sym_name(),
+    b.create<scf::YieldOp>();
   };
 
   auto asyncDispatch = [&](OpBuilder &nestedBuilder, Location loc) {
-    ImplicitLocOpBuilder nb(loc, nestedBuilder);
+    ImplicitLocOpBuilder b(loc, nestedBuilder);
 
     // Create an async.group to wait on all async tokens from the concurrent
     // execution of multiple parallel compute function. First block will be
     // executed synchronously in the caller thread.
-    Value groupSize = nb.create<arith::SubIOp>(blockCount, c1);
-    Value group = nb.create<CreateGroupOp>(GroupType::get(ctx), groupSize);
+    Value groupSize = b.create<arith::SubIOp>(blockCount, c1);
+    Value group = b.create<CreateGroupOp>(GroupType::get(ctx), groupSize);
 
     // Launch async dispatch function for [0, blockCount) range.
     SmallVector<Value> operands = {group, c0, blockCount, blockSize};
     appendBlockComputeOperands(operands);
 
-    nb.create<CallOp>(asyncDispatchFunction.sym_name(),
-                      asyncDispatchFunction.getCallableResults(), operands);
+    b.create<CallOp>(asyncDispatchFunction.sym_name(),
 
     // Wait for the completion of all parallel compute operations.
-    nb.create<AwaitAllOp>(group);
+    b.create<AwaitAllOp>(group);
 
-    nb.create<scf::YieldOp>();
+    b.create<scf::YieldOp>();
   };
 
   // Dispatch either single block compute function, or launch async dispatch.
   b.create<scf::IfOp>(TypeRange(), isSingleBlock, syncDispatch, asyncDispatch);
 }
 
 // Dispatch parallel compute functions by submitting all async compute tasks
 // from a simple for loop in the caller thread.
auto computeFuncOperands = [&](Value blockIndex) -> SmallVector<Value> {
     computeFuncOperands.append(op.getStep().begin(), op.getStep().end());
     computeFuncOperands.append(parallelComputeFunction.captures);
     return computeFuncOperands;
   };
 
   // Induction variable is the index of the block: [0, blockCount).
   LoopBodyBuilder loopBuilder = [&](OpBuilder &loopBuilder, Location loc,
                                     Value iv, ValueRange args) {
-    ImplicitLocOpBuilder nb(loc, loopBuilder);
+    ImplicitLocOpBuilder b(loc, loopBuilder);
 
     // Call parallel compute function inside the async.execute region.
     auto executeBodyBuilder = [&](OpBuilder &executeBuilder,
                                   Location executeLoc, ValueRange executeArgs) {
       executeBuilder.create<CallOp>(executeLoc, compute.sym_name(),
                                     compute.getCallableResults(),
                                     computeFuncOperands(iv));
       executeBuilder.create<async::YieldOp>(executeLoc, ValueRange());
     };
 
     // Create async.execute operation to launch parallel computate function.
-    auto execute = nb.create<ExecuteOp>(TypeRange(), ValueRange(), ValueRange(),
-                                        executeBodyBuilder);
-    nb.create<AddToGroupOp>(rewriter.getIndexType(), execute.token(), group);
-    nb.create<scf::YieldOp>();
+    auto execute = b.create<ExecuteOp>(TypeRange(), ValueRange(), ValueRange(),
+    b.create<AddToGroupOp>(rewriter.getIndexType(), execute.token(), group);
+    b.create<scf::YieldOp>();
   };
 
   // Iterate over all compute blocks and launch parallel compute operations.
   b.create<scf::ForOp>(c1, blockCount, c1, ValueRange(), loopBuilder);
 
   // Call parallel compute function for the first block in the caller thread.
   b.create<CallOp>(compute.sym_name(), compute.getCallableResults(),
                    computeFuncOperands(c0));
AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
   // Do absolutely nothing if the trip count is zero.
   auto noOp = [&](OpBuilder &nestedBuilder, Location loc) {
     nestedBuilder.create<scf::YieldOp>(loc);
   };
 
   // Compute the parallel block size and dispatch concurrent tasks computing
   // results for each block.
   auto dispatch = [&](OpBuilder &nestedBuilder, Location loc) {
-    ImplicitLocOpBuilder nb(loc, nestedBuilder);
+    ImplicitLocOpBuilder b(loc, nestedBuilder);
 
     // Collect statically known constants defining the loop nest in the parallel
     // compute function. LLVM can't always push constants across the non-trivial
     // async dispatch call graph, by providing these values explicitly we can
     // choose to build more efficient loop nest, and rely on a better constant
     // folding, loop unrolling and vectorization.
     ParallelComputeFunctionBounds staticBounds = {
         integerConstants(tripCounts),
auto dispatch = [&](OpBuilder &nestedBuilder, Location loc) {
     // the parallel operation body for a subset of iteration space.
 
     // Compute the number of parallel compute blocks.
     Value blockCount = b.create<arith::CeilDivSIOp>(tripCount, blockSize);
 
     // Unroll when numUnrollableLoops > 0 && blockSize >= maxIterations.
     bool staticShouldUnroll = numUnrollableLoops > 0;
     auto dispatchNotUnrollable = [&](OpBuilder &nestedBuilder, Location loc) {
-      ImplicitLocOpBuilder nb(loc, nestedBuilder);
+      ImplicitLocOpBuilder b(loc, nestedBuilder);
       doDispatch(b, rewriter, notUnrollableParallelComputeFunction, op,
                  blockSize, blockCount, tripCounts);
-      nb.create<scf::YieldOp>();
+      b.create<scf::YieldOp>();
     };
 
     if (staticShouldUnroll) {
       Value dynamicShouldUnroll = b.create<arith::CmpIOp>(
           arith::CmpIPredicate::sge, blockSize,
           b.create<arith::ConstantIndexOp>(maxIterations));
 
       ParallelComputeFunction unrollableParallelComputeFunction =
           createParallelComputeFunction(op, staticBounds, numUnrollableLoops,
                                         rewriter);
 
       auto dispatchUnrollable = [&](OpBuilder &nestedBuilder, Location loc) {
-        ImplicitLocOpBuilder nb(loc, nestedBuilder);
+        ImplicitLocOpBuilder b(loc, nestedBuilder);
         // Align the block size to be a multiple of the statically known
         // number of iterations in the inner loops.
-        Value numIters = nb.create<arith::ConstantIndexOp>(
+        Value numIters = b.create<arith::ConstantIndexOp>(
             numIterations[op.getNumLoops() - numUnrollableLoops]);
-        Value alignedBlockSize = nb.create<arith::MulIOp>(
-            nb.create<arith::CeilDivSIOp>(blockSize, numIters), numIters);
+        Value alignedBlockSize = b.create<arith::MulIOp>(
+            b.create<arith::CeilDivSIOp>(blockSize, numIters), numIters);
         doDispatch(b, rewriter, unrollableParallelComputeFunction, op,
                    alignedBlockSize, blockCount, tripCounts);
-        nb.create<scf::YieldOp>();
+        b.create<scf::YieldOp>();
       };
 
       b.create<scf::IfOp>(TypeRange(), dynamicShouldUnroll, dispatchUnrollable,
                           dispatchNotUnrollable);
-      nb.create<scf::YieldOp>();
+      b.create<scf::YieldOp>();
     } else {
-      dispatchNotUnrollable(nb, loc);
+      dispatchNotUnrollable(b, loc);
     }
   };
 
   // Replace the `scf.parallel` operation with the parallel compute function.
   b.create<scf::IfOp>(TypeRange(), isZeroIterations, noOp, dispatch);
 
   // Parallel operation was replaced with a block iteration loop.
   rewriter.eraseOp(op);