mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir

   1 // RUN: mlir-opt %s -split-input-file -async-to-async-runtime -convert-async-to-llvm | FileCheck %s
   2
   3 // CHECK-LABEL: reference_counting
   4 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
  19 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: 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<i8>
  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
  74 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: 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<i8>
 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<i8>
 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
 132 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<i8>
 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<i8>, %arg1: f32, %arg2: memref<1xf32>)
 160 // CHECK-SAME: -> !llvm.ptr<i8>
 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
 180 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<i8>)
 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
 222 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: %[[ST_F32:.*]] = llvm.bitcast %[[STORAGE]]
 231   // CHECK: %[[LOADED:.*]] = llvm.load %[[ST_F32]] :  !llvm.ptr<f32>
 232   %0 = async.await %result : !async.value<f32>
 233
 234   return %0 : f32
 235 }
 236
 237 // Function outlined from the async.execute operation.
 238 // CHECK-LABEL: func private @async_execute_fn()
 239 // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken()
 240 // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
 241 // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
 242
 243 // Suspend coroutine in the beginning.
 244 // CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]],
 245 // CHECK: llvm.intr.coro.suspend
 246
 247 // Emplace result value.
 248 // CHECK: %[[CST:.*]] = arith.constant 1.230000e+02 : f32
 249 // CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[VALUE]])
 250 // CHECK: %[[ST_F32:.*]] = llvm.bitcast %[[STORAGE]]
 251 // CHECK: llvm.store %[[CST]], %[[ST_F32]] : !llvm.ptr<f32>
 252 // CHECK: call @mlirAsyncRuntimeEmplaceValue(%[[VALUE]])
 253
 254 // Emplace result token.
 255 // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]])
 256
 257 // -----
 258
 259 // CHECK-LABEL: @async_value_operands
 260 func @async_value_operands() {
 261   // CHECK: %[[RET:.*]]:2 = call @async_execute_fn
 262   %token, %result = async.execute -> !async.value<f32> {
 263     %c0 = arith.constant 123.0 : f32
 264     async.yield %c0 : f32
 265   }
 266
 267   // CHECK: %[[TOKEN:.*]] = call @async_execute_fn_0(%[[RET]]#1)
 268   %token0 = async.execute(%result as %value: !async.value<f32>) {
 269     %0 = arith.addf %value, %value : f32
 270     async.yield
 271   }
 272
 273   // CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]])
 274   async.await %token0 : !async.token
 275
 276   return
 277 }
 278
 279 // Function outlined from the first async.execute operation.
 280 // CHECK-LABEL: func private @async_execute_fn()
 281
 282 // Function outlined from the second async.execute operation.
 283 // CHECK-LABEL: func private @async_execute_fn_0(%arg0: !llvm.ptr<i8>)
 284 // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken()
 285 // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
 286
 287 // Suspend coroutine in the beginning.
 288 // CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]],
 289 // CHECK: llvm.intr.coro.suspend
 290
 291 // Suspend coroutine second time waiting for the async operand.
 292 // CHECK: llvm.intr.coro.save
 293 // CHECK: call @mlirAsyncRuntimeAwaitValueAndExecute(%arg0, %[[HDL]],
 294 // CHECK: llvm.intr.coro.suspend
 295
 296 // Get the operand value storage, cast to f32 and add the value.
 297 // CHECK: %[[STORAGE:.*]] = call @mlirAsyncRuntimeGetValueStorage(%arg0)
 298 // CHECK: %[[ST_F32:.*]] = llvm.bitcast %[[STORAGE]]
 299 // CHECK: %[[LOADED:.*]] = llvm.load %[[ST_F32]] :  !llvm.ptr<f32>
 300 // CHECK: arith.addf %[[LOADED]], %[[LOADED]] : f32
 301
 302 // Emplace result token.
 303 // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]])
 304
 305