1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
|
// RUN: mlir-opt %s -split-input-file -convert-async-to-llvm | FileCheck %s
// CHECK-LABEL: execute_no_async_args
func @execute_no_async_args(%arg0: f32, %arg1: memref<1xf32>) {
// CHECK: %[[TOKEN:.*]] = call @async_execute_fn(%arg0, %arg1)
%token = async.execute {
%c0 = constant 0 : index
store %arg0, %arg1[%c0] : memref<1xf32>
async.yield
}
// CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]])
// CHECK-NEXT: return
async.await %token : !async.token
return
}
// Function outlined from the async.execute operation.
// CHECK-LABEL: func private @async_execute_fn(%arg0: f32, %arg1: memref<1xf32>)
// CHECK-SAME: -> !llvm.ptr<i8>
// Create token for return op, and mark a function as a coroutine.
// CHECK: %[[RET:.*]] = call @mlirAsyncRuntimeCreateToken()
// CHECK: %[[HDL:.*]] = llvm.call @llvm.coro.begin
// Pass a suspended coroutine to the async runtime.
// CHECK: %[[RESUME:.*]] = llvm.mlir.addressof @__resume
// CHECK: %[[STATE:.*]] = llvm.call @llvm.coro.save
// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL]], %[[RESUME]])
// CHECK: %[[SUSPENDED:.*]] = llvm.call @llvm.coro.suspend(%[[STATE]]
// Decide the next block based on the code returned from suspend.
// CHECK: %[[ZERO:.*]] = llvm.mlir.constant(0 : i8)
// CHECK: %[[NONE:.*]] = llvm.mlir.constant(-1 : i8)
// CHECK: %[[IS_NONE:.*]] = llvm.icmp "eq" %[[SUSPENDED]], %[[NONE]]
// CHECK: llvm.cond_br %[[IS_NONE]], ^[[SUSPEND:.*]], ^[[RESUME_OR_CLEANUP:.*]]
// Decide if branch to resume or cleanup block.
// CHECK: ^[[RESUME_OR_CLEANUP]]:
// CHECK: %[[IS_ZERO:.*]] = llvm.icmp "eq" %[[SUSPENDED]], %[[ZERO]]
// CHECK: llvm.cond_br %[[IS_ZERO]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]
// Resume coroutine after suspension.
// CHECK: ^[[RESUME]]:
// CHECK: store %arg0, %arg1[%c0] : memref<1xf32>
// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET]])
// Delete coroutine.
// CHECK: ^[[CLEANUP]]:
// CHECK: %[[MEM:.*]] = llvm.call @llvm.coro.free
// CHECK: llvm.call @free(%[[MEM]])
// Suspend coroutine, and also a return statement for ramp function.
// CHECK: ^[[SUSPEND]]:
// CHECK: llvm.call @llvm.coro.end
// CHECK: return %[[RET]]
// -----
// CHECK-LABEL: nested_async_execute
func @nested_async_execute(%arg0: f32, %arg1: f32, %arg2: memref<1xf32>) {
// CHECK: %[[TOKEN:.*]] = call @async_execute_fn_0(%arg0, %arg2, %arg1)
%token0 = async.execute {
%c0 = constant 0 : index
%token1 = async.execute {
%c1 = constant 1: index
store %arg0, %arg2[%c0] : memref<1xf32>
async.yield
}
async.await %token1 : !async.token
store %arg1, %arg2[%c0] : memref<1xf32>
async.yield
}
// CHECK: call @mlirAsyncRuntimeAwaitToken(%[[TOKEN]])
// CHECK-NEXT: return
async.await %token0 : !async.token
return
}
// Function outlined from the inner async.execute operation.
// CHECK-LABEL: func private @async_execute_fn(%arg0: f32, %arg1: memref<1xf32>, %arg2: index)
// CHECK-SAME: -> !llvm.ptr<i8>
// CHECK: %[[RET_0:.*]] = call @mlirAsyncRuntimeCreateToken()
// CHECK: %[[HDL_0:.*]] = llvm.call @llvm.coro.begin
// CHECK: call @mlirAsyncRuntimeExecute
// CHECK: llvm.call @llvm.coro.suspend
// CHECK: store %arg0, %arg1[%arg2] : memref<1xf32>
// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_0]])
// Function outlined from the outer async.execute operation.
// CHECK-LABEL: func private @async_execute_fn_0(%arg0: f32, %arg1: memref<1xf32>, %arg2: f32)
// CHECK-SAME: -> !llvm.ptr<i8>
// CHECK: %[[RET_1:.*]] = call @mlirAsyncRuntimeCreateToken()
// CHECK: %[[HDL_1:.*]] = llvm.call @llvm.coro.begin
// Suspend coroutine in the beginning.
// CHECK: call @mlirAsyncRuntimeExecute
// CHECK: llvm.call @llvm.coro.suspend
// Suspend coroutine second time waiting for the completion of inner execute op.
// CHECK: %[[TOKEN_1:.*]] = call @async_execute_fn
// CHECK: llvm.call @llvm.coro.save
// CHECK: call @mlirAsyncRuntimeAwaitTokenAndExecute(%[[TOKEN_1]], %[[HDL_1]]
// CHECK: llvm.call @llvm.coro.suspend
// Emplace result token after second resumption.
// CHECK: store %arg2, %arg1[%c0] : memref<1xf32>
// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_1]])
// -----
// CHECK-LABEL: async_execute_token_dependency
func @async_execute_token_dependency(%arg0: f32, %arg1: memref<1xf32>) {
// CHECK: %0 = call @async_execute_fn(%arg0, %arg1)
%token = async.execute {
%c0 = constant 0 : index
store %arg0, %arg1[%c0] : memref<1xf32>
async.yield
}
// CHECK: %1 = call @async_execute_fn_0(%0, %arg0, %arg1)
%token_0 = async.execute [%token] {
%c0 = constant 0 : index
store %arg0, %arg1[%c0] : memref<1xf32>
async.yield
}
return
}
// Function outlined from the first async.execute operation.
// CHECK-LABEL: func private @async_execute_fn(%arg0: f32, %arg1: memref<1xf32>)
// CHECK-SAME: -> !llvm.ptr<i8>
// CHECK: %[[RET_0:.*]] = call @mlirAsyncRuntimeCreateToken()
// CHECK: %[[HDL_0:.*]] = llvm.call @llvm.coro.begin
// CHECK: call @mlirAsyncRuntimeExecute
// CHECK: llvm.call @llvm.coro.suspend
// CHECK: store %arg0, %arg1[%c0] : memref<1xf32>
// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_0]])
// Function outlined from the second async.execute operation with dependency.
// CHECK-LABEL: func private @async_execute_fn_0(%arg0: !llvm.ptr<i8>, %arg1: f32, %arg2: memref<1xf32>)
// CHECK-SAME: -> !llvm.ptr<i8>
// CHECK: %[[RET_1:.*]] = call @mlirAsyncRuntimeCreateToken()
// CHECK: %[[HDL_1:.*]] = llvm.call @llvm.coro.begin
// Suspend coroutine in the beginning.
// CHECK: call @mlirAsyncRuntimeExecute(%[[HDL_1]],
// CHECK: llvm.call @llvm.coro.suspend
// Suspend coroutine second time waiting for the completion of token dependency.
// CHECK: llvm.call @llvm.coro.save
// CHECK: call @mlirAsyncRuntimeAwaitTokenAndExecute(%arg0, %[[HDL_1]],
// CHECK: llvm.call @llvm.coro.suspend
// Emplace result token after second resumption.
// CHECK: store %arg1, %arg2[%c0] : memref<1xf32>
// CHECK: call @mlirAsyncRuntimeEmplaceToken(%[[RET_1]])
|
