[mlir:Async] Submit accidentally omitted changes

[ezhulenev]

Summary

Accidentally pushed old branches that did not include all the changes discussed in the PRs.

https://reviews.llvm.org/rGd43b23608ad664f02f56e965ca78916bde220950
https://reviews.llvm.org/rG86ad0af87054c3cccd68d32e103a6f1f6c6194c7

[ezhulenev]

Update Async/Passes.h

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

 #define 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);
+
 std::unique_ptr<OperationPass<ModuleOp>> createAsyncToAsyncRuntimePass();
 
 std::unique_ptr<Pass> createAsyncRuntimeRefCountingPass();
 
 std::unique_ptr<Pass> createAsyncRuntimeRefCountingOptPass();
 
 //===----------------------------------------------------------------------===//
 // Registration

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

class RuntimeCreateGroupOpLowering
     : public OpConversionPattern<RuntimeCreateGroupOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(RuntimeCreateGroupOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     TypeConverter *converter = getTypeConverter();
-    Type resultType = op->getResultTypes()[0];
+    Type resultType = op.getResult().getType();
 
     rewriter.replaceOpWithNewOp<CallOp>(
         op, kCreateGroup, converter->convertType(resultType), operands);
     return success();
   }
 };
 } // namespace
 

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

 //
 //   scf.for %block_index = %c0 to %block_count {
 //      call @parallel_compute_fn(%block_index, %block_size, ...)
 //   }
 //
 struct AsyncParallelForPass
     : public AsyncParallelForBase<AsyncParallelForPass> {
   AsyncParallelForPass() = default;
+
+  AsyncParallelForPass(bool asyncDispatch, int32_t numWorkerThreads,
+                       int32_t targetBlockSize) {
+    this->asyncDispatch = asyncDispatch;
+    this->numWorkerThreads = numWorkerThreads;
+    this->targetBlockSize = targetBlockSize;
+  }
+
   void runOnOperation() override;
 };
 
 struct AsyncParallelForRewrite : public OpRewritePattern<scf::ParallelOp> {
 public:
   AsyncParallelForRewrite(MLIRContext *ctx, bool asyncDispatch,
                           int32_t numWorkerThreads, int32_t targetBlockSize)
       : OpRewritePattern(ctx), asyncDispatch(asyncDispatch),
struct ParallelComputeFunction {
   llvm::SmallVector<Value> captures;
 };
 
 } // namespace
 
 // Converts one-dimensional iteration index in the [0, tripCount) interval
 // into multidimensional iteration coordinate.
 static SmallVector<Value> delinearize(ImplicitLocOpBuilder &b, Value index,
-                                      const SmallVector<Value> &tripCounts) {
+                                      ArrayRef<Value> tripCounts) {
   SmallVector<Value> coords(tripCounts.size());
   assert(!tripCounts.empty() && "tripCounts must be not empty");
 
   for (ssize_t i = tripCounts.size() - 1; i >= 0; --i) {
     coords[i] = b.create<SignedRemIOp>(index, tripCounts[i]);
     index = b.create<SignedDivIOp>(index, tripCounts[i]);
   }
 

 
 // Create a parallel compute fuction from the parallel operation.
 static ParallelComputeFunction
 createParallelComputeFunction(scf::ParallelOp op, PatternRewriter &rewriter) {
   OpBuilder::InsertionGuard guard(rewriter);
   ImplicitLocOpBuilder b(op.getLoc(), rewriter);
 
   ModuleOp module = op->getParentOfType<ModuleOp>();
-  b.setInsertionPointToStart(&module->getRegion(0).front());
 
   ParallelComputeFunctionType computeFuncType =
       getParallelComputeFunctionType(op, rewriter);
 
   FunctionType type = computeFuncType.type;
   FuncOp func = FuncOp::create(op.getLoc(), "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.
   Block *block = b.createBlock(&func.getBody(), func.begin(), type.getInputs());
   b.setInsertionPointToEnd(block);
 
   unsigned offset = 0; // argument offset for arguments decoding
 
-  // Load multiple arguments into values vector.
-  auto getArguments = [&](unsigned num_arguments) -> SmallVector<Value> {
-    SmallVector<Value> values(num_arguments);
-    for (unsigned i = 0; i < num_arguments; ++i)
-      values[i] = block->getArgument(offset++);
-    return values;
+  // Returns `numArguments` arguments starting from `offset` and updates offset
+  // by moving forward to the next argument.
+  auto getArguments = [&](unsigned numArguments) -> ArrayRef<Value> {
+    auto args = block->getArguments();
+    auto slice = args.drop_front(offset).take_front(numArguments);
+    offset += numArguments;
+    return {slice.begin(), slice.end()};
   };
 
   // Block iteration position defined by the block index and size.
   Value blockIndex = block->getArgument(offset++);
   Value blockSize = block->getArgument(offset++);
 
   // Constants used below.
-  Value c0 = b.create<ConstantOp>(b.getIndexAttr(0));
-  Value c1 = b.create<ConstantOp>(b.getIndexAttr(1));
+  Value c0 = b.create<ConstantIndexOp>(0);
+  Value c1 = b.create<ConstantIndexOp>(1);
 
   // Multi-dimensional parallel iteration space defined by the loop trip counts.
-  SmallVector<Value> tripCounts = getArguments(op.getNumLoops());
+  ArrayRef<Value> tripCounts = getArguments(op.getNumLoops());
 
   // Compute a product of trip counts to get the size of the flattened
   // one-dimensional iteration space.
   Value tripCount = tripCounts[0];
   for (unsigned i = 1; i < tripCounts.size(); ++i)
     tripCount = b.create<MulIOp>(tripCount, tripCounts[i]);
 
-  // Parallel operation lower bound, upper bound and step.
-  SmallVector<Value> lowerBound = getArguments(op.getNumLoops());
-  SmallVector<Value> upperBound = getArguments(op.getNumLoops());
-  SmallVector<Value> step = getArguments(op.getNumLoops());
+  // Parallel operation lower bound and step.
+  ArrayRef<Value> lowerBound = getArguments(op.getNumLoops());
+  offset += op.getNumLoops(); // skip upper bound arguments
+  ArrayRef<Value> step = getArguments(op.getNumLoops());
 
   // Remaining arguments are implicit captures of the parallel operation.
-  SmallVector<Value> captures = getArguments(block->getNumArguments() - offset);
+  ArrayRef<Value> captures = getArguments(block->getNumArguments() - offset);
 
   // Find one-dimensional iteration bounds: [blockFirstIndex, blockLastIndex]:
   //   blockFirstIndex = blockIndex * blockSize
   Value blockFirstIndex = b.create<MulIOp>(blockIndex, blockSize);
 
   // The last one-dimensional index in the block defined by the `blockIndex`:
-  //   blockLastIndex = max((blockIndex + 1) * blockSize, tripCount) - 1
-  Value blockEnd0 = b.create<AddIOp>(blockIndex, c1);
-  Value blockEnd1 = b.create<MulIOp>(blockEnd0, blockSize);
-  Value blockEnd2 = b.create<CmpIOp>(CmpIPredicate::sge, blockEnd1, tripCount);
-  Value blockEnd3 = b.create<SelectOp>(blockEnd2, tripCount, blockEnd1);
-  Value blockLastIndex = b.create<SubIOp>(blockEnd3, c1);
+  //   blockLastIndex = max(blockFirstIndex + blockSize, tripCount) - 1
+  Value blockEnd0 = b.create<AddIOp>(blockFirstIndex, blockSize);
+  Value blockEnd1 = b.create<CmpIOp>(CmpIPredicate::sge, blockEnd0, tripCount);
+  Value blockEnd2 = b.create<SelectOp>(blockEnd1, tripCount, blockEnd0);
+  Value blockLastIndex = b.create<SubIOp>(blockEnd2, c1);
 
   // Convert one-dimensional indices to multi-dimensional coordinates.
   auto blockFirstCoord = delinearize(b, blockFirstIndex, tripCounts);
   auto blockLastCoord = delinearize(b, blockLastIndex, tripCounts);
 
-  // Compute compute loops upper bounds from the block last coordinates:
+  // Compute loops upper bounds derived from the block last coordinates:
   //   blockEndCoord[i] = blockLastCoord[i] + 1
   //
   // Block first and last coordinates can be the same along the outer compute
-  // dimension when inner compute dimension containts multple blocks.
+  // dimension when inner compute dimension contains multiple blocks.
   SmallVector<Value> blockEndCoord(op.getNumLoops());
   for (size_t i = 0; i < blockLastCoord.size(); ++i)
     blockEndCoord[i] = b.create<AddIOp>(blockLastCoord[i], c1);
 
   // Construct a loop nest out of scf.for operations that will iterate over
   // all coordinates in [blockFirstCoord, blockLastCoord] range.
   using LoopBodyBuilder =
       std::function<void(OpBuilder &, Location, Value, ValueRange)>;
return [&, loopIdx](OpBuilder &nestedBuilder, Location loc, Value iv,
           lowerBound[loopIdx], nb.create<MulIOp>(iv, step[loopIdx]));
 
       // Check if we are inside first or last iteration of the loop.
       isBlockFirstCoord[loopIdx] =
           nb.create<CmpIOp>(CmpIPredicate::eq, iv, blockFirstCoord[loopIdx]);
       isBlockLastCoord[loopIdx] =
           nb.create<CmpIOp>(CmpIPredicate::eq, iv, blockLastCoord[loopIdx]);
 
-      // Check if the previous loop is in its first of last iteration.
+      // Check if the previous loop is in its first or last iteration.
       if (loopIdx > 0) {
         isBlockFirstCoord[loopIdx] = nb.create<AndOp>(
             isBlockFirstCoord[loopIdx], isBlockFirstCoord[loopIdx - 1]);
         isBlockLastCoord[loopIdx] = nb.create<AndOp>(
             isBlockLastCoord[loopIdx], isBlockLastCoord[loopIdx - 1]);
       }
 
       // Keep building loop nest.

 //
 static FuncOp createAsyncDispatchFunction(ParallelComputeFunction &computeFunc,
                                           PatternRewriter &rewriter) {
   OpBuilder::InsertionGuard guard(rewriter);
   Location loc = computeFunc.func.getLoc();
   ImplicitLocOpBuilder b(loc, rewriter);
 
   ModuleOp module = computeFunc.func->getParentOfType<ModuleOp>();
-  b.setInsertionPointToStart(&module->getRegion(0).front());
 
   ArrayRef<Type> computeFuncInputTypes =
       computeFunc.func.type().cast<FunctionType>().getInputs();
 
   // Compared to the parallel compute function async dispatch function takes
   // additional !async.group argument. Also instead of a single `blockIndex` it
   // takes `blockStart` and `blockEnd` arguments to define the range of
   // dispatched blocks.
static FuncOp createAsyncDispatchFunction(ParallelComputeFunction &computeFunc,
   symbolTable.insert(func);
   rewriter.getListener()->notifyOperationInserted(func);
 
   // Create function entry block.
   Block *block = b.createBlock(&func.getBody(), func.begin(), type.getInputs());
   b.setInsertionPointToEnd(block);
 
   Type indexTy = b.getIndexType();
-  Value c1 = b.create<ConstantOp>(b.getIndexAttr(1));
-  Value c2 = b.create<ConstantOp>(b.getIndexAttr(2));
+  Value c1 = b.create<ConstantIndexOp>(1);
+  Value c2 = b.create<ConstantIndexOp>(2);
 
   // Get the async group that will track async dispatch completion.
   Value group = block->getArgument(0);
 
   // Get the block iteration range: [blockStart, blockEnd)
   Value blockStart = block->getArgument(1);
   Value blockEnd = block->getArgument(2);
 
// `after` block and launch one more async dispatch.
     Value start = before->getArgument(0);
     Value end = before->getArgument(1);
     Value distance = b.create<SubIOp>(end, start);
     Value dispatch = b.create<CmpIOp>(CmpIPredicate::sgt, distance, c1);
     b.create<scf::ConditionOp>(dispatch, before->getArguments());
   }
 
   // Setup the async dispatch loop body: recursively call dispatch function
-  // for second the half of the original range and go to the next iteration.
+  // for the seconds half of the original range and go to the next iteration.
   {
     b.setInsertionPointToEnd(after);
     Value start = after->getArgument(0);
     Value end = after->getArgument(1);
     Value distance = b.create<SubIOp>(end, start);
     Value halfDistance = b.create<SignedDivIOp>(distance, c2);
-    Value midIndex = b.create<AddIOp>(after->getArgument(0), halfDistance);
+    Value midIndex = b.create<AddIOp>(start, halfDistance);
 
     // Call parallel compute function inside the async.execute region.
     auto executeBodyBuilder = [&](OpBuilder &executeBuilder,
                                   Location executeLoc, ValueRange executeArgs) {
       // Update the original `blockStart` and `blockEnd` with new range.
       SmallVector<Value> operands{block->getArguments().begin(),
                                   block->getArguments().end()};
       operands[1] = midIndex;
       operands[2] = end;
 
       executeBuilder.create<CallOp>(executeLoc, func.sym_name(),
                                     func.getCallableResults(), operands);
       executeBuilder.create<async::YieldOp>(executeLoc, ValueRange());
     };
 
     // Create async.execute operation to dispatch half of the block range.
     auto execute = b.create<ExecuteOp>(TypeRange(), ValueRange(), ValueRange(),
                                        executeBodyBuilder);
     b.create<AddToGroupOp>(indexTy, execute.token(), group);
-    b.create<scf::YieldOp>(ValueRange({after->getArgument(0), midIndex}));
+    b.create<scf::YieldOp>(ValueRange({start, midIndex}));
   }
 
   // After dispatching async operations to process the tail of the block range
   // call the parallel compute function for the first block of the range.
   b.setInsertionPointAfter(whileOp);
 
   // Drop async dispatch specific arguments: async group, block start and end.
   auto forwardedInputs = block->getArguments().drop_front(3);
static void doAsyncDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
                             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);
 
-  Value c0 = b.create<ConstantOp>(b.getIndexAttr(0));
-  Value c1 = b.create<ConstantOp>(b.getIndexAttr(1));
+  Value c0 = b.create<ConstantIndexOp>(0);
+  Value c1 = b.create<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<SubIOp>(blockCount, c1);
   Value group = b.create<CreateGroupOp>(GroupType::get(ctx), groupSize);
 
   // Pack the async dispath function operands to launch the work splitting.

 doSequantialDispatch(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<ConstantOp>(b.getIndexAttr(0));
-  Value c1 = b.create<ConstantOp>(b.getIndexAttr(1));
+  Value c0 = b.create<ConstantIndexOp>(0);
+  Value c1 = b.create<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<SubIOp>(blockCount, c1);
   Value group = b.create<CreateGroupOp>(GroupType::get(ctx), groupSize);
 
   // Call parallel compute function for all blocks.
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]);
 
-  auto indexTy = b.getIndexType();
-
   // Do not overload worker threads with too many compute blocks.
-  Value maxComputeBlocks = b.create<ConstantOp>(
-      indexTy, b.getIndexAttr(numWorkerThreads * kMaxOversharding));
+  Value maxComputeBlocks =
+      b.create<ConstantIndexOp>(numWorkerThreads * kMaxOversharding);
 
   // Target block size from the pass parameters.
-  Value targetComputeBlockSize =
-      b.create<ConstantOp>(indexTy, b.getIndexAttr(targetBlockSize));
+  Value targetComputeBlockSize = b.create<ConstantIndexOp>(targetBlockSize);
 
   // Compute parallel block size from the parallel problem size:
   //   blockSize = min(tripCount,
-  //                   max(divup(tripCount, maxComputeBlocks),
+  //                   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 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,
                          blockCount, tripCounts);
   }
 
-  // Parallel operation was replaces with a block iteration loop.
+  // 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);
 
   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) {
+  return std::make_unique<AsyncParallelForPass>(asyncDispatch, numWorkerThreads,
+                                                targetBlockSize);
+}

mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir

func @loop_1d(%arg0: index, %arg1: index, %arg2: index, %arg3: memref<?xf32>) {
   return
 }
 
 // CHECK-LABEL: func private @parallel_compute_fn
 // CHECK:       scf.for
 // CHECK:         memref.store
 
 // CHECK-LABEL: func private @async_dispatch_fn
+// CHECK-SAME:  (
 // CHECK-SAME:    %[[GROUP:arg0]]: !async.group,
 // CHECK-SAME:    %[[BLOCK_START:arg1]]: index
 // CHECK-SAME:    %[[BLOCK_END:arg2]]: index
-
-// CHECK:         scf.while (%[[S:.*]] = %[[BLOCK_START]],
-// CHECK-SAME:               %[[E:.*]] = %[[BLOCK_END]])
+// CHECK-SAME:  )
+// CHECK:         %[[C1:.*]] = constant 1 : index
+// CHECK:         %[[C2:.*]] = constant 2 : index
+// CHECK:         scf.while (%[[S0:.*]] = %[[BLOCK_START]],
+// CHECK-SAME:               %[[E0:.*]] = %[[BLOCK_END]])
+// While loop `before` block decides if we need to dispatch more tasks.
+// CHECK:         {
+// CHECK:           %[[DIFF0:.*]] = subi %[[E0]], %[[S0]]
+// CHECK:           %[[COND:.*]] = cmpi sgt, %[[DIFF0]], %[[C1]]
+// CHECK:           scf.condition(%[[COND]])
+// While loop `after` block splits the range in half and submits async task
+// to process the second half using the call to the same dispatch function.
 // CHECK:         } do {
