mlir/test/mlir-cpu-runner/async-func.mlir

   1 // RUN:   mlir-opt %s -pass-pipeline="builtin.module(async-func-to-async-runtime,async-to-async-runtime,func.func(async-runtime-ref-counting,async-runtime-ref-counting-opt),convert-async-to-llvm,test-lower-to-llvm)" \
   2 // RUN: | mlir-cpu-runner                                                      \
   3 // RUN:     -e main -entry-point-result=void -O0                               \
   4 // RUN:     -shared-libs=%mlir_c_runner_utils  \
   5 // RUN:     -shared-libs=%mlir_runner_utils    \
   6 // RUN:     -shared-libs=%mlir_async_runtime   \
   7 // RUN: | FileCheck %s --dump-input=always
   8
   9 // FIXME: https://github.com/llvm/llvm-project/issues/57231
  10 // UNSUPPORTED: hwasan
  11 // FIXME: Windows does not have aligned_alloc
  12 // UNSUPPORTED: system-windows
  13
  14 async.func @async_func_empty() -> !async.token {
  15   return
  16 }
  17
  18 async.func @async_func_assert() -> !async.token {
  19   %false = arith.constant 0 : i1
  20   cf.assert %false, "error"
  21   return
  22 }
  23
  24 async.func @async_func_nested_assert() -> !async.token {
  25   %token0 = async.call @async_func_assert() : () -> !async.token
  26   async.await %token0 : !async.token
  27   return
  28 }
  29
  30 async.func @async_func_value_assert() -> !async.value<f32> {
  31   %false = arith.constant 0 : i1
  32   cf.assert %false, "error"
  33   %0 = arith.constant 123.45 : f32
  34   return %0 : f32
  35 }
  36
  37 async.func @async_func_value_nested_assert() -> !async.value<f32> {
  38   %value0 = async.call @async_func_value_assert() : () -> !async.value<f32>
  39   %ret = async.await %value0 : !async.value<f32>
  40   return %ret : f32
  41 }
  42
  43 async.func @async_func_return_value() -> !async.value<f32> {
  44   %0 = arith.constant 456.789 : f32
  45   return %0 : f32
  46 }
  47
  48 async.func @async_func_non_blocking_await() -> !async.value<f32> {
  49   %value0 = async.call @async_func_return_value() : () -> !async.value<f32>
  50   %1 = async.await %value0 : !async.value<f32>
  51   return  %1 : f32
  52 }
  53
  54 async.func @async_func_inside_memref() -> !async.value<memref<f32>> {
  55   %0 = memref.alloc() : memref<f32>
  56   %c0 = arith.constant 0.25 : f32
  57   memref.store %c0, %0[] : memref<f32>
  58   return %0 : memref<f32>
  59 }
  60
  61 async.func @async_func_passed_memref(%arg0 : !async.value<memref<f32>>) -> !async.token {
  62   %unwrapped = async.await %arg0 : !async.value<memref<f32>>
  63   %0 = memref.load %unwrapped[] : memref<f32>
  64   %1 = arith.addf %0, %0 : f32
  65   memref.store %1, %unwrapped[] : memref<f32>
  66   return
  67 }
  68
  69 async.func @async_execute_in_async_func(%arg0 : !async.value<memref<f32>>) -> !async.token {
  70   %token0 = async.execute {
  71     %unwrapped = async.await %arg0 : !async.value<memref<f32>>
  72     %0 = memref.load %unwrapped[] : memref<f32>
  73     %1 = arith.addf %0, %0 : f32
  74     memref.store %1, %unwrapped[] : memref<f32>
  75     async.yield
  76   }
  77
  78   async.await %token0 : !async.token
  79   return
  80 }
  81
  82
  83 func.func @main() {
  84   %false = arith.constant 0 : i1
  85
  86   // ------------------------------------------------------------------------ //
  87   // Check that simple async.func completes without errors.
  88   // ------------------------------------------------------------------------ //
  89   %token0 = async.call @async_func_empty() : () -> !async.token
  90   async.runtime.await %token0 : !async.token
  91
  92   // CHECK: 0
  93   %err0 = async.runtime.is_error %token0 : !async.token
  94   vector.print %err0 : i1
  95
  96   // ------------------------------------------------------------------------ //
  97   // Check that assertion in the async.func converted to async error.
  98   // ------------------------------------------------------------------------ //
  99   %token1 = async.call @async_func_assert() : () -> !async.token
 100   async.runtime.await %token1 : !async.token
 101
 102   // CHECK: 1
 103   %err1 = async.runtime.is_error %token1 : !async.token
 104   vector.print %err1 : i1
 105
 106   // ------------------------------------------------------------------------ //
 107   // Check error propagation from the nested async.func.
 108   // ------------------------------------------------------------------------ //
 109   %token2 = async.call @async_func_nested_assert() : () -> !async.token
 110   async.runtime.await %token2 : !async.token
 111
 112   // CHECK: 1
 113   %err2 = async.runtime.is_error %token2 : !async.token
 114   vector.print %err2 : i1
 115
 116   // ------------------------------------------------------------------------ //
 117   // Check error propagation from the nested async.func with async values.
 118   // ------------------------------------------------------------------------ //
 119   %value3 = async.call @async_func_value_nested_assert() : () -> !async.value<f32>
 120   async.runtime.await %value3 : !async.value<f32>
 121
 122   // CHECK: 1
 123   %err3_0 = async.runtime.is_error %value3 : !async.value<f32>
 124   vector.print %err3_0 : i1
 125
 126   // ------------------------------------------------------------------------ //
 127   // Non-blocking async.await inside the async.func
 128   // ------------------------------------------------------------------------ //
 129   %result0 = async.call @async_func_non_blocking_await() : () -> !async.value<f32>
 130   %4 = async.await %result0 : !async.value<f32>
 131
 132   // CHECK: 456.789
 133   vector.print %4 : f32
 134
 135   // ------------------------------------------------------------------------ //
 136   // Memref allocated inside async.func.
 137   // ------------------------------------------------------------------------ //
 138   %result1 = async.call @async_func_inside_memref() : () -> !async.value<memref<f32>>
 139   %5 = async.await %result1 : !async.value<memref<f32>>
 140   %6 = memref.cast %5 :  memref<f32> to memref<*xf32>
 141
 142   // CHECK: Unranked Memref
 143   // CHECK-SAME: rank = 0 offset = 0 sizes = [] strides = []
 144   // CHECK-NEXT: [0.25]
 145   call @printMemrefF32(%6) : (memref<*xf32>) -> ()
 146
 147   // ------------------------------------------------------------------------ //
 148   // Memref passed as async.func parameter
 149   // ------------------------------------------------------------------------ //
 150   %token3 = async.call @async_func_passed_memref(%result1) : (!async.value<memref<f32>>) -> !async.token
 151   async.await %token3 : !async.token
 152
 153   // CHECK: Unranked Memref
 154   // CHECK-SAME: rank = 0 offset = 0 sizes = [] strides = []
 155   // CHECK-NEXT: [0.5]
 156   call @printMemrefF32(%6) : (memref<*xf32>) -> ()
 157
 158   // ------------------------------------------------------------------------ //
 159   // async.execute inside async.func
 160   // ------------------------------------------------------------------------ //
 161   %token4 = async.call @async_execute_in_async_func(%result1) : (!async.value<memref<f32>>) -> !async.token
 162   async.await %token4 : !async.token
 163
 164   // CHECK: Unranked Memref
 165   // CHECK-SAME: rank = 0 offset = 0 sizes = [] strides = []
 166   // CHECK-NEXT: [1]
 167   call @printMemrefF32(%6) : (memref<*xf32>) -> ()
 168
 169   memref.dealloc %5 : memref<f32>
 170
 171   return
 172 }
 173
 174 func.func private @printMemrefF32(memref<*xf32>)
 175   attributes { llvm.emit_c_interface }