[llvm] 17631ac - [Coroutines] Store the index for final suspend point if there is unwind coro end

Chuanqi Xu via llvm-commits llvm-commits at lists.llvm.org
Thu Aug 25 23:07:06 PDT 2022


Author: Chuanqi Xu
Date: 2022-08-26T14:05:46+08:00
New Revision: 17631ac676bc44b7fd5c79abc66844efb4e1b533

URL: https://github.com/llvm/llvm-project/commit/17631ac676bc44b7fd5c79abc66844efb4e1b533
DIFF: https://github.com/llvm/llvm-project/commit/17631ac676bc44b7fd5c79abc66844efb4e1b533.diff

LOG: [Coroutines] Store the index for final suspend point if there is unwind coro end

Closing https://github.com/llvm/llvm-project/issues/57339

The root cause for this issue is an pre-mature optimization to eliminate
the index for the final suspend point since we feel like we can judge
if a coroutine is suspended at the final suspend by if resume_fn_addr is
null. However this is not true if the coroutine exists via an exception
in promise.unhandled_exception(). According to
[dcl.fct.def.coroutine]p14:

> If the evaluation of the expression promise.unhandled_exception()
> exits via an exception, the coroutine is considered suspended at the
> final suspend point.

But from the perspective of the implementation, we can't set the coro
index to the final suspend point directly since it breaks the states.

To fix the issue, we block the optimization if we find there is any
unwind coro end, which indicates that it is possible that the coroutine
exists via an exception from promise.unhandled_exception().

Test Plan: folly

Added: 
    clang/test/CodeGenCoroutines/coro-destructor-of-final_suspend.cpp
    llvm/test/Transforms/Coroutines/coro-split-final-suspend.ll

Modified: 
    clang/test/CodeGenCoroutines/pr56919.cpp
    llvm/lib/Transforms/Coroutines/CoroInternal.h
    llvm/lib/Transforms/Coroutines/CoroSplit.cpp
    llvm/lib/Transforms/Coroutines/Coroutines.cpp

Removed: 
    


