[llvm] e190b7c - [Coroutines] Maintain the position of final suspend

Chuanqi Xu via llvm-commits llvm-commits at lists.llvm.org
Thu Aug 11 22:05:38 PDT 2022


Author: Chuanqi Xu
Date: 2022-08-12T13:05:08+08:00
New Revision: e190b7cc90ca5ee712ca3982bf476afa9e8acb3b

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

LOG: [Coroutines] Maintain the position of final suspend

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

The problem happens when we try to simplify the suspend points. We might
break the assumption that the final suspend lives in the last slot of
Shape.CoroSuspends. This patch tries to main the assumption and fixes
the problem.

Added: 
    clang/test/CodeGenCoroutines/pr56329.cpp
    llvm/test/Transforms/Coroutines/coro-preserve-final.ll

Modified: 
    llvm/lib/Transforms/Coroutines/CoroSplit.cpp

Removed: 
    


################################################################################
diff  --git a/clang/test/CodeGenCoroutines/pr56329.cpp b/clang/test/CodeGenCoroutines/pr56329.cpp
new file mode 100644
index 0000000000000..3918acae0f08f
--- /dev/null
+++ b/clang/test/CodeGenCoroutines/pr56329.cpp
@@ -0,0 +1,117 @@
+// Test for PR56919. Tests the we won't contain the resumption of final suspend point.
+//
+// RUN: %clang_cc1 -triple %itanium_abi_triple -std=c++20 %s -O3 -S -emit-llvm -o - | FileCheck %s
+
+#include "Inputs/coroutine.h"
+
+void _exit(int status) __attribute__ ((__noreturn__));
+
+class Promise;
+
+// An object that can be co_awaited, but we always resume immediately from
+// await_suspend.
+struct ResumeFromAwaitSuspend{};
+
+struct Task {
+  using promise_type = Promise;
+  Promise& promise;
+};
+
+struct Promise {
+  static std::coroutine_handle<Promise> GetHandle(Promise& promise) {
+    return std::coroutine_handle<Promise>::from_promise(promise);
+  }
+
+  void unhandled_exception() {}
+  Task get_return_object() { return Task{*this}; }
+  void return_void() {}
+
+  // Always suspend before starting the coroutine body. We actually run the body
+  // when we are co_awaited.
+  std::suspend_always initial_suspend() { return {}; }
+
+  // We support awaiting tasks. We do so by configuring them to resume us when
+  // they are finished, and then resuming them from their initial suspend.
+  auto await_transform(Task&& task) {
+    struct Awaiter {
+      bool await_ready() { return false; }
+
+      std::coroutine_handle<> await_suspend(
+          const std::coroutine_handle<> handle) {
+        // Tell the child to resume the parent once it finishes.
+        child.resume_at_final_suspend = GetHandle(parent);
+
+        // Run the child.
+        return GetHandle(child);
+      }
+
+      void await_resume() {
+        // The child is now at its final suspend point, and can be destroyed.
+        return GetHandle(child).destroy();
+      }
+
+      Promise& parent;
+      Promise& child;
+    };
+
+    return Awaiter{
+        .parent = *this,
+        .child = task.promise,
+    };
+  }
+
+  // Make evaluation of `co_await ResumeFromAwaitSuspend{}` go through the
+  // await_suspend path, but cause it to resume immediately by returning our own
+  // handle to resume.
+  auto await_transform(ResumeFromAwaitSuspend) {
+    struct Awaiter {
+      bool await_ready() { return false; }
+
+      std::coroutine_handle<> await_suspend(const std::coroutine_handle<> h) {
+        return h;
+      }
+
+      void await_resume() {}
+    };
+
+    return Awaiter{};
+  }
+
+  // Always suspend at the final suspend point, transferring control back to our
+  // caller. We expect never to be resumed from the final suspend.
+  auto final_suspend() noexcept {
+    struct FinalSuspendAwaitable final {
+      bool await_ready() noexcept { return false; }
+
+      std::coroutine_handle<> await_suspend(std::coroutine_handle<>) noexcept {
+        return promise.resume_at_final_suspend;
+      }
+
+      void await_resume() noexcept {
+        _exit(1);
+      }
+
+      Promise& promise;
+    };
+
+    return FinalSuspendAwaitable{.promise = *this};
+  }
+
+  // The handle we will resume once we hit final suspend.
+  std::coroutine_handle<> resume_at_final_suspend;
+};
+
+Task Inner();
+
+Task Outer() {
+  co_await ResumeFromAwaitSuspend();
+  co_await Inner();
+}
+
+// CHECK: define{{.*}}@_Z5Outerv.resume(
+// CHECK-NOT: }
+// CHECK-NOT: _exit
+// CHECK: musttail call
+// CHECK: musttail call
+// CHECK-NEXT: ret void
+// CHEKC-NEXT: }

