[mlir:Async] Add the size parameter to the async.group

[ezhulenev]

Summary

Specify the !async.group size (the number of tokens that will be added to it) at construction time. async.await_all operation can potentially race with async.execute operations that keep updating the group, for this reason it is required to know upfront how many tokens will be added to the group.

[ezhulenev]

Update create_grop documentation in the td file

[ftynse]

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

599

Nit: op.result().getType() is slightly more readable than magic 0.

[ezhulenev]

Resolve comments

[ezhulenev]

Rebase

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

def Async_AwaitOp : Async_Op<"await"> {
 
   let verifier = [{ return ::verify(*this); }];
 }
 
 def Async_CreateGroupOp : Async_Op<"create_group", [NoSideEffect]> {
   let summary = "creates an empty async group";
   let description = [{
     The `async.create_group` allocates an empty async group. Async tokens or
-    values can be added to this group later.
+    values can be added to this group later. The size of the group must be
+    specified at construction time, and `await_all` operation will first
+    wait until the number of added tokens or values reaches the group size.
 
     Example:
 
     ```mlir
-    %0 = async.create_group
+    %size = ... : index
+    %group = async.create_group %size : !async.group
     ...
-    async.await_all %0
+    async.await_all %group
     ```
   }];
 
+  let arguments = (ins Index:$size);
   let results = (outs Async_GroupType:$result);
 
-  let assemblyFormat = "attr-dict";
+  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
     for the group lifetime.
 
     Example:
 
     ```mlir
-    %0 = async.create_group
+    %0 = async.create_group %size : !async.group
     %1 = ... : !async.token
     %2 = async.add_to_group %1, %0 : !async.token
     ```
   }];
 
   let arguments = (ins Async_AnyValueOrTokenType:$operand,
                        Async_GroupType:$group);
   let results = (outs Index:$rank);
 
   let assemblyFormat = "$operand `,` $group `:` type($operand) attr-dict";
 }
 
 def Async_AwaitAllOp : Async_Op<"await_all", []> {
   let summary = "waits for the all async tokens or values in the group to "
                 "become ready";
   let description = [{
     The `async.await_all` operation waits until all the tokens or values in the
     group become ready.
 
     Example:
 
     ```mlir
-    %0 = async.create_group
+    %0 = async.create_group %size : !async.group
 
     %1 = ... : !async.token
     %2 = async.add_to_group %1, %0 : !async.token
 
     %3 = ... : !async.token
     %4 = async.add_to_group %2, %0 : !async.token
 
     async.await_all %0

 // Async Dialect Runtime Operations.
 //===----------------------------------------------------------------------===//
 
 // The following operations are intermediate async dialect operations to help
 // lowering from high level async operation like `async.execute` to the Async
 // Runtime API defined in the `ExecutionEngine/AsyncRuntime.h`.
 
 def Async_RuntimeCreateOp : Async_Op<"runtime.create"> {
-  let summary = "creates an async runtime value (token, value or group)";
+  let summary = "creates an async runtime token or value";
   let description = [{
-    The `async.runtime.create` operation creates an async dialect value
-    (token, value or group). Tokens and values are created in non-ready state.
-    Groups are created in empty state.
+    The `async.runtime.create` operation creates an async dialect token or
+    value. Tokens and values are created in the non-ready state.
   }];
 
-  let results = (outs Async_AnyAsyncType:$result);
+  let results = (outs Async_AnyValueOrTokenType:$result);
   let assemblyFormat = "attr-dict `:` type($result)";
 }
 
+def Async_RuntimeCreateGroupOp : Async_Op<"runtime.create_group"> {
+  let summary = "creates an async runtime group";
+  let description = [{
+    The `async.runtime.create_group` operation creates an async dialect group
+    of the given size. Group created in the empty state.
+  }];
+
+  let arguments = (ins Index:$size);
+  let results = (outs Async_GroupType:$result);
+  let assemblyFormat = "$size `:` type($result) attr-dict ";
+}
+
 def Async_RuntimeSetAvailableOp : Async_Op<"runtime.set_available"> {
   let summary = "switches token or value to available state";
   let description = [{
     The `async.runtime.set_available` operation switches async token or value
     state to available.
   }];
 
   let arguments = (ins Async_AnyValueOrTokenType:$operand);

mlir/include/mlir/ExecutionEngine/AsyncRuntime.h

 
 // Create a new `async.value` in not-ready state. Size parameter specifies the
 // number of bytes that will be allocated for the async value storage. Storage
 // is owned by the `async.value` and deallocated when the async value is
 // destructed (reference count drops to zero).
 extern "C" AsyncValue *mlirAsyncRuntimeCreateValue(int32_t);
 
 // Create a new `async.group` in empty state.
-extern "C" AsyncGroup *mlirAsyncRuntimeCreateGroup();
+extern "C" AsyncGroup *mlirAsyncRuntimeCreateGroup(int64_t size);
 
 extern "C" int64_t mlirAsyncRuntimeAddTokenToGroup(AsyncToken *, AsyncGroup *);
 
 // Switches `async.token` to ready state and runs all awaiters.
 extern "C" void mlirAsyncRuntimeEmplaceToken(AsyncToken *);
 
 // Switches `async.value` to ready state and runs all awaiters.
 extern "C" void mlirAsyncRuntimeEmplaceValue(AsyncValue *);

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

struct AsyncAPI {
 
   static FunctionType createValueFunctionType(MLIRContext *ctx) {
     auto i32 = IntegerType::get(ctx, 32);
     auto value = opaquePointerType(ctx);
     return FunctionType::get(ctx, {i32}, {value});
   }
 
   static FunctionType createGroupFunctionType(MLIRContext *ctx) {
-    return FunctionType::get(ctx, {}, {GroupType::get(ctx)});
+    auto i64 = IntegerType::get(ctx, 64);
+    return FunctionType::get(ctx, {i64}, {GroupType::get(ctx)});
   }
 
   static FunctionType getValueStorageFunctionType(MLIRContext *ctx) {
     auto value = opaquePointerType(ctx);
     auto storage = opaquePointerType(ctx);
     return FunctionType::get(ctx, {value}, {storage});
   }
 
public:
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(RuntimeCreateOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     TypeConverter *converter = getTypeConverter();
     Type resultType = op->getResultTypes()[0];
 
-    // Tokens and Groups lowered to function calls without arguments.
-    if (resultType.isa<TokenType>() || resultType.isa<GroupType>()) {
-      rewriter.replaceOpWithNewOp<CallOp>(
-          op, resultType.isa<TokenType>() ? kCreateToken : kCreateGroup,
+    // Tokens creation maps to a simple function call.
+    if (resultType.isa<TokenType>()) {
+      rewriter.replaceOpWithNewOp<CallOp>(op, kCreateToken,
                                           converter->convertType(resultType));
       return success();
     }
 
     // To create a value we need to compute the storage requirement.
     if (auto value = resultType.dyn_cast<ValueType>()) {
       // Returns the size requirements for the async value storage.
       auto sizeOf = [&](ValueType valueType) -> Value {
         auto loc = op->getLoc();
matchAndRewrite(RuntimeCreateOp op, ArrayRef<Value> operands,
     }
 
     return rewriter.notifyMatchFailure(op, "unsupported async type");
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
+// Convert async.runtime.create_group to the corresponding runtime API call.
+//===----------------------------------------------------------------------===//
+
+namespace {
+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];
+
+    rewriter.replaceOpWithNewOp<CallOp>(
+        op, kCreateGroup, converter->convertType(resultType), operands);
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
 // Convert async.runtime.set_available to the corresponding runtime API call.
 //===----------------------------------------------------------------------===//
 
 namespace {
 class RuntimeSetAvailableOpLowering
     : public OpConversionPattern<RuntimeSetAvailableOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
void ConvertAsyncToLLVMPass::runOnOperation() {
   patterns.add<RuntimeSetAvailableOpLowering, RuntimeSetErrorOpLowering,
                RuntimeIsErrorOpLowering, RuntimeAwaitOpLowering,
                RuntimeAwaitAndResumeOpLowering, RuntimeResumeOpLowering,
                RuntimeAddToGroupOpLowering, RuntimeAddRefOpLowering,
                RuntimeDropRefOpLowering>(converter, ctx);
 
   // Lower async.runtime operations that rely on LLVM type converter to convert
   // from async value payload type to the LLVM type.
-  patterns.add<RuntimeCreateOpLowering, RuntimeStoreOpLowering,
-               RuntimeLoadOpLowering>(llvmConverter, ctx);
+  patterns.add<RuntimeCreateOpLowering, RuntimeCreateGroupOpLowering,
+               RuntimeStoreOpLowering, RuntimeLoadOpLowering>(llvmConverter,
+                                                              ctx);
 
   // Lower async coroutine operations to LLVM coroutine intrinsics.
   patterns
       .add<CoroIdOpConversion, CoroBeginOpConversion, CoroFreeOpConversion,
            CoroEndOpConversion, CoroSaveOpConversion, CoroSuspendOpConversion>(
           converter, ctx);
 
   ConversionTarget target(*ctx);

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

AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
 
   // Assign remaining available concurrency to other induction variables.
   for (size_t i = 1; i < op.getNumLoops(); ++i) {
     targetNumBlocks[i] = divup(targetNumBlocks[i - 1], numBlocks[i - 1]);
     blockSize[i] = divup(tripCounts[i], targetNumBlocks[i]);
     numBlocks[i] = divup(tripCounts[i], blockSize[i]);
   }
 
+  // Total number of async compute blocks.
+  Value totalBlocks = numBlocks[0];
+  for (size_t i = 1; i < op.getNumLoops(); ++i)
+    totalBlocks = rewriter.create<MulIOp>(loc, totalBlocks, numBlocks[i]);
+
   // Create an async.group to wait on all async tokens from async execute ops.
-  auto group = rewriter.create<CreateGroupOp>(loc, GroupType::get(ctx));
+  auto group =
+      rewriter.create<CreateGroupOp>(loc, GroupType::get(ctx), totalBlocks);
 
   // Build a scf.for loop nest from the parallel operation.
 
   // Lower/upper bounds for nest block level computations.
   SmallVector<Value, 4> blockLowerBounds(op.getNumLoops());
   SmallVector<Value, 4> blockUpperBounds(op.getNumLoops());
   SmallVector<Value, 4> blockInductionVars(op.getNumLoops());
 

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

auto callOutlinedFunc = callBuilder.create<CallOp>(
       func.getName(), execute.getResultTypes(), functionInputs.getArrayRef());
   execute.replaceAllUsesWith(callOutlinedFunc.getResults());
   execute.erase();
 
   return {func, coro};
 }
 
 //===----------------------------------------------------------------------===//
-// Convert async.create_group operation to async.runtime.create
+// Convert async.create_group operation to async.runtime.create_group
 //===----------------------------------------------------------------------===//
 
 namespace {
 class CreateGroupOpLowering : public OpConversionPattern<CreateGroupOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(CreateGroupOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<RuntimeCreateOp>(
-        op, GroupType::get(op->getContext()));
+    rewriter.replaceOpWithNewOp<RuntimeCreateGroupOp>(
+        op, GroupType::get(op->getContext()), operands);
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Convert async.add_to_group operation to async.runtime.add_to_group.
 //===----------------------------------------------------------------------===//

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

struct AsyncValue : public RefCounted {
   std::condition_variable cv;
   std::vector<std::function<void()>> awaiters;
 };
 
 // Async group provides a mechanism to group together multiple async tokens or
 // values to await on all of them together (wait for the completion of all
 // tokens or values added to the group).
 struct AsyncGroup : public RefCounted {
-  AsyncGroup(AsyncRuntime *runtime)
-      : RefCounted(runtime), pendingTokens(0), numErrors(0), rank(0) {}
+  AsyncGroup(AsyncRuntime *runtime, int64_t size)
+      : RefCounted(runtime), pendingTokens(size), numErrors(0), rank(0) {}
 
   std::atomic<int> pendingTokens;
   std::atomic<int> numErrors;
   std::atomic<int> rank;
 
   // Pending awaiters are guarded by a mutex.
   std::mutex mu;
   std::condition_variable cv;

 
 // Creates a new `async.value` in not-ready state.
 extern "C" AsyncValue *mlirAsyncRuntimeCreateValue(int32_t size) {
   AsyncValue *value = new AsyncValue(getDefaultAsyncRuntime(), size);
   return value;
 }
 
 // Create a new `async.group` in empty state.
-extern "C" AsyncGroup *mlirAsyncRuntimeCreateGroup() {
-  AsyncGroup *group = new AsyncGroup(getDefaultAsyncRuntime());
+extern "C" AsyncGroup *mlirAsyncRuntimeCreateGroup(int64_t size) {
+  AsyncGroup *group = new AsyncGroup(getDefaultAsyncRuntime(), size);
   return group;
 }
 
 extern "C" int64_t mlirAsyncRuntimeAddTokenToGroup(AsyncToken *token,
                                                    AsyncGroup *group) {
   std::unique_lock<std::mutex> lockToken(token->mu);
   std::unique_lock<std::mutex> lockGroup(group->mu);
 
   // Get the rank of the token inside the group before we drop the reference.
   int rank = group->rank.fetch_add(1);
-  group->pendingTokens.fetch_add(1);
 
   auto onTokenReady = [group, token]() {
     // Increment the number of errors in the group.
     if (State(token->state).isError())
       group->numErrors.fetch_add(1);
 
+    // If pending tokens go below zero it means that more tokens than the group
+    // size were added to this group.
+    assert(group->pendingTokens > 0 && "wrong group size");
+
     // Run all group awaiters if it was the last token in the group.
     if (group->pendingTokens.fetch_sub(1) == 1) {
       group->cv.notify_all();
       for (auto &awaiter : group->awaiters)
         awaiter();
     }
   };
 

mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm.mlir

-// RUN: mlir-opt %s -convert-async-to-llvm | FileCheck %s
+// RUN: mlir-opt %s -convert-async-to-llvm | FileCheck %s --dump-input=always
 
 // CHECK-LABEL: @create_token
 func @create_token() {
   // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
   %0 = async.runtime.create : !async.token
   return
 }
 
 // CHECK-LABEL: @create_value
 func @create_value() {
   // CHECK: %[[NULL:.*]] = llvm.mlir.null : !llvm.ptr<f32>
   // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : i32
   // CHECK: %[[OFFSET:.*]] = llvm.getelementptr %[[NULL]][%[[ONE]]]
   // CHECK: %[[SIZE:.*]] = llvm.ptrtoint %[[OFFSET]]
   // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue(%[[SIZE]])
   %0 = async.runtime.create : !async.value<f32>
   return
 }
 
 // CHECK-LABEL: @create_group
 func @create_group() {
-  // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup
-  %0 = async.runtime.create : !async.group
+  // CHECK: %[[C:.*]] = constant 1 : index
+  // CHECK: %[[S:.*]] = llvm.mlir.cast %[[C]] : index to i64
+  %c = constant 1 : index
+  // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup(%[[S]])
+  %0 = async.runtime.create_group  %c: !async.group
   return
 }
 
 // CHECK-LABEL: @set_token_available
 func @set_token_available() {
   // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
   %0 = async.runtime.create : !async.token
   // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]])
func @await_value() {
   %0 = async.runtime.create : !async.value<f32>
   // CHECK: call @mlirAsyncRuntimeAwaitValue(%[[VALUE]])
   async.runtime.await %0 : !async.value<f32>
   return
 }
 
 // CHECK-LABEL: @await_group
 func @await_group() {
+  %c = constant 1 : index
   // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup
-  %0 = async.runtime.create : !async.group
+  %0 = async.runtime.create_group %c: !async.group
   // CHECK: call @mlirAsyncRuntimeAwaitAllInGroup(%[[GROUP]])
   async.runtime.await %0 : !async.group
   return
 }
 
 // CHECK-LABEL: @await_and_resume_token
 func @await_and_resume_token() {
   %0 = async.coro.id
func @await_and_resume_value() {
   // CHECK: call @mlirAsyncRuntimeAwaitValueAndExecute
   // CHECK-SAME: (%[[VALUE]], %[[HDL]], %[[RESUME]])
   async.runtime.await_and_resume %2, %1 : !async.value<f32>
   return
 }
 
 // CHECK-LABEL: @await_and_resume_group
 func @await_and_resume_group() {
+  %c = constant 1 : index
   %0 = async.coro.id
   // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
   %1 = async.coro.begin %0
   // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateGroup
-  %2 = async.runtime.create : !async.group
+  %2 = async.runtime.create_group %c : !async.group
   // CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume
   // CHECK: call @mlirAsyncRuntimeAwaitAllInGroupAndExecute
   // CHECK-SAME: (%[[TOKEN]], %[[HDL]], %[[RESUME]])
   async.runtime.await_and_resume %2, %1 : !async.group
   return
 }
 
 // CHECK-LABEL: @resume
func @load() -> f32 {
   // CHECK: %[[VALUE:.*]] = llvm.load %[[P1]]
   %1 = async.runtime.load %0 : !async.value<f32>
   // CHECK: return %[[VALUE]] : f32
   return %1 : f32
 }
 
 // CHECK-LABEL: @add_token_to_group
 func @add_token_to_group() {
+  %c = constant 1 : index
   // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
   %0 = async.runtime.create : !async.token
   // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup
-  %1 = async.runtime.create : !async.group
+  %1 = async.runtime.create_group %c : !async.group
   // CHECK: call @mlirAsyncRuntimeAddTokenToGroup(%[[TOKEN]], %[[GROUP]])
   async.runtime.add_to_group %0, %1 : !async.token
   return
 }

mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir

 // Emplace result token after second resumption.
 // CHECK: memref.store %arg1, %arg2[%c0] : memref<1xf32>
 // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_1]])
 
 // -----
 
 // CHECK-LABEL: async_group_await_all
 func @async_group_await_all(%arg0: f32, %arg1: memref<1xf32>) {
-  // CHECK: %0 = call @mlirAsyncRuntimeCreateGroup()
-  %0 = async.create_group
+  %c = constant 1 : index
+  // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup
+  %0 = async.create_group %c : !async.group
 
   // CHECK: %[[TOKEN:.*]] = call @async_execute_fn
   %token = async.execute { async.yield }
-  // CHECK: call @mlirAsyncRuntimeAddTokenToGroup(%[[TOKEN]], %0)
+  // CHECK: call @mlirAsyncRuntimeAddTokenToGroup(%[[TOKEN]], %[[GROUP]])
   async.add_to_group %token, %0 : !async.token
 
   // CHECK: call @async_execute_fn_0
   async.execute {
     async.await_all %0
     async.yield
   }
 
-  // CHECK: call @mlirAsyncRuntimeAwaitAllInGroup(%0)
+  // CHECK: call @mlirAsyncRuntimeAwaitAllInGroup(%[[GROUP]])
   async.await_all %0
 
   return
 }
 
 // Function outlined from the async.execute operation.
 // CHECK: func private @async_execute_fn_0(%arg0: !llvm.ptr<i8>)
 // CHECK: %[[RET_1:.*]] = call @mlirAsyncRuntimeCreateToken()

mlir/test/Dialect/Async/async-to-async-runtime.mlir

 
 // CHECK: ^[[CLEANUP]]:
 // CHECK: ^[[SUSPEND]]:
 
 // -----
 
 // CHECK-LABEL: @async_group_await_all
 func @async_group_await_all(%arg0: f32, %arg1: memref<1xf32>) {
-  // CHECK: %[[GROUP:.*]] = async.runtime.create : !async.group
-  %0 = async.create_group
+  // CHECK: %[[C:.*]] = constant 1 : index
+  %c = constant 1 : index
+  // CHECK: %[[GROUP:.*]] = async.runtime.create_group %[[C]] : !async.group
+  %0 = async.create_group %c : !async.group
 
   // CHECK: %[[TOKEN:.*]] = call @async_execute_fn
   %token = async.execute { async.yield }
   // CHECK: async.runtime.add_to_group %[[TOKEN]], %[[GROUP]]
   async.add_to_group %token, %0 : !async.token
 
   // CHECK: call @async_execute_fn_0
   async.execute {

mlir/test/Dialect/Async/ops.mlir

 // CHECK-LABEL: @await_value
 func @await_value(%arg0: !async.value<f32>) -> f32 {
   // CHECK: async.await %arg0
   %0 = async.await %arg0 : !async.value<f32>
   return %0 : f32
 }
 
 // CHECK-LABEL: @create_group_and_await_all
-func @create_group_and_await_all(%arg0: !async.token, %arg1: !async.value<f32>) -> index {
-  %0 = async.create_group
+func @create_group_and_await_all(%arg0: !async.token,
+                                 %arg1: !async.value<f32>) -> index {
+  %c = constant 2 : index
+  %0 = async.create_group %c : !async.group
 
   // CHECK: async.add_to_group %arg0
   // CHECK: async.add_to_group %arg1
   %1 = async.add_to_group %arg0, %0 : !async.token
   %2 = async.add_to_group %arg1, %0 : !async.value<f32>
   async.await_all %0
 
   %3 = addi %1, %2 : index
   return %3 : index
 }

mlir/test/Dialect/Async/runtime.mlir

func @create_value() -> !async.value<f32> {
   // CHECK: %0 = async.runtime.create : !async.value<f32>
   %0 = async.runtime.create : !async.value<f32>
   // CHECK: return %0 : !async.value<f32>
   return %0 : !async.value<f32>
 }
 
 // CHECK-LABEL: @create_group
 func @create_group() -> !async.group {
-  // CHECK: %0 = async.runtime.create : !async.group
-  %0 = async.runtime.create : !async.group
-  // CHECK: return %0 : !async.group
+  // CHECK: %[[C:.*]] = constant 10 : index
+  %c = constant 10 : index
+  // CHECK: %[[V:.*]] = async.runtime.create_group %[[C]] : !async.group
+  %0 = async.runtime.create_group %c : !async.group
+  // CHECK: return %[[V]] : !async.group
   return %0 : !async.group
 }
 
 // CHECK-LABEL: @set_token_available
 func @set_token_available(%arg0: !async.token) {
   // CHECK: async.runtime.set_available %arg0 : !async.token
   async.runtime.set_available %arg0 : !async.token
   return

mlir/test/mlir-cpu-runner/async-error.mlir

func @main() {
   %err3_1 = async.runtime.is_error %value3 : !async.value<f32>
   vector.print %err3_0 : i1
   vector.print %err3_1 : i1
 
   // ------------------------------------------------------------------------ //
   // Check error propagation from a token to the group.
   // ------------------------------------------------------------------------ //
 
-  %group0 = async.create_group
+  %c2 = constant 2 : index
+  %group0 = async.create_group %c2 : !async.group
 
   %token4 = async.execute {
     async.yield
   }
 
   %token5 = async.execute {
     assert %false, "error"
     async.yield

mlir/test/mlir-cpu-runner/async-group.mlir

 // RUN:   mlir-opt %s -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
 // RUN:               -async-runtime-ref-counting-opt                          \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-std-to-llvm                                     \
 // RUN: | mlir-cpu-runner                                                      \
 // RUN:     -e main -entry-point-result=void -O0                               \
 // RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_c_runner_utils%shlibext  \
 // RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext    \
 // RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_async_runtime%shlibext   \
 // RUN: | FileCheck %s
 
 func @main() {
-  %group = async.create_group
+  %c1 = constant 1 : index
+  %c5 = constant 5 : index
+
+  %group = async.create_group %c5 : !async.group
 
   %token0 = async.execute { async.yield }
   %token1 = async.execute { async.yield }
   %token2 = async.execute { async.yield }
   %token3 = async.execute { async.yield }
   %token4 = async.execute { async.yield }
 
   %0 = async.add_to_group %token0, %group : !async.token
   %1 = async.add_to_group %token1, %group : !async.token
   %2 = async.add_to_group %token2, %group : !async.token
   %3 = async.add_to_group %token3, %group : !async.token
   %4 = async.add_to_group %token4, %group : !async.token
 
   %token5 = async.execute {
     async.await_all %group
     async.yield
   }
 
-  %group0 = async.create_group
+  %group0 = async.create_group %c1 : !async.group
   %5 = async.add_to_group %token5, %group0 : !async.token
   async.await_all %group0
 
   // CHECK: Current thread id: [[THREAD:.*]]
   call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
 
   return
 }
 
 func private @mlirAsyncRuntimePrintCurrentThreadId() -> ()