[clang] [Clang][Coroutines] Improve GRO handling to better fit scopes and cleanups (PR #66706)

via cfe-commits cfe-commits at lists.llvm.org
Mon Sep 18 14:53:59 PDT 2023


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-clang-codegen

<details>
<summary>Changes</summary>

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
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]]
 }

``````````

</details>


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


More information about the cfe-commits mailing list