llvm · NewSigma · Apr 1, 2025 · Apr 2, 2025 · Apr 3, 2025 · Apr 3, 2025
diff --git 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"
@@ -971,6 +972,8 @@ struct DSEState {
   // function due to removing a store causing a previously captured pointer to
   // no longer be captured.
   bool ShouldIterateEndOfFunctionDSE;
+  // Pre-split coroutine promise is specially handled
+  AllocaInst *PresplitCoroPromise = nullptr;
 
   /// Dead instructions to be removed at the end of DSE.
   SmallVector<Instruction *> ToRemove;
@@ -987,9 +990,17 @@ struct DSEState {
     // Collect blocks with throwing instructions not modeled in MemorySSA and
     // alloc-like objects.
     unsigned PO = 0;
+    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 (IsPresplitCoroutine && PresplitCoroPromise == nullptr &&
+            II != nullptr) {
+          if (auto *CoroId = getPresplitCoroId(II))
+            PresplitCoroPromise = CoroId->getPromise();
+        }
+
         MemoryAccess *MA = MSSA.getMemoryAccess(&I);
         if (I.mayThrow() && !MA)
           ThrowingBlocks.insert(I.getParent());
@@ -1194,7 +1205,8 @@ struct DSEState {
 
   bool isInvisibleToCallerAfterRet(const Value *V) {
     if (isa<AllocaInst>(V))
-      return true;
+      return V != PresplitCoroPromise;
+
     auto I = InvisibleToCallerAfterRet.insert({V, false});
     if (I.second) {
       if (!isInvisibleToCallerOnUnwind(V)) {
@@ -1223,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;

diff --git a/llvm/test/Transforms/DeadStoreElimination/coro-alloca.ll b/llvm/test/Transforms/DeadStoreElimination/coro-alloca.ll
@@ -0,0 +1,33 @@
+; Test that store-load operation that crosses suspension point will not be eliminated by DSE before CoroSplit
+; RUN: opt < %s -passes='dse' -S | FileCheck %s
+
+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
+
+resume:
+  call void @llvm.lifetime.end.p0(i64 8, ptr %promise)
+  br label %suspend
+
+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 %promise, align 8
+  store ptr null, ptr %promise, align 8
+  ret void
+}
+
+declare ptr @malloc(i64)
+declare void @await_suspend_wrapper_void(ptr, ptr)