[mlir] AsyncRuntime: use int64_t for ref counting operations

[ezhulenev]

Summary

Workaround for SystemZ ABI problem: https://bugs.llvm.org/show_bug.cgi?id=51898

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

 def Async_RuntimeAddRefOp : Async_Op<"runtime.add_ref"> {
   let summary = "adds a reference to async value";
   let description = [{
     The `async.runtime.add_ref` operation adds a reference(s) to async value
     (token, value or group).
   }];
 
   let arguments = (ins Async_AnyAsyncType:$operand,
-                       Confined<I32Attr, [IntPositive]>:$count);
+                       Confined<I64Attr, [IntPositive]>:$count);
 
   let assemblyFormat = [{
     $operand attr-dict `:` type($operand)
   }];
 }
 
 def Async_RuntimeDropRefOp : Async_Op<"runtime.drop_ref"> {
   let summary = "drops a reference to async value";
   let description = [{
     The `async.runtime.drop_ref` operation drops a reference(s) to async value
     (token, value or group).
   }];
 
   let arguments = (ins Async_AnyAsyncType:$operand,
-                       Confined<I32Attr, [IntPositive]>:$count);
+                       Confined<I64Attr, [IntPositive]>:$count);
 
   let assemblyFormat = [{
     $operand attr-dict `:` type($operand)
   }];
 }
 
 #endif // ASYNC_OPS

mlir/include/mlir/ExecutionEngine/AsyncRuntime.h

 using CoroHandle = void *;           // coroutine handle
 using CoroResume = void (*)(void *); // coroutine resume function
 
 // Async runtime uses reference counting to manage the lifetime of async values
 // (values of async types like tokens, values and groups).
 using RefCountedObjPtr = void *;
 
 // Adds references to reference counted runtime object.
-extern "C" void mlirAsyncRuntimeAddRef(RefCountedObjPtr, int32_t);
+extern "C" void mlirAsyncRuntimeAddRef(RefCountedObjPtr, int64_t);
 
 // Drops references from reference counted runtime object.
-extern "C" void mlirAsyncRuntimeDropRef(RefCountedObjPtr, int32_t);
+extern "C" void mlirAsyncRuntimeDropRef(RefCountedObjPtr, int64_t);
 
 // Create a new `async.token` in not-ready state.
 extern "C" AsyncToken *mlirAsyncRuntimeCreateToken();
 
 // 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);
+extern "C" AsyncValue *mlirAsyncRuntimeCreateValue(int64_t);
 
 // Create a new `async.group` in empty state.
 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 *);

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

struct AsyncAPI {
   }
 
   static LLVM::LLVMTokenType tokenType(MLIRContext *ctx) {
     return LLVM::LLVMTokenType::get(ctx);
   }
 
   static FunctionType addOrDropRefFunctionType(MLIRContext *ctx) {
     auto ref = opaquePointerType(ctx);
-    auto count = IntegerType::get(ctx, 32);
+    auto count = IntegerType::get(ctx, 64);
     return FunctionType::get(ctx, {ref, count}, {});
   }
 
   static FunctionType createTokenFunctionType(MLIRContext *ctx) {
     return FunctionType::get(ctx, {}, {TokenType::get(ctx)});
   }
 
   static FunctionType createValueFunctionType(MLIRContext *ctx) {
-    auto i32 = IntegerType::get(ctx, 32);
+    auto i64 = IntegerType::get(ctx, 64);
     auto value = opaquePointerType(ctx);
-    return FunctionType::get(ctx, {i32}, {value});
+    return FunctionType::get(ctx, {i64}, {value});
   }
 
   static FunctionType createGroupFunctionType(MLIRContext *ctx) {
     auto i64 = IntegerType::get(ctx, 64);
     return FunctionType::get(ctx, {i64}, {GroupType::get(ctx)});
   }
 
   static FunctionType getValueStorageFunctionType(MLIRContext *ctx) {
if (resultType.isa<TokenType>()) {
       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();
-        auto i32 = rewriter.getI32Type();
+        auto i64 = rewriter.getI64Type();
 
         auto storedType = converter->convertType(valueType.getValueType());
         auto storagePtrType = LLVM::LLVMPointerType::get(storedType);
 
         // %Size = getelementptr %T* null, int 1
-        // %SizeI = ptrtoint %T* %Size to i32
+        // %SizeI = ptrtoint %T* %Size to i64
         auto nullPtr = rewriter.create<LLVM::NullOp>(loc, storagePtrType);
         auto one = rewriter.create<LLVM::ConstantOp>(
-            loc, i32, rewriter.getI32IntegerAttr(1));
+            loc, i64, rewriter.getI64IntegerAttr(1));
         auto gep = rewriter.create<LLVM::GEPOp>(loc, storagePtrType, nullPtr,
                                                 one.getResult());
-        return rewriter.create<LLVM::PtrToIntOp>(loc, i32, gep);
+        return rewriter.create<LLVM::PtrToIntOp>(loc, i64, gep);
       };
 
       rewriter.replaceOpWithNewOp<CallOp>(op, kCreateValue, resultType,
                                           sizeOf(value));
 
       return success();
     }
 