+// CHECK:         ^bb0(%[[S1:.*]]: index, %[[E1:.*]]: index):
+// CHECK:           %[[DIFF1:.*]] = subi %[[E1]], %[[S1]]
+// CHECK:           %[[HALF:.*]] = divi_signed %[[DIFF1]], %[[C2]]
+// CHECK:           %[[MID:.*]] = addi %[[S1]], %[[HALF]]
 // CHECK:           %[[TOKEN:.*]] = async.execute
 // CHECK:             call @async_dispatch_fn
-// CHECK:             async.add_to_group
+// CHECK:           async.add_to_group
+// CHECK:           scf.yield %[[S1]], %[[MID]]
 // CHECK:         }
-
+// After async dispatch the first block processed in the caller thread.
 // CHECK:         call @parallel_compute_fn(%[[BLOCK_START]]
 
 // -----
 
 // CHECK-LABEL: @loop_2d
 func @loop_2d(%arg0: index, %arg1: index, %arg2: index, // lb, ub, step
               %arg3: index, %arg4: index, %arg5: index, // lb, ub, step
               %arg6: memref<?x?xf32>) {

mlir/test/Dialect/Async/async-parallel-for-seq-dispatch.mlir

 // RUN: mlir-opt %s -split-input-file -async-parallel-for=async-dispatch=false  \
 // RUN: | FileCheck %s --dump-input=always
 
+// The structure of @parallel_compute_fn checked in the async dispatch test.
+// Here we only check the structure of the sequential dispatch loop.
+
 // CHECK-LABEL: @loop_1d
 func @loop_1d(%arg0: index, %arg1: index, %arg2: index, %arg3: memref<?xf32>) {
   // CHECK: %[[GROUP:.*]] = async.create_group
   // CHECK: scf.for
   // CHECK:   %[[TOKEN:.*]] = async.execute
   // CHECK:     call @parallel_compute_fn
   // CHECK:     async.yield
   // CHECK:   async.add_to_group %[[TOKEN]], %[[GROUP]]