| // RUN: mlir-opt %s -async-runtime-ref-counting | FileCheck %s |
| |
| // CHECK-LABEL: @token |
| func.func private @token() -> !async.token |
| |
| // CHECK-LABEL: @cond |
| func.func private @cond() -> i1 |
| |
| // CHECK-LABEL: @take_token |
| func.func private @take_token(%arg0: !async.token) |
| |
| // CHECK-LABEL: @token_arg_no_uses |
| // CHECK: %[[TOKEN:.*]]: !async.token |
| func.func @token_arg_no_uses(%arg0: !async.token) { |
| // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} |
| return |
| } |
| |
| // CHECK-LABEL: @token_value_no_uses |
| func.func @token_value_no_uses() { |
| // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token |
| // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} |
| %0 = async.runtime.create : !async.token |
| return |
| } |
| |
| // CHECK-LABEL: @token_returned_no_uses |
| func.func @token_returned_no_uses() { |
| // CHECK: %[[TOKEN:.*]] = call @token |
| // 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.func @token_arg_to_func(%arg0: !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 : i64} : !async.token |
| return |
| } |
| |
| // CHECK-LABEL: @token_value_to_func |
| 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 : i64} : !async.token |
| call @take_token(%0): (!async.token) -> () |
| // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} |
| return |
| } |
| |
| // CHECK-LABEL: @token_arg_cond_br_await_with_fallthough |
| // CHECK: %[[TOKEN:.*]]: !async.token |
| func.func @token_arg_cond_br_await_with_fallthough(%arg0: !async.token, %arg1: i1) { |
| // CHECK: cf.cond_br |
| // CHECK-SAME: ^[[BB1:.*]], ^[[BB2:.*]] |
| cf.cond_br %arg1, ^bb1, ^bb2 |
| ^bb1: |
| // CHECK: ^[[BB1]]: |
| // CHECK: cf.br ^[[BB2]] |
| cf.br ^bb2 |
| ^bb2: |
| // CHECK: ^[[BB2]]: |
| // CHECK: async.runtime.await %[[TOKEN]] |
| // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} |
| async.runtime.await %arg0 : !async.token |
| return |
| } |
| |
| // CHECK-LABEL: @token_simple_return |
| func.func @token_simple_return() -> !async.token { |
| // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token |
| %token = async.runtime.create : !async.token |
| // CHECK: return %[[TOKEN]] |
| return %token : !async.token |
| } |
| |
| // CHECK-LABEL: @token_coro_return |
| // CHECK-NOT: async.runtime.drop_ref |
| // CHECK-NOT: async.runtime.add_ref |
| func.func @token_coro_return() -> !async.token { |
| %token = async.runtime.create : !async.token |
| %id = async.coro.id |
| %hdl = async.coro.begin %id |
| %saved = async.coro.save %hdl |
| async.runtime.resume %hdl |
| async.coro.suspend %saved, ^suspend, ^resume, ^cleanup |
| ^resume: |
| cf.br ^cleanup |
| ^cleanup: |
| async.coro.free %id, %hdl |
| cf.br ^suspend |
| ^suspend: |
| async.coro.end %hdl |
| return %token : !async.token |
| } |
| |
| // CHECK-LABEL: @token_coro_await_and_resume |
| // CHECK: %[[TOKEN:.*]]: !async.token |
| func.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 : i64} |
| async.coro.suspend %saved, ^suspend, ^resume, ^cleanup |
| ^resume: |
| cf.br ^cleanup |
| ^cleanup: |
| async.coro.free %id, %hdl |
| cf.br ^suspend |
| ^suspend: |
| async.coro.end %hdl |
| return %token : !async.token |
| } |
| |
| // CHECK-LABEL: @value_coro_await_and_resume |
| // CHECK: %[[VALUE:.*]]: !async.value<f32> |
| func.func @value_coro_await_and_resume(%arg0: !async.value<f32>) -> !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 %[[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 : i64} |
| %0 = async.runtime.load %arg0 : !async.value<f32> |
| // CHECK: arith.addf %[[LOADED]], %[[LOADED]] |
| %1 = arith.addf %0, %0 : f32 |
| cf.br ^cleanup |
| ^cleanup: |
| async.coro.free %id, %hdl |
| cf.br ^suspend |
| ^suspend: |
| async.coro.end %hdl |
| return %token : !async.token |
| } |
| |
| // CHECK-LABEL: @outlined_async_execute |
| // CHECK: %[[TOKEN:.*]]: !async.token |
| func.func private @outlined_async_execute(%arg0: !async.token) -> !async.token { |
| %0 = async.runtime.create : !async.token |
| %1 = async.coro.id |
| %2 = async.coro.begin %1 |
| %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 : 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 |
| cf.br ^cleanup |
| ^cleanup: |
| // CHECK: ^[[CLEANUP:.*]]: |
| // CHECK: async.coro.free |
| async.coro.free %1, %2 |
| cf.br ^suspend |
| ^suspend: |
| // CHECK: ^[[SUSPEND:.*]]: |
| // CHECK: async.coro.end |
| async.coro.end %2 |
| return %0 : !async.token |
| } |
| |
| // CHECK-LABEL: @token_await_inside_nested_region |
| // CHECK: %[[ARG:.*]]: i1 |
| func.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 : i64} |
| // CHECK: return |
| return |
| } |
| |
| // CHECK-LABEL: @token_defined_in_the_loop |
| func.func @token_defined_in_the_loop() { |
| cf.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 : i64} |
| async.runtime.await %token : !async.token |
| %0 = call @cond(): () -> (i1) |
| cf.cond_br %0, ^bb1, ^bb2 |
| ^bb2: |
| // CHECK: ^[[BB2:.*]]: |
| // CHECK: return |
| return |
| } |
| |
| // CHECK-LABEL: @divergent_liveness_one_token |
| func.func @divergent_liveness_one_token(%arg0 : i1) { |
| // CHECK: %[[TOKEN:.*]] = call @token() |
| %token = call @token() : () -> !async.token |
| // CHECK: cf.cond_br %arg0, ^[[LIVE_IN:.*]], ^[[REF_COUNTING:.*]] |
| cf.cond_br %arg0, ^bb1, ^bb2 |
| ^bb1: |
| // CHECK: ^[[LIVE_IN]]: |
| // CHECK: async.runtime.await %[[TOKEN]] |
| // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} |
| // CHECK: cf.br ^[[RETURN:.*]] |
| async.runtime.await %token : !async.token |
| cf.br ^bb2 |
| // CHECK: ^[[REF_COUNTING:.*]]: |
| // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} |
| // CHECK: cf.br ^[[RETURN:.*]] |
| ^bb2: |
| // CHECK: ^[[RETURN]]: |
| // CHECK: return |
| return |
| } |
| |
| // CHECK-LABEL: @divergent_liveness_unique_predecessor |
| func.func @divergent_liveness_unique_predecessor(%arg0 : i1) { |
| // CHECK: %[[TOKEN:.*]] = call @token() |
| %token = call @token() : () -> !async.token |
| // CHECK: cf.cond_br %arg0, ^[[LIVE_IN:.*]], ^[[NO_LIVE_IN:.*]] |
| cf.cond_br %arg0, ^bb2, ^bb1 |
| ^bb1: |
| // CHECK: ^[[NO_LIVE_IN]]: |
| // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} |
| // CHECK: cf.br ^[[RETURN:.*]] |
| cf.br ^bb3 |
| ^bb2: |
| // CHECK: ^[[LIVE_IN]]: |
| // CHECK: async.runtime.await %[[TOKEN]] |
| // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} |
| // CHECK: cf.br ^[[RETURN]] |
| async.runtime.await %token : !async.token |
| cf.br ^bb3 |
| ^bb3: |
| // CHECK: ^[[RETURN]]: |
| // CHECK: return |
| return |
| } |
| |
| // CHECK-LABEL: @divergent_liveness_two_tokens |
| func.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: cf.cond_br %arg0, ^[[AWAIT0:.*]], ^[[AWAIT1:.*]] |
| cf.cond_br %arg0, ^await0, ^await1 |
| ^await0: |
| // CHECK: ^[[AWAIT0]]: |
| // CHECK: async.runtime.drop_ref %[[TOKEN1]] {count = 1 : i64} |
| // CHECK: async.runtime.await %[[TOKEN0]] |
| // CHECK: async.runtime.drop_ref %[[TOKEN0]] {count = 1 : i64} |
| // CHECK: cf.br ^[[RETURN:.*]] |
| async.runtime.await %token0 : !async.token |
| cf.br ^ret |
| ^await1: |
| // CHECK: ^[[AWAIT1]]: |
| // CHECK: async.runtime.drop_ref %[[TOKEN0]] {count = 1 : i64} |
| // CHECK: async.runtime.await %[[TOKEN1]] |
| // CHECK: async.runtime.drop_ref %[[TOKEN1]] {count = 1 : i64} |
| // CHECK: cf.br ^[[RETURN]] |
| async.runtime.await %token1 : !async.token |
| cf.br ^ret |
| ^ret: |
| // CHECK: ^[[RETURN]]: |
| // CHECK: return |
| return |
| } |