explicit RefCountingOpLowering(TypeConverter &converter, MLIRContext *ctx,
                                  StringRef apiFunctionName)
       : OpConversionPattern<RefCountingOp>(converter, ctx),
         apiFunctionName(apiFunctionName) {}
 
   LogicalResult
   matchAndRewrite(RefCountingOp op, typename RefCountingOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto count =
-        rewriter.create<ConstantOp>(op->getLoc(), rewriter.getI32Type(),
-                                    rewriter.getI32IntegerAttr(op.count()));
+        rewriter.create<ConstantOp>(op->getLoc(), rewriter.getI64Type(),
+                                    rewriter.getI64IntegerAttr(op.count()));
 
     auto operand = adaptor.operand();
     rewriter.replaceOpWithNewOp<CallOp>(op, TypeRange(), apiFunctionName,
                                         ValueRange({operand, count}));
 
     return success();
   }
 

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

static LogicalResult dropRefIfNoUses(Value value, unsigned count = 1) {
 
   // Set insertion point after the operation producing a value, or at the
   // beginning of the block if the value defined by the block argument.
   if (Operation *op = value.getDefiningOp())
     b.setInsertionPointAfter(op);
   else
     b.setInsertionPointToStart(value.getParentBlock());
 
-  b.create<RuntimeDropRefOp>(value.getLoc(), value, b.getI32IntegerAttr(1));
+  b.create<RuntimeDropRefOp>(value.getLoc(), value, b.getI64IntegerAttr(1));
   return success();
 }
 
 // Calls `addRefCounting` for every reference counted value defined by the
 // operation `op` (block arguments and values defined in nested regions).
 static LogicalResult walkReferenceCountedValues(
     Operation *op, llvm::function_ref<LogicalResult(Value)> addRefCounting) {
   // Check that we do not have high level async operations in the IR because
for (Operation *lastUser : lastUsers) {
 
     // We can't currently handle other types of terminators.
     if (lastUser->hasTrait<OpTrait::IsTerminator>())
       return lastUser->emitError() << "async reference counting can't handle "
                                       "terminators that are not ReturnLike";
 
     // Add a drop_ref immediately after the last user.
     builder.setInsertionPointAfter(lastUser);
-    builder.create<RuntimeDropRefOp>(loc, value, builder.getI32IntegerAttr(1));
+    builder.create<RuntimeDropRefOp>(loc, value, builder.getI64IntegerAttr(1));
   }
 
   return success();
 }
 
 LogicalResult
 AsyncRuntimeRefCountingPass::addAddRefBeforeFunctionCall(Value value) {
   OpBuilder builder(value.getContext());
   Location loc = value.getLoc();
 
   for (Operation *user : value.getUsers()) {
     if (!isa<CallOp>(user))
       continue;
 
     // Add a reference before the function call to pass the value at `+1`
     // reference to the function entry block.
     builder.setInsertionPoint(user);
-    builder.create<RuntimeAddRefOp>(loc, value, builder.getI32IntegerAttr(1));
+    builder.create<RuntimeAddRefOp>(loc, value, builder.getI64IntegerAttr(1));
   }
 
   return success();
 }
 
 LogicalResult
 AsyncRuntimeRefCountingPass::addDropRefInDivergentLivenessSuccessor(
     Value value) {
for (Block *successor : successors) {
         refCountingBlock = &successor->getParent()->emplaceBlock();
         refCountingBlock->moveBefore(successor);
         OpBuilder builder = OpBuilder::atBlockEnd(refCountingBlock);
         builder.create<BranchOp>(value.getLoc(), successor);
       }
 
       OpBuilder builder = OpBuilder::atBlockBegin(refCountingBlock);
       builder.create<RuntimeDropRefOp>(value.getLoc(), value,
-                                       builder.getI32IntegerAttr(1));
+                                       builder.getI64IntegerAttr(1));
 
       // No need to update the terminator operation.
       if (successor == refCountingBlock)
         continue;
 
       // Update terminator `successor` block to `refCountingBlock`.
       for (auto pair : llvm::enumerate(terminator->getSuccessors()))
         if (pair.value() == successor)
