diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp --- a/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp +++ b/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp @@ -615,13 +615,13 @@ }; 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(blockCount, c1); - Value group = b.create(GroupType::get(ctx), groupSize); - - ImplicitLocOpBuilder nb(loc, nestedBuilder); + Value groupSize = nb.create(blockCount, c1); + Value group = nb.create(GroupType::get(ctx), groupSize); // Launch async dispatch function for [0, blockCount) range. SmallVector operands = {group, c0, blockCount, blockSize}; @@ -631,7 +631,7 @@ asyncDispatchFunction.getCallableResults(), operands); // Wait for the completion of all parallel compute operations. - b.create(group); + nb.create(group); nb.create(); }; diff --git a/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir b/mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir --- 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 @@ // 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