Rename target block size to min task size for clarity.

[ezhulenev]

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

98

variable_ is not compatible with the style guide, does the name shadowing warnings show up anywhere

687

minTaskSizeCst is more consistent with other constants in the mlir code base

[bakhtiyarneyman]

Appease clang.

mlir/include/mlir/Dialect/Async/Passes.td

Option<"asyncDispatch", "async-dispatch",
       "Dispatch async compute tasks using recursive work splitting. If `false` "
       "async compute tasks will be launched using simple for loop in the "
       "caller thread.">,
 
     Option<"numWorkerThreads", "num-workers",
       "int32_t", /*default=*/"8",
       "The number of available workers to execute async operations.">,
 
-    Option<"targetBlockSize", "target-block-size",
+    Option<"minTaskSize", "min-task-size",
       "int32_t", /*default=*/"1000",
-      "The target block size for sharding parallel operation.">
+      "The minimum task size for sharding parallel operation.">
   ];
 
   let dependentDialects = ["async::AsyncDialect", "scf::SCFDialect"];
 }
 
 def AsyncToAsyncRuntime : Pass<"async-to-async-runtime", "ModuleOp"> {
   let summary = "Lower high level async operations (e.g. async.execute) to the"
                 "explicit async.runtime and async.coro operations";

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

 //      call @parallel_compute_fn(%block_index, %block_size, ...)
 //   }
 //
 struct AsyncParallelForPass
     : public AsyncParallelForBase<AsyncParallelForPass> {
   AsyncParallelForPass() = default;
 
   AsyncParallelForPass(bool asyncDispatch, int32_t numWorkerThreads,
-                       int32_t targetBlockSize) {
+                       int32_t minTaskSize) {
     this->asyncDispatch = asyncDispatch;
     this->numWorkerThreads = numWorkerThreads;
-    this->targetBlockSize = targetBlockSize;
+    this->minTaskSize = minTaskSize;
   }
 
   void runOnOperation() override;
 };
 
 struct AsyncParallelForRewrite : public OpRewritePattern<scf::ParallelOp> {
 public:
   AsyncParallelForRewrite(MLIRContext *ctx, bool asyncDispatch,
-                          int32_t numWorkerThreads, int32_t targetBlockSize)
+                          int32_t numWorkerThreads, int32_t minTaskSize)
       : OpRewritePattern(ctx), asyncDispatch(asyncDispatch),
-        numWorkerThreads(numWorkerThreads), targetBlockSize(targetBlockSize) {}
+        numWorkerThreads(numWorkerThreads), minTaskSize(minTaskSize) {}
 
   LogicalResult matchAndRewrite(scf::ParallelOp op,
                                 PatternRewriter &rewriter) const override;
 
 private:
   bool asyncDispatch;
   int32_t numWorkerThreads;
-  int32_t targetBlockSize;
+  int32_t minTaskSize;
 };
 
 struct ParallelComputeFunctionType {
   FunctionType type;
   llvm::SmallVector<Value> captures;
 };
 
 struct ParallelComputeFunction {
static void doAsyncDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
 
   // Wait for the completion of all parallel compute operations.
   b.create<AwaitAllOp>(group);
 }
 
 // Dispatch parallel compute functions by submitting all async compute tasks
 // from a simple for loop in the caller thread.
 static void
-doSequantialDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
+doSequentialDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
                      ParallelComputeFunction &parallelComputeFunction,
                      scf::ParallelOp op, Value blockSize, Value blockCount,
                      const SmallVector<Value> &tripCounts) {
   MLIRContext *ctx = op->getContext();
 
   FuncOp compute = parallelComputeFunction.func;
 
   Value c0 = b.create<ConstantIndexOp>(0);
float overshardingFactor = numWorkerThreads <= 4    ? 8.0
                                : numWorkerThreads <= 64 ? 0.8
                                                         : 0.6;
 
     // Do not overload worker threads with too many compute blocks.
     Value maxComputeBlocks = b.create<ConstantIndexOp>(
         std::max(1, static_cast<int>(numWorkerThreads * overshardingFactor)));
 
     // Target block size from the pass parameters.
-    Value targetComputeBlock = b.create<ConstantIndexOp>(targetBlockSize);
+    Value minTaskSizeCst = b.create<ConstantIndexOp>(minTaskSize);
 
     // Compute parallel block size from the parallel problem size:
     //   blockSize = min(tripCount,
     //                   max(ceil_div(tripCount, maxComputeBlocks),
-    //                       targetComputeBlock))
+    //                       ceil_div(minTaskSize, bodySize)))
     Value bs0 = b.create<SignedCeilDivIOp>(tripCount, maxComputeBlocks);
-    Value bs1 = b.create<CmpIOp>(CmpIPredicate::sge, bs0, targetComputeBlock);
-    Value bs2 = b.create<SelectOp>(bs1, bs0, targetComputeBlock);
+    Value bs1 = b.create<CmpIOp>(CmpIPredicate::sge, bs0, minTaskSizeCst);
+    Value bs2 = b.create<SelectOp>(bs1, bs0, minTaskSizeCst);
     Value bs3 = b.create<CmpIOp>(CmpIPredicate::sle, tripCount, bs2);
     Value blockSize0 = b.create<SelectOp>(bs3, tripCount, bs2);
     Value blockCount0 = b.create<SignedCeilDivIOp>(tripCount, blockSize0);
 
     // Compute balanced block size for the estimated block count.
     Value blockSize = b.create<SignedCeilDivIOp>(tripCount, blockCount0);
     Value blockCount = b.create<SignedCeilDivIOp>(tripCount, blockSize);
 
     // 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, 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 {
-      doSequantialDispatch(b, rewriter, parallelComputeFunction, op, blockSize,
+      doSequentialDispatch(b, rewriter, parallelComputeFunction, op, blockSize,
                            blockCount, 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,
-                                        targetBlockSize);
+                                        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 targetBlockSize) {
+std::unique_ptr<Pass> mlir::createAsyncParallelForPass(bool asyncDispatch,
+                                                       int32_t numWorkerThreads,
+                                                       int32_t minTaskSize) {
   return std::make_unique<AsyncParallelForPass>(asyncDispatch, numWorkerThreads,
-                                                targetBlockSize);
+                                                minTaskSize);
 }

mlir/test/Integration/Dialect/Async/CPU/test-async-parallel-for-1d.mlir

 // RUN: | mlir-cpu-runner                                                      \
 // RUN:  -e entry -entry-point-result=void -O0                                 \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_async_runtime%shlibext\
 // RUN: | FileCheck %s --dump-input=always
 
 // RUN:   mlir-opt %s -async-parallel-for="async-dispatch=false                \
 // RUN:                                    num-workers=20                      \
-// RUN:                                    target-block-size=1"                \
+// RUN:                                    min-task-size=1"                \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
 // RUN:               -async-runtime-ref-counting-opt                          \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-scf-to-std                                      \
 // RUN:               -convert-memref-to-llvm                                  \
 // RUN:               -std-expand                                              \
 // RUN:               -convert-std-to-llvm                                     \

mlir/test/Integration/Dialect/Async/CPU/test-async-parallel-for-2d.mlir

 // RUN: | mlir-cpu-runner                                                      \
 // RUN:  -e entry -entry-point-result=void -O0                                 \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_async_runtime%shlibext\
 // RUN: | FileCheck %s --dump-input=always
 
 // RUN:   mlir-opt %s -async-parallel-for="async-dispatch=false                \
 // RUN:                                    num-workers=20                      \
-// RUN:                                    target-block-size=1"                \
+// RUN:                                    min-task-size=1"                \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
 // RUN:               -async-runtime-ref-counting-opt                          \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-scf-to-std                                      \
 // RUN:               -convert-memref-to-llvm                                  \
 // RUN:               -convert-std-to-llvm                                     \
 // RUN:               -reconcile-unrealized-casts                              \