llvm · ChuanqiXu9 · May 16, 2025 · May 9, 2025 · May 12, 2025 · May 12, 2025
diff --git a/llvm/docs/Coroutines.rst b/llvm/docs/Coroutines.rst
@@ -2121,10 +2121,11 @@ Coroutine Transformation Passes
 ===============================
 CoroEarly
 ---------
-The pass CoroEarly lowers coroutine intrinsics that hide the details of the
-structure of the coroutine frame, but, otherwise not needed to be preserved to
-help later coroutine passes. This pass lowers `coro.frame`_, `coro.done`_,
-and `coro.promise`_ intrinsics.
+The CoroEarly pass ensures later middle end passes correctly interpret coroutine 
+semantics and lowers coroutine intrinsics that not needed to be preserved to 
+help later coroutine passes. This pass lowers `coro.promise`_, `coro.frame`_ and 
+`coro.done`_ intrinsics. Afterwards, it replace uses of promise alloca with 
+`coro.promise`_ intrinsic.

 .. _CoroSplit:


diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroShape.h b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
@@ -79,15 +79,17 @@ struct Shape {

  // Scan the function and collect the above intrinsics for later processing
  void analyze(Function &F, SmallVectorImpl<CoroFrameInst *> &CoroFrames,
-               SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves);
+               SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
+               SmallVectorImpl<CoroPromiseInst *> &CoroPromises);
  // If for some reason, we were not able to find coro.begin, bailout.
  void invalidateCoroutine(Function &F,
                           SmallVectorImpl<CoroFrameInst *> &CoroFrames);
  // Perform ABI related initial transformation
  void initABI();
  // Remove orphaned and unnecessary intrinsics
  void cleanCoroutine(SmallVectorImpl<CoroFrameInst *> &CoroFrames,
-                      SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves);
+                      SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
+                      SmallVectorImpl<CoroPromiseInst *> &CoroPromises);

  // Field indexes for special fields in the switch lowering.
  struct SwitchFieldIndex {
@@ -265,13 +267,14 @@ struct Shape {
  explicit Shape(Function &F) {
    SmallVector<CoroFrameInst *, 8> CoroFrames;
    SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
+    SmallVector<CoroPromiseInst *, 2> CoroPromises;

-    analyze(F, CoroFrames, UnusedCoroSaves);
+    analyze(F, CoroFrames, UnusedCoroSaves, CoroPromises);
    if (!CoroBegin) {
      invalidateCoroutine(F, CoroFrames);
      return;
    }
-    cleanCoroutine(CoroFrames, UnusedCoroSaves);
+    cleanCoroutine(CoroFrames, UnusedCoroSaves, CoroPromises);
  }
 };


diff --git a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
@@ -165,6 +165,7 @@ static void setCannotDuplicate(CoroIdInst *CoroId) {

 void Lowerer::lowerEarlyIntrinsics(Function &F) {
  CoroIdInst *CoroId = nullptr;
+  CoroBeginInst *CoroBegin = nullptr;
  SmallVector<CoroFreeInst *, 4> CoroFrees;
  bool HasCoroSuspend = false;
  for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
@@ -175,6 +176,23 @@ void Lowerer::lowerEarlyIntrinsics(Function &F) {
    switch (CB->getIntrinsicID()) {
      default:
        continue;
+      case Intrinsic::coro_begin:
+      case Intrinsic::coro_begin_custom_abi: {
+        auto CBI = cast<CoroBeginInst>(&I);
+
+        // Ignore coro id's that aren't pre-split.
+        auto Id = dyn_cast<CoroIdInst>(CBI->getId());
+        if (Id && !Id->getInfo().isPreSplit())
+          break;
+
+        if (CoroBegin)
+          report_fatal_error(
+              "coroutine should have exactly one defining @llvm.coro.begin");
+        CBI->addRetAttr(Attribute::NonNull);
+        CBI->addRetAttr(Attribute::NoAlias);
+        CoroBegin = CBI;
+        break;
+      }
      case Intrinsic::coro_free:
        CoroFrees.push_back(cast<CoroFreeInst>(&I));
        break;
@@ -227,13 +245,31 @@ void Lowerer::lowerEarlyIntrinsics(Function &F) {
    }
  }

-  // Make sure that all CoroFree reference the coro.id intrinsic.
-  // Token type is not exposed through coroutine C/C++ builtins to plain C, so
-  // we allow specifying none and fixing it up here.
-  if (CoroId)
+  if (CoroId) {
+    // Make sure that all CoroFree reference the coro.id intrinsic.
+    // Token type is not exposed through coroutine C/C++ builtins to plain C, so
+    // we allow specifying none and fixing it up here.
    for (CoroFreeInst *CF : CoroFrees)
      CF->setArgOperand(0, CoroId);

+    auto *PA = CoroId->getPromise();
+    if (PA && CoroBegin) {
+      assert(isa<AllocaInst>(PA) && "Must pass alloca to coro.id");
+      Builder.SetInsertPoint(*CoroBegin->getInsertionPointAfterDef());
+
+      auto *Alignment = Builder.getInt32(PA->getAlign().value());
+      auto *FromPromise = Builder.getInt1(false);
+      SmallVector<Value *, 3> Arg{CoroBegin, Alignment, FromPromise};
+      auto *PI = Builder.CreateIntrinsic(
+          Builder.getPtrTy(), Intrinsic::coro_promise, Arg, {}, "promise.addr");
+      PA->replaceUsesWithIf(PI, [CoroId](Use &U) {
+        bool IsBitcast = U == U.getUser()->stripPointerCasts();
+        bool IsCoroId = U.getUser() == CoroId;
+        return !IsBitcast && !IsCoroId;
+      });
+    }
+  }
+
  // Coroutine suspention could potentially lead to any argument modified
  // outside of the function, hence arguments should not have noalias
  // attributes.

diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -192,7 +192,8 @@ static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
 // Collect "interesting" coroutine intrinsics.
 void coro::Shape::analyze(Function &F,
                          SmallVectorImpl<CoroFrameInst *> &CoroFrames,
-                          SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves) {
+                          SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
+                          SmallVectorImpl<CoroPromiseInst *> &CoroPromises) {
  clear();

  bool HasFinalSuspend = false;
@@ -286,6 +287,9 @@ void coro::Shape::analyze(Function &F,
          }
        }
        break;
+      case Intrinsic::coro_promise:
+        CoroPromises.push_back(cast<CoroPromiseInst>(II));
+        break;
      }
    }
  }
@@ -477,7 +481,8 @@ void coro::AnyRetconABI::init() {

 void coro::Shape::cleanCoroutine(
    SmallVectorImpl<CoroFrameInst *> &CoroFrames,
-    SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves) {
+    SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
+    SmallVectorImpl<CoroPromiseInst *> &CoroPromises) {
  // The coro.frame intrinsic is always lowered to the result of coro.begin.
  for (CoroFrameInst *CF : CoroFrames) {
    CF->replaceAllUsesWith(CoroBegin);
@@ -489,6 +494,13 @@ void coro::Shape::cleanCoroutine(
  for (CoroSaveInst *CoroSave : UnusedCoroSaves)
    CoroSave->eraseFromParent();
  UnusedCoroSaves.clear();
+
+  auto *AI = getPromiseAlloca();
+  for (auto *PI : CoroPromises) {
+    PI->replaceAllUsesWith(PI->isFromPromise() ? cast<Value>(CoroBegin)
+                                               : cast<Value>(AI));
+    PI->eraseFromParent();
+  }
 }

 static void propagateCallAttrsFromCallee(CallInst *Call, Function *Callee) {

diff --git a/llvm/test/Transforms/Coroutines/gh105595.ll b/llvm/test/Transforms/Coroutines/gh105595.ll
@@ -0,0 +1,32 @@
+; Test that store-load operation that crosses suspension point will not be eliminated by DSE
+; Coro result conversion function that attempts to modify promise shall produce this pattern
+; RUN: opt < %s -passes='coro-early,dse' -S | FileCheck %s
+
+define void @fn() presplitcoroutine {
+  %__promise = alloca ptr, align 8
+  %id = call token @llvm.coro.id(i32 16, ptr %__promise, ptr @fn, ptr null)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
+; CHECK: %promise.addr = call ptr @llvm.coro.promise(ptr %hdl, i32 8, i1 false)
+  %save = call token @llvm.coro.save(ptr null)
+  %sp = call i8 @llvm.coro.suspend(token %save, i1 false)
+  %flag = icmp ule i8 %sp, 1
+  br i1 %flag, label %resume, label %suspend
+
+resume:
+; CHECK: call void @use_value(ptr %promise.addr)
+  call void @use_value(ptr %__promise)
+  br label %suspend
+
+suspend:
+  call i1 @llvm.coro.end(ptr null, i1 false, token none)
+; load value when resume
+; CHECK: %null = load ptr, ptr %promise.addr, align 8
+  %null = load ptr, ptr %__promise, align 8
+  call void @use_value(ptr %null)
+; store value when suspend
+; CHECK: store ptr null, ptr %promise.addr, align 8
+  store ptr null, ptr %__promise, align 8
+  ret void
+}
+
+declare void @use_value(ptr)