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