[Mlir-commits] [mlir] 92db09c - [mlir] AsyncRuntime: use int64_t for ref counting operations
Eugene Zhulenev
llvmlistbot at llvm.org
Mon Sep 27 07:55:08 PDT 2021
Author: Eugene Zhulenev
Date: 2021-09-27T07:55:01-07:00
New Revision: 92db09cde04933d2e703a590054fe709dcfa88c0
URL: https://github.com/llvm/llvm-project/commit/92db09cde04933d2e703a590054fe709dcfa88c0
DIFF: https://github.com/llvm/llvm-project/commit/92db09cde04933d2e703a590054fe709dcfa88c0.diff
LOG: [mlir] AsyncRuntime: use int64_t for ref counting operations
Workaround for SystemZ ABI problem: https://bugs.llvm.org/show_bug.cgi?id=51898
Reviewed By: ftynse
Differential Revision: https://reviews.llvm.org/D110550
Added:
Modified:
mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
mlir/include/mlir/ExecutionEngine/AsyncRuntime.h
mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp
mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCountingOpt.cpp
mlir/lib/ExecutionEngine/AsyncRuntime.cpp
mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm.mlir
mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir
mlir/test/Dialect/Async/async-runtime-policy-based-ref-counting.mlir
mlir/test/Dialect/Async/async-runtime-ref-counting-opt.mlir
mlir/test/Dialect/Async/async-runtime-ref-counting.mlir
mlir/test/Dialect/Async/runtime.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
index ca7aa3264605..a7fc743864d0 100644
--- a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
+++ b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
@@ -508,7 +508,7 @@ def Async_RuntimeAddRefOp : Async_Op<"runtime.add_ref"> {
}];
let arguments = (ins Async_AnyAsyncType:$operand,
- Confined<I32Attr, [IntPositive]>:$count);
+ Confined<I64Attr, [IntPositive]>:$count);
let assemblyFormat = [{
$operand attr-dict `:` type($operand)
@@ -523,7 +523,7 @@ def Async_RuntimeDropRefOp : Async_Op<"runtime.drop_ref"> {
}];
let arguments = (ins Async_AnyAsyncType:$operand,
- Confined<I32Attr, [IntPositive]>:$count);
+ Confined<I64Attr, [IntPositive]>:$count);
let assemblyFormat = [{
$operand attr-dict `:` type($operand)
diff --git a/mlir/include/mlir/ExecutionEngine/AsyncRuntime.h b/mlir/include/mlir/ExecutionEngine/AsyncRuntime.h
index a101b28bb428..af4f98e154d8 100644
--- a/mlir/include/mlir/ExecutionEngine/AsyncRuntime.h
+++ b/mlir/include/mlir/ExecutionEngine/AsyncRuntime.h
@@ -51,10 +51,10 @@ using CoroResume = void (*)(void *); // coroutine resume function
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 @@ extern "C" AsyncToken *mlirAsyncRuntimeCreateToken();
// 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
index 8e0c1a23f2e4..e4b31faedd49 100644
--- a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
+++ b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
@@ -76,7 +76,7 @@ struct AsyncAPI {
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 @@ struct AsyncAPI {
}
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 @@ class RuntimeCreateOpLowering : public OpConversionPattern<RuntimeCreateOp> {
// 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,
@@ -904,8 +904,8 @@ class RefCountingOpLowering : public OpConversionPattern<RefCountingOp> {
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,
diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp
index 2dc6cfe9625e..4b4217a94616 100644
--- a/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp
+++ b/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp
@@ -44,7 +44,7 @@ static LogicalResult dropRefIfNoUses(Value value, unsigned count = 1) {
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();
}
@@ -304,7 +304,7 @@ LogicalResult AsyncRuntimeRefCountingPass::addDropRefAfterLastUse(Value value) {
// 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();
@@ -322,7 +322,7 @@ AsyncRuntimeRefCountingPass::addAddRefBeforeFunctionCall(Value value) {
// 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();
@@ -411,7 +411,7 @@ AsyncRuntimeRefCountingPass::addDropRefInDivergentLivenessSuccessor(
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)
@@ -502,13 +502,13 @@ AsyncRuntimePolicyBasedRefCountingPass::addRefCounting(Value value) {
// 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));
}
}
}
diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCountingOpt.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCountingOpt.cpp
index 063c2050e37a..025dab515bc5 100644
--- a/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCountingOpt.cpp
+++ b/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCountingOpt.cpp
@@ -118,11 +118,11 @@ LogicalResult AsyncRuntimeRefCountingOptPass::optimizeReferenceCounting(
//
// %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
index a8aeaec60d3a..d38967c4c258 100644
--- a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp
+++ b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp
@@ -55,7 +55,7 @@ class AsyncRuntime {
"all ref counted objects must be destroyed");
}
- int32_t getNumRefCountedObjects() {
+ int64_t getNumRefCountedObjects() {
return numRefCountedObjects.load(std::memory_order_relaxed);
}
@@ -73,7 +73,7 @@ class AsyncRuntime {
numRefCountedObjects.fetch_sub(1, std::memory_order_relaxed);
}
- std::atomic<int32_t> numRefCountedObjects;
+ std::atomic<int64_t> numRefCountedObjects;
llvm::ThreadPool threadPool;
};
@@ -123,7 +123,7 @@ class State {
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 @@ class RefCounted {
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 @@ class RefCounted {
private:
AsyncRuntime *runtime;
- std::atomic<int32_t> refCount;
+ std::atomic<int64_t> refCount;
};
} // namespace
@@ -192,7 +192,7 @@ struct AsyncToken : public RefCounted {
// 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 @@ struct AsyncGroup : public RefCounted {
};
// 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);
}
@@ -243,7 +243,7 @@ extern "C" AsyncToken *mlirAsyncRuntimeCreateToken() {
}
// 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
index ec35dab19c9a..6c40a083e635 100644
--- a/mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm.mlir
+++ b/mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm.mlir
@@ -10,7 +10,7 @@ func @create_token() {
// 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]])
diff --git a/mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir b/mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir
index cfb762035040..112cd515d3a2 100644
--- 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
index 34cfb84bee35..1456e1828b8f 100644
--- 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 @@ func @group_await(%arg0: !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
}
@@ -32,7 +32,7 @@ func @add_token_to_group(%arg0: !async.group, %arg1: !async.token) {
// 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
}
@@ -41,7 +41,7 @@ func @value_load(%arg0: !async.value<f32>) {
// 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
index 5d32201e9b91..6da28758f06c 100644
--- 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 private @consume_token(%arg0: !async.token)
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 @@ func @cancellable_operations_0(%arg0: !async.token) {
// 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 @@ func @cancellable_operations_2(%arg0: !async.token) {
// 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 @@ func @cancellable_operations_3(%arg0: !async.token) {
// 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
index cb3f04ff91ca..561696f6304c 100644
--- a/mlir/test/Dialect/Async/async-runtime-ref-counting.mlir
+++ b/mlir/test/Dialect/Async/async-runtime-ref-counting.mlir
@@ -12,14 +12,14 @@ 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
}
@@ -27,7 +27,7 @@ func @token_value_no_uses() {
// 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 @@ func @token_returned_no_uses() {
// 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_arg_to_func(%arg0: !async.token) {
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 @@ func @token_arg_cond_br_await_with_fallthough(%arg0: !async.token, %arg1: i1) {
^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 @@ func @token_coro_await_and_resume(%arg0: !async.token) -> !async.token {
%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 @@ func @value_coro_await_and_resume(%arg0: !async.value<f32>) -> !async.token {
^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
@@ -160,7 +160,7 @@ func private @outlined_async_execute(%arg0: !async.token) -> !async.token {
// 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 @@ func @token_await_inside_nested_region(%arg0: i1) {
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 @@ func @token_defined_in_the_loop() {
// 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 @@ func @divergent_liveness_one_token(%arg0 : i1) {
^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 @@ func @divergent_liveness_unique_predecessor(%arg0 : i1) {
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 @@ func @divergent_liveness_two_tokens(%arg0 : i1) {
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
index 1b39e6420b87..085ac7849bb0 100644
--- a/mlir/test/Dialect/Async/runtime.mlir
+++ b/mlir/test/Dialect/Async/runtime.mlir
@@ -154,14 +154,14 @@ func @add_to_group(%arg0: !async.token, %arg1: !async.value<f32>,
// 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
}
More information about the Mlir-commits
mailing list