501 lines · plain
1// RUN: mlir-opt %s -split-input-file -async-func-to-async-runtime \2// RUN: -async-to-async-runtime | FileCheck %s --dump-input=always3 4// CHECK-LABEL: @execute_no_async_args5func.func @execute_no_async_args(%arg0: f32, %arg1: memref<1xf32>) {6 %token = async.execute {7 %c0 = arith.constant 0 : index8 memref.store %arg0, %arg1[%c0] : memref<1xf32>9 async.yield10 }11 async.await %token : !async.token12 return13}14 15// Function outlined from the async.execute operation.16// CHECK-LABEL: func private @async_execute_fn17// CHECK-SAME: -> !async.token18 19// Create token for return op, and mark a function as a coroutine.20// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token21// CHECK: %[[ID:.*]] = async.coro.id22// CHECK: %[[HDL:.*]] = async.coro.begin23 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.store33// 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_execute52func.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 : index56 57 %token1 = async.execute {58 %c1 = arith.constant 1: index59 memref.store %arg0, %arg2[%c0] : memref<1xf32>60 async.yield61 }62 async.await %token1 : !async.token63 64 memref.store %arg1, %arg2[%c0] : memref<1xf32>65 async.yield66 }67 // CHECK: async.runtime.await %[[TOKEN]]68 // CHECK: %[[IS_ERROR:.*]] = async.runtime.is_error %[[TOKEN]]69 // CHECK: %[[TRUE:.*]] = arith.constant true70 // CHECK: %[[NOT_ERROR:.*]] = arith.xori %[[IS_ERROR]], %[[TRUE]] : i171 // CHECK: cf.assert %[[NOT_ERROR]]72 // CHECK-NEXT: return73 async.await %token0 : !async.token74 return75}76 77// Function outlined from the inner async.execute operation.78// CHECK-LABEL: func private @async_execute_fn79// CHECK-SAME: -> !async.token80 81// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token82// CHECK: %[[ID:.*]] = async.coro.id83// CHECK: %[[HDL:.*]] = async.coro.begin84 85// CHECK: async.runtime.resume %[[HDL]]86// CHECK: async.coro.suspend87// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[DESTROY:.*]]88 89// CHECK: ^[[RESUME]]:90// CHECK: memref.store91// 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_099// CHECK-SAME: -> !async.token100 101// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token102// CHECK: %[[ID:.*]] = async.coro.id103// CHECK: %[[HDL:.*]] = async.coro.begin104 105// Suspend coroutine in the beginning.106// CHECK: async.runtime.resume %[[HDL]]107// CHECK: async.coro.suspend108// 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_fn113// 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.store126// 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_dependency137func.func @async_execute_token_dependency(%arg0: f32, %arg1: memref<1xf32>) {138 // CHECK: %[[TOKEN:.*]] = call @async_execute_fn139 %token = async.execute {140 %c0 = arith.constant 0 : index141 memref.store %arg0, %arg1[%c0] : memref<1xf32>142 async.yield143 }144 // CHECK: call @async_execute_fn_0(%[[TOKEN]], %arg0, %arg1)145 %token_0 = async.execute [%token] {146 %c0 = arith.constant 0 : index147 memref.store %arg0, %arg1[%c0] : memref<1xf32>148 async.yield149 }150 return151}152 153// Function outlined from the first async.execute operation.154// CHECK-LABEL: func private @async_execute_fn155// CHECK-SAME: -> !async.token156// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token157// CHECK: return %[[TOKEN]] : !async.token158 159// Function outlined from the second async.execute operation with dependency.160// CHECK-LABEL: func private @async_execute_fn_0161// CHECK-SAME: %[[ARG0:.*]]: !async.token162// CHECK-SAME: %[[ARG1:.*]]: f32163// CHECK-SAME: %[[ARG2:.*]]: memref<1xf32>164// CHECK-SAME: -> !async.token165// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token166// CHECK: %[[HDL:.*]] = async.coro.begin167 168// Suspend coroutine in the beginning.169// CHECK: async.runtime.resume %[[HDL]]170// CHECK: async.coro.suspend171// 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.store188// CHECK: async.runtime.set_available %[[TOKEN]]189 190// CHECK: ^[[CLEANUP]]:191// CHECK: ^[[SUSPEND]]:192 193// -----194 195// CHECK-LABEL: @async_group_await_all196func.func @async_group_await_all(%arg0: f32, %arg1: memref<1xf32>) {197 // CHECK: %[[C:.*]] = arith.constant 1 : index198 %c = arith.constant 1 : index199 // CHECK: %[[GROUP:.*]] = async.runtime.create_group %[[C]] : !async.group200 %0 = async.create_group %c : !async.group201 202 // CHECK: %[[TOKEN:.*]] = call @async_execute_fn203 %token = async.execute { async.yield }204 // CHECK: async.runtime.add_to_group %[[TOKEN]], %[[GROUP]]205 async.add_to_group %token, %0 : !async.token206 207 // CHECK: call @async_execute_fn_0208 async.execute {209 async.await_all %0210 async.yield211 }212 213 // CHECK: async.runtime.await %[[GROUP]] : !async.group214 async.await_all %0215 return216}217 218// Function outlined from the second async.execute operation.219// CHECK-LABEL: func private @async_execute_fn_0220// CHECK-SAME: (%[[ARG:.*]]: !async.group) -> !async.token221 222// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token223// CHECK: %[[HDL:.*]] = async.coro.begin224 225// Suspend coroutine in the beginning.226// CHECK: async.runtime.resume %[[HDL]]227// CHECK: async.coro.suspend228// 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.suspend234// 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_f32251func.func @execute_and_return_f32() -> f32 {252 // CHECK: %[[RET:.*]]:2 = call @async_execute_fn253 %token, %result = async.execute -> !async.value<f32> {254 %c0 = arith.constant 123.0 : f32255 async.yield %c0 : f32256 }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 : f32264}265 266// Function outlined from the async.execute operation.267// CHECK-LABEL: func private @async_execute_fn()268// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token269// CHECK: %[[VALUE:.*]] = async.runtime.create : !async.value<f32>270// CHECK: %[[HDL:.*]] = async.coro.begin271 272// Suspend coroutine in the beginning.273// CHECK: async.runtime.resume %[[HDL]]274// CHECK: async.coro.suspend275// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]276 277// Emplace result value.278// CHECK: ^[[RESUME]]:279// CHECK: %[[CST:.*]] = arith.constant 1.230000e+02 : f32280// 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_operands290func.func @async_value_operands() {291 // CHECK: %[[RET:.*]]:2 = call @async_execute_fn292 %token, %result = async.execute -> !async.value<f32> {293 %c0 = arith.constant 123.0 : f32294 async.yield %c0 : f32295 }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 : f32300 async.yield301 }302 303 // CHECK: async.runtime.await %[[TOKEN]] : !async.token304 async.await %token0 : !async.token305 306 return307}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_0314// CHECK-SAME: (%[[ARG:.*]]: !async.value<f32>) -> !async.token315 316// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token317// CHECK: %[[HDL:.*]] = async.coro.begin318 319// Suspend coroutine in the beginning.320// CHECK: async.runtime.resume %[[HDL]]321// CHECK: async.coro.suspend322// 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.suspend328// 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]] : <f32338// CHECK: arith.addf %[[LOADED]], %[[LOADED]] : f32339// CHECK: async.runtime.set_available %[[TOKEN]]340 341// CHECK: ^[[CLEANUP]]:342// CHECK: ^[[SUSPEND]]:343 344// -----345 346// CHECK-LABEL: @execute_assertion347func.func @execute_assertion(%arg0: i1) {348 %token = async.execute {349 cf.assert %arg0, "error"350 async.yield351 }352 async.await %token : !async.token353 return354}355 356// Function outlined from the async.execute operation.357// CHECK-LABEL: func private @async_execute_fn(358// CHECK-SAME: %[[ARG0:.*]]: i1359// CHECK-SAME: -> !async.token360 361// Create token for return op, and mark a function as a coroutine.362// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token363// CHECK: %[[ID:.*]] = async.coro.id364// CHECK: %[[HDL:.*]] = async.coro.begin365 366// Initial coroutine suspension.367// CHECK: async.coro.suspend368// 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 be399// lowered to branch-based control flow to enable coroutine CFG rewrite.400 401// CHECK-LABEL: @lower_scf_to_cfg402func.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.token407 } else {408 async.await %token0 : !async.token409 }410 async.yield411 }412 return413}414 415// Function outlined from the first async.execute operation.416// CHECK-LABEL: func private @async_execute_fn(417// CHECK-SAME: -> !async.token418 419// Function outlined from the second async.execute operation.420// CHECK-LABEL: func private @async_execute_fn_0(421// CHECK: %[[TOKEN:.*]]: !async.token422// CHECK: %[[FLAG:.*]]: i1423// CHECK-SAME: -> !async.token424 425// Check that structured control flow lowered to CFG.426// CHECK-NOT: scf.if427// CHECK: cf.cond_br %[[FLAG]]428 429// -----430// Constants captured by the async.execute region should be cloned into the431// outline async execute function.432 433// CHECK-LABEL: @clone_constants434func.func @clone_constants(%arg0: f32, %arg1: memref<1xf32>) {435 %c0 = arith.constant 0 : index436 %token = async.execute {437 memref.store %arg0, %arg1[%c0] : memref<1xf32>438 async.yield439 }440 async.await %token : !async.token441 return442}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.token449// CHECK: %[[CST:.*]] = arith.constant 0 : index450// CHECK: memref.store %[[VALUE]], %[[MEMREF]][%[[CST]]]451 452// -----453// Async Functions should be none blocking454 455// CHECK-LABEL: @async_func_await456async.func @async_func_await(%arg0: f32, %arg1: !async.value<f32>)457 -> !async.token {458 %0 = async.await %arg1 : !async.value<f32>459 return460}461// Create token for return op, and mark a function as a coroutine.462// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token463// CHECK: %[[ID:.*]] = async.coro.id464// CHECK: %[[HDL:.*]] = async.coro.begin465// 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 func476 477// CHECK-LABEL: @execute_in_async_func478async.func @execute_in_async_func(%arg0: f32, %arg1: memref<1xf32>)479 -> !async.token {480 %token = async.execute {481 %c0 = arith.constant 0 : index482 memref.store %arg0, %arg1[%c0] : memref<1xf32>483 async.yield484 }485 async.await %token : !async.token486 return487}488// Call outlind async execute Function489// 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.token493 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.token499// CHECK: %[[CST:.*]] = arith.constant 0 : index500// CHECK: memref.store %[[VALUE]], %[[MEMREF]][%[[CST]]]501