[CIR] Emit ready and suspend branches for cir.await (#168814)

This PR adds codegen for `cir.await` ready and suspend. One notable
difference from the classic codegen is that, in the suspend branch, it
emits an `AwaitSuspendWrapper`(`.__await_suspend_wrapper__init`)
function that is always inlined. This function wraps the suspend logic
inside an internal wrapper that gets inlined. Example here:
https://godbolt.org/z/rWYGcaaG4
This commit is contained in:
Andres-Salamanca
2025-11-25 17:16:21 -05:00
committed by GitHub
parent 5f777b2c8f
commit 97023fba55
5 changed files with 66 additions and 9 deletions

View File

@@ -151,7 +151,6 @@ struct MissingFeatures {
// Coroutines
static bool coroEndBuiltinCall() { return false; }
static bool coroutineFrame() { return false; }
static bool emitBodyAndFallthrough() { return false; }
static bool coroOutsideFrameMD() { return false; }

View File

@@ -502,9 +502,7 @@ RValue CIRGenFunction::emitBuiltinExpr(const GlobalDecl &gd, unsigned builtinID,
return getUndefRValue(e->getType());
case Builtin::BI__builtin_coro_frame: {
cgm.errorNYI(e->getSourceRange(), "BI__builtin_coro_frame NYI");
assert(!cir::MissingFeatures::coroutineFrame());
return getUndefRValue(e->getType());
return emitCoroutineFrame();
}
case Builtin::BI__builtin_coro_free:
case Builtin::BI__builtin_coro_size: {

View File

@@ -97,6 +97,15 @@ struct ParamReferenceReplacerRAII {
}
};
} // namespace
RValue CIRGenFunction::emitCoroutineFrame() {
if (curCoro.data && curCoro.data->coroBegin) {
return RValue::get(curCoro.data->coroBegin);
}
cgm.errorNYI("NYI");
return RValue();
}
static void createCoroData(CIRGenFunction &cgf,
CIRGenFunction::CGCoroInfo &curCoro,
cir::CallOp coroId) {
@@ -302,11 +311,24 @@ emitSuspendExpression(CIRGenFunction &cgf, CGCoroData &coro,
builder, cgf.getLoc(s.getSourceRange()), kind,
/*readyBuilder=*/
[&](mlir::OpBuilder &b, mlir::Location loc) {
builder.createCondition(
cgf.createDummyValue(loc, cgf.getContext().BoolTy));
Expr *condExpr = s.getReadyExpr()->IgnoreParens();
builder.createCondition(cgf.evaluateExprAsBool(condExpr));
},
/*suspendBuilder=*/
[&](mlir::OpBuilder &b, mlir::Location loc) {
// Note that differently from LLVM codegen we do not emit coro.save
// and coro.suspend here, that should be done as part of lowering this
// to LLVM dialect (or some other MLIR dialect)
// A invalid suspendRet indicates "void returning await_suspend"
mlir::Value suspendRet = cgf.emitScalarExpr(s.getSuspendExpr());
// Veto suspension if requested by bool returning await_suspend.
if (suspendRet) {
cgf.cgm.errorNYI("Veto await_suspend");
}
// Signals the parent that execution flows to next region.
cir::YieldOp::create(builder, loc);
},
/*resumeBuilder=*/

View File

@@ -1418,6 +1418,7 @@ public:
cir::CallOp emitCoroAllocBuiltinCall(mlir::Location loc);
cir::CallOp emitCoroBeginBuiltinCall(mlir::Location loc,
mlir::Value coroframeAddr);
RValue emitCoroutineFrame();
void emitDestroy(Address addr, QualType type, Destroyer *destroyer);

View File

@@ -111,6 +111,9 @@ co_invoke_fn co_invoke;
// CIR-DAG: ![[VoidPromisse:.*]] = !cir.record<struct "folly::coro::Task<void>::promise_type" padded {!u8i}>
// CIR-DAG: ![[IntPromisse:.*]] = !cir.record<struct "folly::coro::Task<int>::promise_type" padded {!u8i}>
// CIR-DAG: ![[StdString:.*]] = !cir.record<struct "std::string" padded {!u8i}>
// CIR-DAG: ![[CoroHandleVoid:.*]] = !cir.record<struct "std::coroutine_handle<void>" padded {!u8i}>
// CIR-DAG: ![[CoroHandlePromiseVoid:rec_.*]] = !cir.record<struct "std::coroutine_handle<folly::coro::Task<void>::promise_type>" padded {!u8i}>
// CIR-DAG: ![[CoroHandlePromiseInt:rec_.*]] = !cir.record<struct "std::coroutine_handle<folly::coro::Task<int>::promise_type>" padded {!u8i}>
// CIR-DAG: ![[SuspendAlways:.*]] = !cir.record<struct "std::suspend_always" padded {!u8i}>
// CIR: module {{.*}} {
@@ -160,6 +163,8 @@ VoidTask silly_task() {
// CIR: cir.scope {
// CIR: %[[SuspendAlwaysAddr:.*]] = cir.alloca ![[SuspendAlways]], {{.*}} ["ref.tmp0"] {alignment = 1 : i64}
// CIR: %[[CoroHandleVoidAddr:.*]] = cir.alloca ![[CoroHandleVoid]], {{.*}} ["agg.tmp0"] {alignment = 1 : i64}
// CIR: %[[CoroHandlePromiseAddr:.*]] = cir.alloca ![[CoroHandlePromiseVoid]], {{.*}} ["agg.tmp1"] {alignment = 1 : i64}
// Effectively execute `coawait promise_type::initial_suspend()` by calling initial_suspend() and getting
// the suspend_always struct to use for cir.await. Note that we return by-value since we defer ABI lowering
@@ -175,8 +180,28 @@ VoidTask silly_task() {
// First regions `ready` has a special cir.yield code to veto suspension.
// CIR: cir.await(init, ready : {
// CIR: cir.condition({{.*}})
// CIR: %[[ReadyVeto:.*]] = cir.scope {
// CIR: %[[TmpCallRes:.*]] = cir.call @_ZNSt14suspend_always11await_readyEv(%[[SuspendAlwaysAddr]])
// CIR: cir.yield %[[TmpCallRes:.*]] : !cir.bool
// CIR: }
// CIR: cir.condition(%[[ReadyVeto]])
// Second region `suspend` contains the actual suspend logic.
//
// - Start by getting the coroutine handle using from_address().
// - Implicit convert coroutine handle from task specific promisse
// specialization to a void one.
// - Call suspend_always::await_suspend() passing the handle.
//
// FIXME: add veto support for non-void await_suspends.
// CIR: }, suspend : {
// CIR: %[[FromAddrRes:.*]] = cir.call @_ZNSt16coroutine_handleIN5folly4coro4TaskIvE12promise_typeEE12from_addressEPv(%[[CoroFrameAddr]])
// CIR: cir.store{{.*}} %[[FromAddrRes]], %[[CoroHandlePromiseAddr]] : ![[CoroHandlePromiseVoid]]
// CIR: %[[CoroHandlePromiseReload:.*]] = cir.load{{.*}} %[[CoroHandlePromiseAddr]]
// CIR: cir.call @_ZNSt16coroutine_handleIvEC1IN5folly4coro4TaskIvE12promise_typeEEES_IT_E(%[[CoroHandleVoidAddr]], %[[CoroHandlePromiseReload]])
// CIR: %[[CoroHandleVoidReload:.*]] = cir.load{{.*}} %[[CoroHandleVoidAddr]] : !cir.ptr<![[CoroHandleVoid]]>, ![[CoroHandleVoid]]
// CIR: cir.call @_ZNSt14suspend_always13await_suspendESt16coroutine_handleIvE(%[[SuspendAlwaysAddr]], %[[CoroHandleVoidReload]])
// CIR: cir.yield
// CIR: }, resume : {
// CIR: cir.yield
@@ -203,11 +228,23 @@ folly::coro::Task<int> byRef(const std::string& s) {
// CIR: cir.store {{.*}} %[[RetObj]], %[[IntTaskAddr]] : ![[IntTask]]
// CIR: cir.scope {
// CIR: %[[SuspendAlwaysAddr:.*]] = cir.alloca ![[SuspendAlways]], {{.*}} ["ref.tmp0"] {alignment = 1 : i64}
// CIR: %[[CoroHandleVoidAddr:.*]] = cir.alloca ![[CoroHandleVoid]], {{.*}} ["agg.tmp0"] {alignment = 1 : i64}
// CIR: %[[CoroHandlePromiseAddr:.*]] = cir.alloca ![[CoroHandlePromiseInt]], {{.*}} ["agg.tmp1"] {alignment = 1 : i64}
// CIR: %[[Tmp0:.*]] = cir.call @_ZN5folly4coro4TaskIiE12promise_type15initial_suspendEv(%[[IntPromisseAddr]])
// CIR: cir.await(init, ready : {
// CIR: cir.condition({{.*}})
// CIR: %[[ReadyVeto:.*]] = cir.scope {
// CIR: %[[TmpCallRes:.*]] = cir.call @_ZNSt14suspend_always11await_readyEv(%[[SuspendAlwaysAddr]])
// CIR: cir.yield %[[TmpCallRes:.*]] : !cir.bool
// CIR: }
// CIR: cir.condition(%[[ReadyVeto]])
// CIR: }, suspend : {
// CIR: cir.yield
// CIR: %[[FromAddrRes:.*]] = cir.call @_ZNSt16coroutine_handleIN5folly4coro4TaskIiE12promise_typeEE12from_addressEPv(%[[CoroFrameAddr:.*]])
// CIR: cir.store{{.*}} %[[FromAddrRes]], %[[CoroHandlePromiseAddr]] : ![[CoroHandlePromiseInt]]
// CIR: %[[CoroHandlePromiseReload:.*]] = cir.load{{.*}} %[[CoroHandlePromiseAddr]]
// CIR: cir.call @_ZNSt16coroutine_handleIvEC1IN5folly4coro4TaskIiE12promise_typeEEES_IT_E(%[[CoroHandleVoidAddr]], %[[CoroHandlePromiseReload]])
// CIR: %[[CoroHandleVoidReload:.*]] = cir.load{{.*}} %[[CoroHandleVoidAddr]] : !cir.ptr<![[CoroHandleVoid]]>, ![[CoroHandleVoid]]
// CIR: cir.call @_ZNSt14suspend_always13await_suspendESt16coroutine_handleIvE(%[[SuspendAlwaysAddr]], %[[CoroHandleVoidReload]])
// CIR: cir.yield
// CIR: }, resume : {
// CIR: cir.yield
// CIR: },)