1// RUN: mlir-opt %s -async-runtime-ref-counting | FileCheck %s 2 3// CHECK-LABEL: @token 4func.func private @token() -> !async.token 5 6// CHECK-LABEL: @cond 7func.func private @cond() -> i1 8 9// CHECK-LABEL: @take_token 10func.func private @take_token(%arg0: !async.token) 11 12// CHECK-LABEL: @token_arg_no_uses 13// CHECK: %[[TOKEN:.*]]: !async.token 14func.func @token_arg_no_uses(%arg0: !async.token) { 15 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} 16 return 17} 18 19// CHECK-LABEL: @token_value_no_uses 20func.func @token_value_no_uses() { 21 // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 22 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} 23 %0 = async.runtime.create : !async.token 24 return 25} 26 27// CHECK-LABEL: @token_returned_no_uses 28func.func @token_returned_no_uses() { 29 // CHECK: %[[TOKEN:.*]] = call @token 30 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} 31 %0 = call @token() : () -> !async.token 32 return 33} 34 35// CHECK-LABEL: @token_arg_to_func 36// CHECK: %[[TOKEN:.*]]: !async.token 37func.func @token_arg_to_func(%arg0: !async.token) { 38 // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i64} : !async.token 39 call @take_token(%arg0): (!async.token) -> () 40 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} : !async.token 41 return 42} 43 44// CHECK-LABEL: @token_value_to_func 45func.func @token_value_to_func() { 46 // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 47 %0 = async.runtime.create : !async.token 48 // CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i64} : !async.token 49 call @take_token(%0): (!async.token) -> () 50 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} 51 return 52} 53 54// CHECK-LABEL: @token_arg_cond_br_await_with_fallthough 55// CHECK: %[[TOKEN:.*]]: !async.token 56func.func @token_arg_cond_br_await_with_fallthough(%arg0: !async.token, %arg1: i1) { 57 // CHECK: cf.cond_br 58 // CHECK-SAME: ^[[BB1:.*]], ^[[BB2:.*]] 59 cf.cond_br %arg1, ^bb1, ^bb2 60^bb1: 61 // CHECK: ^[[BB1]]: 62 // CHECK: cf.br ^[[BB2]] 63 cf.br ^bb2 64^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.token 69 return 70} 71 72// CHECK-LABEL: @token_simple_return 73func.func @token_simple_return() -> !async.token { 74 // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 75 %token = async.runtime.create : !async.token 76 // CHECK: return %[[TOKEN]] 77 return %token : !async.token 78} 79 80// CHECK-LABEL: @token_coro_return 81// CHECK-NOT: async.runtime.drop_ref 82// CHECK-NOT: async.runtime.add_ref 83func.func @token_coro_return() -> !async.token { 84 %token = async.runtime.create : !async.token 85 %id = async.coro.id 86 %hdl = async.coro.begin %id 87 %saved = async.coro.save %hdl 88 async.runtime.resume %hdl 89 async.coro.suspend %saved, ^suspend, ^resume, ^cleanup 90^resume: 91 cf.br ^cleanup 92^cleanup: 93 async.coro.free %id, %hdl 94 cf.br ^suspend 95^suspend: 96 async.coro.end %hdl 97 return %token : !async.token 98} 99 100// CHECK-LABEL: @token_coro_await_and_resume 101// CHECK: %[[TOKEN:.*]]: !async.token 102func.func @token_coro_await_and_resume(%arg0: !async.token) -> !async.token { 103 %token = async.runtime.create : !async.token 104 %id = async.coro.id 105 %hdl = async.coro.begin %id 106 %saved = async.coro.save %hdl 107 // CHECK: async.runtime.await_and_resume %[[TOKEN]] 108 async.runtime.await_and_resume %arg0, %hdl : !async.token 109 // CHECK-NEXT: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} 110 async.coro.suspend %saved, ^suspend, ^resume, ^cleanup 111^resume: 112 cf.br ^cleanup 113^cleanup: 114 async.coro.free %id, %hdl 115 cf.br ^suspend 116^suspend: 117 async.coro.end %hdl 118 return %token : !async.token 119} 120 121// CHECK-LABEL: @value_coro_await_and_resume 122// CHECK: %[[VALUE:.*]]: !async.value<f32> 123func.func @value_coro_await_and_resume(%arg0: !async.value<f32>) -> !async.token { 124 %token = async.runtime.create : !async.token 125 %id = async.coro.id 126 %hdl = async.coro.begin %id 127 %saved = async.coro.save %hdl 128 // CHECK: async.runtime.await_and_resume %[[VALUE]] 129 async.runtime.await_and_resume %arg0, %hdl : !async.value<f32> 130 // CHECK: async.coro.suspend 131 // CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]] 132 async.coro.suspend %saved, ^suspend, ^resume, ^cleanup 133^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 : f32 140 cf.br ^cleanup 141^cleanup: 142 async.coro.free %id, %hdl 143 cf.br ^suspend 144^suspend: 145 async.coro.end %hdl 146 return %token : !async.token 147} 148 149// CHECK-LABEL: @outlined_async_execute 150// CHECK: %[[TOKEN:.*]]: !async.token 151func.func private @outlined_async_execute(%arg0: !async.token) -> !async.token { 152 %0 = async.runtime.create : !async.token 153 %1 = async.coro.id 154 %2 = async.coro.begin %1 155 %3 = async.coro.save %2 156 async.runtime.resume %2 157 // CHECK: async.coro.suspend 158 async.coro.suspend %3, ^suspend, ^resume, ^cleanup 159^resume: 160 // CHECK: ^[[RESUME:.*]]: 161 %4 = async.coro.save %2 162 async.runtime.await_and_resume %arg0, %2 : !async.token 163 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} 164 // CHECK: async.coro.suspend 165 async.coro.suspend %4, ^suspend, ^resume_1, ^cleanup 166^resume_1: 167 // CHECK: ^[[RESUME_1:.*]]: 168 // CHECK: async.runtime.set_available 169 async.runtime.set_available %0 : !async.token 170 cf.br ^cleanup 171^cleanup: 172 // CHECK: ^[[CLEANUP:.*]]: 173 // CHECK: async.coro.free 174 async.coro.free %1, %2 175 cf.br ^suspend 176^suspend: 177 // CHECK: ^[[SUSPEND:.*]]: 178 // CHECK: async.coro.end 179 async.coro.end %2 180 return %0 : !async.token 181} 182 183// CHECK-LABEL: @token_await_inside_nested_region 184// CHECK: %[[ARG:.*]]: i1 185func.func @token_await_inside_nested_region(%arg0: i1) { 186 // CHECK: %[[TOKEN:.*]] = call @token() 187 %token = call @token() : () -> !async.token 188 // CHECK: scf.if %[[ARG]] { 189 scf.if %arg0 { 190 // CHECK: async.runtime.await %[[TOKEN]] 191 async.runtime.await %token : !async.token 192 } 193 // CHECK: } 194 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} 195 // CHECK: return 196 return 197} 198 199// CHECK-LABEL: @token_defined_in_the_loop 200func.func @token_defined_in_the_loop() { 201 cf.br ^bb1 202^bb1: 203 // CHECK: ^[[BB1:.*]]: 204 // CHECK: %[[TOKEN:.*]] = call @token() 205 %token = call @token() : () -> !async.token 206 // CHECK: async.runtime.await %[[TOKEN]] 207 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} 208 async.runtime.await %token : !async.token 209 %0 = call @cond(): () -> (i1) 210 cf.cond_br %0, ^bb1, ^bb2 211^bb2: 212 // CHECK: ^[[BB2:.*]]: 213 // CHECK: return 214 return 215} 216 217// CHECK-LABEL: @divergent_liveness_one_token 218func.func @divergent_liveness_one_token(%arg0 : i1) { 219 // CHECK: %[[TOKEN:.*]] = call @token() 220 %token = call @token() : () -> !async.token 221 // CHECK: cf.cond_br %arg0, ^[[LIVE_IN:.*]], ^[[REF_COUNTING:.*]] 222 cf.cond_br %arg0, ^bb1, ^bb2 223^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.token 229 cf.br ^bb2 230 // 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: return 236 return 237} 238 239// CHECK-LABEL: @divergent_liveness_unique_predecessor 240func.func @divergent_liveness_unique_predecessor(%arg0 : i1) { 241 // CHECK: %[[TOKEN:.*]] = call @token() 242 %token = call @token() : () -> !async.token 243 // CHECK: cf.cond_br %arg0, ^[[LIVE_IN:.*]], ^[[NO_LIVE_IN:.*]] 244 cf.cond_br %arg0, ^bb2, ^bb1 245^bb1: 246 // CHECK: ^[[NO_LIVE_IN]]: 247 // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i64} 248 // CHECK: cf.br ^[[RETURN:.*]] 249 cf.br ^bb3 250^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.token 256 cf.br ^bb3 257^bb3: 258 // CHECK: ^[[RETURN]]: 259 // CHECK: return 260 return 261} 262 263// CHECK-LABEL: @divergent_liveness_two_tokens 264func.func @divergent_liveness_two_tokens(%arg0 : i1) { 265 // CHECK: %[[TOKEN0:.*]] = call @token() 266 // CHECK: %[[TOKEN1:.*]] = call @token() 267 %token0 = call @token() : () -> !async.token 268 %token1 = call @token() : () -> !async.token 269 // CHECK: cf.cond_br %arg0, ^[[AWAIT0:.*]], ^[[AWAIT1:.*]] 270 cf.cond_br %arg0, ^await0, ^await1 271^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.token 278 cf.br ^ret 279^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.token 286 cf.br ^ret 287^ret: 288 // CHECK: ^[[RETURN]]: 289 // CHECK: return 290 return 291} 292