[Clang][LLVM][Coroutines] Prevent __coro_gro from outliving __promise

[bcardosolopes]

When dealing with short-circuiting coroutines (e.g. expected), the deferred calls that resolve the get_return_object are currently being emitted after we delete the coroutine frame.

This was caught by ASAN when using optimizations -O1 and above: optimizations after inlining would place the __coro_gro in the heap, and subsequent delete of the coroframe followed by the conversion -> BOOM.

This patch forbids the GRO to be placed in the coroutine frame, by adding a new metadata node that can be attached to alloca instructions.

Fix #49843.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms
@llvm/pr-subscribers-llvm-ir
@llvm/pr-subscribers-clang
@llvm/pr-subscribers-coroutines

@llvm/pr-subscribers-clang-codegen

Changes

When dealing with short-circuiting coroutines (e.g. expected), the deferred calls that resolve the get_return_object are currently being emitted after we delete the coroutine frame.

This was caught by ASAN when using optimizations -O1 and above: optimizations after inlining would place the __coro_gro in the heap, and subsequent delete of the coroframe followed by the conversion -> BOOM.

This patch fixes GRO placement into scopes and cleanups such that:

1. Emit the return block while within CallCoroDelete cleanup scope. This makes the conversion call to happen right after coro.end but before deleting the coroutine frame.

2. Change gro allocation to happen after __promise is allocated. The effect is to bound __coro_gro within the lifetime of __promise, which should make destruction of the conversion result to happen before we also delete the coroframe.

---

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

4 Files Affected:

• (modified) clang/lib/CodeGen/CGCoroutine.cpp (+21-17)
• (modified) clang/test/CodeGenCoroutines/coro-dest-slot.cpp (+16-2)
• (added) clang/test/CodeGenCoroutines/coro-expected.cpp (+223)
• (modified) clang/test/CodeGenCoroutines/coro-gro.cpp (+20-10)

diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index 448943084acedf3..741a8f9990c1ed7 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -657,8 +657,6 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
   CurCoro.Data->CoroBegin = CoroBegin;
 
   GetReturnObjectManager GroManager(*this, S);
-  GroManager.EmitGroAlloca();
-
   CurCoro.Data->CleanupJD = getJumpDestInCurrentScope(RetBB);
   {
     CGDebugInfo *DI = getDebugInfo();
@@ -696,6 +694,12 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
     // promise local variable was not emitted yet.
     CoroId->setArgOperand(1, PromiseAddrVoidPtr);
 
+    // Emit the alloca for ` __coro_gro`  *after* it's done for the promise.
+    // This guarantees that while looking at deferred results from calls to
+    // `get_return_object`, the lifetime of ` __coro_gro` is enclosed by the
+    // `__promise` lifetime, and cleanup order is properly respected.
+    GroManager.EmitGroAlloca();
+
     // Now we have the promise, initialize the GRO
     GroManager.EmitGroInit();
 
@@ -752,22 +756,22 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
       // We don't need FinalBB. Emit it to make sure the block is deleted.
       EmitBlock(FinalBB, /*IsFinished=*/true);
     }
-  }
 
-  EmitBlock(RetBB);
-  // Emit coro.end before getReturnStmt (and parameter destructors), since
-  // resume and destroy parts of the coroutine should not include them.
-  llvm::Function *CoroEnd = CGM.getIntrinsic(llvm::Intrinsic::coro_end);
-  Builder.CreateCall(CoroEnd,
-                     {NullPtr, Builder.getFalse(),
-                      llvm::ConstantTokenNone::get(CoroEnd->getContext())});
-
-  if (Stmt *Ret = S.getReturnStmt()) {
-    // Since we already emitted the return value above, so we shouldn't
-    // emit it again here.
-    if (GroManager.DirectEmit)
-      cast<ReturnStmt>(Ret)->setRetValue(nullptr);
-    EmitStmt(Ret);
+    EmitBlock(RetBB);
+    // Emit coro.end before getReturnStmt (and parameter destructors), since
+    // resume and destroy parts of the coroutine should not include them.
+    llvm::Function *CoroEnd = CGM.getIntrinsic(llvm::Intrinsic::coro_end);
+    Builder.CreateCall(CoroEnd,
+                       {NullPtr, Builder.getFalse(),
+                        llvm::ConstantTokenNone::get(CoroEnd->getContext())});
+
+    if (Stmt *Ret = S.getReturnStmt()) {
+      // Since we already emitted the return value above, so we shouldn't
+      // emit it again here.
+      if (GroManager.DirectEmit)
+        cast<ReturnStmt>(Ret)->setRetValue(nullptr);
+      EmitStmt(Ret);
+    }
   }
 
   // LLVM require the frontend to mark the coroutine.
