[CIR] Upstream CIR await op (#168133)

This PR upstreams `cir.await` and adds initial codegen for emitting a
skeleton of the ready, suspend, and resume branches. Codegen for these
branches is left for a future PR. It also adds a test for the invalid
case where a `cir.func` is marked as a coroutine but does not contain a
`cir.await` op in its body.

4 files changed, 88 insertions, 0 deletions

diff --git a/clang/test/CIR/CodeGen/coro-task.cpp b/clang/test/CIR/CodeGen/coro-task.cpp
index 5738c815909e..01e0786fbda7 100644
--- a/clang/test/CIR/CodeGen/coro-task.cpp
+++ b/clang/test/CIR/CodeGen/coro-task.cpp
@@ -111,6 +111,8 @@ 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: ![[SuspendAlways:.*]] = !cir.record<struct "std::suspend_always" padded {!u8i}>
+
 // CIR: module {{.*}} {
 // CIR-NEXT: cir.global external @_ZN5folly4coro9co_invokeE = #cir.zero : !rec_folly3A3Acoro3A3Aco_invoke_fn
 
@@ -153,6 +155,33 @@ VoidTask silly_task() {
 
 // CIR: %[[RetObj:.*]] = cir.call @_ZN5folly4coro4TaskIvE12promise_type17get_return_objectEv(%[[VoidPromisseAddr]]) nothrow : {{.*}} -> ![[VoidTask]]
 // CIR: cir.store{{.*}} %[[RetObj]], %[[VoidTaskAddr]] : ![[VoidTask]]
+// Start a new scope for the actual codegen for co_await, create temporary allocas for
+// holding coroutine handle and the suspend_always struct.
+
+// CIR: cir.scope {
+// CIR:   %[[SuspendAlwaysAddr:.*]] = cir.alloca ![[SuspendAlways]], {{.*}} ["ref.tmp0"] {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
+// to later passes, same is done elsewhere.
+
+// CIR:   %[[Tmp0:.*]] = cir.call @_ZN5folly4coro4TaskIvE12promise_type15initial_suspendEv(%[[VoidPromisseAddr]])
+// CIR:   cir.store{{.*}} %[[Tmp0:.*]], %[[SuspendAlwaysAddr]]
+
+//
+// Here we start mapping co_await to cir.await.
+//
+
+// First regions `ready` has a special cir.yield code to veto suspension.
+
+// CIR:   cir.await(init, ready : {
+// CIR:     cir.condition({{.*}})
+// CIR:   }, suspend : {
+// CIR:     cir.yield
+// CIR:   }, resume : {
+// CIR:     cir.yield
+// CIR:   },)
+// CIR: }
 
 folly::coro::Task<int> byRef(const std::string& s) {
   co_return s.size();
@@ -172,3 +201,13 @@ folly::coro::Task<int> byRef(const std::string& s) {
 // CIR:    cir.store {{.*}} %[[LOAD]], %[[AllocaFnUse]] : !cir.ptr<![[StdString]]>, !cir.ptr<!cir.ptr<![[StdString]]>>
 // CIR:    %[[RetObj:.*]] = cir.call @_ZN5folly4coro4TaskIiE12promise_type17get_return_objectEv(%4) nothrow : {{.*}} -> ![[IntTask]]
 // CIR:    cir.store {{.*}} %[[RetObj]], %[[IntTaskAddr]] : ![[IntTask]]
+// CIR:    cir.scope {
+// CIR:      %[[SuspendAlwaysAddr:.*]] = cir.alloca ![[SuspendAlways]], {{.*}} ["ref.tmp0"] {alignment = 1 : i64}
+// CIR:      %[[Tmp0:.*]] = cir.call @_ZN5folly4coro4TaskIiE12promise_type15initial_suspendEv(%[[IntPromisseAddr]])
+// CIR:      cir.await(init, ready : {
+// CIR:        cir.condition({{.*}})
+// CIR:      }, suspend : {
+// CIR:        cir.yield
+// CIR:      }, resume : {
+// CIR:        cir.yield
+// CIR:      },)
diff --git a/clang/test/CIR/IR/await.cir b/clang/test/CIR/IR/await.cir
new file mode 100644
index 000000000000..c1fb0d6d7c57
--- /dev/null
+++ b/clang/test/CIR/IR/await.cir
@@ -0,0 +1,21 @@
+// RUN: cir-opt %s --verify-roundtrip | FileCheck %s
+
+cir.func coroutine @checkPrintParse(%arg0 : !cir.bool) {
+  cir.await(user, ready : {
+    cir.condition(%arg0)
+  }, suspend : {
+    cir.yield
+  }, resume : {
+    cir.yield
+  },)
+  cir.return
+}
+
+// CHECK:  cir.func coroutine @checkPrintParse
+// CHECK:  cir.await(user, ready : {
+// CHECK:    cir.condition(%arg0)
+// CHECK:  }, suspend : {
+// CHECK:    cir.yield
+// CHECK:  }, resume : {
+// CHECK:    cir.yield
+// CHECK:  },)
diff --git a/clang/test/CIR/IR/func.cir b/clang/test/CIR/IR/func.cir
index 6e91898a3b45..4e79149315a9 100644
--- a/clang/test/CIR/IR/func.cir
+++ b/clang/test/CIR/IR/func.cir
@@ -101,6 +101,15 @@ cir.func @ullfunc() -> !u64i {
 // CHECK: }
 
 cir.func coroutine @coro() {
+  cir.await(init, ready : {
+    %0 = cir.alloca !cir.bool, !cir.ptr<!cir.bool>, [""] {alignment = 1 : i64}
+    %1 = cir.load align(1) %0 : !cir.ptr<!cir.bool>, !cir.bool
+    cir.condition(%1)
+  }, suspend : {
+    cir.yield
+  }, resume : {
+    cir.yield
+  },)
   cir.return
 }
 // CHECK: cir.func{{.*}} coroutine @coro()
diff --git a/clang/test/CIR/IR/invalid-await.cir b/clang/test/CIR/IR/invalid-await.cir
new file mode 100644
index 000000000000..5f422ca7e954
--- /dev/null
+++ b/clang/test/CIR/IR/invalid-await.cir
@@ -0,0 +1,19 @@
+// RUN: cir-opt %s -verify-diagnostics -split-input-file
+cir.func coroutine @bad_task() { // expected-error {{coroutine body must use at least one cir.await op}}
+  cir.return
+}
+
+// -----
+
+cir.func coroutine @missing_condition() {
+  cir.scope {
+    cir.await(user, ready : { // expected-error {{ready region must end with cir.condition}}
+      cir.yield
+    }, suspend : {
+      cir.yield
+    }, resume : {
+      cir.yield
+    },)
+  }
+  cir.return
+}