You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

254 lines
6.6 KiB

// RUN: mlir-opt %s -async-ref-counting | FileCheck %s
// CHECK-LABEL: @cond
func private @cond() -> i1
// CHECK-LABEL: @token_arg_no_uses
func @token_arg_no_uses(%arg0: !async.token) {
// CHECK: async.drop_ref %arg0 {count = 1 : i32}
return
}
// CHECK-LABEL: @token_arg_conditional_await
func @token_arg_conditional_await(%arg0: !async.token, %arg1: i1) {
cond_br %arg1, ^bb1, ^bb2
^bb1:
// CHECK: async.drop_ref %arg0 {count = 1 : i32}
return
^bb2:
// CHECK: async.await %arg0
// CHECK: async.drop_ref %arg0 {count = 1 : i32}
async.await %arg0 : !async.token
return
}
// CHECK-LABEL: @token_no_uses
func @token_no_uses() {
// CHECK: %[[TOKEN:.*]] = async.execute
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
%token = async.execute {
async.yield
}
return
}
// CHECK-LABEL: @token_return
func @token_return() -> !async.token {
// CHECK: %[[TOKEN:.*]] = async.execute
%token = async.execute {
async.yield
}
// CHECK: return %[[TOKEN]]
return %token : !async.token
}
// CHECK-LABEL: @token_await
func @token_await() {
// CHECK: %[[TOKEN:.*]] = async.execute
%token = async.execute {
async.yield
}
// CHECK: async.await %[[TOKEN]]
async.await %token : !async.token
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
// CHECK: return
return
}
// CHECK-LABEL: @token_await_and_return
func @token_await_and_return() -> !async.token {
// CHECK: %[[TOKEN:.*]] = async.execute
%token = async.execute {
async.yield
}
// CHECK: async.await %[[TOKEN]]
// CHECK-NOT: async.drop_ref
async.await %token : !async.token
// CHECK: return %[[TOKEN]]
return %token : !async.token
}
// CHECK-LABEL: @token_await_inside_scf_if
func @token_await_inside_scf_if(%arg0: i1) {
// CHECK: %[[TOKEN:.*]] = async.execute
%token = async.execute {
async.yield
}
// CHECK: scf.if %arg0 {
scf.if %arg0 {
// CHECK: async.await %[[TOKEN]]
async.await %token : !async.token
}
// CHECK: }
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
// CHECK: return
return
}
// CHECK-LABEL: @token_conditional_await
func @token_conditional_await(%arg0: i1) {
// CHECK: %[[TOKEN:.*]] = async.execute
%token = async.execute {
async.yield
}
cond_br %arg0, ^bb1, ^bb2
^bb1:
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
return
^bb2:
// CHECK: async.await %[[TOKEN]]
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
async.await %token : !async.token
return
}
// CHECK-LABEL: @token_await_in_the_loop
func @token_await_in_the_loop() {
// CHECK: %[[TOKEN:.*]] = async.execute
%token = async.execute {
async.yield
}
br ^bb1
^bb1:
// CHECK: async.await %[[TOKEN]]
async.await %token : !async.token
%0 = call @cond(): () -> (i1)
cond_br %0, ^bb1, ^bb2
^bb2:
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
return
}
// CHECK-LABEL: @token_defined_in_the_loop
func @token_defined_in_the_loop() {
br ^bb1
^bb1:
// CHECK: %[[TOKEN:.*]] = async.execute
%token = async.execute {
async.yield
}
// CHECK: async.await %[[TOKEN]]
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
async.await %token : !async.token
%0 = call @cond(): () -> (i1)
cond_br %0, ^bb1, ^bb2
^bb2:
return
}
// CHECK-LABEL: @token_capture
func @token_capture() {
// CHECK: %[[TOKEN:.*]] = async.execute
%token = async.execute {
async.yield
}
// CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
// CHECK: %[[TOKEN_0:.*]] = async.execute
%token_0 = async.execute {
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
// CHECK-NEXT: async.yield
async.await %token : !async.token
async.yield
}
// CHECK: async.drop_ref %[[TOKEN_0]] {count = 1 : i32}
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
// CHECK: return
return
}
// CHECK-LABEL: @token_nested_capture
func @token_nested_capture() {
// CHECK: %[[TOKEN:.*]] = async.execute
%token = async.execute {
async.yield
}
// CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
// CHECK: %[[TOKEN_0:.*]] = async.execute
%token_0 = async.execute {
// CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
// CHECK: %[[TOKEN_1:.*]] = async.execute
%token_1 = async.execute {
// CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
// CHECK: %[[TOKEN_2:.*]] = async.execute
%token_2 = async.execute {
// CHECK: async.await %[[TOKEN]]
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
async.await %token : !async.token
async.yield
}
// CHECK: async.drop_ref %[[TOKEN_2]] {count = 1 : i32}
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
async.yield
}
// CHECK: async.drop_ref %[[TOKEN_1]] {count = 1 : i32}
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
async.yield
}
// CHECK: async.drop_ref %[[TOKEN_0]] {count = 1 : i32}
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
// CHECK: return
return
}
// CHECK-LABEL: @token_dependency
func @token_dependency() {
// CHECK: %[[TOKEN:.*]] = async.execute
%token = async.execute {
async.yield
}
// CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
// CHECK: %[[TOKEN_0:.*]] = async.execute
%token_0 = async.execute[%token] {
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
// CHECK-NEXT: async.yield
async.yield
}
// CHECK: async.await %[[TOKEN]]
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
async.await %token : !async.token
// CHECK: async.await %[[TOKEN_0]]
// CHECK: async.drop_ref %[[TOKEN_0]] {count = 1 : i32}
async.await %token_0 : !async.token
// CHECK: return
return
}
// CHECK-LABEL: @value_operand
func @value_operand() -> f32 {
// CHECK: %[[TOKEN:.*]], %[[RESULTS:.*]] = async.execute
%token, %results = async.execute -> !async.value<f32> {
%0 = constant 0.0 : f32
async.yield %0 : f32
}
// CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
// CHECK: async.add_ref %[[RESULTS]] {count = 1 : i32}
// CHECK: %[[TOKEN_0:.*]] = async.execute
%token_0 = async.execute[%token](%results as %arg0 : !async.value<f32>) {
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
// CHECK: async.drop_ref %[[RESULTS]] {count = 1 : i32}
// CHECK: async.yield
async.yield
}
// CHECK: async.await %[[TOKEN]]
// CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
async.await %token : !async.token
// CHECK: async.await %[[TOKEN_0]]
// CHECK: async.drop_ref %[[TOKEN_0]] {count = 1 : i32}
async.await %token_0 : !async.token
// CHECK: async.await %[[RESULTS]]
// CHECK: async.drop_ref %[[RESULTS]] {count = 1 : i32}
%0 = async.await %results : !async.value<f32>
// CHECK: return
return %0 : f32
}