[Coro] Distinguish frame allocation functions.

Doing so enables implementations to allocate extra space in a header.

Author: nate-chandler

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.
 

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

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) {

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,

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);

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());
 

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");
@@ -537,7 +539,13 @@ Value *coro::Shape::emitAlloc(IRBuilder<> &Builder, Value *Size,
       sizeParamIndex = 1;
       Args.push_back(RetconLowering.Allocator);
     }
-    auto Alloc = RetconLowering.Alloc;
+    Function *Alloc = nullptr;
+    if (isa_and_nonnull<CoroAllocaAllocFrameInst>(AI)) {
+      assert(ABI == coro::ABI::RetconOnceDynamic);
+      Alloc = RetconLowering.AllocFrame;
+    } else {
+      Alloc = RetconLowering.Alloc;
+    }
     Size = Builder.CreateIntCast(
         Size, Alloc->getFunctionType()->getParamType(sizeParamIndex),
         /*is signed*/ false);
@@ -559,15 +567,21 @@ 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;
+    }
     SmallVector<Value *, 2> Args;
     unsigned allocationParamIndex = 0;
     if (ABI == coro::ABI::RetconOnceDynamic) {

llvm/lib/Transforms/Coroutines/SpillUtils.cpp

@@ -55,12 +55,12 @@ static bool isSuspendReachableFrom(BasicBlock *From,
 
 /// Is the given alloca "local", i.e. bounded in lifetime to not cross a
 /// suspend point?
-static bool isLocalAlloca(CoroAllocaAllocInst *AI) {
+static bool isLocalAlloca(AnyCoroAllocaAllocInst *AI) {
   // Seed the visited set with all the basic blocks containing a free
   // so that we won't pass them up.
   VisitedBlocksSet VisitedOrFreeBBs;
   for (auto *User : AI->users()) {
-    if (auto FI = dyn_cast<CoroAllocaFreeInst>(User))
+    if (auto FI = dyn_cast<AnyCoroAllocaFreeInst>(User))
       VisitedOrFreeBBs.insert(FI->getParent());
   }
 
@@ -71,18 +71,18 @@ static bool isLocalAlloca(CoroAllocaAllocInst *AI) {
 /// This happens during the all-instructions iteration, so it must not
 /// delete the call.
 static Instruction *
-lowerNonLocalAlloca(CoroAllocaAllocInst *AI, const coro::Shape &Shape,
+lowerNonLocalAlloca(AnyCoroAllocaAllocInst *AI, const coro::Shape &Shape,
                     SmallVectorImpl<Instruction *> &DeadInsts) {
   IRBuilder<> Builder(AI);
-  auto Alloc = Shape.emitAlloc(Builder, AI->getSize(), nullptr);
+  auto Alloc = Shape.emitAlloc(Builder, AI->getSize(), AI, nullptr);
 
   for (User *U : AI->users()) {
     if (isa<CoroAllocaGetInst>(U)) {
       U->replaceAllUsesWith(Alloc);
     } else {
-      auto FI = cast<CoroAllocaFreeInst>(U);
+      auto FI = cast<AnyCoroAllocaFreeInst>(U);
       Builder.SetInsertPoint(FI);
-      Shape.emitDealloc(Builder, Alloc, nullptr);
+      Shape.emitDealloc(Builder, Alloc, AI, nullptr);
     }
     DeadInsts.push_back(cast<Instruction>(U));
   }
@@ -460,7 +460,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) {
 
@@ -471,7 +471,7 @@ void collectSpillsAndAllocasFromInsts(
       continue;
 
     // Handle alloca.alloc specially here.
-    if (auto AI = dyn_cast<CoroAllocaAllocInst>(&I)) {
+    if (auto AI = dyn_cast<AnyCoroAllocaAllocInst>(&I)) {
       // Check whether the alloca's lifetime is bounded by suspend points.
       if (isLocalAlloca(AI)) {
         LocalAllocas.push_back(AI);