Merge pull request #11671 from nate-chandler/general-coro/20251009/1

nate-chandler · web-flow · commit e52b2dddd693 · 2025-10-23T20:17:56.000-07:00
stable/21.x: [Coro] RetconOnceDynamic: Distinguish frame allocs and pass frame pointers.
diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
@@ -1771,7 +1771,7 @@ def int_coro_id_retcon_once : Intrinsic<[llvm_token_ty],
     []>;
 def int_coro_id_retcon_once_dynamic : Intrinsic<[llvm_token_ty],
     [llvm_i32_ty, llvm_i32_ty, llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty,
-     llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
+     llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty],
     []>;
 def int_coro_alloc : Intrinsic<[llvm_i1_ty], [llvm_token_ty], []>;
 def int_coro_id_async : Intrinsic<[llvm_token_ty],
@@ -1818,8 +1818,11 @@ def int_coro_prepare_retcon : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty],
                                         [IntrNoMem]>;
 def int_coro_alloca_alloc : Intrinsic<[llvm_token_ty],
                                       [llvm_anyint_ty, llvm_i32_ty], []>;
+def int_coro_alloca_alloc_frame : Intrinsic<[llvm_token_ty],
+                                      [llvm_anyint_ty, llvm_i32_ty], []>;
 def int_coro_alloca_get : Intrinsic<[llvm_ptr_ty], [llvm_token_ty], []>;
 def int_coro_alloca_free : Intrinsic<[], [llvm_token_ty], []>;
+def int_coro_alloca_free_frame : Intrinsic<[], [llvm_token_ty], []>;
 
 // Coroutine Manipulation Intrinsics.
 
diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h b/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h
@@ -327,7 +327,9 @@ class LLVM_LIBRARY_VISIBILITY CoroIdRetconOnceDynamicInst
     StorageArg,
     PrototypeArg,
     AllocArg,
-    DeallocArg
+    DeallocArg,
+    AllocFrameArg,
+    DeallocFrameArg,
   };
 
 public:
@@ -371,6 +373,16 @@ class LLVM_LIBRARY_VISIBILITY CoroIdRetconOnceDynamicInst
     return cast<Function>(getArgOperand(DeallocArg)->stripPointerCasts());
   }
 
+  /// Return the function to use for allocating memory.
+  Function *getAllocFrameFunction() const {
+    return cast<Function>(getArgOperand(AllocFrameArg)->stripPointerCasts());
+  }
+
+  /// Return the function to use for deallocating memory.
+  Function *getDeallocFrameFunction() const {
+    return cast<Function>(getArgOperand(DeallocFrameArg)->stripPointerCasts());
+  }
+
   Value *getAllocator() const { return getArgOperand(AllocatorArg); }
 
   // Methods to support type inquiry through isa, cast, and dyn_cast:
@@ -822,8 +834,7 @@ class CoroAsyncEndInst : public AnyCoroEndInst {
   }
 };
 
-/// This represents the llvm.coro.alloca.alloc instruction.
-class CoroAllocaAllocInst : public IntrinsicInst {
+class AnyCoroAllocaAllocInst : public IntrinsicInst {
   enum { SizeArg, AlignArg };
 
 public:
@@ -832,6 +843,21 @@ class CoroAllocaAllocInst : public IntrinsicInst {
     return cast<ConstantInt>(getArgOperand(AlignArg))->getAlignValue();
   }
 
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const IntrinsicInst *I) {
+    auto ID = I->getIntrinsicID();
+    return ID == Intrinsic::coro_alloca_alloc ||
+           ID == Intrinsic::coro_alloca_alloc_frame;
+  }
+
+  static bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
+/// This represents the llvm.coro.alloca.alloc instruction.
+class CoroAllocaAllocInst : public AnyCoroAllocaAllocInst {
+public:
   // Methods to support type inquiry through isa, cast, and dyn_cast:
   static bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_alloca_alloc;
@@ -841,13 +867,25 @@ class CoroAllocaAllocInst : public IntrinsicInst {
   }
 };
 
+/// This represents the llvm.coro.alloca.alloc.frame instruction.
+class CoroAllocaAllocFrameInst : public AnyCoroAllocaAllocInst {
+public:
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_alloca_alloc_frame;
+  }
+  static bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
 /// This represents the llvm.coro.alloca.get instruction.
 class CoroAllocaGetInst : public IntrinsicInst {
   enum { AllocArg };
 
 public:
-  CoroAllocaAllocInst *getAlloc() const {
-    return cast<CoroAllocaAllocInst>(getArgOperand(AllocArg));
+  AnyCoroAllocaAllocInst *getAlloc() const {
+    return cast<AnyCoroAllocaAllocInst>(getArgOperand(AllocArg));
   }
 
   // Methods to support type inquiry through isa, cast, and dyn_cast:
@@ -859,15 +897,29 @@ class CoroAllocaGetInst : public IntrinsicInst {
   }
 };
 
-/// This represents the llvm.coro.alloca.free instruction.
-class CoroAllocaFreeInst : public IntrinsicInst {
+class AnyCoroAllocaFreeInst : public IntrinsicInst {
   enum { AllocArg };
 
 public:
-  CoroAllocaAllocInst *getAlloc() const {
-    return cast<CoroAllocaAllocInst>(getArgOperand(AllocArg));
+  AnyCoroAllocaAllocInst *getAlloc() const {
+    return cast<AnyCoroAllocaAllocInst>(getArgOperand(AllocArg));
   }
 
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const IntrinsicInst *I) {
+    auto ID = I->getIntrinsicID();
+    return ID == Intrinsic::coro_alloca_free ||
+           ID == Intrinsic::coro_alloca_free_frame;
+  }
+
+  static bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
+/// This represents the llvm.coro.alloca.free instruction.
+class CoroAllocaFreeInst : public IntrinsicInst {
+public:
   // Methods to support type inquiry through isa, cast, and dyn_cast:
   static bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_alloca_free;
@@ -877,6 +929,18 @@ class CoroAllocaFreeInst : public IntrinsicInst {
   }
 };
 
+/// This represents the llvm.coro.alloca.free instruction.
+class CoroAllocaFreeFrameInst : public IntrinsicInst {
+public:
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_alloca_free_frame;
+  }
+  static bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
 } // End namespace llvm.
 
 #endif // LLVM_TRANSFORMS_COROUTINES_COROINSTR_H
diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroShape.h b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
@@ -135,7 +135,9 @@ struct Shape {
   struct RetconLoweringStorage {
     Function *ResumePrototype;
     Function *Alloc;
+    Function *AllocFrame;
     Function *Dealloc;
+    Function *DeallocFrame;
     Value *Allocator;
     BasicBlock *ReturnBlock;
     bool IsFrameInlineInStorage;
@@ -278,13 +280,13 @@ struct Shape {
   ///
   /// \param CG - if non-null, will be updated for the new call
   LLVM_ABI Value *emitAlloc(IRBuilder<> &Builder, Value *Size,
-                            CallGraph *CG) const;
+                            AnyCoroAllocaAllocInst *AI, CallGraph *CG) const;
 
   /// Deallocate memory according to the rules of the active lowering.
   ///
   /// \param CG - if non-null, will be updated for the new call
   LLVM_ABI void emitDealloc(IRBuilder<> &Builder, Value *Ptr,
-                            CallGraph *CG) const;
+                            AnyCoroAllocaAllocInst *AI, CallGraph *CG) const;
 
   Shape() = delete;
   explicit Shape(Function &F) {
diff --git a/llvm/include/llvm/Transforms/Coroutines/SpillUtils.h b/llvm/include/llvm/Transforms/Coroutines/SpillUtils.h
@@ -35,7 +35,7 @@ void collectSpillsFromArgs(SpillInfo &Spills, Function &F,
 void collectSpillsAndAllocasFromInsts(
     SpillInfo &Spills, SmallVector<AllocaInfo, 8> &Allocas,
     SmallVector<Instruction *, 4> &DeadInstructions,
-    SmallVector<CoroAllocaAllocInst *, 4> &LocalAllocas, Function &F,
+    SmallVector<AnyCoroAllocaAllocInst *, 4> &LocalAllocas, Function &F,
     const SuspendCrossingInfo &Checker, const DominatorTree &DT,
     const coro::Shape &Shape);
 void collectSpillsFromDbgInfo(SpillInfo &Spills, Function &F,
diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -1598,12 +1598,12 @@ static bool willLeaveFunctionImmediatelyAfter(BasicBlock *BB,
   return true;
 }
 
-static bool localAllocaNeedsStackSave(CoroAllocaAllocInst *AI) {
+static bool localAllocaNeedsStackSave(AnyCoroAllocaAllocInst *AI) {
   // Look for a free that isn't sufficiently obviously followed by
   // either a suspend or a termination, i.e. something that will leave
   // the coro resumption frame.
   for (auto *U : AI->users()) {
-    auto FI = dyn_cast<CoroAllocaFreeInst>(U);
+    auto FI = dyn_cast<AnyCoroAllocaFreeInst>(U);
     if (!FI) continue;
 
     if (!willLeaveFunctionImmediatelyAfter(FI->getParent()))
@@ -1616,8 +1616,8 @@ static bool localAllocaNeedsStackSave(CoroAllocaAllocInst *AI) {
 
 /// Turn each of the given local allocas into a normal (dynamic) alloca
 /// instruction.
-static void lowerLocalAllocas(ArrayRef<CoroAllocaAllocInst*> LocalAllocas,
-                              SmallVectorImpl<Instruction*> &DeadInsts) {
+static void lowerLocalAllocas(ArrayRef<AnyCoroAllocaAllocInst *> LocalAllocas,
+                              SmallVectorImpl<Instruction *> &DeadInsts) {
   for (auto *AI : LocalAllocas) {
     IRBuilder<> Builder(AI);
 
@@ -1640,7 +1640,7 @@ static void lowerLocalAllocas(ArrayRef<CoroAllocaAllocInst*> LocalAllocas,
       // alloca.alloc is required to obey a stack discipline, although we
       // don't enforce that structurally.
       } else {
-        auto FI = cast<CoroAllocaFreeInst>(U);
+        auto FI = cast<AnyCoroAllocaFreeInst>(U);
         if (StackSave) {
           Builder.SetInsertPoint(FI);
           Builder.CreateStackRestore(StackSave);
@@ -2146,7 +2146,7 @@ void coro::BaseABI::buildCoroutineFrame(bool OptimizeFrame) {
   // Collect the spills for arguments and other not-materializable values.
   coro::collectSpillsFromArgs(Spills, F, Checker);
   SmallVector<Instruction *, 4> DeadInstructions;
-  SmallVector<CoroAllocaAllocInst *, 4> LocalAllocas;
+  SmallVector<AnyCoroAllocaAllocInst *, 4> LocalAllocas;
   coro::collectSpillsAndAllocasFromInsts(Spills, Allocas, DeadInstructions,
                                          LocalAllocas, F, Checker, DT, Shape);
   coro::collectSpillsFromDbgInfo(Spills, F, Checker);
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -163,7 +163,7 @@ static void maybeFreeRetconStorage(IRBuilder<> &Builder,
   if (Shape.RetconLowering.IsFrameInlineInStorage)
     return;
 
-  Shape.emitDealloc(Builder, FramePtr, CG);
+  Shape.emitDealloc(Builder, FramePtr, nullptr, CG);
 }
 
 /// Replace an llvm.coro.end.async.
@@ -1903,7 +1903,8 @@ void coro::AnyRetconABI::splitCoroutine(Function &F, coro::Shape &Shape,
     // Allocate.  We don't need to update the call graph node because we're
     // going to recompute it from scratch after splitting.
     // FIXME: pass the required alignment
-    RawFramePtr = Shape.emitAlloc(Builder, Builder.getInt64(Size), nullptr);
+    RawFramePtr =
+        Shape.emitAlloc(Builder, Builder.getInt64(Size), nullptr, nullptr);
     RawFramePtr =
         Builder.CreateBitCast(RawFramePtr, Shape.CoroBegin->getType());
 
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -325,7 +325,9 @@ void coro::Shape::analyze(Function &F,
     auto Prototype = ContinuationId->getPrototype();
     RetconLowering.ResumePrototype = Prototype;
     RetconLowering.Alloc = ContinuationId->getAllocFunction();
+    RetconLowering.AllocFrame = ContinuationId->getAllocFrameFunction();
     RetconLowering.Dealloc = ContinuationId->getDeallocFunction();
+    RetconLowering.DeallocFrame = ContinuationId->getDeallocFrameFunction();
     RetconLowering.Storage = ContinuationId->getStorage();
     RetconLowering.Allocator = ContinuationId->getAllocator();
     RetconLowering.ReturnBlock = nullptr;
@@ -523,7 +525,7 @@ static void addCallToCallGraph(CallGraph *CG, CallInst *Call, Function *Callee){
 }
 
 Value *coro::Shape::emitAlloc(IRBuilder<> &Builder, Value *Size,
-                              CallGraph *CG) const {
+                              AnyCoroAllocaAllocInst *AI, CallGraph *CG) const {
   switch (ABI) {
   case coro::ABI::Switch:
     llvm_unreachable("can't allocate memory in coro switch-lowering");
@@ -532,12 +534,21 @@ Value *coro::Shape::emitAlloc(IRBuilder<> &Builder, Value *Size,
   case coro::ABI::RetconOnce:
   case coro::ABI::RetconOnceDynamic: {
     unsigned sizeParamIndex = 0;
+    Function *Alloc = nullptr;
+    if (isa_and_nonnull<CoroAllocaAllocFrameInst>(AI)) {
+      assert(ABI == coro::ABI::RetconOnceDynamic);
+      Alloc = RetconLowering.AllocFrame;
+    } else {
+      Alloc = RetconLowering.Alloc;
+    }
     SmallVector<Value *, 2> Args;
     if (ABI == coro::ABI::RetconOnceDynamic) {
-      sizeParamIndex = 1;
+      Args.push_back(RetconLowering.Storage);
       Args.push_back(RetconLowering.Allocator);
+      sizeParamIndex = 2;
+    } else {
+      Alloc = RetconLowering.Alloc;
     }
-    auto Alloc = RetconLowering.Alloc;
     Size = Builder.CreateIntCast(
         Size, Alloc->getFunctionType()->getParamType(sizeParamIndex),
         /*is signed*/ false);
@@ -548,6 +559,9 @@ Value *coro::Shape::emitAlloc(IRBuilder<> &Builder, Value *Size,
         Args.push_back(TypeId);
     }
     auto *Call = Builder.CreateCall(Alloc, Args);
+    if (ABI == coro::ABI::RetconOnceDynamic) {
+      Call->addParamAttr(1, Attribute::SwiftCoro);
+    }
     propagateCallAttrsFromCallee(Call, Alloc);
     addCallToCallGraph(CG, Call, Alloc);
     return Call;
@@ -559,25 +573,37 @@ Value *coro::Shape::emitAlloc(IRBuilder<> &Builder, Value *Size,
 }
 
 void coro::Shape::emitDealloc(IRBuilder<> &Builder, Value *Ptr,
-                              CallGraph *CG) const {
+                              AnyCoroAllocaAllocInst *AI, CallGraph *CG) const {
   switch (ABI) {
   case coro::ABI::Switch:
     llvm_unreachable("can't allocate memory in coro switch-lowering");
 
   case coro::ABI::Retcon:
   case coro::ABI::RetconOnce:
   case coro::ABI::RetconOnceDynamic: {
-    auto Dealloc = RetconLowering.Dealloc;
+    Function *Dealloc = nullptr;
+    if (isa_and_nonnull<CoroAllocaAllocFrameInst>(AI)) {
+      assert(ABI == coro::ABI::RetconOnceDynamic);
+      Dealloc = RetconLowering.DeallocFrame;
+    } else {
+      Dealloc = RetconLowering.Dealloc;
+    }
+    unsigned allocationParamIndex = 0;
     SmallVector<Value *, 2> Args;
-    unsigned sizeParamIndex = 0;
     if (ABI == coro::ABI::RetconOnceDynamic) {
-      sizeParamIndex = 1;
+      allocationParamIndex = 2;
+      Args.push_back(RetconLowering.Storage);
       Args.push_back(RetconLowering.Allocator);
+    } else {
+      Dealloc = RetconLowering.Dealloc;
     }
     Ptr = Builder.CreateBitCast(
-        Ptr, Dealloc->getFunctionType()->getParamType(sizeParamIndex));
+        Ptr, Dealloc->getFunctionType()->getParamType(allocationParamIndex));
     Args.push_back(Ptr);
     auto *Call = Builder.CreateCall(Dealloc, Args);
+    if (ABI == coro::ABI::RetconOnceDynamic) {
+      Call->addParamAttr(1, Attribute::SwiftCoro);
+    }
     propagateCallAttrsFromCallee(Call, Dealloc);
     addCallToCallGraph(CG, Call, Dealloc);
     return;
diff --git a/llvm/lib/Transforms/Coroutines/SpillUtils.cpp b/llvm/lib/Transforms/Coroutines/SpillUtils.cpp
diff --git a/llvm/test/Transforms/Coroutines/coro-retcon-once-dynamic-nocleanup.ll b/llvm/test/Transforms/Coroutines/coro-retcon-once-dynamic-nocleanup.ll
diff --git a/llvm/test/Transforms/Coroutines/coro-retcon-once-dynamic.ll b/llvm/test/Transforms/Coroutines/coro-retcon-once-dynamic.ll