for (auto &func : policy) {
       if (failed(refCount))
         return failure();
 
       int cnt = refCount.getValue();
 
       // Create `add_ref` operation before the operand owner.
       if (cnt > 0) {
         b.setInsertionPoint(operand.getOwner());
-        b.create<RuntimeAddRefOp>(loc, value, b.getI32IntegerAttr(cnt));
+        b.create<RuntimeAddRefOp>(loc, value, b.getI64IntegerAttr(cnt));
       }
 
       // Create `drop_ref` operation after the operand owner.
       if (cnt < 0) {
         b.setInsertionPointAfter(operand.getOwner());
-        b.create<RuntimeDropRefOp>(loc, value, b.getI32IntegerAttr(-cnt));
+        b.create<RuntimeDropRefOp>(loc, value, b.getI64IntegerAttr(-cnt));
       }
     }
   }
 
   return success();
 }
 
 void AsyncRuntimePolicyBasedRefCountingPass::initializeDefaultPolicy() {

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

for (RuntimeAddRefOp addRef : info.addRefs) {
         // When reference counted value passed to a function as an argument,
         // function takes ownership of +1 reference and it will drop it before
         // returning.
         //
         // Example:
         //
         //   %token = ... : !async.token
         //
-        //   async.runtime.add_ref %token {count = 1 : i32} : !async.token
+        //   async.runtime.add_ref %token {count = 1 : i64} : !async.token
         //   call @pass_token(%token: !async.token, ...)
         //
         //   async.await %token : !async.token
-        //   async.runtime.drop_ref %token {count = 1 : i32} : !async.token
+        //   async.runtime.drop_ref %token {count = 1 : i64} : !async.token
         //
         // In this example if we'll cancel a pair of reference counting
         // operations we might end up with a deallocated token when we'll
         // reach `async.await` operation.
         Operation *firstFunctionCallUser = nullptr;
         Operation *lastNonFunctionCallUser = nullptr;
 
         for (Operation *user : info.users) {

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

public:
   AsyncRuntime() : numRefCountedObjects(0) {}
 
   ~AsyncRuntime() {
     threadPool.wait(); // wait for the completion of all async tasks
     assert(getNumRefCountedObjects() == 0 &&
            "all ref counted objects must be destroyed");
   }
 
-  int32_t getNumRefCountedObjects() {
+  int64_t getNumRefCountedObjects() {
     return numRefCountedObjects.load(std::memory_order_relaxed);
   }
 
   llvm::ThreadPool &getThreadPool() { return threadPool; }
 
 private:
   friend class RefCounted;
 
   // Count the total number of reference counted objects in this instance
   // of an AsyncRuntime. For debugging purposes only.
   void addNumRefCountedObjects() {
     numRefCountedObjects.fetch_add(1, std::memory_order_relaxed);
   }
   void dropNumRefCountedObjects() {
     numRefCountedObjects.fetch_sub(1, std::memory_order_relaxed);
   }
 
-  std::atomic<int32_t> numRefCountedObjects;
+  std::atomic<int64_t> numRefCountedObjects;
   llvm::ThreadPool threadPool;
 };
 
 // -------------------------------------------------------------------------- //
 // A state of the async runtime value (token, value or group).
 // -------------------------------------------------------------------------- //
 
 class State {

 };
 
 // -------------------------------------------------------------------------- //
 // A base class for all reference counted objects created by the async runtime.
 // -------------------------------------------------------------------------- //
 
 class RefCounted {
 public:
-  RefCounted(AsyncRuntime *runtime, int32_t refCount = 1)
+  RefCounted(AsyncRuntime *runtime, int64_t refCount = 1)
       : runtime(runtime), refCount(refCount) {
     runtime->addNumRefCountedObjects();
   }
 
   virtual ~RefCounted() {
     assert(refCount.load() == 0 && "reference count must be zero");
     runtime->dropNumRefCountedObjects();
   }
 
   RefCounted(const RefCounted &) = delete;
   RefCounted &operator=(const RefCounted &) = delete;
 
-  void addRef(int32_t count = 1) { refCount.fetch_add(count); }
+  void addRef(int64_t count = 1) { refCount.fetch_add(count); }
 
-  void dropRef(int32_t count = 1) {
-    int32_t previous = refCount.fetch_sub(count);
+  void dropRef(int64_t count = 1) {
+    int64_t previous = refCount.fetch_sub(count);
     assert(previous >= count && "reference count should not go below zero");
     if (previous == count)
       destroy();
   }
 
 protected:
   virtual void destroy() { delete this; }
 
 private:
   AsyncRuntime *runtime;
-  std::atomic<int32_t> refCount;
+  std::atomic<int64_t> refCount;
 };
 
 } // namespace
 
 // Returns the default per-process instance of an async runtime.
 static std::unique_ptr<AsyncRuntime> &getDefaultAsyncRuntimeInstance() {
   static auto runtime = std::make_unique<AsyncRuntime>();
   return runtime;
struct AsyncToken : public RefCounted {
   std::vector<std::function<void()>> awaiters;
 };
 
 // Async value provides a mechanism to access the result of asynchronous
 // operations. It owns the storage that is used to store/load the value of the
 // underlying type, and a flag to signal if the value is ready or not.
 struct AsyncValue : public RefCounted {
   // AsyncValue similar to an AsyncToken created with a reference count of 2.
-  AsyncValue(AsyncRuntime *runtime, int32_t size)
+  AsyncValue(AsyncRuntime *runtime, int64_t size)
       : RefCounted(runtime, /*refCount=*/2), state(State::kUnavailable),
         storage(size) {}
 
   std::atomic<State::StateEnum> state;
 
   // Use vector of bytes to store async value payload.
   std::vector<int8_t> storage;
 
struct AsyncGroup : public RefCounted {
 
   // Pending awaiters are guarded by a mutex.
   std::mutex mu;
   std::condition_variable cv;
   std::vector<std::function<void()>> awaiters;
 };
 
 // Adds references to reference counted runtime object.
-extern "C" void mlirAsyncRuntimeAddRef(RefCountedObjPtr ptr, int32_t count) {
+extern "C" void mlirAsyncRuntimeAddRef(RefCountedObjPtr ptr, int64_t count) {
   RefCounted *refCounted = static_cast<RefCounted *>(ptr);
   refCounted->addRef(count);
 }
 
 // Drops references from reference counted runtime object.
-extern "C" void mlirAsyncRuntimeDropRef(RefCountedObjPtr ptr, int32_t count) {
+extern "C" void mlirAsyncRuntimeDropRef(RefCountedObjPtr ptr, int64_t count) {
   RefCounted *refCounted = static_cast<RefCounted *>(ptr);
   refCounted->dropRef(count);
 }
 
 // Creates a new `async.token` in not-ready state.
 extern "C" AsyncToken *mlirAsyncRuntimeCreateToken() {
   AsyncToken *token = new AsyncToken(getDefaultAsyncRuntime());
   return token;
 }
 
 // Creates a new `async.value` in not-ready state.
-extern "C" AsyncValue *mlirAsyncRuntimeCreateValue(int32_t size) {
+extern "C" AsyncValue *mlirAsyncRuntimeCreateValue(int64_t size) {
   AsyncValue *value = new AsyncValue(getDefaultAsyncRuntime(), size);
   return value;
 }
 
 // Create a new `async.group` in empty state.
 extern "C" AsyncGroup *mlirAsyncRuntimeCreateGroup(int64_t size) {
   AsyncGroup *group = new AsyncGroup(getDefaultAsyncRuntime(), size);
   return group;

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

 // 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: %[[ONE:.*]] = llvm.mlir.constant(1 : i64) : i64
   // 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

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

 // RUN: mlir-opt %s -split-input-file -async-to-async-runtime -convert-async-to-llvm | FileCheck %s
 
 // CHECK-LABEL: reference_counting
 func @reference_counting(%arg0: !async.token) {
-  // CHECK: %[[C2:.*]] = constant 2 : i32
+  // CHECK: %[[C2:.*]] = constant 2 : i64
   // CHECK: call @mlirAsyncRuntimeAddRef(%arg0, %[[C2]])
-  async.runtime.add_ref %arg0 {count = 2 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 2 : i64} : !async.token
 
-  // CHECK: %[[C1:.*]] = constant 1 : i32
+  // CHECK: %[[C1:.*]] = constant 1 : i64
   // CHECK: call @mlirAsyncRuntimeDropRef(%arg0, %[[C1]])
-  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
 
   return
 }
 
 // -----
 
 // CHECK-LABEL: execute_no_async_args
 func @execute_no_async_args(%arg0: f32, %arg1: memref<1xf32>) {

mlir/test/Dialect/Async/async-runtime-policy-based-ref-counting.mlir

func @group_await(%arg0: !async.group) {
   return
 }
 
 // CHECK-LABEL: @add_token_to_group
 // CHECK:         %[[GROUP:.*]]: !async.group
 // CHECK:         %[[TOKEN:.*]]: !async.token
 func @add_token_to_group(%arg0: !async.group, %arg1: !async.token) {
   // CHECK: async.runtime.add_to_group %[[TOKEN]], %[[GROUP]]
-  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   async.runtime.add_to_group %arg1, %arg0 : !async.token
   return
 }
 
 // CHECK-LABEL: @value_load
 // CHECK:         %[[VALUE:.*]]: !async.value<f32>
 func @value_load(%arg0: !async.value<f32>) {
   // CHECK: async.runtime.load %[[VALUE]]
-  // CHECK: async.runtime.drop_ref %[[VALUE]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[VALUE]] {count = 1 : i64}
   %0 = async.runtime.load %arg0 : !async.value<f32>
   return
 }
 
 // CHECK-LABEL: @error_check
 // CHECK:         %[[TOKEN:.*]]: !async.token
 func @error_check(%arg0: !async.token) {
   // CHECK: async.runtime.is_error %[[TOKEN]]
-  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   %0 = async.runtime.is_error %arg0 : !async.token
   return
 }

mlir/test/Dialect/Async/async-runtime-ref-counting-opt.mlir

 // RUN: mlir-opt %s -async-runtime-ref-counting-opt | FileCheck %s
 
 func private @consume_token(%arg0: !async.token)
 
 // CHECK-LABEL: @cancellable_operations_0
 func @cancellable_operations_0(%arg0: !async.token) {
   // CHECK-NOT: async.runtime.add_ref
   // CHECK-NOT: async.runtime.drop_ref
-  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
-  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
+  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @cancellable_operations_1
 func @cancellable_operations_1(%arg0: !async.token) {
   // CHECK-NOT: async.runtime.add_ref
-  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
   // CHECK: call @consume_toke
   call @consume_token(%arg0): (!async.token) -> ()
   // CHECK-NOT: async.runtime.drop_ref
-  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @cancellable_operations_2
 func @cancellable_operations_2(%arg0: !async.token) {
   // CHECK: async.runtime.await
   // CHECK-NEXT: async.runtime.await
   // CHECK-NEXT: async.runtime.await
   // CHECK-NEXT: return
-  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
   async.runtime.await %arg0 : !async.token
-  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
   async.runtime.await %arg0 : !async.token
-  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
   async.runtime.await %arg0 : !async.token
-  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
   return
 }
 
 // CHECK-LABEL: @cancellable_operations_3
 func @cancellable_operations_3(%arg0: !async.token) {
   // CHECK-NOT: add_ref
-  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
   // CHECK: call @consume_toke
   call @consume_token(%arg0): (!async.token) -> ()
   // CHECK-NOT: async.runtime.drop_ref
-  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
   // CHECK: async.runtime.await
   async.runtime.await %arg0 : !async.token
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @not_cancellable_operations_0
 func @not_cancellable_operations_0(%arg0: !async.token) {
   // CHECK: add_ref
-  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
   // CHECK: call @consume_toke
   call @consume_token(%arg0): (!async.token) -> ()
   // CHECK: async.runtime.await
   async.runtime.await %arg0 : !async.token
   // CHECK: async.runtime.drop_ref
-  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
+  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
   // CHECK: return
   return
 }

mlir/test/Dialect/Async/async-runtime-ref-counting.mlir

 // RUN: mlir-opt %s -async-runtime-ref-counting | FileCheck %s
 
 // CHECK-LABEL: @token
 func private @token() -> !async.token
 
 // CHECK-LABEL: @cond
 func private @cond() -> i1
 
 // CHECK-LABEL: @take_token
 func private @take_token(%arg0: !async.token)
 
 // CHECK-LABEL: @token_arg_no_uses
 // CHECK: %[[TOKEN:.*]]: !async.token
 func @token_arg_no_uses(%arg0: !async.token) {
-  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   return
 }
 
 // CHECK-LABEL: @token_value_no_uses
 func @token_value_no_uses() {
   // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
-  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   %0 = async.runtime.create : !async.token
   return
 }
 
 // CHECK-LABEL: @token_returned_no_uses
 func @token_returned_no_uses() {
   // CHECK: %[[TOKEN:.*]] = call @token
-  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   %0 = call @token() : () -> !async.token
   return
 }
 
 // CHECK-LABEL: @token_arg_to_func
 // CHECK: %[[TOKEN:.*]]: !async.token
 func @token_arg_to_func(%arg0: !async.token) {
-  // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32} : !async.token
+  // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i64} : !async.token
   call @take_token(%arg0): (!async.token) -> ()
-  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32} : !async.token
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} : !async.token
   return
 }
 
 // CHECK-LABEL: @token_value_to_func
 func @token_value_to_func() {
   // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
   %0 = async.runtime.create : !async.token
-  // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32} : !async.token
+  // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i64} : !async.token
   call @take_token(%0): (!async.token) -> ()
-  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   return
 }
 
 // CHECK-LABEL: @token_arg_cond_br_await_with_fallthough
 // CHECK: %[[TOKEN:.*]]: !async.token
 func @token_arg_cond_br_await_with_fallthough(%arg0: !async.token, %arg1: i1) {
   // CHECK: cond_br
   // CHECK-SAME: ^[[BB1:.*]], ^[[BB2:.*]]
   cond_br %arg1, ^bb1, ^bb2
 ^bb1:
   // CHECK: ^[[BB1]]:
   // CHECK:   br ^[[BB2]]
   br ^bb2
 ^bb2:
   // CHECK: ^[[BB2]]:
   // CHECK:   async.runtime.await %[[TOKEN]]
-  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   async.runtime.await %arg0 : !async.token
   return
 }
 
 // CHECK-LABEL: @token_simple_return
 func @token_simple_return() -> !async.token {
   // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
   %token = async.runtime.create : !async.token

 // CHECK: %[[TOKEN:.*]]: !async.token
 func @token_coro_await_and_resume(%arg0: !async.token) -> !async.token {
   %token = async.runtime.create : !async.token
   %id = async.coro.id
   %hdl = async.coro.begin %id
   %saved = async.coro.save %hdl
   // CHECK: async.runtime.await_and_resume %[[TOKEN]]
   async.runtime.await_and_resume %arg0, %hdl : !async.token
-  // CHECK-NEXT: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK-NEXT: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   async.coro.suspend %saved, ^suspend, ^resume, ^cleanup
 ^resume:
   br ^cleanup
 ^cleanup:
   async.coro.free %id, %hdl
   br ^suspend
 ^suspend:
   async.coro.end %hdl
func @value_coro_await_and_resume(%arg0: !async.value<f32>) -> !async.token {
   // CHECK: async.runtime.await_and_resume %[[VALUE]]
   async.runtime.await_and_resume %arg0, %hdl : !async.value<f32>
   // CHECK: async.coro.suspend
   // CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]
   async.coro.suspend %saved, ^suspend, ^resume, ^cleanup
 ^resume:
   // CHECK: ^[[RESUME]]:
   // CHECK:   %[[LOADED:.*]] = async.runtime.load %[[VALUE]]
-  // CHECK:   async.runtime.drop_ref %[[VALUE]] {count = 1 : i32}
+  // CHECK:   async.runtime.drop_ref %[[VALUE]] {count = 1 : i64}
   %0 = async.runtime.load %arg0 : !async.value<f32>
   // CHECK:  addf %[[LOADED]], %[[LOADED]]
   %1 = addf %0, %0 : f32
   br ^cleanup
 ^cleanup:
   async.coro.free %id, %hdl
   br ^suspend
 ^suspend:
func private @outlined_async_execute(%arg0: !async.token) -> !async.token {
   %3 = async.coro.save %2
   async.runtime.resume %2
   // CHECK: async.coro.suspend
   async.coro.suspend %3, ^suspend, ^resume, ^cleanup
 ^resume:
   // CHECK: ^[[RESUME:.*]]:
   %4 = async.coro.save %2
   async.runtime.await_and_resume %arg0, %2 : !async.token
-  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   // CHECK: async.coro.suspend
   async.coro.suspend %4, ^suspend, ^resume_1, ^cleanup
 ^resume_1:
   // CHECK: ^[[RESUME_1:.*]]:
   // CHECK:   async.runtime.set_available
   async.runtime.set_available %0 : !async.token
   br ^cleanup
 ^cleanup:
func @token_await_inside_nested_region(%arg0: i1) {
   // CHECK: %[[TOKEN:.*]] = call @token()
   %token = call @token() : () -> !async.token
   // CHECK: scf.if %[[ARG]] {
   scf.if %arg0 {
     // CHECK: async.runtime.await %[[TOKEN]]
     async.runtime.await %token : !async.token
   }
   // CHECK: }
-  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   // CHECK: return
   return
 }
 
 // CHECK-LABEL: @token_defined_in_the_loop
 func @token_defined_in_the_loop() {
   br ^bb1
 ^bb1:
   // CHECK: ^[[BB1:.*]]:
   // CHECK:   %[[TOKEN:.*]] = call @token()
   %token = call @token() : () -> !async.token
   // CHECK:   async.runtime.await %[[TOKEN]]
-  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   async.runtime.await %token : !async.token
   %0 = call @cond(): () -> (i1)
   cond_br %0, ^bb1, ^bb2
 ^bb2:
   // CHECK: ^[[BB2:.*]]:
   // CHECK:   return
   return
 }
 
 // CHECK-LABEL: @divergent_liveness_one_token
 func @divergent_liveness_one_token(%arg0 : i1) {
   // CHECK: %[[TOKEN:.*]] = call @token()
   %token = call @token() : () -> !async.token
   // CHECK: cond_br %arg0, ^[[LIVE_IN:.*]], ^[[REF_COUNTING:.*]]
   cond_br %arg0, ^bb1, ^bb2
 ^bb1:
   // CHECK: ^[[LIVE_IN]]:
   // CHECK:   async.runtime.await %[[TOKEN]]
-  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   // CHECK:   br ^[[RETURN:.*]]
   async.runtime.await %token : !async.token
   br ^bb2
   // CHECK: ^[[REF_COUNTING:.*]]:
-  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   // CHECK:   br ^[[RETURN:.*]]
 ^bb2:
   // CHECK: ^[[RETURN]]:
   // CHECK:   return
   return
 }
 
 // CHECK-LABEL: @divergent_liveness_unique_predecessor
 func @divergent_liveness_unique_predecessor(%arg0 : i1) {
   // CHECK: %[[TOKEN:.*]] = call @token()
   %token = call @token() : () -> !async.token
   // CHECK: cond_br %arg0, ^[[LIVE_IN:.*]], ^[[NO_LIVE_IN:.*]]
   cond_br %arg0, ^bb2, ^bb1
 ^bb1:
   // CHECK: ^[[NO_LIVE_IN]]:
-  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   // CHECK:   br ^[[RETURN:.*]]
   br ^bb3
 ^bb2:
   // CHECK: ^[[LIVE_IN]]:
   // CHECK:   async.runtime.await %[[TOKEN]]
-  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}
   // CHECK:   br ^[[RETURN]]
   async.runtime.await %token : !async.token
   br ^bb3
 ^bb3:
   // CHECK: ^[[RETURN]]:
   // CHECK:  return
   return
 }
 
 // CHECK-LABEL: @divergent_liveness_two_tokens
 func @divergent_liveness_two_tokens(%arg0 : i1) {
   // CHECK: %[[TOKEN0:.*]] = call @token()
   // CHECK: %[[TOKEN1:.*]] = call @token()
   %token0 = call @token() : () -> !async.token
   %token1 = call @token() : () -> !async.token
   // CHECK: cond_br %arg0, ^[[AWAIT0:.*]], ^[[AWAIT1:.*]]
   cond_br %arg0, ^await0, ^await1
 ^await0:
   // CHECK: ^[[AWAIT0]]:
-  // CHECK:   async.runtime.drop_ref %[[TOKEN1]] {count = 1 : i32}
+  // CHECK:   async.runtime.drop_ref %[[TOKEN1]] {count = 1 : i64}
   // CHECK:   async.runtime.await %[[TOKEN0]]
-  // CHECK:   async.runtime.drop_ref %[[TOKEN0]] {count = 1 : i32}
+  // CHECK:   async.runtime.drop_ref %[[TOKEN0]] {count = 1 : i64}
   // CHECK:   br ^[[RETURN:.*]]
   async.runtime.await %token0 : !async.token
   br ^ret
 ^await1:
   // CHECK: ^[[AWAIT1]]:
-  // CHECK:   async.runtime.drop_ref %[[TOKEN0]] {count = 1 : i32}
+  // CHECK:   async.runtime.drop_ref %[[TOKEN0]] {count = 1 : i64}
   // CHECK:   async.runtime.await %[[TOKEN1]]
-  // CHECK:   async.runtime.drop_ref %[[TOKEN1]] {count = 1 : i32}
+  // CHECK:   async.runtime.drop_ref %[[TOKEN1]] {count = 1 : i64}
   // CHECK:   br ^[[RETURN]]
   async.runtime.await %token1 : !async.token
   br ^ret
 ^ret:
   // CHECK: ^[[RETURN]]:
   // CHECK:   return
   return
 }

mlir/test/Dialect/Async/runtime.mlir

func @add_to_group(%arg0: !async.token, %arg1: !async.value<f32>,
   async.runtime.add_to_group %arg0, %arg2 : !async.token
   // CHECK: async.runtime.add_to_group %arg1, %arg2 : !async.value<f32>
   async.runtime.add_to_group %arg1, %arg2 : !async.value<f32>
   return
 }
 
 // CHECK-LABEL: @add_ref
 func @add_ref(%arg0: !async.token) {
-  // CHECK: async.runtime.add_ref %arg0 {count = 1 : i32}
-  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
+  // CHECK: async.runtime.add_ref %arg0 {count = 1 : i64}
+  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
   return
 }
 
 // CHECK-LABEL: @drop_ref
 func @drop_ref(%arg0: !async.token) {
-  // CHECK: async.runtime.drop_ref %arg0 {count = 1 : i32}
-  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
+  // CHECK: async.runtime.drop_ref %arg0 {count = 1 : i64}
+  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
   return
 }