Make AsyncParallelForRewrite parameterizable with a cost model which drives deciding the parallelization granularity.

[mehdi_amini]

The revision title mentions "AsyncParallelFor pass" but this isn't touching the pass itself right?

mlir/include/mlir/Dialect/Async/Transforms.h

35

Can you improve the doc?

[bakhtiyarneyman]

Improve docs.

[bakhtiyarneyman]

In D115423#3181902, @mehdi_amini wrote:

The revision title mentions "AsyncParallelFor pass" but this isn't touching the pass itself right?

Correct.

[mehdi_amini]

mlir/include/mlir/Dialect/Async/Transforms.h

23

[bakhtiyarneyman]

Change the comment.

[bakhtiyarneyman]

Rebase. Note that this made some of the loop unrolling logic dynamic which seeped all the way into loop building.

[ezhulenev]

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

867

Computing numUnrollableLoops as Value (compute function argument) + runtime scf.if in the loop nest prevents LLVM from loop unrolling and vectorization, and it leads to large regressions in: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tfrt/benchmarks/compute_function_benchmark.cc

name                                 old cpu/op  new cpu/op  delta
BM_cpurt_Fresh2/4/process_time       12.3µs ± 9%  12.3µs ±10%      ~     (p=0.961 n=17+18)
BM_cpurt_Fresh2/8/process_time       12.3µs ± 8%  12.5µs ± 7%      ~     (p=0.245 n=18+17)
BM_cpurt_Factorized0/0/process_time  21.1µs ± 4%  21.2µs ± 4%      ~     (p=0.499 n=19+18)
BM_cpurt_Factorized0/4/process_time  21.5µs ± 7%  35.9µs ± 6%   +66.93%  (p=0.000 n=19+20)
BM_cpurt_Factorized0/8/process_time  21.0µs ± 3%  35.7µs ± 5%   +69.52%  (p=0.000 n=17+20)
BM_cpurt_Factorized1/0/process_time  9.35µs ± 6%  9.31µs ± 3%      ~     (p=0.732 n=18+17)
BM_cpurt_Factorized1/4/process_time  35.9µs ± 4%  46.4µs ± 5%   +29.04%  (p=0.000 n=18+19)
BM_cpurt_Factorized1/8/process_time  36.2µs ± 3%  46.3µs ± 6%   +28.17%  (p=0.000 n=18+19)
BM_cpurt_Factorized2/0/process_time  86.7µs ± 8%  86.3µs ± 8%      ~     (p=0.832 n=18+17)
BM_cpurt_Factorized2/4/process_time   124µs ± 7%   442µs ± 3%  +255.67%  (p=0.000 n=18+18)
BM_cpurt_Factorized2/8/process_time   159µs ± 6%   459µs ± 6%  +187.73%  (p=0.000 n=17+18)

numUnrollableLoops should be known at compiled time, with dynamic minTaskSize the structure should look like this:

ParallelComputeFunction unrollableParallelComputeFunction =
        createParallelComputeFunction(op, staticBounds, numUnrollableLoops,
                                      rewriter);

ParallelComputeFunction defaultParallelComputeFunction =
        createParallelComputeFunction(op, staticBounds, numUnrollableLoops,
                                      rewriter);

b.create<scf::IfOp>(..) {
  dispatch unrollableParallelComputeFunction
} else {
  dispatch defaultParallelComputeFunction
}

There is no real need of benefit of computing numUnrollableLoop as Value, because unrolling/vectorization can only happen in trip counts (loop bounds) are know at compile time.

[bakhtiyarneyman]

Address comments: lift the scf.if out of the loops.

[bakhtiyarneyman]

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

867

Done, PTAL.

[mehdi_amini]

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

882

You're creating a value dynamicShouldUnroll here which isn't used in the else branch below, can you sink this in the then branch?

[bakhtiyarneyman]

Elide dynamic check if possible.

[ezhulenev]

