9c53b8e52e
[NFC] No new functionality, mostly a cleanup and one more abstraction level between Async and LLVM IR. Instead of lowering from Async to LLVM coroutines and Async Runtime API in one shot, do it progressively via async.coro and async.runtime operations. 1. Lower from async to async.runtime/coro (e.g. async.execute to function with coro setup and runtime calls) 2. Lower from async.runtime/coro to LLVM intrinsics and runtime API calls Intermediate coro/runtime operations will allow to run transformations on a higher level IR and do not try to match IR based on the LLVM::CallOp properties. Although async.coro is very close to LLVM coroutines, it is not exactly the same API, instead it is optimized for usability in async lowering, and misses a lot of details that are present in @llvm.coro intrinsic. Reviewed By: ftynse Differential Revision: https://reviews.llvm.org/D94923
254 lines
6.9 KiB
MLIR
254 lines
6.9 KiB
MLIR
// RUN: mlir-opt %s -async-ref-counting | FileCheck %s
|
|
|
|
// CHECK-LABEL: @cond
|
|
func private @cond() -> i1
|
|
|
|
// CHECK-LABEL: @token_arg_no_uses
|
|
func @token_arg_no_uses(%arg0: !async.token) {
|
|
// CHECK: async.runtime.drop_ref %arg0 {count = 1 : i32}
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @token_arg_conditional_await
|
|
func @token_arg_conditional_await(%arg0: !async.token, %arg1: i1) {
|
|
cond_br %arg1, ^bb1, ^bb2
|
|
^bb1:
|
|
// CHECK: async.runtime.drop_ref %arg0 {count = 1 : i32}
|
|
return
|
|
^bb2:
|
|
// CHECK: async.await %arg0
|
|
// CHECK: async.runtime.drop_ref %arg0 {count = 1 : i32}
|
|
async.await %arg0 : !async.token
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @token_no_uses
|
|
func @token_no_uses() {
|
|
// CHECK: %[[TOKEN:.*]] = async.execute
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
%token = async.execute {
|
|
async.yield
|
|
}
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @token_return
|
|
func @token_return() -> !async.token {
|
|
// CHECK: %[[TOKEN:.*]] = async.execute
|
|
%token = async.execute {
|
|
async.yield
|
|
}
|
|
// CHECK: return %[[TOKEN]]
|
|
return %token : !async.token
|
|
}
|
|
|
|
// CHECK-LABEL: @token_await
|
|
func @token_await() {
|
|
// CHECK: %[[TOKEN:.*]] = async.execute
|
|
%token = async.execute {
|
|
async.yield
|
|
}
|
|
// CHECK: async.await %[[TOKEN]]
|
|
async.await %token : !async.token
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK: return
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @token_await_and_return
|
|
func @token_await_and_return() -> !async.token {
|
|
// CHECK: %[[TOKEN:.*]] = async.execute
|
|
%token = async.execute {
|
|
async.yield
|
|
}
|
|
// CHECK: async.await %[[TOKEN]]
|
|
// CHECK-NOT: async.runtime.drop_ref
|
|
async.await %token : !async.token
|
|
// CHECK: return %[[TOKEN]]
|
|
return %token : !async.token
|
|
}
|
|
|
|
// CHECK-LABEL: @token_await_inside_scf_if
|
|
func @token_await_inside_scf_if(%arg0: i1) {
|
|
// CHECK: %[[TOKEN:.*]] = async.execute
|
|
%token = async.execute {
|
|
async.yield
|
|
}
|
|
// CHECK: scf.if %arg0 {
|
|
scf.if %arg0 {
|
|
// CHECK: async.await %[[TOKEN]]
|
|
async.await %token : !async.token
|
|
}
|
|
// CHECK: }
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK: return
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @token_conditional_await
|
|
func @token_conditional_await(%arg0: i1) {
|
|
// CHECK: %[[TOKEN:.*]] = async.execute
|
|
%token = async.execute {
|
|
async.yield
|
|
}
|
|
cond_br %arg0, ^bb1, ^bb2
|
|
^bb1:
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
return
|
|
^bb2:
|
|
// CHECK: async.await %[[TOKEN]]
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
async.await %token : !async.token
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @token_await_in_the_loop
|
|
func @token_await_in_the_loop() {
|
|
// CHECK: %[[TOKEN:.*]] = async.execute
|
|
%token = async.execute {
|
|
async.yield
|
|
}
|
|
br ^bb1
|
|
^bb1:
|
|
// CHECK: async.await %[[TOKEN]]
|
|
async.await %token : !async.token
|
|
%0 = call @cond(): () -> (i1)
|
|
cond_br %0, ^bb1, ^bb2
|
|
^bb2:
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @token_defined_in_the_loop
|
|
func @token_defined_in_the_loop() {
|
|
br ^bb1
|
|
^bb1:
|
|
// CHECK: %[[TOKEN:.*]] = async.execute
|
|
%token = async.execute {
|
|
async.yield
|
|
}
|
|
// CHECK: async.await %[[TOKEN]]
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
async.await %token : !async.token
|
|
%0 = call @cond(): () -> (i1)
|
|
cond_br %0, ^bb1, ^bb2
|
|
^bb2:
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @token_capture
|
|
func @token_capture() {
|
|
// CHECK: %[[TOKEN:.*]] = async.execute
|
|
%token = async.execute {
|
|
async.yield
|
|
}
|
|
|
|
// CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK: %[[TOKEN_0:.*]] = async.execute
|
|
%token_0 = async.execute {
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK-NEXT: async.yield
|
|
async.await %token : !async.token
|
|
async.yield
|
|
}
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN_0]] {count = 1 : i32}
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK: return
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @token_nested_capture
|
|
func @token_nested_capture() {
|
|
// CHECK: %[[TOKEN:.*]] = async.execute
|
|
%token = async.execute {
|
|
async.yield
|
|
}
|
|
|
|
// CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK: %[[TOKEN_0:.*]] = async.execute
|
|
%token_0 = async.execute {
|
|
// CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK: %[[TOKEN_1:.*]] = async.execute
|
|
%token_1 = async.execute {
|
|
// CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK: %[[TOKEN_2:.*]] = async.execute
|
|
%token_2 = async.execute {
|
|
// CHECK: async.await %[[TOKEN]]
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
async.await %token : !async.token
|
|
async.yield
|
|
}
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN_2]] {count = 1 : i32}
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
async.yield
|
|
}
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN_1]] {count = 1 : i32}
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
async.yield
|
|
}
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN_0]] {count = 1 : i32}
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK: return
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @token_dependency
|
|
func @token_dependency() {
|
|
// CHECK: %[[TOKEN:.*]] = async.execute
|
|
%token = async.execute {
|
|
async.yield
|
|
}
|
|
|
|
// CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK: %[[TOKEN_0:.*]] = async.execute
|
|
%token_0 = async.execute[%token] {
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK-NEXT: async.yield
|
|
async.yield
|
|
}
|
|
|
|
// CHECK: async.await %[[TOKEN]]
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
async.await %token : !async.token
|
|
// CHECK: async.await %[[TOKEN_0]]
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN_0]] {count = 1 : i32}
|
|
async.await %token_0 : !async.token
|
|
|
|
// CHECK: return
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @value_operand
|
|
func @value_operand() -> f32 {
|
|
// CHECK: %[[TOKEN:.*]], %[[RESULTS:.*]] = async.execute
|
|
%token, %results = async.execute -> !async.value<f32> {
|
|
%0 = constant 0.0 : f32
|
|
async.yield %0 : f32
|
|
}
|
|
|
|
// CHECK: async.runtime.add_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK: async.runtime.add_ref %[[RESULTS]] {count = 1 : i32}
|
|
// CHECK: %[[TOKEN_0:.*]] = async.execute
|
|
%token_0 = async.execute[%token](%results as %arg0 : !async.value<f32>) {
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
// CHECK: async.runtime.drop_ref %[[RESULTS]] {count = 1 : i32}
|
|
// CHECK: async.yield
|
|
async.yield
|
|
}
|
|
|
|
// CHECK: async.await %[[TOKEN]]
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
|
|
async.await %token : !async.token
|
|
|
|
// CHECK: async.await %[[TOKEN_0]]
|
|
// CHECK: async.runtime.drop_ref %[[TOKEN_0]] {count = 1 : i32}
|
|
async.await %token_0 : !async.token
|
|
|
|
// CHECK: async.await %[[RESULTS]]
|
|
// CHECK: async.runtime.drop_ref %[[RESULTS]] {count = 1 : i32}
|
|
%0 = async.await %results : !async.value<f32>
|
|
|
|
// CHECK: return
|
|
return %0 : f32
|
|
}
|