maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[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
This commit is contained in:
Eugene Zhulenev 2021-09-27 07:06:54 -07:00
parent 94ea36649e
commit 92db09cde0
12 changed files with 78 additions and 78 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
View File

@ -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)

6
mlir/include/mlir/ExecutionEngine/AsyncRuntime.h
View File

@ -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);

18
mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
View File

@ -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 @@ public:
// 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 @@ public:
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,

12
mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp
View File

@ -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));
}
}
}

4
mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCountingOpt.cpp
View File

@ -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

22
mlir/lib/ExecutionEngine/AsyncRuntime.cpp
View File

@ -55,7 +55,7 @@ public:
"all ref counted objects must be destroyed");
}
int32_t getNumRefCountedObjects() {
int64_t getNumRefCountedObjects() {
return numRefCountedObjects.load(std::memory_order_relaxed);
}
@ -73,7 +73,7 @@ private:
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 @@ private:
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 @@ public:
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 @@ protected:
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;
}

2
mlir/test/Conversion/AsyncToLLVM/convert-runtime-to-llvm.mlir
View File

@ -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]])

8
mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir
View File

@ -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
}

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

@ -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
}

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

@ -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
}

42
mlir/test/Dialect/Async/async-runtime-ref-counting.mlir
View File

@ -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

8
mlir/test/Dialect/Async/runtime.mlir
View File

@ -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
}