Build is failing: https://buildkite.com/llvm-project/premerge-checks/builds/70123#b97ba1d4-10ba-4287-899a-0bcf5260aefe

[ezhulenev]

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

896

nit: I'd move this lambda close to the dispatchNotUnrollable to reduce the nesting, and put similar things together, although it's used only inside one branch of the if. I think it's ok to put createParallelComputeFunction inside lamdba, so you don't create aligned compute function if you'll not need it.

[bakhtiyarneyman]

Address nit.

[bakhtiyarneyman]

Fix minor C++ compilation issue.

[bakhtiyarneyman]

Same thing for another lambda.

[bakhtiyarneyman]

Move minTaskSize definition earlier to avoid contaminating the analysis results with the partially executed rewrite logic.

mlir/include/mlir/Dialect/Async/Passes.h

 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 std::unique_ptr<Pass> createAsyncParallelForPass();
 
 std::unique_ptr<Pass> createAsyncParallelForPass(bool asyncDispatch,
                                                  int32_t numWorkerThreads,
-                                                 int32_t targetBlockSize);
+                                                 int32_t minTaskSize);
 
 std::unique_ptr<OperationPass<ModuleOp>> createAsyncToAsyncRuntimePass();
 
 std::unique_ptr<Pass> createAsyncRuntimeRefCountingPass();
 
 std::unique_ptr<Pass> createAsyncRuntimeRefCountingOptPass();
 
 std::unique_ptr<Pass> createAsyncRuntimePolicyBasedRefCountingPass();

mlir/include/mlir/Dialect/Async/Transforms.h

+//===- Transforms.h - Async dialect transformation utilities ----*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines transformations on Async operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ASYNC_TRANSFORMS_H_
+#define MLIR_DIALECT_ASYNC_TRANSFORMS_H_
+
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/IR/ImplicitLocOpBuilder.h"
+
+namespace mlir {
+namespace async {
+
+/// Emit the IR to compute the minimum number of iterations of scf.parallel body
+/// that would be viable for a single parallel task. Allows the user to avoid
+/// incurring the overheads of spawning costly parallel tasks in absence of
+/// sufficient amount of parallelizable work.
+///
+/// Must return an index type.
+using AsyncMinTaskSizeComputationFunction =
+    std::function<Value(ImplicitLocOpBuilder, scf::ParallelOp)>;
+
+/// Add a pattern to the given pattern list to lower scf.parallel to async
+/// operations.
+void populateAsyncParallelForPatterns(
+    RewritePatternSet &patterns, bool asyncDispatch, int32_t numWorkerThreads,
+    AsyncMinTaskSizeComputationFunction computeMinTaskSize);
+
+} // namespace async
+} // namespace mlir
+
+#endif // MLIR_DIALECT_ASYNC_TRANSFORMS_H_

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

 //===- AsyncParallelFor.cpp - Implementation of Async Parallel For --------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements scf.parallel to scf.for + async.execute conversion pass.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Async/Passes.h"
