292 lines · plain
1// RUN: mlir-opt %s -async-runtime-ref-counting | FileCheck %s2 3// CHECK-LABEL: @token4func.func private @token() -> !async.token5 6// CHECK-LABEL: @cond7func.func private @cond() -> i18 9// CHECK-LABEL: @take_token10func.func private @take_token(%arg0: !async.token)11 12// CHECK-LABEL: @token_arg_no_uses13// CHECK: %[[TOKEN:.*]]: !async.token14func.func @token_arg_no_uses(%arg0: !async.token) {15 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}16 return17}18 19// CHECK-LABEL: @token_value_no_uses20func.func @token_value_no_uses() {21 // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token22 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}23 %0 = async.runtime.create : !async.token24 return25}26 27// CHECK-LABEL: @token_returned_no_uses28func.func @token_returned_no_uses() {29 // CHECK: %[[TOKEN:.*]] = call @token30 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}31 %0 = call @token() : () -> !async.token32 return33}34 35// CHECK-LABEL: @token_arg_to_func36// CHECK: %[[TOKEN:.*]]: !async.token37func.func @token_arg_to_func(%arg0: !async.token) {38 // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i64} : !async.token39 call @take_token(%arg0): (!async.token) -> ()40 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} : !async.token41 return42}43 44// CHECK-LABEL: @token_value_to_func45func.func @token_value_to_func() {46 // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token47 %0 = async.runtime.create : !async.token48 // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i64} : !async.token49 call @take_token(%0): (!async.token) -> ()50 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}51 return52}53 54// CHECK-LABEL: @token_arg_cond_br_await_with_fallthough55// CHECK: %[[TOKEN:.*]]: !async.token56func.func @token_arg_cond_br_await_with_fallthough(%arg0: !async.token, %arg1: i1) {57 // CHECK: cf.cond_br58 // CHECK-SAME: ^[[BB1:.*]], ^[[BB2:.*]]59 cf.cond_br %arg1, ^bb1, ^bb260^bb1:61 // CHECK: ^[[BB1]]:62 // CHECK: cf.br ^[[BB2]]63 cf.br ^bb264^bb2:65 // CHECK: ^[[BB2]]:66 // CHECK: async.runtime.await %[[TOKEN]]67 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}68 async.runtime.await %arg0 : !async.token69 return70}71 72// CHECK-LABEL: @token_simple_return73func.func @token_simple_return() -> !async.token {74 // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token75 %token = async.runtime.create : !async.token76 // CHECK: return %[[TOKEN]]77 return %token : !async.token78}79 80// CHECK-LABEL: @token_coro_return81// CHECK-NOT: async.runtime.drop_ref82// CHECK-NOT: async.runtime.add_ref83func.func @token_coro_return() -> !async.token {84 %token = async.runtime.create : !async.token85 %id = async.coro.id86 %hdl = async.coro.begin %id87 %saved = async.coro.save %hdl88 async.runtime.resume %hdl89 async.coro.suspend %saved, ^suspend, ^resume, ^cleanup90^resume:91 cf.br ^cleanup92^cleanup:93 async.coro.free %id, %hdl94 cf.br ^suspend95^suspend:96 async.coro.end %hdl97 return %token : !async.token98}99 100// CHECK-LABEL: @token_coro_await_and_resume101// CHECK: %[[TOKEN:.*]]: !async.token102func.func @token_coro_await_and_resume(%arg0: !async.token) -> !async.token {103 %token = async.runtime.create : !async.token104 %id = async.coro.id105 %hdl = async.coro.begin %id106 %saved = async.coro.save %hdl107 // CHECK: async.runtime.await_and_resume %[[TOKEN]]108 async.runtime.await_and_resume %arg0, %hdl : !async.token109 // CHECK-NEXT: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}110 async.coro.suspend %saved, ^suspend, ^resume, ^cleanup111^resume:112 cf.br ^cleanup113^cleanup:114 async.coro.free %id, %hdl115 cf.br ^suspend116^suspend:117 async.coro.end %hdl118 return %token : !async.token119}120 121// CHECK-LABEL: @value_coro_await_and_resume122// CHECK: %[[VALUE:.*]]: !async.value<f32>123func.func @value_coro_await_and_resume(%arg0: !async.value<f32>) -> !async.token {124 %token = async.runtime.create : !async.token125 %id = async.coro.id126 %hdl = async.coro.begin %id127 %saved = async.coro.save %hdl128 // CHECK: async.runtime.await_and_resume %[[VALUE]]129 async.runtime.await_and_resume %arg0, %hdl : !async.value<f32>130 // CHECK: async.coro.suspend131 // CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]132 async.coro.suspend %saved, ^suspend, ^resume, ^cleanup133^resume:134 // CHECK: ^[[RESUME]]:135 // CHECK: %[[LOADED:.*]] = async.runtime.load %[[VALUE]]136 // CHECK: async.runtime.drop_ref %[[VALUE]] {count = 1 : i64}137 %0 = async.runtime.load %arg0 : !async.value<f32>138 // CHECK: arith.addf %[[LOADED]], %[[LOADED]]139 %1 = arith.addf %0, %0 : f32140 cf.br ^cleanup141^cleanup:142 async.coro.free %id, %hdl143 cf.br ^suspend144^suspend:145 async.coro.end %hdl146 return %token : !async.token147}148 149// CHECK-LABEL: @outlined_async_execute150// CHECK: %[[TOKEN:.*]]: !async.token151func.func private @outlined_async_execute(%arg0: !async.token) -> !async.token {152 %0 = async.runtime.create : !async.token153 %1 = async.coro.id154 %2 = async.coro.begin %1155 %3 = async.coro.save %2156 async.runtime.resume %2157 // CHECK: async.coro.suspend158 async.coro.suspend %3, ^suspend, ^resume, ^cleanup159^resume:160 // CHECK: ^[[RESUME:.*]]:161 %4 = async.coro.save %2162 async.runtime.await_and_resume %arg0, %2 : !async.token163 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}164 // CHECK: async.coro.suspend165 async.coro.suspend %4, ^suspend, ^resume_1, ^cleanup166^resume_1:167 // CHECK: ^[[RESUME_1:.*]]:168 // CHECK: async.runtime.set_available169 async.runtime.set_available %0 : !async.token170 cf.br ^cleanup171^cleanup:172 // CHECK: ^[[CLEANUP:.*]]:173 // CHECK: async.coro.free174 async.coro.free %1, %2175 cf.br ^suspend176^suspend:177 // CHECK: ^[[SUSPEND:.*]]:178 // CHECK: async.coro.end179 async.coro.end %2180 return %0 : !async.token181}182 183// CHECK-LABEL: @token_await_inside_nested_region184// CHECK: %[[ARG:.*]]: i1185func.func @token_await_inside_nested_region(%arg0: i1) {186 // CHECK: %[[TOKEN:.*]] = call @token()187 %token = call @token() : () -> !async.token188 // CHECK: scf.if %[[ARG]] {189 scf.if %arg0 {190 // CHECK: async.runtime.await %[[TOKEN]]191 async.runtime.await %token : !async.token192 }193 // CHECK: }194 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}195 // CHECK: return196 return197}198 199// CHECK-LABEL: @token_defined_in_the_loop200func.func @token_defined_in_the_loop() {201 cf.br ^bb1202^bb1:203 // CHECK: ^[[BB1:.*]]:204 // CHECK: %[[TOKEN:.*]] = call @token()205 %token = call @token() : () -> !async.token206 // CHECK: async.runtime.await %[[TOKEN]]207 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}208 async.runtime.await %token : !async.token209 %0 = call @cond(): () -> (i1)210 cf.cond_br %0, ^bb1, ^bb2211^bb2:212 // CHECK: ^[[BB2:.*]]:213 // CHECK: return214 return215}216 217// CHECK-LABEL: @divergent_liveness_one_token218func.func @divergent_liveness_one_token(%arg0 : i1) {219 // CHECK: %[[TOKEN:.*]] = call @token()220 %token = call @token() : () -> !async.token221 // CHECK: cf.cond_br %arg0, ^[[LIVE_IN:.*]], ^[[REF_COUNTING:.*]]222 cf.cond_br %arg0, ^bb1, ^bb2223^bb1:224 // CHECK: ^[[LIVE_IN]]:225 // CHECK: async.runtime.await %[[TOKEN]]226 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}227 // CHECK: cf.br ^[[RETURN:.*]]228 async.runtime.await %token : !async.token229 cf.br ^bb2230 // CHECK: ^[[REF_COUNTING:.*]]:231 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}232 // CHECK: cf.br ^[[RETURN:.*]]233^bb2:234 // CHECK: ^[[RETURN]]:235 // CHECK: return236 return237}238 239// CHECK-LABEL: @divergent_liveness_unique_predecessor240func.func @divergent_liveness_unique_predecessor(%arg0 : i1) {241 // CHECK: %[[TOKEN:.*]] = call @token()242 %token = call @token() : () -> !async.token243 // CHECK: cf.cond_br %arg0, ^[[LIVE_IN:.*]], ^[[NO_LIVE_IN:.*]]244 cf.cond_br %arg0, ^bb2, ^bb1245^bb1:246 // CHECK: ^[[NO_LIVE_IN]]:247 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}248 // CHECK: cf.br ^[[RETURN:.*]]249 cf.br ^bb3250^bb2:251 // CHECK: ^[[LIVE_IN]]:252 // CHECK: async.runtime.await %[[TOKEN]]253 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64}254 // CHECK: cf.br ^[[RETURN]]255 async.runtime.await %token : !async.token256 cf.br ^bb3257^bb3:258 // CHECK: ^[[RETURN]]:259 // CHECK: return260 return261}262 263// CHECK-LABEL: @divergent_liveness_two_tokens264func.func @divergent_liveness_two_tokens(%arg0 : i1) {265 // CHECK: %[[TOKEN0:.*]] = call @token()266 // CHECK: %[[TOKEN1:.*]] = call @token()267 %token0 = call @token() : () -> !async.token268 %token1 = call @token() : () -> !async.token269 // CHECK: cf.cond_br %arg0, ^[[AWAIT0:.*]], ^[[AWAIT1:.*]]270 cf.cond_br %arg0, ^await0, ^await1271^await0:272 // CHECK: ^[[AWAIT0]]:273 // CHECK: async.runtime.drop_ref %[[TOKEN1]] {count = 1 : i64}274 // CHECK: async.runtime.await %[[TOKEN0]]275 // CHECK: async.runtime.drop_ref %[[TOKEN0]] {count = 1 : i64}276 // CHECK: cf.br ^[[RETURN:.*]]277 async.runtime.await %token0 : !async.token278 cf.br ^ret279^await1:280 // CHECK: ^[[AWAIT1]]:281 // CHECK: async.runtime.drop_ref %[[TOKEN0]] {count = 1 : i64}282 // CHECK: async.runtime.await %[[TOKEN1]]283 // CHECK: async.runtime.drop_ref %[[TOKEN1]] {count = 1 : i64}284 // CHECK: cf.br ^[[RETURN]]285 async.runtime.await %token1 : !async.token286 cf.br ^ret287^ret:288 // CHECK: ^[[RETURN]]:289 // CHECK: return290 return291}292