diff --git a/clang/test/CodeGenCoroutines/coro-dest-slot.cpp b/clang/test/CodeGenCoroutines/coro-dest-slot.cpp
index d794c74cd52d61a..0d2416e7913da0d 100644
--- a/clang/test/CodeGenCoroutines/coro-dest-slot.cpp
+++ b/clang/test/CodeGenCoroutines/coro-dest-slot.cpp
@@ -17,6 +17,7 @@ struct coro {
 extern "C" coro f(int) { co_return; }
 // Verify that cleanup.dest.slot is eliminated in a coroutine.
 // CHECK-LABEL: f(
+// CHECK: %[[PROMISE:.+]] = alloca %"struct.coro::promise_type"
 // CHECK: %[[INIT_SUSPEND:.+]] = call i8 @llvm.coro.suspend(
 // CHECK-NEXT: switch i8 %[[INIT_SUSPEND]], label
 // CHECK-NEXT:   i8 0, label %[[INIT_READY:.+]]
@@ -32,9 +33,22 @@ extern "C" coro f(int) { co_return; }
 
 // CHECK: call void @_ZNSt13suspend_never12await_resumeEv(
 // CHECK: %[[CLEANUP_DEST1:.+]] = phi i32 [ 0, %[[FINAL_READY]] ], [ 2, %[[FINAL_CLEANUP]] ]
-// CHECK: %[[CLEANUP_DEST2:.+]] = phi i32 [ %[[CLEANUP_DEST0]], %{{.+}} ], [ %[[CLEANUP_DEST1]], %{{.+}} ], [ 0, %{{.+}} ]
+// CHECK: switch i32 %[[CLEANUP_DEST1]], label %[[CLEANUP_BB_2:.+]] [
+// CHECK:   i32 0, label %[[CLEANUP_CONT:.+]]
+// CHECK: ]
+
+// CHECK: [[CLEANUP_CONT]]:
+// CHECK:   br label %[[CORO_RET:.+]]
+
+// CHECK: [[CORO_RET]]:
+// CHECK:   call i1 @llvm.coro.end
+// CHECK:   br label %cleanup19
+
+// CHECK: [[CLEANUP_BB_2]]:
+// CHECK: %[[CLEANUP_DEST2:.+]] = phi i32 [ %[[CLEANUP_DEST0]], %{{.+}} ], [ 1, %[[CORO_RET]] ], [ %[[CLEANUP_DEST1]], %{{.+}} ]
+// CHECK: call void @llvm.lifetime.end.p0(i64 1, ptr %[[PROMISE]])
 // CHECK: call ptr @llvm.coro.free(
 // CHECK: switch i32 %[[CLEANUP_DEST2]], label %{{.+}} [
-// CHECK-NEXT: i32 0
 // CHECK-NEXT: i32 2
+// CHECK-NEXT: i32 1
 // CHECK-NEXT: ]
diff --git a/clang/test/CodeGenCoroutines/coro-expected.cpp b/clang/test/CodeGenCoroutines/coro-expected.cpp
new file mode 100644
index 000000000000000..5d0f7d1fc261f35
--- /dev/null
+++ b/clang/test/CodeGenCoroutines/coro-expected.cpp
@@ -0,0 +1,223 @@
+// RUN: %clang_cc1 -std=c++20 -triple=x86_64-unknown-linux-gnu -fcxx-exceptions -emit-llvm -o %t.ll %s -disable-llvm-passes
+// RUN: FileCheck --input-file=%t.ll %s
+
+// Verifies lifetime of __coro_gro variable w.r.t to the promise type,
+// more specifically make sure the coroutine frame isn't destroyed *before*
+// the conversion function that unwraps the promise proxy as part of
+// short-circuiting coroutines delayed behavior.
+
+#include "Inputs/coroutine.h"
+
+using namespace std;
+
+namespace folly {
+
+template <class Error>
+struct unexpected final {
+  Error error;
+
+  constexpr unexpected() = default;
+  constexpr unexpected(unexpected const&) = default;
+  unexpected& operator=(unexpected const&) = default;
+  constexpr /* implicit */ unexpected(Error err) : error(err) {}
+};
+
+template <class Value, class Error>
+class expected final {
+  enum class Which : unsigned char { empty, value, error };
+
+  Which which_;
+  Error error_;
+  Value value_;
+
+ public:
+  struct promise_type;
+
+  constexpr expected() noexcept : which_(Which::empty) {}
+  constexpr expected(expected const& that) = default;
+
+  expected(expected*& pointer) noexcept : which_(Which::empty) {
+    pointer = this;
+  }
+
+  constexpr /* implicit */ expected(Value val) noexcept
+      : which_(Which::value), value_(val) {}
+
+  /* implicit */ expected(Error) = delete;
+
+  constexpr /* implicit */ expected(unexpected<Error> err) noexcept
+      : which_(Which::error), error_(err.error) {}
+
+  expected& operator=(expected const& that) = default;
+
+  expected& operator=(Value val) noexcept {
+    which_ = Which::value;
+    value_ = val;
+    return *this;
+  }
+
+  expected& operator=(unexpected<Error> err) noexcept {
+    which_ = Which::error;
+    error_ = err.error;
+    return *this;
+  }
+
+  /* implicit */ expected& operator=(Error) = delete;
+
+  constexpr bool has_value() const noexcept {
+    return Which::value == this->which_;
+  }
+
+  constexpr bool has_error() const noexcept {
+    return Which::error == this->which_;
+  }
+
+  Value value() const {
+    require_value();
+    return value_;
+  }
+
+  Error error() const {
+    require_error();
+    return error_;
+  }
+
+ private:
+  void require_value() const {
+    if (!has_value()) {
+      // terminate
+    }
+  }
+
+  void require_error() const {
+    if (!has_error()) {
+      // terminate
+    }
+  }
+};
+
+template <typename Value, typename Error>
+struct expected<Value, Error>::promise_type {
+  struct return_object;
+
+  expected<Value, Error>* value_ = nullptr;
+
+  promise_type() = default;
+  promise_type(promise_type const&) = delete;
+  return_object get_return_object() noexcept {
+    return return_object{value_};
+  }
+  std::suspend_never initial_suspend() const noexcept {
+    return {};
+  }
+  std::suspend_never final_suspend() const noexcept {
+    return {};
+  }
+  void return_value(Value u) {
+    *value_ = u;
+  }
+  void unhandled_exception() {
+    throw;
+  }
+};
+
+template <typename Value, typename Error>
+struct expected<Value, Error>::promise_type::return_object {
+  expected<Value, Error> storage_;
+  expected<Value, Error>*& pointer_;
+
+  explicit return_object(expected<Value, Error>*& p) noexcept : pointer_{p} {
+    pointer_ = &storage_;
+  }
+  return_object(return_object const&) = delete;
+  ~return_object() {}
+
+  /* implicit */ operator expected<Value, Error>() {
+    return storage_; // deferred
+  }
+};
+
+template <typename Value, typename Error>
+struct expected_awaitable {
+  using promise_type = typename expected<Value, Error>::promise_type;
+
+  expected<Value, Error> o_;
+
+  explicit expected_awaitable(expected<Value, Error> o) : o_(o) {}
+
+  bool await_ready() const noexcept {
+    return o_.has_value();
+  }
+  Value await_resume() {
+    return o_.value();
+  }
+  void await_suspend(std::coroutine_handle<promise_type> h) {
+    *h.promise().value_ = unexpected<Error>(o_.error());
+    h.destroy();
+  }
+};
+
+template <typename Value, typename Error>
+expected_awaitable<Value, Error>
+/* implicit */ operator co_await(expected<Value, Error> o) {
+  return expected_awaitable<Value, Error>{o};
+}
+
+} // namespace folly
+
+struct err {};
+
+folly::expected<int, err> go(folly::expected<int, err> e) {
+  co_return co_await e;
+}
+
+int main() {
+  return go(0).value();
+}
+
+// CHECK: define {{.*}} @_Z2goN5folly8expectedIi3errEE
+
+// CHECK: %[[RetVal:.+]] = alloca %"class.folly::expected"
+// CHECK: %[[Promise:.+]] = alloca %"struct.folly::expected<int, err>::promise_type"
+// CHECK: %[[GroActive:.+]] = alloca i1
+// CHECK: %[[CoroGro:.+]] = alloca %"struct.folly::expected<int, err>::promise_type::return_object"
+
+// __promise lifetime should enclose __coro_gro's.
+
+// CHECK: call void @llvm.lifetime.start.p0(i64 8, ptr %[[Promise]])
+// CHECK: call void @_ZN5folly8expectedIi3errE12promise_typeC1Ev({{.*}} %[[Promise]])
+// CHECK: store i1 false, ptr %[[GroActive]]
+// CHECK: call void @llvm.lifetime.start.p0(i64 16, ptr %[[CoroGro]])
+// CHECK: call void @_ZN5folly8expectedIi3errE12promise_type17get_return_objectEv({{.*}} %[[CoroGro]],{{.*}} %[[Promise]])
+// CHECK: store i1 true, ptr %[[GroActive]]
+
+// Calls into `folly::expected<int, err>::promise_type::return_object::operator folly::expected<int, err>()` to unwrap
+// the delayed proxied promise.
+
+// CHECK: call i1 @llvm.coro.end
+// CHECK: %[[DelayedConv:.+]] = call i64 @_ZN5folly8expectedIi3errE12promise_type13return_objectcvS2_Ev(ptr noundef nonnull align 8 dereferenceable(16) %[[CoroGro:.+]])
+// CHECK: store i64 %[[DelayedConv]], ptr %[[RetVal]]
+// CHECK: br label %[[Cleanup0:.+]]
+
+// CHECK: [[Cleanup0]]:
+// CHECK:   %[[IsCleanupActive:.+]] = load i1, ptr %[[GroActive]]
+// CHECK:   br i1 %[[IsCleanupActive]], label %[[CleanupAction:.+]], label %[[CleanupDone:.+]]
+
+// Call into `folly::expected<int, err>::promise_type::return_object::~return_object()` while __promise is alive.
+
+// CHECK: [[CleanupAction]]:
+// CHECK:   call void @_ZN5folly8expectedIi3errE12promise_type13return_objectD1Ev(ptr noundef nonnull align 8 dereferenceable(16) %[[CoroGro]])
+// CHECK:   br label %[[CleanupDone]]
+
+// __promise lifetime should end after __coro_gro's.
+
+// CHECK: [[CleanupDone]]:
+// CHECK:   call void @llvm.lifetime.end.p0(i64 16, ptr %[[CoroGro]])
+// CHECK:   call void @llvm.lifetime.end.p0(i64 8, ptr %[[Promise]])
+// CHECK:   call ptr @llvm.coro.free
+// CHECK:   br i1 {{.*}}, label %[[CoroFree:.+]], {{.*}}
+
+// Delete coroutine frame.
+
+// CHECK: [[CoroFree]]:
+// CHECK: call void @_ZdlPv
\ No newline at end of file
diff --git a/clang/test/CodeGenCoroutines/coro-gro.cpp b/clang/test/CodeGenCoroutines/coro-gro.cpp
index b48b769950ae871..f7ed865931dc2ed 100644
--- a/clang/test/CodeGenCoroutines/coro-gro.cpp
+++ b/clang/test/CodeGenCoroutines/coro-gro.cpp
@@ -29,14 +29,21 @@ void doSomething() noexcept;
 // CHECK: define{{.*}} i32 @_Z1fv(
 int f() {
   // CHECK: %[[RetVal:.+]] = alloca i32
+  // CHECK: %[[Promise:.+]] = alloca %"struct.std::coroutine_traits<int>::promise_type"
   // CHECK: %[[GroActive:.+]] = alloca i1
+  // CHECK: %[[CoroGro:.+]] = alloca %struct.GroType
 
   // CHECK: %[[Size:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: call noalias noundef nonnull ptr @_Znwm(i64 noundef %[[Size]])
-  // CHECK: store i1 false, ptr %[[GroActive]]
+
+  // GRO lifetime should be bound within promise's lifetime.
+  //
+  // CHECK: call void @llvm.lifetime.start.p0(i64 1, ptr %[[Promise]])
   // CHECK: call void @_ZNSt16coroutine_traitsIiJEE12promise_typeC1Ev(
+  // CHECK: store i1 false, ptr %[[GroActive]]
+  // CHECK: call void @llvm.lifetime.start.p0(i64 1, ptr %[[CoroGro]])
   // CHECK: call void @_ZNSt16coroutine_traitsIiJEE12promise_type17get_return_objectEv(
-  // CHECK: store i1 true, ptr %[[GroActive]]
+  // CHECK: store i1 true, ptr %[[GroActive]], align 1
 
   Cleanup cleanup;
   doSomething();
@@ -46,12 +53,6 @@ int f() {
   // CHECK: call void @_ZNSt16coroutine_traitsIiJEE12promise_type11return_voidEv(
   // CHECK: call void @_ZN7CleanupD1Ev(
 
-  // Destroy promise and free the memory.
-
-  // CHECK: call void @_ZNSt16coroutine_traitsIiJEE12promise_typeD1Ev(
-  // CHECK: %[[Mem:.+]] = call ptr @llvm.coro.free(
-  // CHECK: call void @_ZdlPv(ptr noundef %[[Mem]])
-
   // Initialize retval from Gro and destroy Gro
   // Note this also tests delaying initialization when Gro and function return
   // types mismatch (see cwg2563).
@@ -63,9 +64,18 @@ int f() {
 
   // CHECK: [[CleanupGro]]:
   // CHECK:   call void @_ZN7GroTypeD1Ev(
-  // CHECK:   br label %[[Done]]
+  // CHECK:   br label %[[CleanupDone:.+]]
+
+  // Destroy promise and free the memory.
+
+  // GRO lifetime should be bound within promise's lifetime.
+  // CHECK: [[CleanupDone]]:
+  // CHECK: call void @llvm.lifetime.end.p0(i64 1, ptr %[[CoroGro]])
+  // CHECK: call void @_ZNSt16coroutine_traitsIiJEE12promise_typeD1Ev(
+  // CHECK: call void @llvm.lifetime.end.p0(i64 1, ptr %[[Promise]])
+  // CHECK: %[[Mem:.+]] = call ptr @llvm.coro.free(
+  // CHECK: call void @_ZdlPv(ptr noundef %[[Mem]])
 
-  // CHECK: [[Done]]:
   // CHECK:   %[[LoadRet:.+]] = load i32, ptr %[[RetVal]]
   // CHECK:   ret i32 %[[LoadRet]]
 }

[ChuanqiXu9]

I remember that there is a defect that we may place the GRO on the coroutine frame.

And my instinct reaction is that would this patch be covered by forcing GRO to not live on the coroutine frame?

[bcardosolopes]

Thanks for the fast reply @ChuanqiXu9

I remember that there is a defect that we may place the GRO on the coroutine frame.

Can you point me to such defect? I had no luck searching for it.

And my instinct reaction is that would this patch be covered by forcing GRO to not live on the coroutine frame?

How do you suggest we do so? Even if we teach the optimizers to not touch the GRO (e.g. -O0) there's still a correctness issue. The generated code is still wrong, given that currently, the lifetime of __coro_gro outlives the __promise, but accessing the proxy for delayed conversion might still depend on the promise.

Also, what do you mean by "patch be covered"? I'm confused.

[bcardosolopes]

Updated the patch to account for failing libcxx tests, and to not change current codegen when GRO allocas are not involved.

[ChuanqiXu9]

Can you point me to such defect? I had no luck searching for it.

It should be #49843. I didn't expect that it was not marked as coroutines. Besides fixing that bug, it is also a minor optimization to not include GRO in coroutine frames.

Also, what do you mean by "patch be covered"? I'm confused.

When I say "patch be covered", I mean the problem may be solved automatically if we can force the GRO not living on the coroutine frames.

How do you suggest we do so? Even if we teach the optimizers to not touch the GRO (e.g. -O0) there's still a correctness issue. The generated code is still wrong, given that currently, the lifetime of __coro_gro outlives the __promise, but accessing the proxy for delayed conversion might still depend on the promise.

By forcing the GRO not living on the coroutine frames, it shouldn't be a problem if the lifetime of __coro_gro outlives the __promise. The only limit is that the initialization of __coro_gro should be in the lifetime of __promise.

[bcardosolopes]

Thanks for the clarifications

By forcing the GRO not living on the coroutine frames, it shouldn't be a problem if the lifetime of __coro_gro outlives the __promise. The only limit is that the initialization of __coro_gro should be in the lifetime of __promise.

Applied your suggestion based on #49843, and edited description + title.

[ChuanqiXu9]

Thanks. It looks good if we can add a llvm test in llvm/test/Transforms/Coroutines for the new metadata.

[bcardosolopes]

Applied suggested changes and updated PR.

[apolloww]

Thanks for the fix! I think originally the deferred conversion is placed after coro.end so that it won't be counted as "use across suspend point", but it doesn't stop optimizer from breaking that assumption.

[ChuanqiXu9]

LGTM. Thanks.