diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp --- a/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp +++ b/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp @@ -373,23 +373,23 @@ 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( - lowerBounds[loopIdx], nb.create(iv, steps[loopIdx])); + computeBlockInductionVars[loopIdx] = b.create( + lowerBounds[loopIdx], b.create(iv, steps[loopIdx])); // Check if we are inside first or last iteration of the loop. - isBlockFirstCoord[loopIdx] = nb.create( + isBlockFirstCoord[loopIdx] = b.create( arith::CmpIPredicate::eq, iv, blockFirstCoord[loopIdx]); - isBlockLastCoord[loopIdx] = nb.create( + isBlockLastCoord[loopIdx] = b.create( 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( + isBlockFirstCoord[loopIdx] = b.create( isBlockFirstCoord[loopIdx], isBlockFirstCoord[loopIdx - 1]); - isBlockLastCoord[loopIdx] = nb.create( + isBlockLastCoord[loopIdx] = b.create( isBlockLastCoord[loopIdx], isBlockLastCoord[loopIdx - 1]); } @@ -398,24 +398,24 @@ if (loopIdx + 1 >= op.getNumLoops() - numBlockAlignedInnerLoops) { // For block aligned loops we always iterate starting from 0 up to // the loop trip counts. - nb.create(c0, tripCounts[loopIdx + 1], c1, ValueRange(), - workLoopBuilder(loopIdx + 1)); + b.create(c0, tripCounts[loopIdx + 1], c1, ValueRange(), + workLoopBuilder(loopIdx + 1)); } else { // Select nested loop lower/upper bounds depending on our position in // the multi-dimensional iteration space. - auto lb = nb.create( - isBlockFirstCoord[loopIdx], blockFirstCoord[loopIdx + 1], c0); + auto lb = b.create(isBlockFirstCoord[loopIdx], + blockFirstCoord[loopIdx + 1], c0); - auto ub = nb.create(isBlockLastCoord[loopIdx], - blockEndCoord[loopIdx + 1], - tripCounts[loopIdx + 1]); + auto ub = b.create(isBlockLastCoord[loopIdx], + blockEndCoord[loopIdx + 1], + tripCounts[loopIdx + 1]); - nb.create(lb, ub, c1, ValueRange(), - workLoopBuilder(loopIdx + 1)); + b.create(lb, ub, c1, ValueRange(), + workLoopBuilder(loopIdx + 1)); } - nb.create(loc); + b.create(loc); return; } @@ -425,7 +425,7 @@ mapping.map(computeFuncType.captures, captures); for (auto &bodyOp : op.getLoopBody().getOps()) - nb.clone(bodyOp, mapping); + b.clone(bodyOp, mapping); }; }; @@ -602,38 +602,38 @@ b.create(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 operands = {c0, blockSize}; appendBlockComputeOperands(operands); - nb.create(parallelComputeFunction.func.sym_name(), - parallelComputeFunction.func.getCallableResults(), - operands); - nb.create(); + b.create(parallelComputeFunction.func.sym_name(), + parallelComputeFunction.func.getCallableResults(), + operands); + b.create(); }; 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(blockCount, c1); - Value group = nb.create(GroupType::get(ctx), groupSize); + Value groupSize = b.create(blockCount, c1); + Value group = b.create(GroupType::get(ctx), groupSize); // Launch async dispatch function for [0, blockCount) range. SmallVector operands = {group, c0, blockCount, blockSize}; appendBlockComputeOperands(operands); - nb.create(asyncDispatchFunction.sym_name(), - asyncDispatchFunction.getCallableResults(), operands); + b.create(asyncDispatchFunction.sym_name(), + asyncDispatchFunction.getCallableResults(), operands); // Wait for the completion of all parallel compute operations. - nb.create(group); + b.create(group); - nb.create(); + b.create(); }; // Dispatch either single block compute function, or launch async dispatch. @@ -680,7 +680,7 @@ // 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, @@ -692,10 +692,10 @@ }; // Create async.execute operation to launch parallel computate function. - auto execute = nb.create(TypeRange(), ValueRange(), ValueRange(), - executeBodyBuilder); - nb.create(rewriter.getIndexType(), execute.token(), group); - nb.create(); + auto execute = b.create(TypeRange(), ValueRange(), ValueRange(), + executeBodyBuilder); + b.create(rewriter.getIndexType(), execute.token(), group); + b.create(); }; // Iterate over all compute blocks and launch parallel compute operations. @@ -758,7 +758,7 @@ // 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 @@ -872,10 +872,10 @@ // 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(); + b.create(); }; if (staticShouldUnroll) { @@ -888,23 +888,23 @@ 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( + Value numIters = b.create( numIterations[op.getNumLoops() - numUnrollableLoops]); - Value alignedBlockSize = nb.create( - nb.create(blockSize, numIters), numIters); + Value alignedBlockSize = b.create( + b.create(blockSize, numIters), numIters); doDispatch(b, rewriter, unrollableParallelComputeFunction, op, alignedBlockSize, blockCount, tripCounts); - nb.create(); + b.create(); }; b.create(TypeRange(), dynamicShouldUnroll, dispatchUnrollable, dispatchNotUnrollable); - nb.create(); + b.create(); } else { - dispatchNotUnrollable(nb, loc); + dispatchNotUnrollable(b, loc); } };