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