################################################################################
diff  --git a/clang/test/CodeGenCoroutines/coro-destructor-of-final_suspend.cpp b/clang/test/CodeGenCoroutines/coro-destructor-of-final_suspend.cpp
new file mode 100644
index 0000000000000..079a0d723eb11
--- /dev/null
+++ b/clang/test/CodeGenCoroutines/coro-destructor-of-final_suspend.cpp
@@ -0,0 +1,90 @@
+// This addresses https://github.com/llvm/llvm-project/issues/57339
+// RUN: %clang_cc1 -triple %itanium_abi_triple -std=c++20 -fcxx-exceptions \
+// RUN:     -fexceptions -S -emit-llvm -o - %s -O1 | FileCheck %s
+
+#include "Inputs/coroutine.h"
+
+struct gen {
+  struct promise_type {
+    gen get_return_object() noexcept {
+      return gen{std::coroutine_handle<promise_type>::from_promise(*this)};
+    }
+    std::suspend_always initial_suspend() noexcept { return {}; }
+
+    struct final_awaiter {
+      ~final_awaiter() noexcept;
+      bool await_ready() noexcept {
+        return false;
+      }
+      void await_suspend(std::coroutine_handle<>) noexcept {}
+      void await_resume() noexcept {}
+    };
+
+    final_awaiter final_suspend() noexcept {
+      return {};
+    }
+
+    void unhandled_exception() {
+      throw;
+    }
+    void return_void() {}
+  };
+
+  gen(std::coroutine_handle<promise_type> coro) noexcept
+  : coro(coro)
+  {
+  }
+
+  ~gen() noexcept {
+    if (coro) {
+      coro.destroy();
+    }
+  }
+
+  gen(gen&& g) noexcept
+  : coro(g.coro)
+  {
+    g.coro = {};
+  }
+
+  std::coroutine_handle<promise_type> coro;
+};
+
+struct X {};
+
+gen maybe_throwing(bool x) {
+  if (x) {
+    throw X{};
+  }
+  co_return;
+}
+
+// CHECK: define{{.*}}@_Z14maybe_throwingb.destroy
+// CHECK: %[[INDEX:.+]] = load i1, ptr %index.addr, align 1
+// CHECK: br i1 %[[INDEX]], label %[[AFTERSUSPEND:.+]], label %[[CORO_FREE:.+]]
+// CHECK: [[AFTERSUSPEND]]:
+// CHECK: call{{.*}}_ZN3gen12promise_type13final_awaiterD1Ev(
+// CHECK: [[CORO_FREE]]:
+// CHECK: call{{.*}}_ZdlPv
+
+void noexcept_call() noexcept;
+
+gen no_throwing() {
+  noexcept_call();
+  co_return;
+}
+
+// CHECK: define{{.*}}@_Z11no_throwingv.resume({{.*}}%[[ARG:.+]])
+// CHECK: resume:
+// CHECK:   call{{.*}}@_Z13noexcept_callv()
+// CHECK:   store ptr null, ptr %[[ARG]]
+// CHECK:   ret void
+
+// CHECK: define{{.*}}@_Z11no_throwingv.destroy({{.*}}%[[ARG:.+]])
+// CHECK: %[[RESUME_FN_ADDR:.+]] = load ptr, ptr %[[ARG]]
+// CHECK: %[[IF_NULL:.+]] = icmp eq ptr %[[RESUME_FN_ADDR]], null
+// CHECK: br i1 %[[IF_NULL]], label %[[AFTERSUSPEND:.+]], label %[[CORO_FREE:.+]]
+// CHECK: [[AFTERSUSPEND]]:
+// CHECK: call{{.*}}_ZN3gen12promise_type13final_awaiterD1Ev(
+// CHECK: [[CORO_FREE]]:
+// CHECK: call{{.*}}_ZdlPv

