From 71baee432642817b69c9e88e4bfc7c3053fc0240 Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma@163.com>
Date: Tue, 1 Apr 2025 22:36:15 +0800
Subject: [PATCH 1/7] [DSE] Defer alloca store elimination for CoroSplit

Allocas are destroyed when returning from functions. However, this is not the case for pre-split coroutines. Any premature elimination will lead to side effects.

Fix 123347
---
 .../Scalar/DeadStoreElimination.cpp           |  9 ++--
 .../DeadStoreElimination/coro-alloca.ll       | 41 +++++++++++++++++++
 2 files changed, 45 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/Transforms/DeadStoreElimination/coro-alloca.ll

diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index 935f21fd484f3..c452f04b29d72 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -1194,14 +1194,13 @@ struct DSEState {
 
   bool isInvisibleToCallerAfterRet(const Value *V) {
     if (isa<AllocaInst>(V))
-      return true;
+      // Defer alloca store elimination, wait for CoroSplit
+      return !F.hasFnAttribute(Attribute::PresplitCoroutine);
+
     auto I = InvisibleToCallerAfterRet.insert({V, false});
     if (I.second) {
-      if (!isInvisibleToCallerOnUnwind(V)) {
-        I.first->second = false;
-      } else if (isNoAliasCall(V)) {
+      if (isInvisibleToCallerOnUnwind(V) && isNoAliasCall(V))
         I.first->second = !PointerMayBeCaptured(V, /*ReturnCaptures=*/true);
-      }
     }
     return I.first->second;
   }
diff --git a/llvm/test/Transforms/DeadStoreElimination/coro-alloca.ll b/llvm/test/Transforms/DeadStoreElimination/coro-alloca.ll
new file mode 100644
index 0000000000000..97883177f89e2
--- /dev/null
+++ b/llvm/test/Transforms/DeadStoreElimination/coro-alloca.ll
@@ -0,0 +1,41 @@
+; Test that store-load operation that crosses suspension point will not be eliminated by DSE before CoroSplit
+; RUN: opt < %s -passes='dse,verify' -S | FileCheck %s
+
+define void @fn(ptr align 8 %0) presplitcoroutine {
+  %2 = alloca ptr, align 8
+  %3 = alloca i8, align 1
+  %4 = call token @llvm.coro.id(i32 16, ptr %2, ptr @fn, ptr null)
+  %5 = call ptr @llvm.coro.begin(token %4, ptr null)
+  %6 = call ptr @malloc(i64 1)
+  call void @llvm.lifetime.start.p0(i64 8, ptr %2)
+  store ptr %6, ptr %2, align 8
+  %7 = call token @llvm.coro.save(ptr null)
+  call void @llvm.coro.await.suspend.void(ptr %3, ptr %5, ptr @await_suspend_wrapper_void)
+  %8 = call i8 @llvm.coro.suspend(token %7, i1 false)
+  %9 = icmp ule i8 %8, 1
+  br i1 %9, label %10, label %11
+
+10:
+  call void @llvm.lifetime.end.p0(i64 8, ptr %2)
+  br label %11
+
+11:
+  %12 = call i1 @llvm.coro.end(ptr null, i1 false, token none)
+  %13 = load ptr, ptr %2, align 8
+  store ptr %13, ptr %0, align 8
+; store when suspend, load when resume
+; CHECK: store ptr null, ptr %2, align 8
+  store ptr null, ptr %2, align 8
+  ret void
+}
+
+declare ptr @malloc(i64)
+declare token @llvm.coro.id(i32, ptr, ptr, ptr)
+declare ptr @llvm.coro.begin(token, ptr)
+declare void @llvm.lifetime.start.p0(i64, ptr)
+declare token @llvm.coro.save(ptr)
+declare void @llvm.lifetime.end.p0(i64, ptr)
+declare void @llvm.coro.await.suspend.void(ptr, ptr, ptr)
+declare i8 @llvm.coro.suspend(token, i1)
+declare i1 @llvm.coro.end(ptr, i1, token)
+declare void @await_suspend_wrapper_void(ptr, ptr)

From 5362a65a08a41fa5e6e8eae5bcd7f0adfd2e9a82 Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma@163.com>
Date: Wed, 2 Apr 2025 22:54:11 +0800
Subject: [PATCH 2/7] Revision 1

---
 .../Scalar/DeadStoreElimination.cpp           |  7 ++-
 .../DeadStoreElimination/coro-alloca.ll       | 54 ++++++++-----------
 2 files changed, 28 insertions(+), 33 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index c452f04b29d72..780b64e70136f 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -1195,12 +1195,15 @@ struct DSEState {
   bool isInvisibleToCallerAfterRet(const Value *V) {
     if (isa<AllocaInst>(V))
       // Defer alloca store elimination, wait for CoroSplit
-      return !F.hasFnAttribute(Attribute::PresplitCoroutine);
+      return !F.isPresplitCoroutine();
 
     auto I = InvisibleToCallerAfterRet.insert({V, false});
     if (I.second) {
-      if (isInvisibleToCallerOnUnwind(V) && isNoAliasCall(V))
+      if (!isInvisibleToCallerOnUnwind(V)) {
+        I.first->second = false;
+      } else if (isNoAliasCall(V)) {
         I.first->second = !PointerMayBeCaptured(V, /*ReturnCaptures=*/true);
+      }
     }
     return I.first->second;
   }
diff --git a/llvm/test/Transforms/DeadStoreElimination/coro-alloca.ll b/llvm/test/Transforms/DeadStoreElimination/coro-alloca.ll
index 97883177f89e2..ec9dc84f2c4ae 100644
--- a/llvm/test/Transforms/DeadStoreElimination/coro-alloca.ll
+++ b/llvm/test/Transforms/DeadStoreElimination/coro-alloca.ll
@@ -1,41 +1,33 @@
 ; Test that store-load operation that crosses suspension point will not be eliminated by DSE before CoroSplit
-; RUN: opt < %s -passes='dse,verify' -S | FileCheck %s
+; RUN: opt < %s -passes='dse' -S | FileCheck %s
 
-define void @fn(ptr align 8 %0) presplitcoroutine {
-  %2 = alloca ptr, align 8
-  %3 = alloca i8, align 1
-  %4 = call token @llvm.coro.id(i32 16, ptr %2, ptr @fn, ptr null)
-  %5 = call ptr @llvm.coro.begin(token %4, ptr null)
-  %6 = call ptr @malloc(i64 1)
-  call void @llvm.lifetime.start.p0(i64 8, ptr %2)
-  store ptr %6, ptr %2, align 8
-  %7 = call token @llvm.coro.save(ptr null)
-  call void @llvm.coro.await.suspend.void(ptr %3, ptr %5, ptr @await_suspend_wrapper_void)
-  %8 = call i8 @llvm.coro.suspend(token %7, i1 false)
-  %9 = icmp ule i8 %8, 1
-  br i1 %9, label %10, label %11
+define void @fn(ptr align 8 %arg) presplitcoroutine {
+  %promise = alloca ptr, align 8
+  %awaiter = alloca i8, align 1
+  %id = call token @llvm.coro.id(i32 16, ptr %promise, ptr @fn, ptr null)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
+  %mem = call ptr @malloc(i64 1)
+  call void @llvm.lifetime.start.p0(i64 8, ptr %promise)
+  store ptr %mem, ptr %promise, align 8
+  %save = call token @llvm.coro.save(ptr null)
+  call void @llvm.coro.await.suspend.void(ptr %awaiter, ptr %hdl, ptr @await_suspend_wrapper_void)
+  %sp = call i8 @llvm.coro.suspend(token %save, i1 false)
+  %flag = icmp ule i8 %sp, 1
+  br i1 %flag, label %resume, label %suspend
 
-10:
-  call void @llvm.lifetime.end.p0(i64 8, ptr %2)
-  br label %11
+resume:
+  call void @llvm.lifetime.end.p0(i64 8, ptr %promise)
+  br label %suspend
 
-11:
-  %12 = call i1 @llvm.coro.end(ptr null, i1 false, token none)
-  %13 = load ptr, ptr %2, align 8
-  store ptr %13, ptr %0, align 8
+suspend:
+  call i1 @llvm.coro.end(ptr null, i1 false, token none)
+  %temp = load ptr, ptr %promise, align 8
+  store ptr %temp, ptr %arg, align 8
 ; store when suspend, load when resume
-; CHECK: store ptr null, ptr %2, align 8
-  store ptr null, ptr %2, align 8
+; CHECK: store ptr null, ptr %promise, align 8
+  store ptr null, ptr %promise, align 8
   ret void
 }
 
 declare ptr @malloc(i64)
-declare token @llvm.coro.id(i32, ptr, ptr, ptr)
-declare ptr @llvm.coro.begin(token, ptr)
-declare void @llvm.lifetime.start.p0(i64, ptr)
-declare token @llvm.coro.save(ptr)
-declare void @llvm.lifetime.end.p0(i64, ptr)
-declare void @llvm.coro.await.suspend.void(ptr, ptr, ptr)
-declare i8 @llvm.coro.suspend(token, i1)
-declare i1 @llvm.coro.end(ptr, i1, token)
 declare void @await_suspend_wrapper_void(ptr, ptr)

From 74bb3f36e8428c1c11d6f3bdf170eed64cb3c945 Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma@163.com>
Date: Thu, 3 Apr 2025 21:29:18 +0800
Subject: [PATCH 3/7] Mark promise of pre-split coroutine visible to caller

---
 llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index 780b64e70136f..e43decfb55a4f 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -75,6 +75,7 @@
 #include "llvm/Support/DebugCounter.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Coroutines/CoroInstr.h"
 #include "llvm/Transforms/Utils/AssumeBundleBuilder.h"
 #include "llvm/Transforms/Utils/BuildLibCalls.h"
 #include "llvm/Transforms/Utils/Local.h"
@@ -972,6 +973,8 @@ struct DSEState {
   // no longer be captured.
   bool ShouldIterateEndOfFunctionDSE;
 
+  CoroIdInst *CoroId = nullptr;
+
   /// Dead instructions to be removed at the end of DSE.
   SmallVector<Instruction *> ToRemove;
 
@@ -990,6 +993,13 @@ struct DSEState {
     for (BasicBlock *BB : post_order(&F)) {
       PostOrderNumbers[BB] = PO++;
       for (Instruction &I : *BB) {
+        if (auto* II = dyn_cast<IntrinsicInst>(&I)) {
+            const auto ID = II->getIntrinsicID();
+            if (ID == Intrinsic::coro_begin ||
+                ID == Intrinsic::coro_begin_custom_abi)
+              CoroId = cast<CoroIdInst>(cast<CoroBeginInst>(II)->getId());
+        }
+
         MemoryAccess *MA = MSSA.getMemoryAccess(&I);
         if (I.mayThrow() && !MA)
           ThrowingBlocks.insert(I.getParent());
@@ -1194,8 +1204,7 @@ struct DSEState {
 
   bool isInvisibleToCallerAfterRet(const Value *V) {
     if (isa<AllocaInst>(V))
-      // Defer alloca store elimination, wait for CoroSplit
-      return !F.isPresplitCoroutine();
+      return !(F.isPresplitCoroutine() && V == CoroId->getPromise());
 
     auto I = InvisibleToCallerAfterRet.insert({V, false});
     if (I.second) {

From c721f76970de7802ea920112e4f24572e6ccf81a Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma@163.com>
Date: Thu, 3 Apr 2025 21:51:16 +0800
Subject: [PATCH 4/7] Format code

---
 llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index e43decfb55a4f..b771616351815 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -994,10 +994,10 @@ struct DSEState {
       PostOrderNumbers[BB] = PO++;
       for (Instruction &I : *BB) {
         if (auto* II = dyn_cast<IntrinsicInst>(&I)) {
-            const auto ID = II->getIntrinsicID();
-            if (ID == Intrinsic::coro_begin ||
-                ID == Intrinsic::coro_begin_custom_abi)
-              CoroId = cast<CoroIdInst>(cast<CoroBeginInst>(II)->getId());
+          const auto ID = II->getIntrinsicID();
+          if (ID == Intrinsic::coro_begin ||
+              ID == Intrinsic::coro_begin_custom_abi)
+            CoroId = cast<CoroIdInst>(cast<CoroBeginInst>(II)->getId());
         }
 
         MemoryAccess *MA = MSSA.getMemoryAccess(&I);

From 9a55f2b54de723094dbcee83d95383a4071dc20e Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma@163.com>
Date: Fri, 4 Apr 2025 10:39:10 +0800
Subject: [PATCH 5/7] Fix ci fail

---
 .../Transforms/Scalar/DeadStoreElimination.cpp  | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index b771616351815..83f2ec28435fb 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -972,8 +972,8 @@ struct DSEState {
   // function due to removing a store causing a previously captured pointer to
   // no longer be captured.
   bool ShouldIterateEndOfFunctionDSE;
-
-  CoroIdInst *CoroId = nullptr;
+  // Pre-split coroutine promise is specially handled
+  AllocaInst *PresplitCoroPromise = nullptr;
 
   /// Dead instructions to be removed at the end of DSE.
   SmallVector<Instruction *> ToRemove;
@@ -993,11 +993,16 @@ struct DSEState {
     for (BasicBlock *BB : post_order(&F)) {
       PostOrderNumbers[BB] = PO++;
       for (Instruction &I : *BB) {
-        if (auto* II = dyn_cast<IntrinsicInst>(&I)) {
+        auto *II = dyn_cast<IntrinsicInst>(&I);
+        if (F.isPresplitCoroutine() && II != nullptr) {
           const auto ID = II->getIntrinsicID();
           if (ID == Intrinsic::coro_begin ||
-              ID == Intrinsic::coro_begin_custom_abi)
-            CoroId = cast<CoroIdInst>(cast<CoroBeginInst>(II)->getId());
+              ID == Intrinsic::coro_begin_custom_abi) {
+            auto *AnyCoroId = cast<CoroBeginInst>(II)->getId();
+            auto *CoroId = dyn_cast_if_present<CoroIdInst>(AnyCoroId);
+            if (CoroId)
+              PresplitCoroPromise = CoroId->getPromise();
+          }
         }
 
         MemoryAccess *MA = MSSA.getMemoryAccess(&I);
@@ -1204,7 +1209,7 @@ struct DSEState {
 
   bool isInvisibleToCallerAfterRet(const Value *V) {
     if (isa<AllocaInst>(V))
-      return !(F.isPresplitCoroutine() && V == CoroId->getPromise());
+      return V != PresplitCoroPromise;
 
     auto I = InvisibleToCallerAfterRet.insert({V, false});
     if (I.second) {

From f3917ffdae3d8e9b43d81be9bca6a34999a17930 Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma@163.com>
Date: Wed, 9 Apr 2025 12:34:14 +0800
Subject: [PATCH 6/7] Check if CoroId corresponds to splitted coroutine

---
 llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index 83f2ec28435fb..b6abd7929bd7e 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -990,18 +990,21 @@ struct DSEState {
     // Collect blocks with throwing instructions not modeled in MemorySSA and
     // alloc-like objects.
     unsigned PO = 0;
+    bool FoundPromise = !F.isPresplitCoroutine();
     for (BasicBlock *BB : post_order(&F)) {
       PostOrderNumbers[BB] = PO++;
       for (Instruction &I : *BB) {
         auto *II = dyn_cast<IntrinsicInst>(&I);
-        if (F.isPresplitCoroutine() && II != nullptr) {
+        if (!FoundPromise && II != nullptr) {
           const auto ID = II->getIntrinsicID();
           if (ID == Intrinsic::coro_begin ||
               ID == Intrinsic::coro_begin_custom_abi) {
             auto *AnyCoroId = cast<CoroBeginInst>(II)->getId();
             auto *CoroId = dyn_cast_if_present<CoroIdInst>(AnyCoroId);
-            if (CoroId)
+            if (CoroId && isa<ConstantPointerNull>(CoroId->getRawInfo())) {
               PresplitCoroPromise = CoroId->getPromise();
+              FoundPromise = true;
+            }
           }
         }
 

From 9f24a22c10871e831703ea7e6dbed21c436ff88e Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma@163.com>
Date: Wed, 9 Apr 2025 13:52:59 +0800
Subject: [PATCH 7/7] Refactor

---
 .../Scalar/DeadStoreElimination.cpp           | 29 +++++++++++--------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index b6abd7929bd7e..f354974471192 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -990,22 +990,15 @@ struct DSEState {
     // Collect blocks with throwing instructions not modeled in MemorySSA and
     // alloc-like objects.
     unsigned PO = 0;
-    bool FoundPromise = !F.isPresplitCoroutine();
+    const bool IsPresplitCoroutine = F.isPresplitCoroutine();
     for (BasicBlock *BB : post_order(&F)) {
       PostOrderNumbers[BB] = PO++;
       for (Instruction &I : *BB) {
         auto *II = dyn_cast<IntrinsicInst>(&I);
-        if (!FoundPromise && II != nullptr) {
-          const auto ID = II->getIntrinsicID();
-          if (ID == Intrinsic::coro_begin ||
-              ID == Intrinsic::coro_begin_custom_abi) {
-            auto *AnyCoroId = cast<CoroBeginInst>(II)->getId();
-            auto *CoroId = dyn_cast_if_present<CoroIdInst>(AnyCoroId);
-            if (CoroId && isa<ConstantPointerNull>(CoroId->getRawInfo())) {
-              PresplitCoroPromise = CoroId->getPromise();
-              FoundPromise = true;
-            }
-          }
+        if (IsPresplitCoroutine && PresplitCoroPromise == nullptr &&
+            II != nullptr) {
+          if (auto *CoroId = getPresplitCoroId(II))
+            PresplitCoroPromise = CoroId->getPromise();
         }
 
         MemoryAccess *MA = MSSA.getMemoryAccess(&I);
@@ -1242,6 +1235,18 @@ struct DSEState {
     return !I.first->second;
   }
 
+  CoroIdInst *getPresplitCoroId(IntrinsicInst *II) {
+    const auto ID = II->getIntrinsicID();
+    if (ID != Intrinsic::coro_begin && ID != Intrinsic::coro_begin_custom_abi)
+      return nullptr;
+
+    auto *AnyCoroId = cast<CoroBeginInst>(II)->getId();
+    auto *CoroId = dyn_cast_if_present<CoroIdInst>(AnyCoroId);
+    if (CoroId && isa<ConstantPointerNull>(CoroId->getRawInfo()))
+      return CoroId;
+    return nullptr;
+  }
+
   std::optional<MemoryLocation> getLocForWrite(Instruction *I) const {
     if (!I->mayWriteToMemory())
       return std::nullopt;