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

   1 // RUN:   mlir-opt %s -async-to-async-runtime                                  \
   2 // RUN:               -async-runtime-ref-counting                              \
   3 // RUN:               -async-runtime-ref-counting-opt                          \
   4 // RUN:               -convert-async-to-llvm                                   \
   5 // RUN:               -convert-linalg-to-loops                                 \
   6 // RUN:               -convert-scf-to-std                                      \
   7 // RUN:               -convert-linalg-to-llvm                                  \
   8 // RUN:               -convert-memref-to-llvm                                  \
   9 // RUN:               -convert-arith-to-llvm                                   \
  10 // RUN:               -convert-std-to-llvm                                     \
  11 // RUN:               -reconcile-unrealized-casts                              \
  12 // RUN: | mlir-cpu-runner                                                      \
  13 // RUN:     -e main -entry-point-result=void -O0                               \
  14 // RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_c_runner_utils%shlibext  \
  15 // RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext    \
  16 // RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_async_runtime%shlibext   \
  17 // RUN: | FileCheck %s
  18
  19 func @main() {
  20   %i0 = arith.constant 0 : index
  21   %i1 = arith.constant 1 : index
  22   %i2 = arith.constant 2 : index
  23   %i3 = arith.constant 3 : index
  24
  25   %c0 = arith.constant 0.0 : f32
  26   %c1 = arith.constant 1.0 : f32
  27   %c2 = arith.constant 2.0 : f32
  28   %c3 = arith.constant 3.0 : f32
  29   %c4 = arith.constant 4.0 : f32
  30
  31   %A = memref.alloc() : memref<4xf32>
  32   linalg.fill(%c0, %A) : f32, memref<4xf32>
  33
  34   // CHECK: [0, 0, 0, 0]
  35   %U = memref.cast %A :  memref<4xf32> to memref<*xf32>
  36   call @print_memref_f32(%U): (memref<*xf32>) -> ()
  37
  38   // CHECK: Current thread id: [[MAIN:.*]]
  39   // CHECK: [1, 0, 0, 0]
  40   memref.store %c1, %A[%i0]: memref<4xf32>
  41   call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
  42   call @print_memref_f32(%U): (memref<*xf32>) -> ()
  43
  44   %outer = async.execute {
  45     // CHECK: Current thread id: [[THREAD0:.*]]
  46     // CHECK: [1, 2, 0, 0]
  47     memref.store %c2, %A[%i1]: memref<4xf32>
  48     call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
  49     call @print_memref_f32(%U): (memref<*xf32>) -> ()
  50
  51     // No op async region to create a token for testing async dependency.
  52     %noop = async.execute {
  53       // CHECK: Current thread id: [[THREAD1:.*]]
  54       call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
  55       async.yield
  56     }
  57
  58     %inner = async.execute [%noop] {
  59       // CHECK: Current thread id: [[THREAD2:.*]]
  60       // CHECK: [1, 2, 3, 0]
  61       memref.store %c3, %A[%i2]: memref<4xf32>
  62       call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
  63       call @print_memref_f32(%U): (memref<*xf32>) -> ()
  64
  65       async.yield
  66     }
  67     async.await %inner : !async.token
  68
  69     // CHECK: Current thread id: [[THREAD3:.*]]
  70     // CHECK: [1, 2, 3, 4]
  71     memref.store %c4, %A[%i3]: memref<4xf32>
  72     call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
  73     call @print_memref_f32(%U): (memref<*xf32>) -> ()
  74
  75     async.yield
  76   }
  77   async.await %outer : !async.token
  78
  79   // CHECK: Current thread id: [[MAIN]]
  80   // CHECK: [1, 2, 3, 4]
  81   call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
  82   call @print_memref_f32(%U): (memref<*xf32>) -> ()
  83
  84   memref.dealloc %A : memref<4xf32>
  85
  86   return
  87 }
  88
  89 func private @mlirAsyncRuntimePrintCurrentThreadId() -> ()
  90
  91 func private @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }