Make use of uniqueness of coro.promise

llvm · ChuanqiXu9 · May 16, 2025 · May 9, 2025 · May 12, 2025 · May 12, 2025
commit dc0bd0af7bf31eb6a4afe0b95a211a14b1f51b75
diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroShape.h b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
@@ -80,7 +80,7 @@ struct Shape {
  // Scan the function and collect the above intrinsics for later processing
  void analyze(Function &F, SmallVectorImpl<CoroFrameInst *> &CoroFrames,
               SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
-               SmallVectorImpl<CoroPromiseInst *> &CoroPromises);
+               CoroPromiseInst *&CoroPromise);
  // If for some reason, we were not able to find coro.begin, bailout.
  void invalidateCoroutine(Function &F,
                           SmallVectorImpl<CoroFrameInst *> &CoroFrames);
@@ -89,7 +89,7 @@ struct Shape {
  // Remove orphaned and unnecessary intrinsics
  void cleanCoroutine(SmallVectorImpl<CoroFrameInst *> &CoroFrames,
                      SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
-                      SmallVectorImpl<CoroPromiseInst *> &CoroPromises);
+                      CoroPromiseInst *CoroPromise);

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

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


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

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

 void coro::Shape::cleanCoroutine(
    SmallVectorImpl<CoroFrameInst *> &CoroFrames,
-    SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves,
-    SmallVectorImpl<CoroPromiseInst *> &CoroPromises) {
+    SmallVectorImpl<CoroSaveInst *> &UnusedCoroSaves, CoroPromiseInst *PI) {
  // The coro.frame intrinsic is always lowered to the result of coro.begin.
  for (CoroFrameInst *CF : CoroFrames) {
    CF->replaceAllUsesWith(CoroBegin);
@@ -495,10 +494,10 @@ void coro::Shape::cleanCoroutine(
    CoroSave->eraseFromParent();
  UnusedCoroSaves.clear();

-  auto *AI = getPromiseAlloca();
-  for (auto *PI : CoroPromises) {
-    PI->replaceAllUsesWith(PI->isFromPromise() ? cast<Value>(CoroBegin)
-                                               : cast<Value>(AI));
+  if (PI) {
+    PI->replaceAllUsesWith(PI->isFromPromise()
+                               ? cast<Value>(CoroBegin)
+                               : cast<Value>(getPromiseAlloca()));
    PI->eraseFromParent();
  }
 }

diff --git a/llvm/test/Transforms/Coroutines/gh105595.ll b/llvm/test/Transforms/Coroutines/gh105595.ll
@@ -18,7 +18,6 @@ resume:
  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