[mlir:Async] Remove async operations if it is statically known that the parallel operation has a single compute block

[ezhulenev]

Summary

Depends On D104850

Add a test that verifies that canonicalization removes all async overheads if it is statically known that the scf.parallel operation will be computed using a single block.

[ezhulenev]

Add more checks to the test

[ezhulenev]

Revert accidental changes

[ezhulenev]

Rebased

[ezhulenev]

Rebased

[bondhugula]

mlir/lib/Dialect/Async/IR/Async.cpp

252–275

This can be in the op's folding hook as opposed to a canonicalization pattern.

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

1516–1525 ↗

(On Diff #354771)

Any reason this is moved up? It's no longer in alphabetical order.

[ezhulenev]

Remove irrelevant files

mlir/lib/Dialect/Async/IR/Async.cpp

252–275

Can it? I think folding hook can only return an attribute, but it doesn't have access to rewriter to erase operations.

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

1516–1525 ↗

(On Diff #354771)

That's the result of out internal automation tool, I removed these changes from this PR.

[bondhugula]

mlir/lib/Dialect/Async/IR/Async.cpp

252–275

Sorry, I had missed the erasure of AwaitAllOps. The folding hook can't of course be used to erase other ops.

[mehdi_amini]

mlir/lib/Dialect/Async/IR/Async.cpp

252–275

The folding hook can return an attribute or an existing SSA value in the IR, it shouldn't create new IR.
This value will be used to replace the existing result, and the original op will be deleted by the folder.

[ezhulenev]

Drop reference if the value has no uses

[ezhulenev]

Rebase

[herhut]

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

568

Nit: I think the comment about async groups is hard to understand outside the scope of this review, where one sees all this in one place. It is also not very relevant here.

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

14

nit: synchronous.
2nd nit: There is no function call anymore. Maybe also mention in the comment that we have a single call that gets inlined.

mlir/include/mlir/Dialect/Async/IR/Async.h

 #define MLIR_DIALECT_ASYNC_IR_ASYNC_H
 
 #include "mlir/Dialect/Async/IR/AsyncTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // Async Dialect
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Async/IR/AsyncOpsDialect.h.inc"

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

let description = [{
     ...
     async.await_all %group
     ```
   }];
 
   let arguments = (ins Index:$size);
   let results = (outs Async_GroupType:$result);
 
+  let hasCanonicalizeMethod = 1;
+
   let assemblyFormat = "$size `:` type($result) attr-dict";
 }
 
 def Async_AddToGroupOp : Async_Op<"add_to_group", []> {
   let summary = "adds and async token or value to the group";
   let description = [{
     The `async.add_to_group` adds an async token or value to the async group.
     Returns the rank of the added element in the group. This rank is fixed

mlir/lib/Dialect/Async/IR/Async.cpp

static LogicalResult verify(ExecuteOp op) {
   if (op.body().getArgumentTypes() != unwrappedTypes)
     return op.emitOpError("async body region argument types do not match the "
                           "execute operation arguments types");
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+/// CreateGroupOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult CreateGroupOp::canonicalize(CreateGroupOp op,
+                                          PatternRewriter &rewriter) {
+  // Find all `await_all` users of the group.
+  llvm::SmallVector<AwaitAllOp> awaitAllUsers;
+
+  auto isAwaitAll = [&](Operation *op) -> bool {
+    if (AwaitAllOp awaitAll = dyn_cast<AwaitAllOp>(op)) {
+      awaitAllUsers.push_back(awaitAll);
+      return true;
+    }
+    return false;
+  };
+
+  // Check if all users of the group are `await_all` operations.
+  if (!llvm::all_of(op->getUsers(), isAwaitAll))
+    return failure();
+
+  // If group is only awaited without adding anything to it, we can safely erase
+  // the create operation and all users.
+  for (AwaitAllOp awaitAll : awaitAllUsers)
+    rewriter.eraseOp(awaitAll);
+  rewriter.eraseOp(op);
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 /// AwaitOp
 //===----------------------------------------------------------------------===//
 
 void AwaitOp::build(OpBuilder &builder, OperationState &result, Value operand,
                     ArrayRef<NamedAttribute> attrs) {
   result.addOperands({operand});
   result.attributes.append(attrs.begin(), attrs.end());
 

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

 //       async.execute { call @async_dispatch(%mid_index, %block_end); }
 //       %block_end = %mid_index
 //     }
 //
 //     // Call parallel compute function for a single block.
 //     call @parallel_compute_fn(%block_start, %block_size, ...);
 //   }
 //
-static FuncOp createAsyncDispatchFunction(ParallelComputeFunction &computeFunc,
-                                          PatternRewriter &rewriter) {
+static FuncOp
+createAsingleBlockDispatchFunction(ParallelComputeFunction &computeFunc,
   OpBuilder::InsertionGuard guard(rewriter);
   Location loc = computeFunc.func.getLoc();
   ImplicitLocOpBuilder b(loc, rewriter);
 
   ModuleOp module = computeFunc.func->getParentOfType<ModuleOp>();
 
   ArrayRef<Type> computeFuncInputTypes =
       computeFunc.func.type().cast<FunctionType>().getInputs();
createAsingleBlockDispatchFunction(ParallelComputeFunction &computeFunc,
   b.create<CallOp>(computeFunc.func.sym_name(),
                    computeFunc.func.getCallableResults(), computeFuncOperands);
   b.create<ReturnOp>(ValueRange());
 
   return func;
 }
 
 // Launch async dispatch of the parallel compute function.
-static void doAsyncDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
-                            ParallelComputeFunction &parallelComputeFunction,
-                            scf::ParallelOp op, Value blockSize,
-                            Value blockCount,
-                            const SmallVector<Value> &tripCounts) {
+static void
+doAsingleBlockDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
+                       scf::ParallelOp op, Value blockSize, Value blockCount,
   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);
+      createAsingleBlockDispatchFunction(parallelComputeFunction, rewriter);
 
   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.
-  SmallVector<Value> asyncDispatchOperands = {group, c0, blockCount, blockSize};
-  asyncDispatchOperands.append(tripCounts);
-  asyncDispatchOperands.append(op.lowerBound().begin(), op.lowerBound().end());
-  asyncDispatchOperands.append(op.upperBound().begin(), op.upperBound().end());
-  asyncDispatchOperands.append(op.step().begin(), op.step().end());
-  asyncDispatchOperands.append(parallelComputeFunction.captures);
+  // Appends operands shared by async dispatch and parallel compute functions to
+  // the given operands vector.
+  auto appendBlockComputeOperands = [&](SmallVector<Value> &operands) {
+    operands.append(tripCounts);
+    operands.append(op.lowerBound().begin(), op.lowerBound().end());
+    operands.append(op.upperBound().begin(), op.upperBound().end());
+    operands.append(op.step().begin(), op.step().end());
+    operands.append(parallelComputeFunction.captures);
+  };
+
+  // Check if the block size is one, in this case we can skip the async dispatch
+  // completely. If this will be known statically, then canonicalization will
+  // erase async group operations.
+  Value isSingleBlock = b.create<CmpIOp>(CmpIPredicate::eq, blockCount, c1);
+
+  auto syncDispatch = [&](OpBuilder &nestedBuilder, Location loc) {
+    ImplicitLocOpBuilder nb(loc, nestedBuilder);
+
+    // Call parallel compute function for the single block.
+    SmallVector<Value> operands = {c0, blockSize};
+    appendBlockComputeOperands(operands);
+
+    nb.create<CallOp>(parallelComputeFunction.func.sym_name(),
+                      parallelComputeFunction.func.getCallableResults(),
+                      operands);
+    nb.create<scf::YieldOp>();
+  };
+
+  auto asyncDispatch = [&](OpBuilder &nestedBuilder, Location loc) {
+    ImplicitLocOpBuilder nb(loc, nestedBuilder);
 
     // Launch async dispatch function for [0, blockCount) range.
-  b.create<CallOp>(asyncDispatchFunction.sym_name(),
-                   asyncDispatchFunction.getCallableResults(),
-                   asyncDispatchOperands);
+    SmallVector<Value> operands = {group, c0, blockCount, blockSize};
+    appendBlockComputeOperands(operands);
+
+    nb.create<CallOp>(asyncDispatchFunction.sym_name(),
+                      asyncDispatchFunction.getCallableResults(), operands);
+    nb.create<scf::YieldOp>();
+  };
+
+  // Dispatch either single block compute function, or launch async dispatch.
+  b.create<scf::IfOp>(TypeRange(), isSingleBlock, syncDispatch, asyncDispatch);
 
   // 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,
                      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);
   Value c1 = b.create<ConstantIndexOp>(1);
AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
   // 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,
+    doAsingleBlockDispatch(b, rewriter, parallelComputeFunction, op, blockSize,
                            blockCount, tripCounts);
   } else {
     doSequantialDispatch(b, rewriter, parallelComputeFunction, op, blockSize,
                          blockCount, tripCounts);
   }
 
   // Parallel operation was replaced with a block iteration loop.
   rewriter.eraseOp(op);
 

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

 // RUN: mlir-opt %s -split-input-file -async-parallel-for=async-dispatch=true  \
 // RUN: | FileCheck %s
 
 // CHECK-LABEL: @loop_1d
 func @loop_1d(%arg0: index, %arg1: index, %arg2: index, %arg3: memref<?xf32>) {
+  // CHECK: %[[C0:.*]] = constant 0 : index
   // CHECK: %[[GROUP:.*]] = async.create_group
-  // CHECK: call @async_dispatch_fn
+  // CHECK: scf.if {{.*}} {
+  // CHECK:   call @parallel_compute_fn(%[[C0]]
+  // CHECK: } else {
+  // CHECK:   call @async_dispatch_fn
+  // CHECK: }
   // CHECK: async.await_all %[[GROUP]]
   scf.parallel (%i) = (%arg0) to (%arg1) step (%arg2) {
     %one = constant 1.0 : f32
     memref.store %one, %arg3[%i] : memref<?xf32>
   }
   return
 }
 

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

+// RUN: mlir-opt %s                                                            \
+// RUN:    -async-parallel-for=async-dispatch=true                             \
+// RUN:    -canonicalize -inline -symbol-dce                                   \
+// RUN: | FileCheck %s
+
+// RUN: mlir-opt %s                                                            \
+// RUN:    -async-parallel-for=async-dispatch=false                            \
+// RUN:    -canonicalize -inline -symbol-dce                                   \
+// RUN: | FileCheck %s
+
+// Check that if we statically know that the parallel operation has a single
+// block then all async operations will be canonicalized away and we will
+// end up with a single synchonous compute function call.
+
+// CHECK-LABEL: @loop_1d(
+// CHECK:       %[[MEMREF:.*]]: memref<?xf32>
+func @loop_1d(%arg0: memref<?xf32>) {
+  // CHECK-DAG: %[[C0:.*]] = constant 0 : index
+  // CHECK-DAG: %[[C1:.*]] = constant 1 : index
+  // CHECK-DAG: %[[C100:.*]] = constant 100 : index
+  // CHECK-DAG: %[[ONE:.*]] = constant 1.000000e+00 : f32
+  // CHECK:     scf.for %[[I:.*]] = %[[C0]] to %[[C100]] step %[[C1]]
+  // CHECK:       memref.store %[[ONE]], %[[MEMREF]][%[[I]]]
+  %lb = constant 0 : index
+  %ub = constant 100 : index
+  %st = constant 1 : index
+  scf.parallel (%i) = (%lb) to (%ub) step (%st) {
+    %one = constant 1.0 : f32
+    memref.store %one, %arg0[%i] : memref<?xf32>
+  }
+
+  return
+}