• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1// RUN: mlir-opt %s -async-ref-counting | FileCheck %s
2
3// CHECK-LABEL: @cond
4func private @cond() -> i1
5
6// CHECK-LABEL: @token_arg_no_uses
7func @token_arg_no_uses(%arg0: !async.token) {
8  // CHECK: async.drop_ref %arg0 {count = 1 : i32}
9  return
10}
11
12// CHECK-LABEL: @token_arg_conditional_await
13func @token_arg_conditional_await(%arg0: !async.token, %arg1: i1) {
14  cond_br %arg1, ^bb1, ^bb2
15^bb1:
16  // CHECK: async.drop_ref %arg0 {count = 1 : i32}
17  return
18^bb2:
19  // CHECK: async.await %arg0
20  // CHECK: async.drop_ref %arg0 {count = 1 : i32}
21  async.await %arg0 : !async.token
22  return
23}
24
25// CHECK-LABEL: @token_no_uses
26func @token_no_uses() {
27  // CHECK: %[[TOKEN:.*]] = async.execute
28  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
29  %token = async.execute {
30    async.yield
31  }
32  return
33}
34
35// CHECK-LABEL: @token_return
36func @token_return() -> !async.token {
37  // CHECK: %[[TOKEN:.*]] = async.execute
38  %token = async.execute {
39    async.yield
40  }
41  // CHECK: return %[[TOKEN]]
42  return %token : !async.token
43}
44
45// CHECK-LABEL: @token_await
46func @token_await() {
47  // CHECK: %[[TOKEN:.*]] = async.execute
48  %token = async.execute {
49    async.yield
50  }
51  // CHECK: async.await %[[TOKEN]]
52  async.await %token : !async.token
53  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
54  // CHECK: return
55  return
56}
57
58// CHECK-LABEL: @token_await_and_return
59func @token_await_and_return() -> !async.token {
60  // CHECK: %[[TOKEN:.*]] = async.execute
61  %token = async.execute {
62    async.yield
63  }
64  // CHECK: async.await %[[TOKEN]]
65  // CHECK-NOT: async.drop_ref
66  async.await %token : !async.token
67  // CHECK: return %[[TOKEN]]
68  return %token : !async.token
69}
70
71// CHECK-LABEL: @token_await_inside_scf_if
72func @token_await_inside_scf_if(%arg0: i1) {
73  // CHECK: %[[TOKEN:.*]] = async.execute
74  %token = async.execute {
75    async.yield
76  }
77  // CHECK: scf.if %arg0 {
78  scf.if %arg0 {
79    // CHECK: async.await %[[TOKEN]]
80    async.await %token : !async.token
81  }
82  // CHECK: }
83  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
84  // CHECK: return
85  return
86}
87
88// CHECK-LABEL: @token_conditional_await
89func @token_conditional_await(%arg0: i1) {
90  // CHECK: %[[TOKEN:.*]] = async.execute
91  %token = async.execute {
92    async.yield
93  }
94  cond_br %arg0, ^bb1, ^bb2
95^bb1:
96  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
97  return
98^bb2:
99  // CHECK: async.await %[[TOKEN]]
100  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
101  async.await %token : !async.token
102  return
103}
104
105// CHECK-LABEL: @token_await_in_the_loop
106func @token_await_in_the_loop() {
107  // CHECK: %[[TOKEN:.*]] = async.execute
108  %token = async.execute {
109    async.yield
110  }
111  br ^bb1
112^bb1:
113  // CHECK: async.await %[[TOKEN]]
114  async.await %token : !async.token
115  %0 = call @cond(): () -> (i1)
116  cond_br %0, ^bb1, ^bb2
117^bb2:
118  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
119  return
120}
121
122// CHECK-LABEL: @token_defined_in_the_loop
123func @token_defined_in_the_loop() {
124  br ^bb1
125^bb1:
126  // CHECK: %[[TOKEN:.*]] = async.execute
127  %token = async.execute {
128    async.yield
129  }
130  // CHECK: async.await %[[TOKEN]]
131  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
132  async.await %token : !async.token
133  %0 = call @cond(): () -> (i1)
134  cond_br %0, ^bb1, ^bb2
135^bb2:
136  return
137}
138
139// CHECK-LABEL: @token_capture
140func @token_capture() {
141  // CHECK: %[[TOKEN:.*]] = async.execute
142  %token = async.execute {
143    async.yield
144  }
145
146  // CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
147  // CHECK: %[[TOKEN_0:.*]] = async.execute
148  %token_0 = async.execute {
149    // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
150    // CHECK-NEXT: async.yield
151    async.await %token : !async.token
152    async.yield
153  }
154  // CHECK: async.drop_ref %[[TOKEN_0]] {count = 1 : i32}
155  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
156  // CHECK: return
157  return
158}
159
160// CHECK-LABEL: @token_nested_capture
161func @token_nested_capture() {
162  // CHECK: %[[TOKEN:.*]] = async.execute
163  %token = async.execute {
164    async.yield
165  }
166
167  // CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
168  // CHECK: %[[TOKEN_0:.*]] = async.execute
169  %token_0 = async.execute {
170    // CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
171    // CHECK: %[[TOKEN_1:.*]] = async.execute
172    %token_1 = async.execute {
173      // CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
174      // CHECK: %[[TOKEN_2:.*]] = async.execute
175      %token_2 = async.execute {
176        // CHECK: async.await %[[TOKEN]]
177        // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
178        async.await %token : !async.token
179        async.yield
180      }
181      // CHECK: async.drop_ref %[[TOKEN_2]] {count = 1 : i32}
182      // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
183      async.yield
184    }
185    // CHECK: async.drop_ref %[[TOKEN_1]] {count = 1 : i32}
186    // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
187    async.yield
188  }
189  // CHECK: async.drop_ref %[[TOKEN_0]] {count = 1 : i32}
190  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
191  // CHECK: return
192  return
193}
194
195// CHECK-LABEL: @token_dependency
196func @token_dependency() {
197  // CHECK: %[[TOKEN:.*]] = async.execute
198  %token = async.execute {
199    async.yield
200  }
201
202  // CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
203  // CHECK: %[[TOKEN_0:.*]] = async.execute
204  %token_0 = async.execute[%token] {
205    // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
206    // CHECK-NEXT: async.yield
207    async.yield
208  }
209
210  // CHECK: async.await %[[TOKEN]]
211  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
212  async.await %token : !async.token
213  // CHECK: async.await %[[TOKEN_0]]
214  // CHECK: async.drop_ref %[[TOKEN_0]] {count = 1 : i32}
215  async.await %token_0 : !async.token
216
217  // CHECK: return
218  return
219}
220
221// CHECK-LABEL: @value_operand
222func @value_operand() -> f32 {
223  // CHECK: %[[TOKEN:.*]], %[[RESULTS:.*]] = async.execute
224  %token, %results = async.execute -> !async.value<f32> {
225    %0 = constant 0.0 : f32
226    async.yield %0 : f32
227  }
228
229  // CHECK: async.add_ref %[[TOKEN]] {count = 1 : i32}
230  // CHECK: async.add_ref %[[RESULTS]] {count = 1 : i32}
231  // CHECK: %[[TOKEN_0:.*]] = async.execute
232  %token_0 = async.execute[%token](%results as %arg0 : !async.value<f32>)  {
233    // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
234    // CHECK: async.drop_ref %[[RESULTS]] {count = 1 : i32}
235    // CHECK: async.yield
236    async.yield
237  }
238
239  // CHECK: async.await %[[TOKEN]]
240  // CHECK: async.drop_ref %[[TOKEN]] {count = 1 : i32}
241  async.await %token : !async.token
242
243  // CHECK: async.await %[[TOKEN_0]]
244  // CHECK: async.drop_ref %[[TOKEN_0]] {count = 1 : i32}
245  async.await %token_0 : !async.token
246
247  // CHECK: async.await %[[RESULTS]]
248  // CHECK: async.drop_ref %[[RESULTS]] {count = 1 : i32}
249  %0 = async.await %results : !async.value<f32>
250
251  // CHECK: return
252  return %0 : f32
253}
254