[Mlir-commits] [mlir] b171583 - [mlir] Async: create async.group inside the scf.if branch

[Eugene Zhulenev]

Author: Eugene Zhulenev
Date: 2022-02-16T14:47:04-08:00
New Revision: b171583ae7d187164b40e8dcd90b9aa2f34c06a2

URL: https://github.com/llvm/llvm-project/commit/b171583ae7d187164b40e8dcd90b9aa2f34c06a2
DIFF: https://github.com/llvm/llvm-project/commit/b171583ae7d187164b40e8dcd90b9aa2f34c06a2.diff

LOG: [mlir] Async: create async.group inside the scf.if branch

Reviewed By: cota

Differential Revision: https://reviews.llvm.org/D119959

Added: 
    

Modified: 
    mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
    mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
index fd98f31e09feb..cdd85e5c5b406 100644
--- a/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
+++ b/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
@@ -615,13 +615,13 @@ static void doAsyncDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
   };
 
   auto asyncDispatch = [&](OpBuilder &nestedBuilder, Location loc) {
+    ImplicitLocOpBuilder nb(loc, nestedBuilder);
+
     // 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<arith::SubIOp>(blockCount, c1);
-    Value group = b.create<CreateGroupOp>(GroupType::get(ctx), groupSize);
-
-    ImplicitLocOpBuilder nb(loc, nestedBuilder);
+    Value groupSize = nb.create<arith::SubIOp>(blockCount, c1);
+    Value group = nb.create<CreateGroupOp>(GroupType::get(ctx), groupSize);
 
     // Launch async dispatch function for [0, blockCount) range.
     SmallVector<Value> operands = {group, c0, blockCount, blockSize};
@@ -631,7 +631,7 @@ static void doAsyncDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
                       asyncDispatchFunction.getCallableResults(), operands);
 
     // Wait for the completion of all parallel compute operations.
-    b.create<AwaitAllOp>(group);
+    nb.create<AwaitAllOp>(group);
 
     nb.create<scf::YieldOp>();
   };

diff  --git a/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir b/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir
index bf2f7404abb1e..3fdbcf0ca153b 100644
--- a/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir
+++ b/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir
@@ -12,7 +12,7 @@ func @loop_1d(%arg0: index, %arg1: index, %arg2: index, %arg3: memref<?xf32>) {
 
   // CHECK:      scf.if %[[IS_NOOP]] {
   // CHECK-NEXT: } else {
-    // CHECK:        scf.if {{.*}} {
+  // CHECK:        scf.if {{.*}} {
   // CHECK:          call @parallel_compute_fn(%[[C0]]
   // CHECK:        } else {
   // CHECK:          %[[GROUP:.*]] = async.create_group