diff  --git a/clang/test/CodeGenCoroutines/pr56919.cpp b/clang/test/CodeGenCoroutines/pr56919.cpp
index dd6df99a56232..c7de08ef72d7f 100644
--- a/clang/test/CodeGenCoroutines/pr56919.cpp
+++ b/clang/test/CodeGenCoroutines/pr56919.cpp
@@ -28,7 +28,7 @@ class Task final {
    public: 
     Task<void> get_return_object() { return Task<void>(std::coroutine_handle<promise_type>::from_promise(*this)); }
 
-    void unhandled_exception();
+    void unhandled_exception() {}
 
     std::suspend_always initial_suspend() { return {}; }
 

diff  --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h
index af35b45c2eaf5..032361c220452 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInternal.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInternal.h
@@ -112,6 +112,7 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
     unsigned IndexAlign;
     unsigned IndexOffset;
     bool HasFinalSuspend;
+    bool HasUnwindCoroEnd;
   };
 
   struct RetconLoweringStorage {

diff  --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 0107d09059c17..8c4d197a28998 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -396,11 +396,22 @@ static void createResumeEntryBlock(Function &F, coro::Shape &Shape) {
       // The coroutine should be marked done if it reaches the final suspend
       // point.
       markCoroutineAsDone(Builder, Shape, FramePtr);
-    } else {
+    }
+
+    // If the coroutine don't have unwind coro end, we could omit the store to
+    // the final suspend point since we could infer the coroutine is suspended
+    // at the final suspend point by the nullness of ResumeFnAddr.
+    // However, we can't skip it if the coroutine have unwind coro end. Since
+    // the coroutine reaches unwind coro end is considered suspended at the
+    // final suspend point (the ResumeFnAddr is null) but in fact the coroutine
+    // didn't complete yet. We need the IndexVal for the final suspend point
+    // to make the states clear.
+    if (!S->isFinal() || Shape.SwitchLowering.HasUnwindCoroEnd) {
       auto *GepIndex = Builder.CreateStructGEP(
           FrameTy, FramePtr, Shape.getSwitchIndexField(), "index.addr");
       Builder.CreateStore(IndexVal, GepIndex);
     }
+
     Save->replaceAllUsesWith(ConstantTokenNone::get(C));
     Save->eraseFromParent();
 
@@ -449,19 +460,22 @@ static void createResumeEntryBlock(Function &F, coro::Shape &Shape) {
   Shape.SwitchLowering.ResumeEntryBlock = NewEntry;
 }
 
-
-// Rewrite final suspend point handling. We do not use suspend index to
-// represent the final suspend point. Instead we zero-out ResumeFnAddr in the
-// coroutine frame, since it is undefined behavior to resume a coroutine
-// suspended at the final suspend point. Thus, in the resume function, we can
-// simply remove the last case (when coro::Shape is built, the final suspend
-// point (if present) is always the last element of CoroSuspends array).
-// In the destroy function, we add a code sequence to check if ResumeFnAddress
-// is Null, and if so, jump to the appropriate label to handle cleanup from the
-// final suspend point.
+// In the resume function, we remove the last case  (when coro::Shape is built,
+// the final suspend point (if present) is always the last element of
+// CoroSuspends array) since it is an undefined behavior to resume a coroutine
+// suspended at the final suspend point.
+// In the destroy function, if it isn't possible that the ResumeFnAddr is NULL
+// and the coroutine doesn't suspend at the final suspend point actually (this
+// is possible since the coroutine is considered suspended at the final suspend
+// point if promise.unhandled_exception() exits via an exception), we can
+// remove the last case.
 void CoroCloner::handleFinalSuspend() {
   assert(Shape.ABI == coro::ABI::Switch &&
          Shape.SwitchLowering.HasFinalSuspend);
+
+  if (isSwitchDestroyFunction() && Shape.SwitchLowering.HasUnwindCoroEnd)
+    return;
+
   auto *Switch = cast<SwitchInst>(VMap[Shape.SwitchLowering.ResumeSwitch]);
   auto FinalCaseIt = std::prev(Switch->case_end());
   BasicBlock *ResumeBB = FinalCaseIt->getCaseSuccessor();

diff  --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index 1742e9319c3b1..ca723a724a5f8 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -171,6 +171,7 @@ static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
 // Collect "interesting" coroutine intrinsics.
 void coro::Shape::buildFrom(Function &F) {
   bool HasFinalSuspend = false;
+  bool HasUnwindCoroEnd = false;
   size_t FinalSuspendIndex = 0;
   clear(*this);
   SmallVector<CoroFrameInst *, 8> CoroFrames;
@@ -242,6 +243,10 @@ void coro::Shape::buildFrom(Function &F) {
         if (auto *AsyncEnd = dyn_cast<CoroAsyncEndInst>(II)) {
           AsyncEnd->checkWellFormed();
         }
+
+        if (CoroEnds.back()->isUnwind())
+          HasUnwindCoroEnd = true;
+
         if (CoroEnds.back()->isFallthrough() && isa<CoroEndInst>(II)) {
           // Make sure that the fallthrough coro.end is the first element in the
           // CoroEnds vector.
@@ -290,6 +295,7 @@ void coro::Shape::buildFrom(Function &F) {
     auto SwitchId = cast<CoroIdInst>(Id);
     this->ABI = coro::ABI::Switch;
     this->SwitchLowering.HasFinalSuspend = HasFinalSuspend;
+    this->SwitchLowering.HasUnwindCoroEnd = HasUnwindCoroEnd;
     this->SwitchLowering.ResumeSwitch = nullptr;
     this->SwitchLowering.PromiseAlloca = SwitchId->getPromise();
     this->SwitchLowering.ResumeEntryBlock = nullptr;

diff  --git a/llvm/test/Transforms/Coroutines/coro-split-final-suspend.ll b/llvm/test/Transforms/Coroutines/coro-split-final-suspend.ll
new file mode 100644
index 0000000000000..73bd3c547404f
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-split-final-suspend.ll
@@ -0,0 +1,115 @@
+; Tests that we'll generate the store to the final suspend index if we see the unwind coro end.
+; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -S | FileCheck %s
+
+define ptr @unwind_coro_end() presplitcoroutine personality i32 3 {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
+  call void @print(i32 0)
+  br label %init
+
+init:
+  %initial_suspend = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %initial_suspend, label %init_suspend [i8 0, label %init_resume
+                                                   i8 1, label %init_suspend]
+
+init_suspend:
+  ret ptr %hdl
+
+init_resume:
+  br label %susp
+
+susp:
+  %0 = call i8 @llvm.coro.suspend(token none, i1 true)
+  switch i8 %0, label %suspend [i8 0, label %resume
+                                i8 1, label %suspend]
+resume:
+  invoke void @print(i32 1) to label %suspend unwind label %lpad
+
+suspend:
+  call i1 @llvm.coro.end(ptr %hdl, i1 false)
+  call void @print(i32 0)
+  ret ptr %hdl
+
+lpad:
+  %lpval = landingpad { ptr, i32 }
+     cleanup
+
+  call void @print(i32 2)
+  %need.resume = call i1 @llvm.coro.end(ptr null, i1 true)
+  br i1 %need.resume, label %eh.resume, label %cleanup.cont
+
+cleanup.cont:
+  call void @print(i32 3)
+  br label %eh.resume
+
+eh.resume:
+  resume { ptr, i32 } %lpval
+}
+
+; Tests that we need to store the final index if we see unwind coro end.
+; CHECK: define{{.*}}@unwind_coro_end.resume
+; CHECK: store i1 true, ptr %index.addr
+
+; Tests the use of final index in the destroy function.
+; CHECK: define{{.*}}@unwind_coro_end.destroy
+; CHECK: %[[INDEX:.+]] = load i1, ptr %index.addr
+; CHECK-NEXT: switch i1 %[[INDEX]],
+
+define ptr @nounwind_coro_end(i1 %val) presplitcoroutine personality i32 3 {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
+  call void @print(i32 0)
+  br label %init
+
+init:
+  %initial_suspend = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %initial_suspend, label %init_suspend [i8 0, label %init_resume
+                                                   i8 1, label %init_suspend]
+
+init_suspend:
+  ret ptr %hdl
+
+init_resume:
+  br label %susp
+
+susp:
+  %0 = call i8 @llvm.coro.suspend(token none, i1 true)
+  switch i8 %0, label %suspend [i8 0, label %resume
+                                i8 1, label %suspend]
+resume:
+  call void @print(i32 1)
+  br label %suspend
+
+suspend:
+  call i1 @llvm.coro.end(ptr %hdl, i1 false)
+  call void @print(i32 0)
+  ret ptr %hdl
+}
+
+; Tests that we can omit to store the final suspend index if we don't
+; see unwind coro end.
+; CHECK: define{{.*}}@nounwind_coro_end.resume
+; CHECK-NOT: store i1 true, ptr %index.addr
+; CHECK: }
+
+; Tests that we judge the final suspend case by the nullness of resume function.
+; CHECK: define{{.*}}@nounwind_coro_end.destroy
+; CHECK:  %[[RESUME_FN:.+]] = load ptr, ptr %hdl, align 8
+; CHECK:  %{{.*}} = icmp eq ptr %[[RESUME_FN]], null
+
+declare ptr @llvm.coro.free(token, ptr)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+declare void @llvm.coro.resume(ptr)
+declare void @llvm.coro.destroy(ptr)
+
+declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+declare ptr @llvm.coro.alloc(token)
+declare ptr @llvm.coro.begin(token, ptr)
+declare i1 @llvm.coro.end(ptr, i1)
+
+declare noalias ptr @malloc(i32)
+declare void @print(i32)
+declare void @free(ptr)


        


More information about the llvm-commits mailing list