[Coro] Hide promise alloca for later passes

NewSigma · NewSigma · commit fa9149a8b066 · 2025-05-09T18:43:07.000+08:00
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`_ that outside the 
+coroutine body, `coro.frame`_ and `coro.done`_ intrinsics. 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
@@ -9,6 +9,7 @@
 #include "llvm/Transforms/Coroutines/CoroEarly.h"
 #include "CoroInternal.h"
 #include "llvm/IR/DIBuilder.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
@@ -165,6 +166,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 +177,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;
@@ -218,22 +237,44 @@ void Lowerer::lowerEarlyIntrinsics(Function &F) {
       case Intrinsic::coro_destroy:
         lowerResumeOrDestroy(*CB, CoroSubFnInst::DestroyIndex);
         break;
-      case Intrinsic::coro_promise:
-        lowerCoroPromise(cast<CoroPromiseInst>(&I));
+      case Intrinsic::coro_promise: {
+        bool OutsideCoro = CoroBegin == nullptr;
+        if (OutsideCoro)
+          lowerCoroPromise(cast<CoroPromiseInst>(&I));
         break;
+      }
       case Intrinsic::coro_done:
         lowerCoroDone(cast<IntrinsicInst>(&I));
         break;
     }
   }
 
-  // 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);
 
+    if (auto *PA = CoroId->getPromise()) {
+      assert(CoroBegin && "Use Switch-Resumed ABI but missing coro.begin");
+
+      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,11 @@ void coro::Shape::cleanCoroutine(
   for (CoroSaveInst *CoroSave : UnusedCoroSaves)
     CoroSave->eraseFromParent();
   UnusedCoroSaves.clear();
+
+  for (auto *PI : CoroPromises) {
+    PI->replaceAllUsesWith(getPromiseAlloca());
+    PI->eraseFromParent();
+  }
 }
 
 static void propagateCallAttrsFromCallee(CallInst *Call, Function *Callee) {
diff --git a/llvm/test/Instrumentation/AddressSanitizer/skip-coro.ll b/llvm/test/Instrumentation/AddressSanitizer/skip-coro.ll
@@ -15,8 +15,9 @@ define ptr @foo() #0 {
 entry:
   %__promise = alloca %struct.Promise, align 8
   %0 = call token @llvm.coro.id(i32 16, ptr nonnull %__promise, ptr null, ptr null)
-  %1 = call ptr @llvm.coro.noop()
-  ret ptr %1
+  %1 = call ptr @llvm.coro.begin(token %0, ptr null)
+  %2 = call ptr @llvm.coro.noop()
+  ret ptr %2
 }
 
 declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr)
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)
