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

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

auto syncDispatch = [&](OpBuilder &nestedBuilder, Location loc) {
 
     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);
+
     // 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};
     appendBlockComputeOperands(operands);
 
     nb.create<CallOp>(asyncDispatchFunction.sym_name(),
                       asyncDispatchFunction.getCallableResults(), operands);
 
     // Wait for the completion of all parallel compute operations.
-    b.create<AwaitAllOp>(group);
+    nb.create<AwaitAllOp>(group);
 
     nb.create<scf::YieldOp>();
   };
 
   // Dispatch either single block compute function, or launch async dispatch.
   b.create<scf::IfOp>(TypeRange(), isSingleBlock, syncDispatch, asyncDispatch);
 }
 

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 --dump-input=always
 
 // CHECK-LABEL: @loop_1d(
 // CHECK-SAME:    %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index
 func @loop_1d(%arg0: index, %arg1: index, %arg2: index, %arg3: memref<?xf32>) {
   // CHECK:      %[[C0:.*]] = arith.constant 0 : index
 
   // CHECK:      %[[RANGE:.*]] = arith.subi %[[UB]], %[[LB]]
   // CHECK:      %[[TRIP_CNT:.*]] = arith.ceildivsi %[[RANGE]], %[[STEP]]
   // CHECK:      %[[IS_NOOP:.*]] = arith.cmpi eq, %[[TRIP_CNT]], %[[C0]] : index
 
   // CHECK:      scf.if %[[IS_NOOP]] {
   // CHECK-NEXT: } else {
-    // CHECK:        scf.if {{.*}} {
   // CHECK:          call @parallel_compute_fn(%[[C0]]
   // CHECK:        } else {
   // CHECK:          %[[GROUP:.*]] = async.create_group
   // CHECK:          call @async_dispatch_fn
   // CHECK:          async.await_all %[[GROUP]]
   // CHECK:        }
   // CHECK:      }
   scf.parallel (%i) = (%arg0) to (%arg1) step (%arg2) {