xref: /llvm-project/mlir/test/Dialect/Async/async-runtime-ref-counting-opt.mlir (revision 5e7dea225be10d3ba0d01e87fb36e80c6764bd83)
1// RUN: mlir-opt %s -async-runtime-ref-counting-opt | FileCheck %s
2
3func.func private @consume_token(%arg0: !async.token)
4
5// CHECK-LABEL: @cancellable_operations_0
6func.func @cancellable_operations_0(%arg0: !async.token) {
7  // CHECK-NOT: async.runtime.add_ref
8  // CHECK-NOT: async.runtime.drop_ref
9  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
10  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
11  // CHECK: return
12  return
13}
14
15// CHECK-LABEL: @cancellable_operations_1
16func.func @cancellable_operations_1(%arg0: !async.token) {
17  // CHECK-NOT: async.runtime.add_ref
18  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
19  // CHECK: call @consume_toke
20  call @consume_token(%arg0): (!async.token) -> ()
21  // CHECK-NOT: async.runtime.drop_ref
22  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
23  // CHECK: return
24  return
25}
26
27// CHECK-LABEL: @cancellable_operations_2
28func.func @cancellable_operations_2(%arg0: !async.token) {
29  // CHECK: async.runtime.await
30  // CHECK-NEXT: async.runtime.await
31  // CHECK-NEXT: async.runtime.await
32  // CHECK-NEXT: return
33  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
34  async.runtime.await %arg0 : !async.token
35  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
36  async.runtime.await %arg0 : !async.token
37  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
38  async.runtime.await %arg0 : !async.token
39  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
40  return
41}
42
43// CHECK-LABEL: @cancellable_operations_3
44func.func @cancellable_operations_3(%arg0: !async.token) {
45  // CHECK-NOT: add_ref
46  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
47  // CHECK: call @consume_toke
48  call @consume_token(%arg0): (!async.token) -> ()
49  // CHECK-NOT: async.runtime.drop_ref
50  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
51  // CHECK: async.runtime.await
52  async.runtime.await %arg0 : !async.token
53  // CHECK: return
54  return
55}
56
57// CHECK-LABEL: @not_cancellable_operations_0
58func.func @not_cancellable_operations_0(%arg0: !async.token) {
59  // CHECK: add_ref
60  async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token
61  // CHECK: call @consume_toke
62  call @consume_token(%arg0): (!async.token) -> ()
63  // CHECK: async.runtime.await
64  async.runtime.await %arg0 : !async.token
65  // CHECK: async.runtime.drop_ref
66  async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token
67  // CHECK: return
68  return
69}
70