[mlir:Async] Change async-parallel-for block size/count calculation

[ezhulenev]

Summary

Depends On D105037

Avoid creating too many tasks when the number of workers is large.

[herhut]

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

658

Can kMaxOversharding be removed?

663

Does this mean that for large number of threads, we do not use all of them? That is because we assume we have too many threads for the available memory bandwidth?

668

Should this be std::max?

[ezhulenev]

Fix min/max bug

[ezhulenev]

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

663

Yes. This is an ugly hack because of an absence of a good cost model stolen from Eigen: https://gitlab.com/libeigen/eigen/-/blob/master/unsupported/Eigen/CXX11/src/Tensor/TensorContractionThreadPool.h#L241-246.

Assumption is that it's rarely beneficial to use all 64+ threads for problem that reads/writes from/to memory, and that you likely have competing tasks that can use idle threads.

I plan to implement something better than this for lowering from linalg.

[herhut]

Thanks for the explanation!

[stella.stamenova]

After this change kMaxOvershading is unused, causing warnings in the clang build:

/usr/bin/clang++ -DGTEST_HAS_RTTI=0 -DMLIR_CUDA_CONVERSIONS_ENABLED=0 -DMLIR_ROCM_CONVERSIONS_ENABLED=0 -D_DEBUG -D_GNU_SOURCE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -Itools/mlir/lib/Dialect/Async/Transforms -I/mnt/vss/_work/2/s/mlir/lib/Dialect/Async/Transforms -Iinclude -I/mnt/vss/_work/2/s/llvm/include -I/mnt/vss/_work/2/s/mlir/include -Itools/mlir/include -fPIC -fvisibility-inlines-hidden -Werror -Werror=date-time -Werror=unguarded-availability-new -Wall -Wextra -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wmissing-field-initializers -pedantic -Wno-long-long -Wc++98-compat-extra-semi -Wimplicit-fallthrough -Wcovered-switch-default -Wno-noexcept-type -Wnon-virtual-dtor -Wdelete-non-virtual-dtor -Wsuggest-override -Wstring-conversion -Wmisleading-indentation -fdiagnostics-color -Werror=global-constructors -g  -fno-exceptions -fno-rtti -std=c++14 -MD -MT tools/mlir/lib/Dialect/Async/Transforms/CMakeFiles/obj.MLIRAsyncTransforms.dir/AsyncParallelFor.cpp.o -MF tools/mlir/lib/Dialect/Async/Transforms/CMakeFiles/obj.MLIRAsyncTransforms.dir/AsyncParallelFor.cpp.o.d -o tools/mlir/lib/Dialect/Async/Transforms/CMakeFiles/obj.MLIRAsyncTransforms.dir/AsyncParallelFor.cpp.o -c /mnt/vss/_work/2/s/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
/mnt/vss/_work/2/s/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp:116:28: error: unused variable 'kMaxOversharding' [-Werror,-Wunused-const-variable]
  static constexpr int32_t kMaxOversharding = 4;

If the variable is no longer needed, can we remove it?

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

AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
   }
 
   // Compute a product of trip counts to get the 1-dimensional iteration space
   // for the scf.parallel operation.
   Value tripCount = tripCounts[0];
   for (size_t i = 1; i < tripCounts.size(); ++i)
     tripCount = b.create<MulIOp>(tripCount, tripCounts[i]);
 
+  // With large number of threads the value of creating many compute blocks
+  // is reduced because the problem typically becomes memory bound. For small
+  // number of threads it helps with stragglers.
+  float overshardingFactor = numWorkerThreads <= 4    ? 8.0
+                             : numWorkerThreads <= 8  ? 4.0
+                             : numWorkerThreads <= 16 ? 2.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<ConstantIndexOp>(numWorkerThreads * kMaxOversharding);
+  Value maxComputeBlocks = b.create<ConstantIndexOp>(
+      std::max(1, static_cast<int>(numWorkerThreads * overshardingFactor)));
 
   // Target block size from the pass parameters.
   Value targetComputeBlockSize = b.create<ConstantIndexOp>(targetBlockSize);
 
   // Compute parallel block size from the parallel problem size:
   //   blockSize = min(tripCount,
   //                   max(ceil_div(tripCount, maxComputeBlocks),
   //                       targetComputeBlockSize))
   Value bs0 = b.create<SignedCeilDivIOp>(tripCount, maxComputeBlocks);
   Value bs1 = b.create<CmpIOp>(CmpIPredicate::sge, bs0, targetComputeBlockSize);
   Value bs2 = b.create<SelectOp>(bs1, bs0, targetComputeBlockSize);
   Value bs3 = b.create<CmpIOp>(CmpIPredicate::sle, tripCount, bs2);
-  Value blockSize = b.create<SelectOp>(bs3, 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