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