diff  --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 722a1c6ec0cee..5efe377f1f938 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -1555,6 +1555,8 @@ static void simplifySuspendPoints(coro::Shape &Shape) {
   size_t I = 0, N = S.size();
   if (N == 0)
     return;
+
+  size_t ChangedFinalIndex = std::numeric_limits<size_t>::max();
   while (true) {
     auto SI = cast<CoroSuspendInst>(S[I]);
     // Leave final.suspend to handleFinalSuspend since it is undefined behavior
@@ -1562,13 +1564,27 @@ static void simplifySuspendPoints(coro::Shape &Shape) {
     if (!SI->isFinal() && simplifySuspendPoint(SI, Shape.CoroBegin)) {
       if (--N == I)
         break;
+
       std::swap(S[I], S[N]);
+
+      if (cast<CoroSuspendInst>(S[I])->isFinal()) {
+        assert(Shape.SwitchLowering.HasFinalSuspend);
+        ChangedFinalIndex = I;
+      }
+
       continue;
     }
     if (++I == N)
       break;
   }
   S.resize(N);
+
+  // Maintain final.suspend in case final suspend was swapped.
+  // Due to we requrie the final suspend to be the last element of CoroSuspends.
+  if (ChangedFinalIndex < N) {
+    assert(cast<CoroSuspendInst>(S[ChangedFinalIndex])->isFinal());
+    std::swap(S[ChangedFinalIndex], S.back());
+  }
 }
 
 static void splitSwitchCoroutine(Function &F, coro::Shape &Shape,

diff  --git a/llvm/test/Transforms/Coroutines/coro-preserve-final.ll b/llvm/test/Transforms/Coroutines/coro-preserve-final.ll
new file mode 100644
index 0000000000000..730868a9f1ccd
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-preserve-final.ll
@@ -0,0 +1,131 @@
+; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -S | FileCheck %s
+
+%"struct.std::__n4861::noop_coroutine_promise" = type { i8 }
+%struct.Promise = type { %"struct.std::__n4861::coroutine_handle" }
+%"struct.std::__n4861::coroutine_handle" = type { ptr }
+
+define dso_local ptr @_Z5Outerv() #1 {
+entry:
+  %__promise = alloca %struct.Promise, align 8
+  %0 = call token @llvm.coro.id(i32 16, ptr nonnull %__promise, ptr nonnull @_Z5Outerv, ptr null)
+  %1 = call i1 @llvm.coro.alloc(token %0)
+  br i1 %1, label %coro.alloc, label %init.suspend
+
+coro.alloc:                                       ; preds = %entry
+  %2 = tail call i64 @llvm.coro.size.i64()
+  %call = call noalias noundef nonnull ptr @_Znwm(i64 noundef %2) #12
+  br label %init.suspend
+
+init.suspend:                                     ; preds = %entry, %coro.alloc
+  %3 = phi ptr [ null, %entry ], [ %call, %coro.alloc ]
+  %4 = call ptr @llvm.coro.begin(token %0, ptr %3) #13
+  call void @llvm.lifetime.start.p0(i64 8, ptr nonnull %__promise) #3
+  store ptr null, ptr %__promise, align 8
+  %5 = call token @llvm.coro.save(ptr null)
+  %6 = call i8 @llvm.coro.suspend(token %5, i1 false)
+  switch i8 %6, label %coro.ret [
+  i8 0, label %await.suspend
+  i8 1, label %cleanup62
+  ]
+
+await.suspend:                                    ; preds = %init.suspend
+  %7 = call token @llvm.coro.save(ptr null)
+  %8 = call ptr @llvm.coro.subfn.addr(ptr %4, i8 0)
+  call fastcc void %8(ptr %4) #3
+  %9 = call i8 @llvm.coro.suspend(token %7, i1 false)
+  switch i8 %9, label %coro.ret [
+  i8 0, label %await2.suspend
+  i8 1, label %cleanup62
+  ]
+
+await2.suspend:                                   ; preds = %await.suspend
+  %call27 = call ptr @_Z5Innerv() #3
+  %10 = call token @llvm.coro.save(ptr null)
+  %11 = getelementptr inbounds i8, ptr %__promise, i64 -16
+  store ptr %11, ptr %call27, align 8
+  %12 = getelementptr inbounds i8, ptr %call27, i64 -16
+  %13 = call ptr @llvm.coro.subfn.addr(ptr nonnull %12, i8 0)
+  call fastcc void %13(ptr nonnull %12) #3
+  %14 = call i8 @llvm.coro.suspend(token %10, i1 false)
+  switch i8 %14, label %coro.ret [
+  i8 0, label %final.suspend
+  i8 1, label %cleanup62
+  ]
+
+final.suspend:                                    ; preds = %await2.suspend
+  %15 = call ptr @llvm.coro.subfn.addr(ptr nonnull %12, i8 1)
+  call fastcc void %15(ptr nonnull %12) #3
+  %16 = call token @llvm.coro.save(ptr null)
+  %retval.sroa.0.0.copyload.i = load ptr, ptr %__promise, align 8
+  %17 = call ptr @llvm.coro.subfn.addr(ptr %retval.sroa.0.0.copyload.i, i8 0)
+  call fastcc void %17(ptr %retval.sroa.0.0.copyload.i) #3
+  %18 = call i8 @llvm.coro.suspend(token %16, i1 true) #13
+  switch i8 %18, label %coro.ret [
+  i8 0, label %final.ready
+  i8 1, label %cleanup62
+  ]
+
+final.ready:                                      ; preds = %final.suspend
+  call void @_Z5_exiti(i32 noundef 1) #14
+  unreachable
+
+cleanup62:                                        ; preds = %await2.suspend, %await.suspend, %init.suspend, %final.suspend
+  call void @llvm.lifetime.end.p0(i64 8, ptr nonnull %__promise) #3
+  %19 = call ptr @llvm.coro.free(token %0, ptr %4)
+  %.not = icmp eq ptr %19, null
+  br i1 %.not, label %coro.ret, label %coro.free
+
+coro.free:                                        ; preds = %cleanup62
+  call void @_ZdlPv(ptr noundef nonnull %19) #3
+  br label %coro.ret
+
+coro.ret:                                         ; preds = %coro.free, %cleanup62, %final.suspend, %await2.suspend, %await.suspend, %init.suspend
+  %20 = call i1 @llvm.coro.end(ptr null, i1 false) #13
+  ret ptr %__promise
+}
+
+declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr) #2
+declare i1 @llvm.coro.alloc(token) #3
+declare dso_local noundef nonnull ptr @_Znwm(i64 noundef) local_unnamed_addr #4
+declare i64 @llvm.coro.size.i64() #5
+declare ptr @llvm.coro.begin(token, ptr writeonly) #3
+declare void @llvm.lifetime.start.p0(i64 immarg, ptr nocapture) #6
+declare token @llvm.coro.save(ptr) #7
+declare void @llvm.lifetime.end.p0(i64 immarg, ptr nocapture) #6
+declare i8 @llvm.coro.suspend(token, i1) #3
+declare dso_local ptr @_Z5Innerv() local_unnamed_addr #8
+declare dso_local void @_ZdlPv(ptr noundef) local_unnamed_addr #9
+declare ptr @llvm.coro.free(token, ptr nocapture readonly) #2
+declare i1 @llvm.coro.end(ptr, i1) #3
+declare dso_local void @_Z5_exiti(i32 noundef) local_unnamed_addr #10
+declare ptr @llvm.coro.subfn.addr(ptr nocapture readonly, i8) #11
+
+attributes #0 = { mustprogress nounwind uwtable "frame-pointer"="none" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #1 = { nounwind presplitcoroutine uwtable "frame-pointer"="none" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #2 = { argmemonly nofree nounwind readonly }
+attributes #3 = { nounwind }
+attributes #4 = { nobuiltin allocsize(0) "frame-pointer"="none" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #5 = { nofree nosync nounwind readnone }
+attributes #6 = { argmemonly mustprogress nocallback nofree nosync nounwind willreturn }
+attributes #7 = { nomerge nounwind }
+attributes #8 = { "frame-pointer"="none" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #9 = { nobuiltin nounwind "frame-pointer"="none" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #10 = { noreturn "frame-pointer"="none" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }
+attributes #11 = { argmemonly nounwind readonly }
+attributes #12 = { nounwind allocsize(0) }
+attributes #13 = { noduplicate }
+attributes #14 = { noreturn nounwind }
+
+; CHECK: define{{.*}}@_Z5Outerv.resume(
+; CHECK: entry.resume:
+; CHECK: switch i2 %index
+; CHECK-NEXT:    i2 0, label %await2.suspend
+; CHECK-NEXT:    i2 1, label %final.suspend
+;
+; CHECK: await2.suspend:
+; CHECK: musttail call
+; CHECK-NEXT: ret void
+;
+; CHECK: final.suspend:
+; CHECK: musttail call
+; CHECK-NEXT: ret void


        


More information about the llvm-commits mailing list