1a6628e59SEugene Zhulenev// RUN: mlir-opt %s -async-runtime-ref-counting-opt | FileCheck %s 2a6628e59SEugene Zhulenev 3*5e7dea22SRiver Riddlefunc.func private @consume_token(%arg0: !async.token) 4a6628e59SEugene Zhulenev 5a6628e59SEugene Zhulenev// CHECK-LABEL: @cancellable_operations_0 6*5e7dea22SRiver Riddlefunc.func @cancellable_operations_0(%arg0: !async.token) { 7a6628e59SEugene Zhulenev // CHECK-NOT: async.runtime.add_ref 8a6628e59SEugene Zhulenev // CHECK-NOT: async.runtime.drop_ref 992db09cdSEugene Zhulenev async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token 1092db09cdSEugene Zhulenev async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token 11a6628e59SEugene Zhulenev // CHECK: return 12a6628e59SEugene Zhulenev return 13a6628e59SEugene Zhulenev} 14a6628e59SEugene Zhulenev 15a6628e59SEugene Zhulenev// CHECK-LABEL: @cancellable_operations_1 16*5e7dea22SRiver Riddlefunc.func @cancellable_operations_1(%arg0: !async.token) { 17a6628e59SEugene Zhulenev // CHECK-NOT: async.runtime.add_ref 1892db09cdSEugene Zhulenev async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token 19a6628e59SEugene Zhulenev // CHECK: call @consume_toke 20a6628e59SEugene Zhulenev call @consume_token(%arg0): (!async.token) -> () 21a6628e59SEugene Zhulenev // CHECK-NOT: async.runtime.drop_ref 2292db09cdSEugene Zhulenev async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token 23a6628e59SEugene Zhulenev // CHECK: return 24a6628e59SEugene Zhulenev return 25a6628e59SEugene Zhulenev} 26a6628e59SEugene Zhulenev 27a6628e59SEugene Zhulenev// CHECK-LABEL: @cancellable_operations_2 28*5e7dea22SRiver Riddlefunc.func @cancellable_operations_2(%arg0: !async.token) { 29a6628e59SEugene Zhulenev // CHECK: async.runtime.await 30a6628e59SEugene Zhulenev // CHECK-NEXT: async.runtime.await 31a6628e59SEugene Zhulenev // CHECK-NEXT: async.runtime.await 32a6628e59SEugene Zhulenev // CHECK-NEXT: return 3392db09cdSEugene Zhulenev async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token 34a6628e59SEugene Zhulenev async.runtime.await %arg0 : !async.token 3592db09cdSEugene Zhulenev async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token 36a6628e59SEugene Zhulenev async.runtime.await %arg0 : !async.token 3792db09cdSEugene Zhulenev async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token 38a6628e59SEugene Zhulenev async.runtime.await %arg0 : !async.token 3992db09cdSEugene Zhulenev async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token 40a6628e59SEugene Zhulenev return 41a6628e59SEugene Zhulenev} 42a6628e59SEugene Zhulenev 43a6628e59SEugene Zhulenev// CHECK-LABEL: @cancellable_operations_3 44*5e7dea22SRiver Riddlefunc.func @cancellable_operations_3(%arg0: !async.token) { 45a6628e59SEugene Zhulenev // CHECK-NOT: add_ref 4692db09cdSEugene Zhulenev async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token 47a6628e59SEugene Zhulenev // CHECK: call @consume_toke 48a6628e59SEugene Zhulenev call @consume_token(%arg0): (!async.token) -> () 49a6628e59SEugene Zhulenev // CHECK-NOT: async.runtime.drop_ref 5092db09cdSEugene Zhulenev async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token 51a6628e59SEugene Zhulenev // CHECK: async.runtime.await 52a6628e59SEugene Zhulenev async.runtime.await %arg0 : !async.token 53a6628e59SEugene Zhulenev // CHECK: return 54a6628e59SEugene Zhulenev return 55a6628e59SEugene Zhulenev} 569ccdaac8SEugene Zhulenev 579ccdaac8SEugene Zhulenev// CHECK-LABEL: @not_cancellable_operations_0 58*5e7dea22SRiver Riddlefunc.func @not_cancellable_operations_0(%arg0: !async.token) { 599ccdaac8SEugene Zhulenev // CHECK: add_ref 6092db09cdSEugene Zhulenev async.runtime.add_ref %arg0 {count = 1 : i64} : !async.token 619ccdaac8SEugene Zhulenev // CHECK: call @consume_toke 629ccdaac8SEugene Zhulenev call @consume_token(%arg0): (!async.token) -> () 639ccdaac8SEugene Zhulenev // CHECK: async.runtime.await 649ccdaac8SEugene Zhulenev async.runtime.await %arg0 : !async.token 659ccdaac8SEugene Zhulenev // CHECK: async.runtime.drop_ref 6692db09cdSEugene Zhulenev async.runtime.drop_ref %arg0 {count = 1 : i64} : !async.token 679ccdaac8SEugene Zhulenev // CHECK: return 689ccdaac8SEugene Zhulenev return 699ccdaac8SEugene Zhulenev} 70