[CIR] Emit ready and suspend branches for cir.await

[Andres-Salamanca]

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

[llvmbot]

@llvm/pr-subscribers-clangir

@llvm/pr-subscribers-clang

Author: None (Andres-Salamanca)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/168814.diff

5 Files Affected:

• (modified) clang/include/clang/CIR/MissingFeatures.h (-1)
• (modified) clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp (+1-3)
• (modified) clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp (+23-2)
• (modified) clang/lib/CIR/CodeGen/CIRGenFunction.h (+1)
• (modified) clang/test/CIR/CodeGen/coro-task.cpp (+40-3)

diff --git a/clang/include/clang/CIR/MissingFeatures.h b/clang/include/clang/CIR/MissingFeatures.h
index 6b5c34d28ce2a..a6bdf11a01a4b 100644
--- a/clang/include/clang/CIR/MissingFeatures.h
+++ b/clang/include/clang/CIR/MissingFeatures.h
@@ -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; }
 
diff --git a/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp b/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
index 77f19343653db..862f25cc4d3c8 100644
--- a/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
@@ -470,9 +470,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: {
diff --git a/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp b/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp
index bb55991d9366a..d105d64ea5d31 100644
--- a/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenCoroutine.cpp
@@ -97,6 +97,14 @@ struct ParamReferenceReplacerRAII {
   }
 };
 } // namespace
+
+RValue CIRGenFunction::emitCoroutineFrame() {
+  if (curCoro.data && curCoro.data->coroBegin) {
+    return RValue::get(curCoro.data->coroBegin);
+  }
+  cgm.errorNYI("NYI");
+}
+
 static void createCoroData(CIRGenFunction &cgf,
                            CIRGenFunction::CGCoroInfo &curCoro,
                            cir::CallOp coroId) {
@@ -302,11 +310,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=*/
diff --git a/clang/lib/CIR/CodeGen/CIRGenFunction.h b/clang/lib/CIR/CodeGen/CIRGenFunction.h
index b426f3389ff1b..d09dfbd84975a 100644
--- a/clang/lib/CIR/CodeGen/CIRGenFunction.h
+++ b/clang/lib/CIR/CodeGen/CIRGenFunction.h
@@ -1415,6 +1415,7 @@ class CIRGenFunction : public CIRGenTypeCache {
   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);
 
diff --git a/clang/test/CIR/CodeGen/coro-task.cpp b/clang/test/CIR/CodeGen/coro-task.cpp
index 01e0786fbda71..4843f2433fa64 100644
--- a/clang/test/CIR/CodeGen/coro-task.cpp
+++ b/clang/test/CIR/CodeGen/coro-task.cpp
@@ -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:      },)

[github-actions]

🐧 Linux x64 Test Results

• 112070 tests passed
• 4073 tests skipped

[bcardosolopes]

Great, thanks! LGTM