[llvm] [Coroutines][NFC] Elide coro.free based on frame instead of coro.id (PR #187627)

[Weibo He via llvm-commits]

https://github.com/NewSigma created https://github.com/llvm/llvm-project/pull/187627

Part 2/4: Implement HALO for coroutines that flow off final suspend. Parent PR approved in https://github.com/llvm/llvm-project/pull/185336, with no change since then

Since `coro.id` is unavailable in resumers, Elide `coro.free` based on frame instead of `coro.id`

>From 0fac973740057fed40794d17da225d5de153bfc2 Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma at 163.com>
Date: Sun, 8 Mar 2026 17:22:40 +0800
Subject: [PATCH] [Coroutines][NFC] Elide coro.free based on frame instead of
 coro.id

---
 llvm/lib/Transforms/Coroutines/CoroElide.cpp  |  2 +-
 llvm/lib/Transforms/Coroutines/CoroInternal.h |  2 +-
 llvm/lib/Transforms/Coroutines/CoroSplit.cpp  | 15 ++++++---------
 llvm/lib/Transforms/Coroutines/Coroutines.cpp | 12 ++++--------
 4 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/llvm/lib/Transforms/Coroutines/CoroElide.cpp b/llvm/lib/Transforms/Coroutines/CoroElide.cpp
index 1c8d4a8592d60..3e260d48a5d88 100644
--- a/llvm/lib/Transforms/Coroutines/CoroElide.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroElide.cpp
@@ -228,6 +228,7 @@ void CoroIdElider::elideHeapAllocations(uint64_t FrameSize, Align FrameAlign) {
       new BitCastInst(Frame, PointerType::getUnqual(C), "vFrame", InsertPt);
 
   for (auto *CB : CoroBegins) {
+    coro::elideCoroFree(CB);
     CB->replaceAllUsesWith(FrameVoidPtr);
     CB->eraseFromParent();
   }
@@ -410,7 +411,6 @@ bool CoroIdElider::attemptElide() {
 
   if (EligibleForElide && FrameSizeAndAlign) {
     elideHeapAllocations(FrameSizeAndAlign->first, FrameSizeAndAlign->second);
-    coro::replaceCoroFree(CoroId, /*Elide=*/true);
     NumOfCoroElided++;
 
 #ifndef NDEBUG
diff --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h
index cc47a557ee5c0..319e600870091 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInternal.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInternal.h
@@ -21,7 +21,7 @@ namespace llvm::coro {
 bool isSuspendBlock(BasicBlock *BB);
 bool declaresAnyIntrinsic(const Module &M);
 bool declaresIntrinsics(const Module &M, ArrayRef<Intrinsic::ID> List);
-void replaceCoroFree(CoroIdInst *CoroId, bool Elide);
+void elideCoroFree(Value *FramePtr);
 
 /// Replaces all @llvm.coro.alloc intrinsics calls associated with a given
 /// call @llvm.coro.id instruction with boolean value false.
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index f780a7bbc8f66..d28004c2d6312 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -1102,8 +1102,7 @@ void coro::SwitchCloner::create() {
 
   // Replacing coro.free with 'null' in cleanup to suppress deallocation code.
   if (FKind == coro::CloneKind::SwitchCleanup)
-    coro::replaceCoroFree(cast<CoroIdInst>(VMap[Shape.CoroBegin->getId()]),
-                          /*Elide=*/FKind == coro::CloneKind::SwitchCleanup);
+    elideCoroFree(NewFramePtr);
 }
 
 static void updateAsyncFuncPointerContextSize(coro::Shape &Shape) {
@@ -1163,10 +1162,9 @@ static void handleNoSuspendCoroutine(coro::Shape &Shape) {
   auto *CoroBegin = Shape.CoroBegin;
   switch (Shape.ABI) {
   case coro::ABI::Switch: {
-    auto SwitchId = Shape.getSwitchCoroId();
-    auto *AllocInst = SwitchId->getCoroAlloc();
-    coro::replaceCoroFree(SwitchId, /*Elide=*/AllocInst != nullptr);
-    if (AllocInst) {
+    if (auto *AllocInst = Shape.getSwitchCoroId()->getCoroAlloc()) {
+      coro::elideCoroFree(CoroBegin);
+
       IRBuilder<> Builder(AllocInst);
       // Create an alloca for a byte array of the frame size
       auto *FrameTy = ArrayType::get(Type::getInt8Ty(Builder.getContext()),
@@ -1441,9 +1439,8 @@ struct SwitchCoroutineSplitter {
     if (Shape.CoroBegin) {
       auto *NewCoroBegin =
           cast_if_present<CoroBeginInst>(VMap[Shape.CoroBegin]);
-      auto *NewCoroId = cast<CoroIdInst>(NewCoroBegin->getId());
-      coro::replaceCoroFree(NewCoroId, /*Elide=*/true);
-      coro::suppressCoroAllocs(NewCoroId);
+      coro::elideCoroFree(NewCoroBegin);
+      coro::suppressCoroAllocs(cast<CoroIdInst>(NewCoroBegin->getId()));
       NewCoroBegin->replaceAllUsesWith(NoAllocF->getArg(FrameIdx));
       NewCoroBegin->eraseFromParent();
     }
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index 2922e39a85e81..a68a5bf1623b5 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -118,11 +118,10 @@ bool coro::declaresIntrinsics(const Module &M, ArrayRef<Intrinsic::ID> List) {
   return false;
 }
 
-// Replace all coro.frees associated with the provided CoroId either with 'null'
-// if Elide is true and with its frame parameter otherwise.
-void coro::replaceCoroFree(CoroIdInst *CoroId, bool Elide) {
+// Replace all coro.frees associated with the provided frame with 'null'
+void coro::elideCoroFree(Value *FramePtr) {
   SmallVector<CoroFreeInst *, 4> CoroFrees;
-  for (User *U : CoroId->users())
+  for (User *U : FramePtr->users())
     if (auto CF = dyn_cast<CoroFreeInst>(U))
       CoroFrees.push_back(CF);
 
@@ -130,10 +129,7 @@ void coro::replaceCoroFree(CoroIdInst *CoroId, bool Elide) {
     return;
 
   Value *Replacement =
-      Elide
-          ? ConstantPointerNull::get(PointerType::get(CoroId->getContext(), 0))
-          : CoroFrees.front()->getFrame();
-
+      ConstantPointerNull::get(PointerType::get(FramePtr->getContext(), 0));
   for (CoroFreeInst *CF : CoroFrees) {
     CF->replaceAllUsesWith(Replacement);
     CF->eraseFromParent();