1// RUN: mlir-opt %s -split-input-file -async-func-to-async-runtime \ 2// RUN: -async-to-async-runtime | FileCheck %s --dump-input=always 3 4// CHECK-LABEL: @execute_no_async_args 5func.func @execute_no_async_args(%arg0: f32, %arg1: memref<1xf32>) { 6 %token = async.execute { 7 %c0 = arith.constant 0 : index 8 memref.store %arg0, %arg1[%c0] : memref<1xf32> 9 async.yield 10 } 11 async.await %token : !async.token 12 return 13} 14 15// Function outlined from the async.execute operation. 16// CHECK-LABEL: func private @async_execute_fn 17// CHECK-SAME: -> !async.token 18 19// Create token for return op, and mark a function as a coroutine. 20// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 21// CHECK: %[[ID:.*]] = async.coro.id 22// CHECK: %[[HDL:.*]] = async.coro.begin 23 24// Pass a suspended coroutine to the async runtime. 25// CHECK: %[[SAVED:.*]] = async.coro.save %[[HDL]] 26// CHECK: async.runtime.resume %[[HDL]] 27// CHECK: async.coro.suspend %[[SAVED]] 28// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[DESTROY:.*]] 29 30// Resume coroutine after suspension. 31// CHECK: ^[[RESUME]]: 32// CHECK: memref.store 33// CHECK: async.runtime.set_available %[[TOKEN]] 34// CHECK: cf.br ^[[CLEANUP:.*]] 35 36// Delete coroutine. 37// CHECK: ^[[CLEANUP]]: 38// CHECK: async.coro.free %[[ID]], %[[HDL]] 39 40// Delete coroutine. 41// CHECK: ^[[DESTROY]]: 42// CHECK: async.coro.free %[[ID]], %[[HDL]] 43 44// Suspend coroutine, and also a return statement for ramp function. 45// CHECK: ^[[SUSPEND]]: 46// CHECK: async.coro.end %[[HDL]] 47// CHECK: return %[[TOKEN]] 48 49// ----- 50 51// CHECK-LABEL: @nested_async_execute 52func.func @nested_async_execute(%arg0: f32, %arg1: f32, %arg2: memref<1xf32>) { 53 // CHECK: %[[TOKEN:.*]] = call @async_execute_fn_0(%arg0, %arg2, %arg1) 54 %token0 = async.execute { 55 %c0 = arith.constant 0 : index 56 57 %token1 = async.execute { 58 %c1 = arith.constant 1: index 59 memref.store %arg0, %arg2[%c0] : memref<1xf32> 60 async.yield 61 } 62 async.await %token1 : !async.token 63 64 memref.store %arg1, %arg2[%c0] : memref<1xf32> 65 async.yield 66 } 67 // CHECK: async.runtime.await %[[TOKEN]] 68 // CHECK: %[[IS_ERROR:.*]] = async.runtime.is_error %[[TOKEN]] 69 // CHECK: %[[TRUE:.*]] = arith.constant true 70 // CHECK: %[[NOT_ERROR:.*]] = arith.xori %[[IS_ERROR]], %[[TRUE]] : i1 71 // CHECK: cf.assert %[[NOT_ERROR]] 72 // CHECK-NEXT: return 73 async.await %token0 : !async.token 74 return 75} 76 77// Function outlined from the inner async.execute operation. 78// CHECK-LABEL: func private @async_execute_fn 79// CHECK-SAME: -> !async.token 80 81// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 82// CHECK: %[[ID:.*]] = async.coro.id 83// CHECK: %[[HDL:.*]] = async.coro.begin 84 85// CHECK: async.runtime.resume %[[HDL]] 86// CHECK: async.coro.suspend 87// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[DESTROY:.*]] 88 89// CHECK: ^[[RESUME]]: 90// CHECK: memref.store 91// CHECK: async.runtime.set_available %[[TOKEN]] 92// CHECK: cf.br ^[[CLEANUP:.*]] 93 94// CHECK: ^[[CLEANUP]]: 95// CHECK: ^[[DESTROY]]: 96 97// Function outlined from the outer async.execute operation. 98// CHECK-LABEL: func private @async_execute_fn_0 99// CHECK-SAME: -> !async.token 100 101// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 102// CHECK: %[[ID:.*]] = async.coro.id 103// CHECK: %[[HDL:.*]] = async.coro.begin 104 105// Suspend coroutine in the beginning. 106// CHECK: async.runtime.resume %[[HDL]] 107// CHECK: async.coro.suspend 108// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_0:.*]], ^[[DESTROY_0:.*]] 109 110// Suspend coroutine second time waiting for the completion of inner execute op. 111// CHECK: ^[[RESUME_0]]: 112// CHECK: %[[INNER_TOKEN:.*]] = call @async_execute_fn 113// CHECK: %[[SAVED:.*]] = async.coro.save %[[HDL]] 114// CHECK: async.runtime.await_and_resume %[[INNER_TOKEN]], %[[HDL]] 115// CHECK: async.coro.suspend %[[SAVED]] 116// CHECK-SAME: ^[[SUSPEND]], ^[[RESUME_1:.*]], ^[[DESTROY_0]] 117 118// Check the error of the awaited token after resumption. 119// CHECK: ^[[RESUME_1]]: 120// CHECK: %[[ERR:.*]] = async.runtime.is_error %[[INNER_TOKEN]] 121// CHECK: cf.cond_br %[[ERR]], ^[[SET_ERROR:.*]], ^[[CONTINUATION:.*]] 122 123// Set token available if the token is not in the error state. 124// CHECK: ^[[CONTINUATION:.*]]: 125// CHECK: memref.store 126// CHECK: async.runtime.set_available %[[TOKEN]] 127// CHECK: cf.br ^[[CLEANUP_0:.*]] 128 129// CHECK: ^[[SET_ERROR]]: 130// CHECK: ^[[CLEANUP_0]]: 131// CHECK: ^[[DESTROY_0]]: 132// CHECK: ^[[SUSPEND]]: 133 134// ----- 135 136// CHECK-LABEL: @async_execute_token_dependency 137func.func @async_execute_token_dependency(%arg0: f32, %arg1: memref<1xf32>) { 138 // CHECK: %[[TOKEN:.*]] = call @async_execute_fn 139 %token = async.execute { 140 %c0 = arith.constant 0 : index 141 memref.store %arg0, %arg1[%c0] : memref<1xf32> 142 async.yield 143 } 144 // CHECK: call @async_execute_fn_0(%[[TOKEN]], %arg0, %arg1) 145 %token_0 = async.execute [%token] { 146 %c0 = arith.constant 0 : index 147 memref.store %arg0, %arg1[%c0] : memref<1xf32> 148 async.yield 149 } 150 return 151} 152 153// Function outlined from the first async.execute operation. 154// CHECK-LABEL: func private @async_execute_fn 155// CHECK-SAME: -> !async.token 156// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 157// CHECK: return %[[TOKEN]] : !async.token 158 159// Function outlined from the second async.execute operation with dependency. 160// CHECK-LABEL: func private @async_execute_fn_0 161// CHECK-SAME: %[[ARG0:.*]]: !async.token 162// CHECK-SAME: %[[ARG1:.*]]: f32 163// CHECK-SAME: %[[ARG2:.*]]: memref<1xf32> 164// CHECK-SAME: -> !async.token 165// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 166// CHECK: %[[HDL:.*]] = async.coro.begin 167 168// Suspend coroutine in the beginning. 169// CHECK: async.runtime.resume %[[HDL]] 170// CHECK: async.coro.suspend 171// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_0:.*]], ^[[CLEANUP:.*]] 172 173// Suspend coroutine second time waiting for the completion of token dependency. 174// CHECK: ^[[RESUME_0]]: 175// CHECK: %[[SAVED:.*]] = async.coro.save %[[HDL]] 176// CHECK: async.runtime.await_and_resume %[[ARG0]], %[[HDL]] 177// CHECK: async.coro.suspend %[[SAVED]] 178// CHECK-SAME: ^[[SUSPEND]], ^[[RESUME_1:.*]], ^[[CLEANUP]] 179 180// Check the error of the awaited token after resumption. 181// CHECK: ^[[RESUME_1]]: 182// CHECK: %[[ERR:.*]] = async.runtime.is_error %[[ARG0]] 183// CHECK: cf.cond_br %[[ERR]], ^[[SET_ERROR:.*]], ^[[CONTINUATION:.*]] 184 185// Emplace result token after second resumption and error checking. 186// CHECK: ^[[CONTINUATION:.*]]: 187// CHECK: memref.store 188// CHECK: async.runtime.set_available %[[TOKEN]] 189 190// CHECK: ^[[CLEANUP]]: 191// CHECK: ^[[SUSPEND]]: 192 193// ----- 194 195// CHECK-LABEL: @async_group_await_all 196func.func @async_group_await_all(%arg0: f32, %arg1: memref<1xf32>) { 197 // CHECK: %[[C:.*]] = arith.constant 1 : index 198 %c = arith.constant 1 : index 199 // CHECK: %[[GROUP:.*]] = async.runtime.create_group %[[C]] : !async.group 200 %0 = async.create_group %c : !async.group 201 202 // CHECK: %[[TOKEN:.*]] = call @async_execute_fn 203 %token = async.execute { async.yield } 204 // CHECK: async.runtime.add_to_group %[[TOKEN]], %[[GROUP]] 205 async.add_to_group %token, %0 : !async.token 206 207 // CHECK: call @async_execute_fn_0 208 async.execute { 209 async.await_all %0 210 async.yield 211 } 212 213 // CHECK: async.runtime.await %[[GROUP]] : !async.group 214 async.await_all %0 215 return 216} 217 218// Function outlined from the second async.execute operation. 219// CHECK-LABEL: func private @async_execute_fn_0 220// CHECK-SAME: (%[[ARG:.*]]: !async.group) -> !async.token 221 222// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 223// CHECK: %[[HDL:.*]] = async.coro.begin 224 225// Suspend coroutine in the beginning. 226// CHECK: async.runtime.resume %[[HDL]] 227// CHECK: async.coro.suspend 228// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_0:.*]], ^[[CLEANUP:.*]] 229 230// Suspend coroutine second time waiting for the group. 231// CHECK: ^[[RESUME_0]]: 232// CHECK: async.runtime.await_and_resume %[[ARG]], %[[HDL]] 233// CHECK: async.coro.suspend 234// CHECK-SAME: ^[[SUSPEND]], ^[[RESUME_1:.*]], ^[[CLEANUP]] 235 236// Check the error of the awaited token after resumption. 237// CHECK: ^[[RESUME_1]]: 238// CHECK: %[[ERR:.*]] = async.runtime.is_error %[[ARG]] 239// CHECK: cf.cond_br %[[ERR]], ^[[SET_ERROR:.*]], ^[[CONTINUATION:.*]] 240 241// Emplace result token after error checking. 242// CHECK: ^[[CONTINUATION:.*]]: 243// CHECK: async.runtime.set_available %[[TOKEN]] 244 245// CHECK: ^[[CLEANUP]]: 246// CHECK: ^[[SUSPEND]]: 247 248// ----- 249 250// CHECK-LABEL: @execute_and_return_f32 251func.func @execute_and_return_f32() -> f32 { 252 // CHECK: %[[RET:.*]]:2 = call @async_execute_fn 253 %token, %result = async.execute -> !async.value<f32> { 254 %c0 = arith.constant 123.0 : f32 255 async.yield %c0 : f32 256 } 257 258 // CHECK: async.runtime.await %[[RET]]#1 : !async.value<f32> 259 // CHECK: %[[VALUE:.*]] = async.runtime.load %[[RET]]#1 : <f32> 260 %0 = async.await %result : !async.value<f32> 261 262 // CHECK: return %[[VALUE]] 263 return %0 : f32 264} 265 266// Function outlined from the async.execute operation. 267// CHECK-LABEL: func private @async_execute_fn() 268// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 269// CHECK: %[[VALUE:.*]] = async.runtime.create : !async.value<f32> 270// CHECK: %[[HDL:.*]] = async.coro.begin 271 272// Suspend coroutine in the beginning. 273// CHECK: async.runtime.resume %[[HDL]] 274// CHECK: async.coro.suspend 275// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]] 276 277// Emplace result value. 278// CHECK: ^[[RESUME]]: 279// CHECK: %[[CST:.*]] = arith.constant 1.230000e+02 : f32 280// CHECK: async.runtime.store %cst, %[[VALUE]] 281// CHECK: async.runtime.set_available %[[VALUE]] 282// CHECK: async.runtime.set_available %[[TOKEN]] 283 284// CHECK: ^[[CLEANUP]]: 285// CHECK: ^[[SUSPEND]]: 286 287// ----- 288 289// CHECK-LABEL: @async_value_operands 290func.func @async_value_operands() { 291 // CHECK: %[[RET:.*]]:2 = call @async_execute_fn 292 %token, %result = async.execute -> !async.value<f32> { 293 %c0 = arith.constant 123.0 : f32 294 async.yield %c0 : f32 295 } 296 297 // CHECK: %[[TOKEN:.*]] = call @async_execute_fn_0(%[[RET]]#1) 298 %token0 = async.execute(%result as %value: !async.value<f32>) { 299 %0 = arith.addf %value, %value : f32 300 async.yield 301 } 302 303 // CHECK: async.runtime.await %[[TOKEN]] : !async.token 304 async.await %token0 : !async.token 305 306 return 307} 308 309// Function outlined from the first async.execute operation. 310// CHECK-LABEL: func private @async_execute_fn() 311 312// Function outlined from the second async.execute operation. 313// CHECK-LABEL: func private @async_execute_fn_0 314// CHECK-SAME: (%[[ARG:.*]]: !async.value<f32>) -> !async.token 315 316// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 317// CHECK: %[[HDL:.*]] = async.coro.begin 318 319// Suspend coroutine in the beginning. 320// CHECK: async.runtime.resume %[[HDL]] 321// CHECK: async.coro.suspend 322// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_0:.*]], ^[[CLEANUP:.*]] 323 324// Suspend coroutine second time waiting for the async operand. 325// CHECK: ^[[RESUME_0]]: 326// CHECK: async.runtime.await_and_resume %[[ARG]], %[[HDL]] 327// CHECK: async.coro.suspend 328// CHECK-SAME: ^[[SUSPEND]], ^[[RESUME_1:.*]], ^[[CLEANUP]] 329 330// Check the error of the awaited token after resumption. 331// CHECK: ^[[RESUME_1]]: 332// CHECK: %[[ERR:.*]] = async.runtime.is_error %[[ARG]] 333// CHECK: cf.cond_br %[[ERR]], ^[[SET_ERROR:.*]], ^[[CONTINUATION:.*]] 334 335// // Load from the async.value argument after error checking. 336// CHECK: ^[[CONTINUATION:.*]]: 337// CHECK: %[[LOADED:.*]] = async.runtime.load %[[ARG]] : <f32 338// CHECK: arith.addf %[[LOADED]], %[[LOADED]] : f32 339// CHECK: async.runtime.set_available %[[TOKEN]] 340 341// CHECK: ^[[CLEANUP]]: 342// CHECK: ^[[SUSPEND]]: 343 344// ----- 345 346// CHECK-LABEL: @execute_assertion 347func.func @execute_assertion(%arg0: i1) { 348 %token = async.execute { 349 cf.assert %arg0, "error" 350 async.yield 351 } 352 async.await %token : !async.token 353 return 354} 355 356// Function outlined from the async.execute operation. 357// CHECK-LABEL: func private @async_execute_fn( 358// CHECK-SAME: %[[ARG0:.*]]: i1 359// CHECK-SAME: -> !async.token 360 361// Create token for return op, and mark a function as a coroutine. 362// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 363// CHECK: %[[ID:.*]] = async.coro.id 364// CHECK: %[[HDL:.*]] = async.coro.begin 365 366// Initial coroutine suspension. 367// CHECK: async.coro.suspend 368// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[DESTROY:.*]] 369 370// Resume coroutine after suspension. 371// CHECK: ^[[RESUME]]: 372// CHECK: cf.cond_br %[[ARG0]], ^[[SET_AVAILABLE:.*]], ^[[SET_ERROR:.*]] 373 374// Set coroutine completion token to available state. 375// CHECK: ^[[SET_AVAILABLE]]: 376// CHECK: async.runtime.set_available %[[TOKEN]] 377// CHECK: cf.br ^[[CLEANUP:.*]] 378 379// Set coroutine completion token to error state. 380// CHECK: ^[[SET_ERROR]]: 381// CHECK: async.runtime.set_error %[[TOKEN]] 382// CHECK: cf.br ^[[CLEANUP]] 383 384// Delete coroutine. 385// CHECK: ^[[CLEANUP]]: 386// CHECK: async.coro.free %[[ID]], %[[HDL]] 387 388// Delete coroutine. 389// CHECK: ^[[DESTROY]]: 390// CHECK: async.coro.free %[[ID]], %[[HDL]] 391 392// Suspend coroutine, and also a return statement for ramp function. 393// CHECK: ^[[SUSPEND]]: 394// CHECK: async.coro.end %[[HDL]] 395// CHECK: return %[[TOKEN]] 396 397// ----- 398// Structured control flow operations with async operations in the body must be 399// lowered to branch-based control flow to enable coroutine CFG rewrite. 400 401// CHECK-LABEL: @lower_scf_to_cfg 402func.func @lower_scf_to_cfg(%arg0: f32, %arg1: memref<1xf32>, %arg2: i1) { 403 %token0 = async.execute { async.yield } 404 %token1 = async.execute { 405 scf.if %arg2 { 406 async.await %token0 : !async.token 407 } else { 408 async.await %token0 : !async.token 409 } 410 async.yield 411 } 412 return 413} 414 415// Function outlined from the first async.execute operation. 416// CHECK-LABEL: func private @async_execute_fn( 417// CHECK-SAME: -> !async.token 418 419// Function outlined from the second async.execute operation. 420// CHECK-LABEL: func private @async_execute_fn_0( 421// CHECK: %[[TOKEN:.*]]: !async.token 422// CHECK: %[[FLAG:.*]]: i1 423// CHECK-SAME: -> !async.token 424 425// Check that structured control flow lowered to CFG. 426// CHECK-NOT: scf.if 427// CHECK: cf.cond_br %[[FLAG]] 428 429// ----- 430// Constants captured by the async.execute region should be cloned into the 431// outline async execute function. 432 433// CHECK-LABEL: @clone_constants 434func.func @clone_constants(%arg0: f32, %arg1: memref<1xf32>) { 435 %c0 = arith.constant 0 : index 436 %token = async.execute { 437 memref.store %arg0, %arg1[%c0] : memref<1xf32> 438 async.yield 439 } 440 async.await %token : !async.token 441 return 442} 443 444// Function outlined from the async.execute operation. 445// CHECK-LABEL: func private @async_execute_fn( 446// CHECK-SAME: %[[VALUE:arg[0-9]+]]: f32, 447// CHECK-SAME: %[[MEMREF:arg[0-9]+]]: memref<1xf32> 448// CHECK-SAME: ) -> !async.token 449// CHECK: %[[CST:.*]] = arith.constant 0 : index 450// CHECK: memref.store %[[VALUE]], %[[MEMREF]][%[[CST]]] 451 452// ----- 453// Async Functions should be none blocking 454 455// CHECK-LABEL: @async_func_await 456async.func @async_func_await(%arg0: f32, %arg1: !async.value<f32>) 457 -> !async.token { 458 %0 = async.await %arg1 : !async.value<f32> 459 return 460} 461// Create token for return op, and mark a function as a coroutine. 462// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token 463// CHECK: %[[ID:.*]] = async.coro.id 464// CHECK: %[[HDL:.*]] = async.coro.begin 465// CHECK: cf.br ^[[ORIGIN_ENTRY:.*]] 466 467// CHECK: ^[[ORIGIN_ENTRY]]: 468// CHECK: %[[SAVED:.*]] = async.coro.save %[[HDL]] 469// CHECK: async.runtime.await_and_resume %[[arg1:.*]], %[[HDL]] : 470// CHECK-SAME: !async.value<f32> 471// CHECK: async.coro.suspend %[[SAVED]] 472// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]] 473 474// ----- 475// Async execute inside async func 476 477// CHECK-LABEL: @execute_in_async_func 478async.func @execute_in_async_func(%arg0: f32, %arg1: memref<1xf32>) 479 -> !async.token { 480 %token = async.execute { 481 %c0 = arith.constant 0 : index 482 memref.store %arg0, %arg1[%c0] : memref<1xf32> 483 async.yield 484 } 485 async.await %token : !async.token 486 return 487} 488// Call outlind async execute Function 489// CHECK: %[[RES:.*]] = call @async_execute_fn( 490// CHECK-SAME: %[[VALUE:arg[0-9]+]], 491// CHECK-SAME: %[[MEMREF:arg[0-9]+]] 492// CHECK-SAME: ) : (f32, memref<1xf32>) -> !async.token 493 494// Function outlined from the async.execute operation. 495// CHECK-LABEL: func private @async_execute_fn( 496// CHECK-SAME: %[[VALUE:arg[0-9]+]]: f32, 497// CHECK-SAME: %[[MEMREF:arg[0-9]+]]: memref<1xf32> 498// CHECK-SAME: ) -> !async.token 499// CHECK: %[[CST:.*]] = arith.constant 0 : index 500// CHECK: memref.store %[[VALUE]], %[[MEMREF]][%[[CST]]] 501