1// RUN: mlir-opt %s -split-input-file -async-to-async-runtime -convert-async-to-llvm | FileCheck %s 2 3// CHECK-LABEL: reference_counting 4func.func @reference_counting(%arg0: !async.token) { 5 // CHECK: %[[C2:.*]] = arith.constant 2 : i64 6 // CHECK: call @mlirAsyncRuntimeAddRef(%arg0, %[[C2]]) 7 async.runtime.add_ref %arg0 {count = 2 : i64} : !async.token 8 9 // CHECK: %[[C1:.*]] = arith.constant 1 : i64 10 // CHECK: call @mlirAsyncRuntimeDropRef(%arg0, %[[C1]]) 11 async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token 12 13 return 14} 15 16// ----- 17 18// CHECK-LABEL: execute_no_async_args 19func.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 : index 23 memref.store %arg0, %arg1[%c0] : memref<1xf32> 24 async.yield 25 } 26 // CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]]) 27 // CHECK: %[[IS_ERROR:.*]] = call @mlirAsyncRuntimeIsTokenError(%[[TOKEN]]) 28 // CHECK: %[[TRUE:.*]] = arith.constant true 29 // CHECK: %[[NOT_ERROR:.*]] = arith.xori %[[IS_ERROR]], %[[TRUE]] : i1 30 // CHECK: cf.assert %[[NOT_ERROR]] 31 // CHECK-NEXT: return 32 async.await %token : !async.token 33 return 34} 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.ptr 39 40// Create token for return op, and mark a function as a coroutine. 41// CHECK: %[[RET:.*]] = call @mlirAsyncRuntimeCreateToken() 42// CHECK: %[[HDL:.*]] = llvm.intr.coro.begin 43 44// Pass a suspended coroutine to the async runtime. 45// CHECK: %[[STATE:.*]] = llvm.intr.coro.save 46// CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume 47// 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 i32 52// 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.free 64// 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.end 69// CHECK: return %[[RET]] 70 71// ----- 72 73// CHECK-LABEL: nested_async_execute 74func.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 : index 78 79 %token1 = async.execute { 80 %c1 = arith.constant 1: index 81 memref.store %arg0, %arg2[%c0] : memref<1xf32> 82 async.yield 83 } 84 async.await %token1 : !async.token 85 86 memref.store %arg1, %arg2[%c0] : memref<1xf32> 87 async.yield 88 } 89 // CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]]) 90 // CHECK: %[[IS_ERROR:.*]] = call @mlirAsyncRuntimeIsTokenError(%[[TOKEN]]) 91 // CHECK: %[[TRUE:.*]] = arith.constant true 92 // CHECK: %[[NOT_ERROR:.*]] = arith.xori %[[IS_ERROR]], %[[TRUE]] : i1 93 // CHECK: cf.assert %[[NOT_ERROR]] 94 async.await %token0 : !async.token 95 return 96} 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.ptr 101// CHECK: %[[RET_0:.*]] = call @mlirAsyncRuntimeCreateToken() 102// CHECK: %[[HDL_0:.*]] = llvm.intr.coro.begin 103// CHECK: call @mlirAsyncRuntimeExecute 104// CHECK: llvm.intr.coro.suspend 105// CHECK: %[[C0:.*]] = arith.constant 0 : index 106// 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.ptr 112// CHECK: %[[RET_1:.*]] = call @mlirAsyncRuntimeCreateToken() 113// CHECK: %[[HDL_1:.*]] = llvm.intr.coro.begin 114 115// Suspend coroutine in the beginning. 116// CHECK: call @mlirAsyncRuntimeExecute 117// CHECK: llvm.intr.coro.suspend 118 119// Suspend coroutine second time waiting for the completion of inner execute op. 120// CHECK: %[[TOKEN_1:.*]] = call @async_execute_fn 121// CHECK: llvm.intr.coro.save 122// CHECK: call @mlirAsyncRuntimeAwaitTokenAndExecute(%[[TOKEN_1]], %[[HDL_1]] 123// CHECK: llvm.intr.coro.suspend 124 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_dependency 132func.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 : index 136 memref.store %arg0, %arg1[%c0] : memref<1xf32> 137 async.yield 138 } 139 // CHECK: %1 = call @async_execute_fn_0(%0, %arg0, %arg1) 140 %token_0 = async.execute [%token] { 141 %c0 = arith.constant 0 : index 142 memref.store %arg0, %arg1[%c0] : memref<1xf32> 143 async.yield 144 } 145 return 146} 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.ptr 151// CHECK: %[[RET_0:.*]] = call @mlirAsyncRuntimeCreateToken() 152// CHECK: %[[HDL_0:.*]] = llvm.intr.coro.begin 153// CHECK: call @mlirAsyncRuntimeExecute 154// CHECK: llvm.intr.coro.suspend 155// 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.ptr 161// CHECK: %[[RET_1:.*]] = call @mlirAsyncRuntimeCreateToken() 162// CHECK: %[[HDL_1:.*]] = llvm.intr.coro.begin 163 164// Suspend coroutine in the beginning. 165// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL_1]], 166// CHECK: llvm.intr.coro.suspend 167 168// Suspend coroutine second time waiting for the completion of token dependency. 169// CHECK: llvm.intr.coro.save 170// CHECK: call @mlirAsyncRuntimeAwaitTokenAndExecute(%arg0, %[[HDL_1]], 171// CHECK: llvm.intr.coro.suspend 172 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_all 180func.func @async_group_await_all(%arg0: f32, %arg1: memref<1xf32>) { 181 %c = arith.constant 1 : index 182 // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup 183 %0 = async.create_group %c : !async.group 184 185 // CHECK: %[[TOKEN:.*]] = call @async_execute_fn 186 %token = async.execute { async.yield } 187 // CHECK: call @mlirAsyncRuntimeAddTokenToGroup(%[[TOKEN]], %[[GROUP]]) 188 async.add_to_group %token, %0 : !async.token 189 190 // CHECK: call @async_execute_fn_0 191 async.execute { 192 async.await_all %0 193 async.yield 194 } 195 196 // CHECK: call @mlirAsyncRuntimeAwaitAllInGroup(%[[GROUP]]) 197 async.await_all %0 198 199 return 200} 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.begin 206 207// Suspend coroutine in the beginning. 208// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL_1]], 209// CHECK: llvm.intr.coro.suspend 210 211// Suspend coroutine second time waiting for the group. 212// CHECK: llvm.intr.coro.save 213// CHECK: call @mlirAsyncRuntimeAwaitAllInGroupAndExecute(%arg0, %[[HDL_1]], 214// CHECK: llvm.intr.coro.suspend 215 216// Emplace result token. 217// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_1]]) 218 219// ----- 220 221// CHECK-LABEL: execute_and_return_f32 222func.func @execute_and_return_f32() -> f32 { 223 // CHECK: %[[RET:.*]]:2 = call @async_execute_fn 224 %token, %result = async.execute -> !async.value<f32> { 225 %c0 = arith.constant 123.0 : f32 226 async.yield %c0 : f32 227 } 228 229 // CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[RET]]#1) 230 // CHECK: %[[LOADED:.*]] = llvm.load %[[STORAGE]] : !llvm.ptr -> f32 231 %0 = async.await %result : !async.value<f32> 232 233 return %0 : f32 234} 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 @mlirAsyncRuntimeCreateValue 240// CHECK: %[[HDL:.*]] = llvm.intr.coro.begin 241 242// Suspend coroutine in the beginning. 243// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]], 244// CHECK: llvm.intr.coro.suspend 245 246// Emplace result value. 247// CHECK: %[[CST:.*]] = arith.constant 1.230000e+02 : f32 248// CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[VALUE]]) 249// CHECK: llvm.store %[[CST]], %[[STORAGE]] : f32, !llvm.ptr 250// CHECK: call @mlirAsyncRuntimeEmplaceValue(%[[VALUE]]) 251 252// Emplace result token. 253// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]]) 254 255// ----- 256 257// CHECK-LABEL: @async_value_operands 258func.func @async_value_operands() { 259 // CHECK: %[[RET:.*]]:2 = call @async_execute_fn 260 %token, %result = async.execute -> !async.value<f32> { 261 %c0 = arith.constant 123.0 : f32 262 async.yield %c0 : f32 263 } 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 : f32 268 async.yield 269 } 270 271 // CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]]) 272 async.await %token0 : !async.token 273 274 return 275} 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.begin 284 285// Suspend coroutine in the beginning. 286// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]], 287// CHECK: llvm.intr.coro.suspend 288 289// Suspend coroutine second time waiting for the async operand. 290// CHECK: llvm.intr.coro.save 291// CHECK: call @mlirAsyncRuntimeAwaitValueAndExecute(%arg0, %[[HDL]], 292// CHECK: llvm.intr.coro.suspend 293 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 -> f32 297// CHECK: arith.addf %[[LOADED]], %[[LOADED]] : f32 298 299// Emplace result token. 300// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]]) 301 302 303