Promote readability by factoring out creation of min/max operation. Remove unnecessary divisions.

[ezhulenev]

We should migrate it to min/max after https://github.com/llvm/llvm-project/commit/9fb57c8c1dd87df36daf7b6f7dee3b7423475afc

[bakhtiyarneyman]

Rebase.

[bakhtiyarneyman]

Use min/max ops.

[ezhulenev]

Can you please sync to head and re-upload, can't apply the patch to the head

[ezhulenev]

There is one failed test

Failed Tests (1):
  MLIR :: Dialect/Async/async-parallel-for-canonicalize.mlir

Also the diff is agains the commit that is not in the llvm repo:

This diff is against commit 8a77386b7f31d4d983972b9b4fc5386817f417d6, but
    the commit is nowhere in the working copy

[bakhtiyarneyman]

Rebase to head.

[ezhulenev]

I have a failing test for this patch: Integration/Dialect/Async/CPU/test-async-parallel-for-2d.mlir

Looks like the problem is that there is no LLVM lowering for minsi:

%14 = builtin.unrealized_conversion_cast %13 : i64 to index
    %17 = builtin.unrealized_conversion_cast %16 : i64 to index
    %18 = minsi %17, %14 : index
    %19 = builtin.unrealized_conversion_cast %18 : index to i64

[bakhtiyarneyman]

Fix the failing test.

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

createParallelComputeFunction(scf::ParallelOp op, PatternRewriter &rewriter) {
   // Remaining arguments are implicit captures of the parallel operation.
   ArrayRef<Value> captures = getArguments(block->getNumArguments() - offset);
 
   // Find one-dimensional iteration bounds: [blockFirstIndex, blockLastIndex]:
   //   blockFirstIndex = blockIndex * blockSize
   Value blockFirstIndex = b.create<arith::MulIOp>(blockIndex, blockSize);
 
   // The last one-dimensional index in the block defined by the `blockIndex`:
-  //   blockLastIndex = max(blockFirstIndex + blockSize, tripCount) - 1
   Value blockEnd0 = b.create<arith::AddIOp>(blockFirstIndex, blockSize);
-  Value blockEnd1 =
-      b.create<arith::CmpIOp>(arith::CmpIPredicate::sge, blockEnd0, tripCount);
-  Value blockEnd2 = b.create<SelectOp>(blockEnd1, tripCount, blockEnd0);
-  Value blockLastIndex = b.create<arith::SubIOp>(blockEnd2, c1);
+  Value blockEnd1 = b.create<arith::MinSIOp>(blockEnd0, tripCount);
+  Value blockLastIndex = b.create<arith::SubIOp>(blockEnd1, c1);
 
   // Convert one-dimensional indices to multi-dimensional coordinates.
   auto blockFirstCoord = delinearize(b, blockFirstIndex, tripCounts);
   auto blockLastCoord = delinearize(b, blockLastIndex, tripCounts);
 
   // Compute loops upper bounds derived from the block last coordinates:
   //   blockEndCoord[i] = blockLastCoord[i] + 1
   //
auto dispatch = [&](OpBuilder &nestedBuilder, Location loc) {
 
     // 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)))
-    Value bs0 = b.create<arith::DivSIOp>(tripCount, maxComputeBlocks);
-    Value bs1 =
-        b.create<arith::CmpIOp>(arith::CmpIPredicate::sge, bs0, minTaskSizeCst);
-    Value bs2 = b.create<SelectOp>(bs1, bs0, minTaskSizeCst);
-    Value bs3 =
-        b.create<arith::CmpIOp>(arith::CmpIPredicate::sle, tripCount, bs2);
-    Value blockSize0 = b.create<SelectOp>(bs3, tripCount, bs2);
-    Value blockCount0 = b.create<arith::CeilDivSIOp>(tripCount, blockSize0);
-
-    // Compute balanced block size for the estimated block count.
-    Value blockSize = b.create<arith::CeilDivSIOp>(tripCount, blockCount0);
+    Value bs0 = b.create<arith::CeilDivSIOp>(tripCount, maxComputeBlocks);
+    Value bs1 = b.create<arith::MaxSIOp>(bs0, minTaskSizeCst);
+    Value blockSize = b.create<arith::MinSIOp>(tripCount, bs1);
     Value blockCount = b.create<arith::CeilDivSIOp>(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,

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

 // RUN:   mlir-opt %s -async-parallel-for                                      \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
 // RUN:               -async-runtime-ref-counting-opt                          \
+// RUN:               -arith-expand                                            \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-scf-to-std                                      \
 // RUN:               -convert-memref-to-llvm                                  \
 // RUN:               -convert-std-to-llvm                                     \
 // RUN:               -reconcile-unrealized-casts                              \
 // 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                                      \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-policy-based-ref-counting                 \
+// RUN:               -arith-expand                                            \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-scf-to-std                                      \
 // RUN:               -convert-memref-to-llvm                                  \
 // RUN:               -convert-std-to-llvm                                     \
 // RUN:               -reconcile-unrealized-casts                              \
 // 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:                                    min-task-size=1"                \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
 // RUN:               -async-runtime-ref-counting-opt                          \
+// RUN:               -arith-expand                                            \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-scf-to-std                                      \
 // RUN:               -convert-memref-to-llvm                                  \
 // RUN:               -convert-std-to-llvm                                     \
 // RUN:               -reconcile-unrealized-casts                              \
 // RUN: | mlir-cpu-runner                                                      \
 // RUN:  -e entry -entry-point-result=void -O0                                 \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \