1// RUN: mlir-opt %s -async-ref-counting-optimization | FileCheck %s 2 3// CHECK-LABEL: @cancellable_operations_0 4func @cancellable_operations_0(%arg0: !async.token) { 5 // CHECK-NOT: async.add_ref 6 // CHECK-NOT: async.drop_ref 7 async.add_ref %arg0 {count = 1 : i32} : !async.token 8 async.drop_ref %arg0 {count = 1 : i32} : !async.token 9 // CHECK: return 10 return 11} 12 13// CHECK-LABEL: @cancellable_operations_1 14func @cancellable_operations_1(%arg0: !async.token) { 15 // CHECK-NOT: async.add_ref 16 // CHECK: async.execute 17 async.add_ref %arg0 {count = 1 : i32} : !async.token 18 async.execute [%arg0] { 19 // CHECK: async.drop_ref 20 async.drop_ref %arg0 {count = 1 : i32} : !async.token 21 // CHECK-NEXT: async.yield 22 async.yield 23 } 24 // CHECK-NOT: async.drop_ref 25 async.drop_ref %arg0 {count = 1 : i32} : !async.token 26 // CHECK: return 27 return 28} 29 30// CHECK-LABEL: @cancellable_operations_2 31func @cancellable_operations_2(%arg0: !async.token) { 32 // CHECK: async.await 33 // CHECK-NEXT: async.await 34 // CHECK-NEXT: async.await 35 // CHECK-NEXT: return 36 async.add_ref %arg0 {count = 1 : i32} : !async.token 37 async.await %arg0 : !async.token 38 async.drop_ref %arg0 {count = 1 : i32} : !async.token 39 async.await %arg0 : !async.token 40 async.add_ref %arg0 {count = 1 : i32} : !async.token 41 async.await %arg0 : !async.token 42 async.drop_ref %arg0 {count = 1 : i32} : !async.token 43 return 44} 45 46// CHECK-LABEL: @cancellable_operations_3 47func @cancellable_operations_3(%arg0: !async.token) { 48 // CHECK-NOT: add_ref 49 async.add_ref %arg0 {count = 1 : i32} : !async.token 50 %token = async.execute { 51 async.await %arg0 : !async.token 52 // CHECK: async.drop_ref 53 async.drop_ref %arg0 {count = 1 : i32} : !async.token 54 async.yield 55 } 56 // CHECK-NOT: async.drop_ref 57 async.drop_ref %arg0 {count = 1 : i32} : !async.token 58 // CHECK: async.await 59 async.await %arg0 : !async.token 60 // CHECK: return 61 return 62} 63 64// CHECK-LABEL: @not_cancellable_operations_0 65func @not_cancellable_operations_0(%arg0: !async.token, %arg1: i1) { 66 // It is unsafe to cancel `add_ref` / `drop_ref` pair because it is possible 67 // that the body of the `async.execute` operation will run before the await 68 // operation in the function body, and will destroy the `%arg0` token. 69 // CHECK: add_ref 70 async.add_ref %arg0 {count = 1 : i32} : !async.token 71 %token = async.execute { 72 // CHECK: async.await 73 async.await %arg0 : !async.token 74 // CHECK: async.drop_ref 75 async.drop_ref %arg0 {count = 1 : i32} : !async.token 76 // CHECK: async.yield 77 async.yield 78 } 79 // CHECK: async.await 80 async.await %arg0 : !async.token 81 // CHECK: drop_ref 82 async.drop_ref %arg0 {count = 1 : i32} : !async.token 83 // CHECK: return 84 return 85} 86 87// CHECK-LABEL: @not_cancellable_operations_1 88func @not_cancellable_operations_1(%arg0: !async.token, %arg1: i1) { 89 // Same reason as above, although `async.execute` is inside the nested 90 // region or "regular" opeation. 91 // 92 // NOTE: This test is not correct w.r.t. reference counting, and at runtime 93 // would leak %arg0 value if %arg1 is false. IR like this will not be 94 // constructed by automatic reference counting pass, because it would 95 // place `async.add_ref` right before the `async.execute` inside `scf.if`. 96 97 // CHECK: async.add_ref 98 async.add_ref %arg0 {count = 1 : i32} : !async.token 99 scf.if %arg1 { 100 %token = async.execute { 101 async.await %arg0 : !async.token 102 // CHECK: async.drop_ref 103 async.drop_ref %arg0 {count = 1 : i32} : !async.token 104 async.yield 105 } 106 } 107 // CHECK: async.await 108 async.await %arg0 : !async.token 109 // CHECK: async.drop_ref 110 async.drop_ref %arg0 {count = 1 : i32} : !async.token 111 // CHECK: return 112 return 113} 114