diff --git a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td --- a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td +++ b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td @@ -508,7 +508,7 @@ }]; let arguments = (ins Async_AnyAsyncType:$operand, - Confined:$count); + Confined:$count); let assemblyFormat = [{ $operand attr-dict `:` type($operand) @@ -523,7 +523,7 @@ }]; let arguments = (ins Async_AnyAsyncType:$operand, - Confined:$count); + Confined:$count); let assemblyFormat = [{ $operand attr-dict `:` type($operand) diff --git a/mlir/include/mlir/ExecutionEngine/AsyncRuntime.h b/mlir/include/mlir/ExecutionEngine/AsyncRuntime.h --- a/mlir/include/mlir/ExecutionEngine/AsyncRuntime.h +++ b/mlir/include/mlir/ExecutionEngine/AsyncRuntime.h @@ -51,10 +51,10 @@ 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(); @@ -63,7 +63,7 @@ // 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); diff --git a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp --- a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp +++ b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp @@ -76,7 +76,7 @@ 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}, {}); } @@ -85,9 +85,9 @@ } 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) { @@ -559,19 +559,19 @@ // 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(loc, storagePtrType); auto one = rewriter.create( - loc, i32, rewriter.getI32IntegerAttr(1)); + loc, i64, rewriter.getI64IntegerAttr(1)); auto gep = rewriter.create(loc, storagePtrType, nullPtr, one.getResult()); - return rewriter.create(loc, i32, gep); + return rewriter.create(loc, i64, gep); }; rewriter.replaceOpWithNewOp(op, kCreateValue, resultType, @@ -904,8 +904,8 @@ matchAndRewrite(RefCountingOp op, typename RefCountingOp::Adaptor adaptor, ConversionPatternRewriter &rewriter) const override { auto count = - rewriter.create(op->getLoc(), rewriter.getI32Type(), - rewriter.getI32IntegerAttr(op.count())); + rewriter.create(op->getLoc(), rewriter.getI64Type(), + rewriter.getI64IntegerAttr(op.count())); auto operand = adaptor.operand(); rewriter.replaceOpWithNewOp(op, TypeRange(), apiFunctionName, diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp --- a/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp +++ b/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp @@ -44,7 +44,7 @@ else b.setInsertionPointToStart(value.getParentBlock()); - b.create(value.getLoc(), value, b.getI32IntegerAttr(1)); + b.create(value.getLoc(), value, b.getI64IntegerAttr(1)); return success(); } @@ -304,7 +304,7 @@ // Add a drop_ref immediately after the last user. builder.setInsertionPointAfter(lastUser); - builder.create(loc, value, builder.getI32IntegerAttr(1)); + builder.create(loc, value, builder.getI64IntegerAttr(1)); } return success(); @@ -322,7 +322,7 @@ // Add a reference before the function call to pass the value at `+1` // reference to the function entry block. builder.setInsertionPoint(user); - builder.create(loc, value, builder.getI32IntegerAttr(1)); + builder.create(loc, value, builder.getI64IntegerAttr(1)); } return success(); @@ -411,7 +411,7 @@ OpBuilder builder = OpBuilder::atBlockBegin(refCountingBlock); builder.create(value.getLoc(), value, - builder.getI32IntegerAttr(1)); + builder.getI64IntegerAttr(1)); // No need to update the terminator operation. if (successor == refCountingBlock) @@ -502,13 +502,13 @@ // Create `add_ref` operation before the operand owner. if (cnt > 0) { b.setInsertionPoint(operand.getOwner()); - b.create(loc, value, b.getI32IntegerAttr(cnt)); + b.create(loc, value, b.getI64IntegerAttr(cnt)); } // Create `drop_ref` operation after the operand owner. if (cnt < 0) { b.setInsertionPointAfter(operand.getOwner()); - b.create(loc, value, b.getI32IntegerAttr(-cnt)); + b.create(loc, value, b.getI64IntegerAttr(-cnt)); } } } diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCountingOpt.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCountingOpt.cpp --- a/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCountingOpt.cpp +++ b/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCountingOpt.cpp @@ -118,11 +118,11 @@ // // %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 diff --git a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp --- a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp +++ b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp @@ -55,7 +55,7 @@ "all ref counted objects must be destroyed"); } - int32_t getNumRefCountedObjects() { + int64_t getNumRefCountedObjects() { return numRefCountedObjects.load(std::memory_order_relaxed); } @@ -73,7 +73,7 @@ numRefCountedObjects.fetch_sub(1, std::memory_order_relaxed); } - std::atomic numRefCountedObjects; + std::atomic numRefCountedObjects; llvm::ThreadPool threadPool; }; @@ -123,7 +123,7 @@ class RefCounted { public: - RefCounted(AsyncRuntime *runtime, int32_t refCount = 1) + RefCounted(AsyncRuntime *runtime, int64_t refCount = 1) : runtime(runtime), refCount(refCount) { runtime->addNumRefCountedObjects(); } @@ -136,10 +136,10 @@ 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(); @@ -150,7 +150,7 @@ private: AsyncRuntime *runtime; - std::atomic refCount; + std::atomic refCount; }; } // namespace @@ -192,7 +192,7 @@ // 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) {} @@ -225,13 +225,13 @@ }; // 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(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(ptr); refCounted->dropRef(count); } @@ -243,7 +243,7 @@ } // 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; } diff --git a/mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm.mlir b/mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm.mlir --- a/mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm.mlir +++ b/mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm.mlir @@ -10,7 +10,7 @@ // CHECK-LABEL: @create_value func @create_value() { // CHECK: %[[NULL:.*]] = llvm.mlir.null : !llvm.ptr - // 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]]) diff --git a/mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir b/mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir --- a/mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir +++ b/mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir @@ -2,13 +2,13 @@ // 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 } diff --git a/mlir/test/Dialect/Async/async-runtime-policy-based-ref-counting.mlir b/mlir/test/Dialect/Async/async-runtime-policy-based-ref-counting.mlir --- a/mlir/test/Dialect/Async/async-runtime-policy-based-ref-counting.mlir +++ b/mlir/test/Dialect/Async/async-runtime-policy-based-ref-counting.mlir @@ -23,7 +23,7 @@ // 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 } @@ -32,7 +32,7 @@ // CHECK: %[[VALUE:.*]]: !async.value func @value_load(%arg0: !async.value) { // 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 return } @@ -41,7 +41,7 @@ // 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 } diff --git a/mlir/test/Dialect/Async/async-runtime-ref-counting-opt.mlir b/mlir/test/Dialect/Async/async-runtime-ref-counting-opt.mlir --- a/mlir/test/Dialect/Async/async-runtime-ref-counting-opt.mlir +++ b/mlir/test/Dialect/Async/async-runtime-ref-counting-opt.mlir @@ -6,8 +6,8 @@ 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 } @@ -15,11 +15,11 @@ // 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 } @@ -30,24 +30,24 @@ // 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 @@ -57,13 +57,13 @@ // 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 } diff --git a/mlir/test/Dialect/Async/async-runtime-ref-counting.mlir b/mlir/test/Dialect/Async/async-runtime-ref-counting.mlir --- a/mlir/test/Dialect/Async/async-runtime-ref-counting.mlir +++ b/mlir/test/Dialect/Async/async-runtime-ref-counting.mlir @@ -12,14 +12,14 @@ // 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 } @@ -27,7 +27,7 @@ // 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 } @@ -35,9 +35,9 @@ // 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 } @@ -45,9 +45,9 @@ 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 } @@ -64,7 +64,7 @@ ^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 } @@ -106,7 +106,7 @@ %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 @@ -133,7 +133,7 @@ ^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 // CHECK: addf %[[LOADED]], %[[LOADED]] %1 = addf %0, %0 : f32 @@ -160,7 +160,7 @@ // 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: @@ -191,7 +191,7 @@ 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 } @@ -204,7 +204,7 @@ // 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 @@ -223,12 +223,12 @@ ^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]]: @@ -244,13 +244,13 @@ 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 @@ -270,17 +270,17 @@ 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 diff --git a/mlir/test/Dialect/Async/runtime.mlir b/mlir/test/Dialect/Async/runtime.mlir --- a/mlir/test/Dialect/Async/runtime.mlir +++ b/mlir/test/Dialect/Async/runtime.mlir @@ -154,14 +154,14 @@ // 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 }