1// RUN: mlir-opt %s | FileCheck %s 2 3// CHECK-LABEL: @identity_token 4func.func @identity_token(%arg0: !async.token) -> !async.token { 5 // CHECK: return %arg0 : !async.token 6 return %arg0 : !async.token 7} 8 9// CHECK-LABEL: @identity_value 10func.func @identity_value(%arg0 : !async.value<f32>) -> !async.value<f32> { 11 // CHECK: return %arg0 : !async.value<f32> 12 return %arg0 : !async.value<f32> 13} 14 15// CHECK-LABEL: @empty_async_execute 16func.func @empty_async_execute() -> !async.token { 17 // CHECK: async.execute 18 %token = async.execute { 19 async.yield 20 } 21 22 // CHECK: return %token : !async.token 23 return %token : !async.token 24} 25 26// CHECK-LABEL: @return_async_value 27func.func @return_async_value() -> !async.value<f32> { 28 // CHECK: async.execute -> !async.value<f32> 29 %token, %bodyResults = async.execute -> !async.value<f32> { 30 %cst = arith.constant 1.000000e+00 : f32 31 async.yield %cst : f32 32 } 33 34 // CHECK: return %bodyResults : !async.value<f32> 35 return %bodyResults : !async.value<f32> 36} 37 38// CHECK-LABEL: @return_captured_value 39func.func @return_captured_value() -> !async.token { 40 %cst = arith.constant 1.000000e+00 : f32 41 // CHECK: async.execute -> !async.value<f32> 42 %token, %results = async.execute -> !async.value<f32> { 43 async.yield %cst : f32 44 } 45 46 // CHECK: return %token : !async.token 47 return %token : !async.token 48} 49 50// CHECK-LABEL: @return_async_values 51func.func @return_async_values() -> (!async.value<f32>, !async.value<f32>) { 52 %token, %bodyResults:2 = async.execute -> (!async.value<f32>, !async.value<f32>) { 53 %cst1 = arith.constant 1.000000e+00 : f32 54 %cst2 = arith.constant 2.000000e+00 : f32 55 async.yield %cst1, %cst2 : f32, f32 56 } 57 58 // CHECK: return %bodyResults#0, %bodyResults#1 : !async.value<f32>, !async.value<f32> 59 return %bodyResults#0, %bodyResults#1 : !async.value<f32>, !async.value<f32> 60} 61 62// CHECK-LABEL: @async_token_dependencies 63func.func @async_token_dependencies(%arg0: !async.token) -> !async.token { 64 // CHECK: async.execute [%arg0] 65 %token = async.execute [%arg0] { 66 async.yield 67 } 68 69 // CHECK: return %token : !async.token 70 return %token : !async.token 71} 72 73// CHECK-LABEL: @async_value_operands 74func.func @async_value_operands(%arg0: !async.value<f32>) -> !async.token { 75 // CHECK: async.execute (%arg0 as %arg1: !async.value<f32>) -> !async.value<f32> 76 %token, %results = async.execute (%arg0 as %arg1: !async.value<f32>) -> !async.value<f32> { 77 async.yield %arg1 : f32 78 } 79 80 // CHECK: return %token : !async.token 81 return %token : !async.token 82} 83 84// CHECK-LABEL: @async_token_and_value_operands 85func.func @async_token_and_value_operands(%arg0: !async.token, %arg1: !async.value<f32>) -> !async.token { 86 // CHECK: async.execute [%arg0] (%arg1 as %arg2: !async.value<f32>) -> !async.value<f32> 87 %token, %results = async.execute [%arg0] (%arg1 as %arg2: !async.value<f32>) -> !async.value<f32> { 88 async.yield %arg2 : f32 89 } 90 91 // CHECK: return %token : !async.token 92 return %token : !async.token 93} 94 95// CHECK-LABEL: @empty_tokens_or_values_operands 96func.func @empty_tokens_or_values_operands() { 97 // CHECK: async.execute { 98 %token0 = async.execute [] () -> () { async.yield } 99 // CHECK: async.execute { 100 %token1 = async.execute () -> () { async.yield } 101 // CHECK: async.execute { 102 %token2 = async.execute -> () { async.yield } 103 // CHECK: async.execute { 104 %token3 = async.execute () { async.yield } 105 // CHECK: async.execute { 106 %token4 = async.execute [] { async.yield } 107 return 108} 109 110// CHECK-LABEL: @await_token 111func.func @await_token(%arg0: !async.token) { 112 // CHECK: async.await %arg0 113 async.await %arg0 : !async.token 114 return 115} 116 117// CHECK-LABEL: @await_value 118func.func @await_value(%arg0: !async.value<f32>) -> f32 { 119 // CHECK: async.await %arg0 120 %0 = async.await %arg0 : !async.value<f32> 121 return %0 : f32 122} 123 124// CHECK-LABEL: @create_group_and_await_all 125func.func @create_group_and_await_all(%arg0: !async.token, 126 %arg1: !async.value<f32>) -> index { 127 %c = arith.constant 2 : index 128 %0 = async.create_group %c : !async.group 129 130 // CHECK: async.add_to_group %arg0 131 // CHECK: async.add_to_group %arg1 132 %1 = async.add_to_group %arg0, %0 : !async.token 133 %2 = async.add_to_group %arg1, %0 : !async.value<f32> 134 async.await_all %0 135 136 %3 = arith.addi %1, %2 : index 137 return %3 : index 138} 139 140// CHECK-LABEL: @async_func_return_token 141async.func @async_func_return_token() -> !async.token { 142 // CHECK: return 143 return 144} 145 146// CHECK-LABEL: @async_func_return_value 147async.func @async_func_return_value() -> !async.value<i32> { 148 %0 = arith.constant 42 : i32 149 // CHECK: return %[[value:.*]] : i32 150 return %0 : i32 151} 152 153// CHECK-LABEL: @async_func_return_optional_token 154async.func @async_func_return_optional_token() -> (!async.token, !async.value<i32>) { 155 %0 = arith.constant 42 : i32 156 // CHECK: return %[[value:.*]] : i32 157 return %0 : i32 158} 159 160// CHECK-LABEL: @async_call 161func.func @async_call() { 162 // CHECK: async.call @async_func_return_token 163 // CHECK: async.call @async_func_return_value 164 // CHECK: async.call @async_func_return_optional_token 165 %0 = async.call @async_func_return_token() : () -> !async.token 166 %1 = async.call @async_func_return_value() : () -> !async.value<i32> 167 %2, %3 = async.call @async_func_return_optional_token() : () -> (!async.token, !async.value<i32>) 168 return 169} 170