+#include "mlir/Dialect/Async/Transforms.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/RegionUtils.h"
AsyncParallelForPass(bool asyncDispatch, int32_t numWorkerThreads,
     this->minTaskSize = minTaskSize;
   }
 
   void runOnOperation() override;
 };
 
 struct AsyncParallelForRewrite : public OpRewritePattern<scf::ParallelOp> {
 public:
-  AsyncParallelForRewrite(MLIRContext *ctx, bool asyncDispatch,
-                          int32_t numWorkerThreads, int32_t minTaskSize)
+  AsyncParallelForRewrite(
+      MLIRContext *ctx, bool asyncDispatch, int32_t numWorkerThreads,
+      AsyncMinTaskSizeComputationFunction computeMinTaskSize)
       : OpRewritePattern(ctx), asyncDispatch(asyncDispatch),
-        numWorkerThreads(numWorkerThreads), minTaskSize(minTaskSize) {}
+        numWorkerThreads(numWorkerThreads),
+        computeMinTaskSize(computeMinTaskSize) {}
 
   LogicalResult matchAndRewrite(scf::ParallelOp op,
                                 PatternRewriter &rewriter) const override;
 
 private:
   bool asyncDispatch;
   int32_t numWorkerThreads;
-  int32_t minTaskSize;
+  AsyncMinTaskSizeComputationFunction computeMinTaskSize;
 };
 
 struct ParallelComputeFunctionType {
   FunctionType type;
   SmallVector<Value> captures;
 };
 
 // Helper struct to parse parallel compute function argument list.
 struct ParallelComputeFunctionArgs {
   BlockArgument blockIndex();
   BlockArgument blockSize();
+  BlockArgument numBlockAlignedInnerLoops();
   ArrayRef<BlockArgument> tripCounts();
   ArrayRef<BlockArgument> lowerBounds();
   ArrayRef<BlockArgument> upperBounds();
   ArrayRef<BlockArgument> steps();
   ArrayRef<BlockArgument> captures();
 
   unsigned numLoops;
   ArrayRef<BlockArgument> args;
struct ParallelComputeFunction {
   FuncOp func;
   llvm::SmallVector<Value> captures;
 };
 
 } // namespace
 
 BlockArgument ParallelComputeFunctionArgs::blockIndex() { return args[0]; }
 BlockArgument ParallelComputeFunctionArgs::blockSize() { return args[1]; }
