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

825

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

825

Done, PTAL.

[mehdi_amini]

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

833

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

847

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.
static ParallelComputeFunction createParallelComputeFunction(
   ImplicitLocOpBuilder b(op.getLoc(), rewriter);
 
   ModuleOp module = op->getParentOfType<ModuleOp>();
 
   ParallelComputeFunctionType computeFuncType =
       getParallelComputeFunctionType(op, rewriter);
 
   FunctionType type = computeFuncType.type;
-  FuncOp func = FuncOp::create(op.getLoc(), "parallel_compute_fn", type);
+  FuncOp func = FuncOp::create(op.getLoc(),
+                               numBlockAlignedInnerLoops > 0
+                                   ? "parallel_compute_fn_with_aligned_loops"
+                                   : "parallel_compute_fn",
+                               type);
   func.setPrivate();
 
   // Insert function into the module symbol table and assign it unique name.
   SymbolTable symbolTable(module);
   symbolTable.insert(func);
   rewriter.getListener()->notifyOperationInserted(func);
 
   // Create function entry block.

 AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
                                          PatternRewriter &rewriter) const {
   // We do not currently support rewrite for parallel op with reductions.
   if (op.getNumReductions() != 0)
     return failure();
 
   ImplicitLocOpBuilder b(op.getLoc(), rewriter);
 
+  // Computing minTaskSize emits IR and can be implemented as executing a cost
+  // model on the body of the scf.parallel. Thus it needs to be computed before
+  // the body of the scf.parallel has been manipulated.
+  Value minTaskSize = computeMinTaskSize(b, op);
+
   // Make sure that all constants will be inside the parallel operation body to
   // reduce the number of parallel compute function arguments.
   cloneConstantsIntoTheRegion(op.getLoopBody(), rewriter);
 
   // Compute trip count for each loop induction variable:
   //   tripCount = ceil_div(upperBound - lowerBound, step);
   SmallVector<Value> tripCounts(op.getNumLoops());
   for (size_t i = 0; i < op.getNumLoops(); ++i) {
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 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>(
-          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;
-    }
+    ParallelComputeFunction notUnrollableParallelComputeFunction =
+        createParallelComputeFunction(op, staticBounds, 0, 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.
+    auto doDispatch = asyncDispatch ? doAsyncDispatch : doSequentialDispatch;
 
     // Create a parallel compute function that takes a block id and computes
     // the parallel operation body for a subset of iteration space.
-    ParallelComputeFunction parallelComputeFunction =
+
+    // 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);
+      doDispatch(b, rewriter, notUnrollableParallelComputeFunction, op,
+                 blockSize, blockCount, tripCounts);
+      nb.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);
 
-    // 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);
-    } else {
-      doSequentialDispatch(b, rewriter, parallelComputeFunction, op, blockSize,
-                           blockCount, tripCounts);
-    }
+      auto dispatchUnrollable = [&](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 alignedBlockSize = nb.create<arith::MulIOp>(
+            nb.create<arith::CeilDivSIOp>(blockSize, numIters), numIters);
+        doDispatch(b, rewriter, unrollableParallelComputeFunction, op,
+                   alignedBlockSize, blockCount, tripCounts);
+        nb.create<scf::YieldOp>();
+      };
 
+      b.create<scf::IfOp>(TypeRange(), dynamicShouldUnroll, dispatchUnrollable,
+                          dispatchNotUnrollable);
       nb.create<scf::YieldOp>();
+    } else {
+      dispatchNotUnrollable(nb, 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);
 
   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

 // CHECK-SAME:   %[[LB0:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[LB1:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[UB0:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[UB1:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[STEP0:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[STEP1:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[MEMREF:arg[0-9]+]]: memref<?x10xf32>
 // CHECK-SAME: ) {
+// CHECK:        scf.for %[[I:arg[0-9]+]]
+// CHECK:          select
+// CHECK:          scf.for %[[J:arg[0-9]+]]
+// CHECK:          memref.store
+
+// CHECK-LABEL: func private @parallel_compute_fn_with_aligned_loops(
+// CHECK-SAME:   %[[BLOCK_INDEX:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[BLOCK_SIZE:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[TRIP_COUNT0:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[TRIP_COUNT1:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[LB0:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[LB1:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[UB0:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[UB1:arg[0-9]+]]: index,
+// CHECK-SAME:   %[[STEP0:arg[0-9]+]]: index,
+// 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",