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

   1 // RUN: mlir-opt %s -convert-async-to-llvm | FileCheck %s --dump-input=always
   2
   3 // CHECK-LABEL: @create_token
   4 func @create_token() {
   5   // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
   6   %0 = async.runtime.create : !async.token
   7   return
   8 }
   9
  10 // CHECK-LABEL: @create_value
  11 func @create_value() {
  12   // CHECK: %[[NULL:.*]] = llvm.mlir.null : !llvm.ptr<f32>
  13   // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : i64) : i64
  14   // CHECK: %[[OFFSET:.*]] = llvm.getelementptr %[[NULL]][%[[ONE]]]
  15   // CHECK: %[[SIZE:.*]] = llvm.ptrtoint %[[OFFSET]]
  16   // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue(%[[SIZE]])
  17   %0 = async.runtime.create : !async.value<f32>
  18   return
  19 }
  20
  21 // CHECK-LABEL: @create_group
  22 func @create_group() {
  23   // CHECK: %[[C:.*]] = arith.constant 1 : index
  24   // CHECK: %[[S:.*]] = builtin.unrealized_conversion_cast %[[C]] : index to i64
  25   %c = arith.constant 1 : index
  26   // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup(%[[S]])
  27   %0 = async.runtime.create_group  %c: !async.group
  28   return
  29 }
  30
  31 // CHECK-LABEL: @set_token_available
  32 func @set_token_available() {
  33   // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
  34   %0 = async.runtime.create : !async.token
  35   // CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[TOKEN]])
  36   async.runtime.set_available %0 : !async.token
  37   return
  38 }
  39
  40 // CHECK-LABEL: @set_value_available
  41 func @set_value_available() {
  42   // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
  43   %0 = async.runtime.create : !async.value<f32>
  44   // CHECK: call @mlirAsyncRuntimeEmplaceValue(%[[VALUE]])
  45   async.runtime.set_available %0 : !async.value<f32>
  46   return
  47 }
  48
  49 // CHECK-LABEL: @is_token_error
  50 func @is_token_error() -> i1 {
  51   // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
  52   %0 = async.runtime.create : !async.token
  53   // CHECK: %[[ERR:.*]] = call @mlirAsyncRuntimeIsTokenError(%[[TOKEN]])
  54   %1 = async.runtime.is_error %0 : !async.token
  55   return %1 : i1
  56 }
  57
  58 // CHECK-LABEL: @is_value_error
  59 func @is_value_error() -> i1 {
  60   // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
  61   %0 = async.runtime.create : !async.value<f32>
  62   // CHECK: %[[ERR:.*]] = call @mlirAsyncRuntimeIsValueError(%[[VALUE]])
  63   %1 = async.runtime.is_error %0 : !async.value<f32>
  64   return %1 : i1
  65 }
  66
  67 // CHECK-LABEL: @await_token
  68 func @await_token() {
  69   // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
  70   %0 = async.runtime.create : !async.token
  71   // CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]])
  72   async.runtime.await %0 : !async.token
  73   return
  74 }
  75
  76 // CHECK-LABEL: @await_value
  77 func @await_value() {
  78   // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
  79   %0 = async.runtime.create : !async.value<f32>
  80   // CHECK: call @mlirAsyncRuntimeAwaitValue(%[[VALUE]])
  81   async.runtime.await %0 : !async.value<f32>
  82   return
  83 }
  84
  85 // CHECK-LABEL: @await_group
  86 func @await_group() {
  87   %c = arith.constant 1 : index
  88   // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup
  89   %0 = async.runtime.create_group %c: !async.group
  90   // CHECK: call @mlirAsyncRuntimeAwaitAllInGroup(%[[GROUP]])
  91   async.runtime.await %0 : !async.group
  92   return
  93 }
  94
  95 // CHECK-LABEL: @await_and_resume_token
  96 func @await_and_resume_token() {
  97   %0 = async.coro.id
  98   // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
  99   %1 = async.coro.begin %0
 100   // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
 101   %2 = async.runtime.create : !async.token
 102   // CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume
 103   // CHECK: call @mlirAsyncRuntimeAwaitTokenAndExecute
 104   // CHECK-SAME: (%[[TOKEN]], %[[HDL]], %[[RESUME]])
 105   async.runtime.await_and_resume %2, %1 : !async.token
 106   return
 107 }
 108
 109 // CHECK-LABEL: @await_and_resume_value
 110 func @await_and_resume_value() {
 111   %0 = async.coro.id
 112   // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
 113   %1 = async.coro.begin %0
 114   // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
 115   %2 = async.runtime.create : !async.value<f32>
 116   // CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume
 117   // CHECK: call @mlirAsyncRuntimeAwaitValueAndExecute
 118   // CHECK-SAME: (%[[VALUE]], %[[HDL]], %[[RESUME]])
 119   async.runtime.await_and_resume %2, %1 : !async.value<f32>
 120   return
 121 }
 122
 123 // CHECK-LABEL: @await_and_resume_group
 124 func @await_and_resume_group() {
 125   %c = arith.constant 1 : index
 126   %0 = async.coro.id
 127   // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
 128   %1 = async.coro.begin %0
 129   // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateGroup
 130   %2 = async.runtime.create_group %c : !async.group
 131   // CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume
 132   // CHECK: call @mlirAsyncRuntimeAwaitAllInGroupAndExecute
 133   // CHECK-SAME: (%[[TOKEN]], %[[HDL]], %[[RESUME]])
 134   async.runtime.await_and_resume %2, %1 : !async.group
 135   return
 136 }
 137
 138 // CHECK-LABEL: @resume
 139 func @resume() {
 140   %0 = async.coro.id
 141   // CHECK: %[[HDL:.*]] = llvm.intr.coro.begin
 142   %1 = async.coro.begin %0
 143   // CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume
 144   // CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]], %[[RESUME]])
 145   async.runtime.resume %1
 146   return
 147 }
 148
 149 // CHECK-LABEL: @store
 150 func @store() {
 151   // CHECK: %[[CST:.*]] = arith.constant 1.0
 152   %0 = arith.constant 1.0 : f32
 153   // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
 154   %1 = async.runtime.create : !async.value<f32>
 155   // CHECK: %[[P0:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[VALUE]])
 156   // CHECK: %[[P1:.*]] = llvm.bitcast %[[P0]] : !llvm.ptr<i8> to !llvm.ptr<f32>
 157   // CHECK: llvm.store %[[CST]], %[[P1]]
 158   async.runtime.store %0, %1 : !async.value<f32>
 159   return
 160 }
 161
 162 // CHECK-LABEL: @load
 163 func @load() -> f32 {
 164   // CHECK: %[[VALUE:.*]] = call @mlirAsyncRuntimeCreateValue
 165   %0 = async.runtime.create : !async.value<f32>
 166   // CHECK: %[[P0:.*]] = call @mlirAsyncRuntimeGetValueStorage(%[[VALUE]])
 167   // CHECK: %[[P1:.*]] = llvm.bitcast %[[P0]] : !llvm.ptr<i8> to !llvm.ptr<f32>
 168   // CHECK: %[[VALUE:.*]] = llvm.load %[[P1]]
 169   %1 = async.runtime.load %0 : !async.value<f32>
 170   // CHECK: return %[[VALUE]] : f32
 171   return %1 : f32
 172 }
 173
 174 // CHECK-LABEL: @add_token_to_group
 175 func @add_token_to_group() {
 176   %c = arith.constant 1 : index
 177   // CHECK: %[[TOKEN:.*]] = call @mlirAsyncRuntimeCreateToken
 178   %0 = async.runtime.create : !async.token
 179   // CHECK: %[[GROUP:.*]] = call @mlirAsyncRuntimeCreateGroup
 180   %1 = async.runtime.create_group %c : !async.group
 181   // CHECK: call @mlirAsyncRuntimeAddTokenToGroup(%[[TOKEN]], %[[GROUP]])
 182   async.runtime.add_to_group %0, %1 : !async.token
 183   return
 184 }