+BlockArgument ParallelComputeFunctionArgs::numBlockAlignedInnerLoops() {
+  return args[2];
+}
 
 ArrayRef<BlockArgument> ParallelComputeFunctionArgs::tripCounts() {
-  return args.drop_front(2).take_front(numLoops);
+  return args.drop_front(3).take_front(numLoops);
 }
 
 ArrayRef<BlockArgument> ParallelComputeFunctionArgs::lowerBounds() {
-  return args.drop_front(2 + 1 * numLoops).take_front(numLoops);
+  return args.drop_front(3 + 1 * numLoops).take_front(numLoops);
 }
 
 ArrayRef<BlockArgument> ParallelComputeFunctionArgs::upperBounds() {
-  return args.drop_front(2 + 2 * numLoops).take_front(numLoops);
+  return args.drop_front(3 + 2 * numLoops).take_front(numLoops);
 }
 
 ArrayRef<BlockArgument> ParallelComputeFunctionArgs::steps() {
-  return args.drop_front(2 + 3 * numLoops).take_front(numLoops);
+  return args.drop_front(3 + 3 * numLoops).take_front(numLoops);
 }
 
 ArrayRef<BlockArgument> ParallelComputeFunctionArgs::captures() {
-  return args.drop_front(2 + 4 * numLoops);
+  return args.drop_front(3 + 4 * numLoops);
 }
 
 template <typename ValueRange>
 static SmallVector<IntegerAttr> integerConstants(ValueRange values) {
   SmallVector<IntegerAttr> attrs(values.size());
   for (unsigned i = 0; i < values.size(); ++i)
     matchPattern(values[i], m_Constant(&attrs[i]));
   return attrs;
getParallelComputeFunctionType(scf::ParallelOp op, PatternRewriter &rewriter) {
   SmallVector<Type> inputs;
   inputs.reserve(2 + 4 * op.getNumLoops() + captures.size());
 
   Type indexTy = rewriter.getIndexType();
 
   // One-dimensional iteration space defined by the block index and size.
   inputs.push_back(indexTy); // blockIndex
   inputs.push_back(indexTy); // blockSize
+  inputs.push_back(indexTy); // numBlockAlignedInnerLoops
 
   // Multi-dimensional parallel iteration space defined by the loop trip counts.
   for (unsigned i = 0; i < op.getNumLoops(); ++i)
     inputs.push_back(indexTy); // loop tripCount
 
   // Parallel operation lower bound, upper bound and step. Lower bound, upper
   // bound and step passed as contiguous arguments:
   //   call @compute(%lb0, %lb1, ..., %ub0, %ub1, ..., %step0, %step1, ...)
   for (unsigned i = 0; i < op.getNumLoops(); ++i) {
     inputs.push_back(indexTy); // lower bound
     inputs.push_back(indexTy); // upper bound
     inputs.push_back(indexTy); // step
   }
 
   // Types of the implicit captures.
   for (Value capture : captures)
     inputs.push_back(capture.getType());
 
   // Convert captures to vector for later convenience.
   SmallVector<Value> capturesVector(captures.begin(), captures.end());
   return {rewriter.getFunctionType(inputs, TypeRange()), capturesVector};
 }
 
 // Create a parallel compute fuction from the parallel operation.
-static ParallelComputeFunction createParallelComputeFunction(
-    scf::ParallelOp op, ParallelComputeFunctionBounds bounds,
-    unsigned numBlockAlignedInnerLoops, PatternRewriter &rewriter) {
+static ParallelComputeFunction
+createParallelComputeFunction(scf::ParallelOp op,
+                              ParallelComputeFunctionBounds bounds,
+                              PatternRewriter &rewriter) {
   OpBuilder::InsertionGuard guard(rewriter);
   ImplicitLocOpBuilder b(op.getLoc(), rewriter);
 
   ModuleOp module = op->getParentOfType<ModuleOp>();
 
   ParallelComputeFunctionType computeFuncType =
       getParallelComputeFunctionType(op, rewriter);
 
createParallelComputeFunction(scf::ParallelOp op,
   Block *block = b.createBlock(&func.getBody(), func.begin(), type.getInputs());
   b.setInsertionPointToEnd(block);
 
   ParallelComputeFunctionArgs args = {op.getNumLoops(), func.getArguments()};
 
   // Block iteration position defined by the block index and size.
   BlockArgument blockIndex = args.blockIndex();
   BlockArgument blockSize = args.blockSize();
+  BlockArgument numBlockAlignedInnerLoops = args.numBlockAlignedInnerLoops();
 
   // Constants used below.
   Value c0 = b.create<arith::ConstantIndexOp>(0);
   Value c1 = b.create<arith::ConstantIndexOp>(1);
 
   // Materialize known constants as constant operation in the function body.
   auto values = [&](ArrayRef<BlockArgument> args, ArrayRef<IntegerAttr> attrs) {
     return llvm::to_vector(
createParallelComputeFunction(scf::ParallelOp op,
   // Value at ith position specifies if all loops in [0, i) range of the loop
   // nest are in the first/last iteration.
   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,
+    return [&, loopIdx](OpBuilder &nestedBuilder, Location nestedLoc, Value iv,
                         ValueRange args) {
-      ImplicitLocOpBuilder nb(loc, nestedBuilder);
+      ImplicitLocOpBuilder nb(nestedLoc, nestedBuilder);
 
       // Compute induction variable for `loopIdx`.
       computeBlockInductionVars[loopIdx] = nb.create<arith::AddIOp>(
           lowerBounds[loopIdx], nb.create<arith::MulIOp>(iv, steps[loopIdx]));
 
       // Check if we are inside first or last iteration of the loop.
       isBlockFirstCoord[loopIdx] = nb.create<arith::CmpIOp>(
           arith::CmpIPredicate::eq, iv, blockFirstCoord[loopIdx]);
       isBlockLastCoord[loopIdx] = nb.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], isBlockFirstCoord[loopIdx - 1]);
         isBlockLastCoord[loopIdx] = nb.create<arith::AndIOp>(
             isBlockLastCoord[loopIdx], isBlockLastCoord[loopIdx - 1]);
       }
 
       // Keep building loop nest.
       if (loopIdx < op.getNumLoops() - 1) {
-        if (loopIdx + 1 >= op.getNumLoops() - numBlockAlignedInnerLoops) {
+        auto inner = nb.create<arith::CmpIOp>(
+            arith::CmpIPredicate::sge, numBlockAlignedInnerLoops,
+            nb.create<arith::ConstantIndexOp>(op.getNumLoops() - loopIdx - 1));
+        nb.create<scf::IfOp>(
+            TypeRange(), inner,
+            [&](OpBuilder &doublyNestedBuilder, Location doublyNestedLoc) {
+              ImplicitLocOpBuilder dnb(doublyNestedLoc, doublyNestedBuilder);
               // 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(),
+              dnb.create<scf::ForOp>(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<SelectOp>(isBlockFirstCoord[loopIdx],
+              return dnb.create<scf::YieldOp>();
+            },
+            [&](OpBuilder &doublyNestedBuilder, Location doublyNestedLoc) {
+              ImplicitLocOpBuilder dnb(doublyNestedLoc, doublyNestedBuilder);
+              // Select nested loop lower/upper bounds depending on our position
+              // in the multi-dimensional iteration space.
+              auto lb = dnb.create<SelectOp>(isBlockFirstCoord[loopIdx],
                                              blockFirstCoord[loopIdx + 1], c0);
 
-          auto ub = nb.create<SelectOp>(isBlockLastCoord[loopIdx],
+              auto ub = dnb.create<SelectOp>(isBlockLastCoord[loopIdx],
                                              blockEndCoord[loopIdx + 1],
                                              tripCounts[loopIdx + 1]);
 
-          nb.create<scf::ForOp>(lb, ub, c1, ValueRange(),
+              dnb.create<scf::ForOp>(lb, ub, c1, ValueRange(),
                                      workLoopBuilder(loopIdx + 1));
-        }
+              return dnb.create<scf::YieldOp>();
+            });
 
-        nb.create<scf::YieldOp>(loc);
+        nb.create<scf::YieldOp>(nestedLoc);
         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);
 
static FuncOp createAsyncDispatchFunction(ParallelComputeFunction &computeFunc,
 
   return func;
 }
 
 // Launch async dispatch of the parallel compute function.
 static void doAsyncDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
                             ParallelComputeFunction &parallelComputeFunction,
                             scf::ParallelOp op, Value blockSize,
-                            Value blockCount,
+                            Value blockCount, Value numBlockAlignedInnerLoops,
                             const SmallVector<Value> &tripCounts) {
   MLIRContext *ctx = op->getContext();
 
   // Add one more level of indirection to dispatch parallel compute functions
   // using async operations and recursive work splitting.
   FuncOp asyncDispatchFunction =
       createAsyncDispatchFunction(parallelComputeFunction, rewriter);
 
static void doAsyncDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
   // 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);
 
     // Call parallel compute function for the single block.
-    SmallVector<Value> operands = {c0, blockSize};
+    SmallVector<Value> operands = {c0, blockSize, numBlockAlignedInnerLoops};
     appendBlockComputeOperands(operands);
 
     nb.create<CallOp>(parallelComputeFunction.func.sym_name(),
                       parallelComputeFunction.func.getCallableResults(),
                       operands);
     nb.create<scf::YieldOp>();
   };
 
   auto asyncDispatch = [&](OpBuilder &nestedBuilder, Location loc) {
     // 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 = b.create<arith::SubIOp>(blockCount, c1);
     Value group = b.create<CreateGroupOp>(GroupType::get(ctx), groupSize);
 
     ImplicitLocOpBuilder nb(loc, nestedBuilder);
 
     // Launch async dispatch function for [0, blockCount) range.
-    SmallVector<Value> operands = {group, c0, blockCount, blockSize};
+    SmallVector<Value> operands = {
+        group, c0, blockCount, blockSize, numBlockAlignedInnerLoops,
+    };
     appendBlockComputeOperands(operands);
 
     nb.create<CallOp>(asyncDispatchFunction.sym_name(),
                       asyncDispatchFunction.getCallableResults(), operands);
 
     // Wait for the completion of all parallel compute operations.
     b.create<AwaitAllOp>(group);
 
     nb.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.
 static void
 doSequentialDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
                      ParallelComputeFunction &parallelComputeFunction,
                      scf::ParallelOp op, Value blockSize, Value blockCount,
+                     Value numBlockAlignedInnerLoops,
                      const SmallVector<Value> &tripCounts) {
   MLIRContext *ctx = op->getContext();
 
   FuncOp compute = parallelComputeFunction.func;
 
   Value c0 = b.create<arith::ConstantIndexOp>(0);
   Value c1 = b.create<arith::ConstantIndexOp>(1);
 
   // 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 = b.create<arith::SubIOp>(blockCount, c1);
   Value group = b.create<CreateGroupOp>(GroupType::get(ctx), groupSize);
 
   // Call parallel compute function for all blocks.
   using LoopBodyBuilder =
       std::function<void(OpBuilder &, Location, Value, ValueRange)>;
 
   // Returns parallel compute function operands to process the given block.
   auto computeFuncOperands = [&](Value blockIndex) -> SmallVector<Value> {
-    SmallVector<Value> computeFuncOperands = {blockIndex, blockSize};
+    SmallVector<Value> computeFuncOperands = {blockIndex, blockSize,
+                                              numBlockAlignedInnerLoops};
     computeFuncOperands.append(tripCounts);
     computeFuncOperands.append(op.lowerBound().begin(), op.lowerBound().end());
     computeFuncOperands.append(op.upperBound().begin(), op.upperBound().end());
     computeFuncOperands.append(op.step().begin(), op.step().end());
     computeFuncOperands.append(parallelComputeFunction.captures);
     return computeFuncOperands;
   };
 
auto dispatch = [&](OpBuilder &nestedBuilder, Location loc) {
     ParallelComputeFunctionBounds staticBounds = {
         integerConstants(tripCounts),
         integerConstants(op.lowerBound()),
         integerConstants(op.upperBound()),
         integerConstants(op.step()),
     };
 
     // Find how many inner iteration dimensions are statically known, and their
-    // product is smaller than the `512`. We aling the parallel compute block
+    // product is smaller than the `512`. We align the parallel compute block
     // size by the product of statically known dimensions, so that we can
     // guarantee that the inner loops executes from 0 to the loop trip counts
     // and we can elide dynamic loop boundaries, and give LLVM an opportunity to
     // unroll the loops. The constant `512` is arbitrary, it should depend on
     // how many iterations LLVM will typically decide to unroll.
     static constexpr int64_t maxIterations = 512;
 
     // The number of inner loops with statically known number of iterations less
float overshardingFactor = numWorkerThreads <= 4    ? 8.0
                                : numWorkerThreads <= 32 ? 1.0
                                : numWorkerThreads <= 64 ? 0.8
                                                         : 0.6;
 
     // Do not overload worker threads with too many compute blocks.
     Value maxComputeBlocks = b.create<arith::ConstantIndexOp>(
         std::max(1, static_cast<int>(numWorkerThreads * overshardingFactor)));
 
-    // Target block size from the pass parameters.
-    Value minTaskSizeCst = b.create<arith::ConstantIndexOp>(minTaskSize);
-
     // Compute parallel block size from the parallel problem size:
     //   blockSize = min(tripCount,
     //                   max(ceil_div(tripCount, maxComputeBlocks),
-    //                       ceil_div(minTaskSize, bodySize)))
+    //                       minTaskSize))
     Value bs0 = b.create<arith::CeilDivSIOp>(tripCount, maxComputeBlocks);
-    Value bs1 = b.create<arith::MaxSIOp>(bs0, minTaskSizeCst);
+    Value minTaskSize = computeMinTaskSize(b, op);
+    Value bs1 = b.create<arith::MaxSIOp>(bs0, minTaskSize);
     Value blockSize = b.create<arith::MinSIOp>(tripCount, bs1);
 
-    // Align the block size to be a multiple of the statically known number
-    // of iterations in the inner loops.
-    if (numUnrollableLoops > 0 && minTaskSize >= maxIterations) {
-      Value numIters = b.create<arith::ConstantIndexOp>(
+    // Unroll when numUnrollableLoops > 0 && minTaskSize >= maxIterations.
+    // Value shouldUnroll = b.create<arith::AndIOp>(
+    //     b.create<arith::CmpIOp>(
+    //         arith::CmpIPredicate::sgt,
+    //         b.create<arith::ConstantIndexOp>(numUnrollableLoops), c0),
+    //     b.create<arith::CmpIOp>(
+    //         arith::CmpIPredicate::sge, minTaskSize,
+    //         b.create<arith::ConstantIndexOp>(maxIterations)));
+    Value shouldUnroll =
+        b.create<arith::ConstantOp>(b.getI1Type(), b.getBoolAttr(false));
+    MLIRContext *ctx = op->getContext();
+    auto numUnrollableLoopsAndBlockSize = b.create<scf::IfOp>(
+        TypeRange{IndexType::get(ctx), IndexType::get(ctx)}, shouldUnroll,
+        [&](OpBuilder &nestedBuilder, Location loc) {
+          ImplicitLocOpBuilder nb(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>(
               numIterations[op.getNumLoops() - numUnrollableLoops]);
-      Value bs2 = b.create<arith::MulIOp>(
-          b.create<arith::CeilDivSIOp>(blockSize, numIters), numIters);
-      blockSize = b.create<arith::MinSIOp>(tripCount, bs2);
-    } else {
-      // Reset the number of unrollable loops if we didn't align the block size.
-      numUnrollableLoops = 0;
-    }
+          Value bs2 = nb.create<arith::MulIOp>(
+              nb.create<arith::CeilDivSIOp>(blockSize, numIters), numIters);
+          return nb.create<scf::YieldOp>(
+              ValueRange{nb.create<arith::ConstantIndexOp>(numUnrollableLoops),
+                         nb.create<arith::MinSIOp>(tripCount, bs2)});
+        },
+        [&](OpBuilder &nestedBuilder, Location loc) {
+          ImplicitLocOpBuilder nb(loc, nestedBuilder);
+          // Reset the number of unrollable loops if we didn't align the block
+          // size.
+          return nb.create<scf::YieldOp>(ValueRange{c0, blockSize});
+        });
 
-    // Compute the number of parallel compute blocks.
-    Value blockCount = b.create<arith::CeilDivSIOp>(tripCount, blockSize);
+    blockSize = numUnrollableLoopsAndBlockSize.getResult(1);
 
     // Create a parallel compute function that takes a block id and computes
     // the parallel operation body for a subset of iteration space.
     ParallelComputeFunction parallelComputeFunction =
-        createParallelComputeFunction(op, staticBounds, numUnrollableLoops,
-                                      rewriter);
+        createParallelComputeFunction(op, staticBounds, rewriter);
+
+    // Compute the number of parallel compute blocks.
+    Value blockCount = b.create<arith::CeilDivSIOp>(tripCount, blockSize);
 
     // Dispatch parallel compute function using async recursive work splitting,
     // or by submitting compute task sequentially from a caller thread.
     if (asyncDispatch) {
       doAsyncDispatch(b, rewriter, parallelComputeFunction, op, blockSize,
-                      blockCount, tripCounts);
+                      blockCount, numUnrollableLoopsAndBlockSize.getResult(0),
+                      tripCounts);
     } else {
-      doSequentialDispatch(b, rewriter, parallelComputeFunction, op, blockSize,
-                           blockCount, tripCounts);
+      doSequentialDispatch(
+          b, rewriter, parallelComputeFunction, op, blockSize, blockCount,
+          numUnrollableLoopsAndBlockSize.getResult(0), tripCounts);
     }
 
     nb.create<scf::YieldOp>();
   };
 
   // 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);
 
   return success();
 }
 
 void AsyncParallelForPass::runOnOperation() {
   MLIRContext *ctx = &getContext();
 
   RewritePatternSet patterns(ctx);
-  patterns.add<AsyncParallelForRewrite>(ctx, asyncDispatch, numWorkerThreads,
-                                        minTaskSize);
-
+  populateAsyncParallelForPatterns(
+      patterns, asyncDispatch, numWorkerThreads,
+      [&](ImplicitLocOpBuilder builder, scf::ParallelOp op) {
+        return builder.create<arith::ConstantIndexOp>(minTaskSize);
+      });
   if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
     signalPassFailure();
 }
 
 std::unique_ptr<Pass> mlir::createAsyncParallelForPass() {
   return std::make_unique<AsyncParallelForPass>();
 }
 
 std::unique_ptr<Pass> mlir::createAsyncParallelForPass(bool asyncDispatch,
                                                        int32_t numWorkerThreads,
                                                        int32_t minTaskSize) {
   return std::make_unique<AsyncParallelForPass>(asyncDispatch, numWorkerThreads,
                                                 minTaskSize);
 }
+
+void mlir::async::populateAsyncParallelForPatterns(
+    RewritePatternSet &patterns, bool asyncDispatch, int32_t numWorkerThreads,
+    AsyncMinTaskSizeComputationFunction computeMinTaskSize) {
+  MLIRContext *ctx = patterns.getContext();
+  patterns.add<AsyncParallelForRewrite>(ctx, asyncDispatch, numWorkerThreads,
+                                        computeMinTaskSize);
+}

mlir/test/Dialect/Async/async-parallel-for-compute-fn.mlir

func @sink_constant_step(%arg0: memref<?xf32>, %lb: index, %ub: index) {
 
   return
 }
 
 // CHECK-LABEL: func private @parallel_compute_fn(
 // CHECK-SAME:   %[[BLOCK_INDEX:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[BLOCK_SIZE:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[TRIP_COUNT:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[NUM_BLOCK_ALIGNED_INNER_LOOPS:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[LB:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[UB:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[STEP:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[MEMREF:arg[0-9]+]]: memref<?xf32>
 // CHECK-SAME: ) {
 // CHECK:        %[[CSTEP:.*]] = arith.constant 123 : index
 // CHECK-NOT:    %[[STEP]]
 // CHECK:        scf.for %[[I:arg[0-9]+]]

 // CHECK-SAME:   %[[STEP1:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[MEMREF:arg[0-9]+]]: memref<?x10xf32>
 // CHECK-SAME: ) {
 // CHECK:        %[[C0:.*]] = arith.constant 0 : index
 // CHECK:        %[[C1:.*]] = arith.constant 1 : index
 // CHECK:        %[[C10:.*]] = arith.constant 10 : index
 // CHECK:        scf.for %[[I:arg[0-9]+]]
 // CHECK-NOT:      select
-// CHECK:          scf.for %[[J:arg[0-9]+]] = %c0 to %c10 step %c1
+// CHECK:          scf.for %[[J:arg[0-9]+]] = %c0 to %c10 step %c1
 

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

 )
 
 cc_library(
     name = "AsyncTransforms",
     srcs = glob([
         "lib/Dialect/Async/Transforms/*.cpp",
         "lib/Dialect/Async/Transforms/*.h",
     ]),
-    hdrs = ["include/mlir/Dialect/Async/Passes.h"],
+    hdrs = [
+        "include/mlir/Dialect/Async/Passes.h",
+        "include/mlir/Dialect/Async/Transforms.h",
+    ],
     includes = ["include"],
     deps = [
         ":Analysis",
         ":ArithmeticDialect",
         ":Async",
         ":AsyncPassIncGen",
         ":IR",
         ":Pass",