303 lines · plain
1// RUN: mlir-opt %s -split-input-file -async-to-async-runtime -convert-async-to-llvm | FileCheck %s2 3// CHECK-LABEL: reference_counting4func.func @reference_counting(%arg0: !async.token) {5 // CHECK: %[[C2:.*]] = arith.constant 2 : i646 // CHECK: call @mlirAsyncRuntimeAddRef(%arg0, %[[C2]])7 async.runtime.add_ref %arg0 {count = 2 : i64} : !async.token8 9 // CHECK: %[[C1:.*]] = arith.constant 1 : i6410 // CHECK: call @mlirAsyncRuntimeDropRef(%arg0, %[[C1]])11 async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token12 13 return14}15 16// -----17 18// CHECK-LABEL: execute_no_async_args19func.func @execute_no_async_args(%arg0: f32, %arg1: memref<1xf32>) {20 // CHECK: %[[TOKEN:.*]] = call @async_execute_fn(%arg0, %arg1)21 %token = async.execute {22 %c0 = arith.constant 0 : index23 memref.store %arg0, %arg1[%c0] : memref<1xf32>24 async.yield25 }26 // CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]])27 // CHECK: %[[IS_ERROR:.*]] = call @mlirAsyncRuntimeIsTokenError(%[[TOKEN]])28 // CHECK: %[[TRUE:.*]] = arith.constant true29 // CHECK: %[[NOT_ERROR:.*]] = arith.xori %[[IS_ERROR]], %[[TRUE]] : i130 // CHECK: cf.assert %[[NOT_ERROR]]31 // CHECK-NEXT: return32 async.await %token : !async.token33 return34}35 36// Function outlined from the async.execute operation.37// CHECK-LABEL: func private @async_execute_fn(%arg0: f32, %arg1: memref<1xf32>)38// CHECK-SAME: -> !llvm.ptr39 40// Create token for return op, and mark a function as a coroutine.41// CHECK: %[[RET:.*]] = call @mlirAsyncRuntimeCreateToken()42// CHECK: %[[HDL:.*]] = llvm.intr.coro.begin43 44// Pass a suspended coroutine to the async runtime.45// CHECK: %[[STATE:.*]] = llvm.intr.coro.save46// CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume47// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]], %[[RESUME]])48// CHECK: %[[SUSPENDED:.*]] = llvm.intr.coro.suspend %[[STATE]]49 50// Decide the next block based on the code returned from suspend.51// CHECK: %[[SEXT:.*]] = llvm.sext %[[SUSPENDED]] : i8 to i3252// CHECK: llvm.switch %[[SEXT]] : i32, ^[[SUSPEND:[b0-9]+]]53// CHECK-NEXT: 0: ^[[RESUME:[b0-9]+]]54// CHECK-NEXT: 1: ^[[CLEANUP:[b0-9]+]]55 56// Resume coroutine after suspension.57// CHECK: ^[[RESUME]]:58// CHECK: memref.store %arg0, %arg1[%c0] : memref<1xf32>59// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET]])60 61// Delete coroutine.62// CHECK: ^[[CLEANUP]]:63// CHECK: %[[MEM:.*]] = llvm.intr.coro.free64// CHECK: llvm.call @free(%[[MEM]])65 66// Suspend coroutine, and also a return statement for ramp function.67// CHECK: ^[[SUSPEND]]:68// CHECK: llvm.intr.coro.end69// CHECK: return %[[RET]]70 71// -----72 73// CHECK-LABEL: nested_async_execute74func.func @nested_async_execute(%arg0: f32, %arg1: f32, %arg2: memref<1xf32>) {75 // CHECK: %[[TOKEN:.*]] = call @async_execute_fn_0(%arg0, %arg2, %arg1)76 %token0 = async.execute {77 %c0 = arith.constant 0 : index78 79 %token1 = async.execute {80 %c1 = arith.constant 1: index81 memref.store %arg0, %arg2[%c0] : memref<1xf32>82 async.yield83 }84 async.await %token1 : !async.token85 86 memref.store %arg1, %arg2[%c0] : memref<1xf32>87 async.yield88 }89 // CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]])90 // CHECK: %[[IS_ERROR:.*]] = call @mlirAsyncRuntimeIsTokenError(%[[TOKEN]])91 // CHECK: %[[TRUE:.*]] = arith.constant true92 // CHECK: %[[NOT_ERROR:.*]] = arith.xori %[[IS_ERROR]], %[[TRUE]] : i193 // CHECK: cf.assert %[[NOT_ERROR]]94 async.await %token0 : !async.token95 return96}97 98// Function outlined from the inner async.execute operation.99// CHECK-LABEL: func private @async_execute_fn(%arg0: f32, %arg1: memref<1xf32>)100// CHECK-SAME: -> !llvm.ptr101// CHECK: %[[RET_0:.*]] = call @mlirAsyncRuntimeCreateToken()102// CHECK: %[[HDL_0:.*]] = llvm.intr.coro.begin103// CHECK: call @mlirAsyncRuntimeExecute104// CHECK: llvm.intr.coro.suspend105// CHECK: %[[C0:.*]] = arith.constant 0 : index106// CHECK: memref.store %arg0, %arg1[%[[C0]]] : memref<1xf32>107// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_0]])108 109// Function outlined from the outer async.execute operation.110// CHECK-LABEL: func private @async_execute_fn_0(%arg0: f32, %arg1: memref<1xf32>, %arg2: f32)111// CHECK-SAME: -> !llvm.ptr112// CHECK: %[[RET_1:.*]] = call @mlirAsyncRuntimeCreateToken()113// CHECK: %[[HDL_1:.*]] = llvm.intr.coro.begin114 115// Suspend coroutine in the beginning.116// CHECK: call @mlirAsyncRuntimeExecute117// CHECK: llvm.intr.coro.suspend118 119// Suspend coroutine second time waiting for the completion of inner execute op.120// CHECK: %[[TOKEN_1:.*]] = call @async_execute_fn121// CHECK: llvm.intr.coro.save122// CHECK: call @mlirAsyncRuntimeAwaitTokenAndExecute(%[[TOKEN_1]], %[[HDL_1]]123// CHECK: llvm.intr.coro.suspend124 125// Emplace result token after second resumption.126// CHECK: memref.store %arg2, %arg1[%c0] : memref<1xf32>127// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_1]])128 129// -----130 131// CHECK-LABEL: async_execute_token_dependency132func.func @async_execute_token_dependency(%arg0: f32, %arg1: memref<1xf32>) {133 // CHECK: %0 = call @async_execute_fn(%arg0, %arg1)134 %token = async.execute {135 %c0 = arith.constant 0 : index136 memref.store %arg0, %arg1[%c0] : memref<1xf32>137 async.yield138 }139 // CHECK: %1 = call @async_execute_fn_0(%0, %arg0, %arg1)140 %token_0 = async.execute [%token] {141 %c0 = arith.constant 0 : index142 memref.store %arg0, %arg1[%c0] : memref<1xf32>143 async.yield144 }145 return146}147 148// Function outlined from the first async.execute operation.149// CHECK-LABEL: func private @async_execute_fn(%arg0: f32, %arg1: memref<1xf32>)150// CHECK-SAME: -> !llvm.ptr151// CHECK: %[[RET_0:.*]] = call @mlirAsyncRuntimeCreateToken()152// CHECK: %[[HDL_0:.*]] = llvm.intr.coro.begin153// CHECK: call @mlirAsyncRuntimeExecute154// CHECK: llvm.intr.coro.suspend155// CHECK: memref.store %arg0, %arg1[%c0] : memref<1xf32>156// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_0]])157 158// Function outlined from the second async.execute operation with dependency.159// CHECK-LABEL: func private @async_execute_fn_0(%arg0: !llvm.ptr, %arg1: f32, %arg2: memref<1xf32>)160// CHECK-SAME: -> !llvm.ptr161// CHECK: %[[RET_1:.*]] = call @mlirAsyncRuntimeCreateToken()162// CHECK: %[[HDL_1:.*]] = llvm.intr.coro.begin163 164// Suspend coroutine in the beginning.165// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL_1]],166// CHECK: llvm.intr.coro.suspend167 168// Suspend coroutine second time waiting for the completion of token dependency.169// CHECK: llvm.intr.coro.save170// CHECK: call @mlirAsyncRuntimeAwaitTokenAndExecute(%arg0, %[[HDL_1]],171// CHECK: llvm.intr.coro.suspend172 173// Emplace result token after second resumption.174// CHECK: memref.store %arg1, %arg2[%c0] : memref<1xf32>175// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_1]])176 177// -----178 179// CHECK-LABEL: async_group_await_all180func.func @async_group_await_all(%arg0: f32, %arg1: memref<1xf32>) {181 %c = arith.constant 1 : index182 // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup183 %0 = async.create_group %c : !async.group184 185 // CHECK: %[[TOKEN:.*]] = call @async_execute_fn186 %token = async.execute { async.yield }187 // CHECK: call @mlirAsyncRuntimeAddTokenToGroup(%[[TOKEN]], %[[GROUP]])188 async.add_to_group %token, %0 : !async.token189 190 // CHECK: call @async_execute_fn_0191 async.execute {192 async.await_all %0193 async.yield194 }195 196 // CHECK: call @mlirAsyncRuntimeAwaitAllInGroup(%[[GROUP]])197 async.await_all %0198 199 return200}201 202// Function outlined from the async.execute operation.203// CHECK: func private @async_execute_fn_0(%arg0: !llvm.ptr)204// CHECK: %[[RET_1:.*]] = call @mlirAsyncRuntimeCreateToken()205// CHECK: %[[HDL_1:.*]] = llvm.intr.coro.begin206 207// Suspend coroutine in the beginning.208// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL_1]],209// CHECK: llvm.intr.coro.suspend210 211// Suspend coroutine second time waiting for the group.212// CHECK: llvm.intr.coro.save213// CHECK: call @mlirAsyncRuntimeAwaitAllInGroupAndExecute(%arg0, %[[HDL_1]],214// CHECK: llvm.intr.coro.suspend215 216// Emplace result token.217// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_1]])218 219// -----220 221// CHECK-LABEL: execute_and_return_f32222func.func @execute_and_return_f32() -> f32 {223 // CHECK: %[[RET:.*]]:2 = call @async_execute_fn224 %token, %result = async.execute -> !async.value<f32> {225 %c0 = arith.constant 123.0 : f32226 async.yield %c0 : f32227 }228 229 // CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[RET]]#1)230 // CHECK: %[[LOADED:.*]] = llvm.load %[[STORAGE]] : !llvm.ptr -> f32231 %0 = async.await %result : !async.value<f32>232 233 return %0 : f32234}235 236// Function outlined from the async.execute operation.237// CHECK-LABEL: func private @async_execute_fn()238// CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken()239// CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue240// CHECK: %[[HDL:.*]] = llvm.intr.coro.begin241 242// Suspend coroutine in the beginning.243// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]],244// CHECK: llvm.intr.coro.suspend245 246// Emplace result value.247// CHECK: %[[CST:.*]] = arith.constant 1.230000e+02 : f32248// CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[VALUE]])249// CHECK: llvm.store %[[CST]], %[[STORAGE]] : f32, !llvm.ptr250// CHECK: call @mlirAsyncRuntimeEmplaceValue(%[[VALUE]])251 252// Emplace result token.253// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]])254 255// -----256 257// CHECK-LABEL: @async_value_operands258func.func @async_value_operands() {259 // CHECK: %[[RET:.*]]:2 = call @async_execute_fn260 %token, %result = async.execute -> !async.value<f32> {261 %c0 = arith.constant 123.0 : f32262 async.yield %c0 : f32263 }264 265 // CHECK: %[[TOKEN:.*]] = call @async_execute_fn_0(%[[RET]]#1)266 %token0 = async.execute(%result as %value: !async.value<f32>) {267 %0 = arith.addf %value, %value : f32268 async.yield269 }270 271 // CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]])272 async.await %token0 : !async.token273 274 return275}276 277// Function outlined from the first async.execute operation.278// CHECK-LABEL: func private @async_execute_fn()279 280// Function outlined from the second async.execute operation.281// CHECK-LABEL: func private @async_execute_fn_0(%arg0: !llvm.ptr)282// CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken()283// CHECK: %[[HDL:.*]] = llvm.intr.coro.begin284 285// Suspend coroutine in the beginning.286// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]],287// CHECK: llvm.intr.coro.suspend288 289// Suspend coroutine second time waiting for the async operand.290// CHECK: llvm.intr.coro.save291// CHECK: call @mlirAsyncRuntimeAwaitValueAndExecute(%arg0, %[[HDL]],292// CHECK: llvm.intr.coro.suspend293 294// Get the operand value storage, cast to f32 and add the value.295// CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%arg0)296// CHECK: %[[LOADED:.*]] = llvm.load %[[STORAGE]] : !llvm.ptr -> f32297// CHECK: arith.addf %[[LOADED]], %[[LOADED]] : f32298 299// Emplace result token.300// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]])301 302 303