[Coro] Add variant of retcon.once ABI.

nate-chandler · bnbarham · commit f1d437f0f901 · 2025-04-21T14:53:20.000-07:00
Like async coroutines, it's fixed-per-function-size frame is caller allocated--the size is stored in a global "coro function pointer". Like retcon coroutines, dynamic allocations are performed via intrinsic-provided allocation and deallocation functions. Unlike both, it takes an allocator struct as an argument which is forwarded to the allocation/deallocation functions. (cherry picked from commit b0e1dc9)
diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
@@ -1742,6 +1742,10 @@ def int_coro_id_retcon_once : Intrinsic<[llvm_token_ty],
     [llvm_i32_ty, llvm_i32_ty, llvm_ptr_ty,
      llvm_ptr_ty, llvm_ptr_ty, llvm_ptr_ty, llvm_vararg_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],
+    []>;
 def int_coro_alloc : Intrinsic<[llvm_i1_ty], [llvm_token_ty], []>;
 def int_coro_id_async : Intrinsic<[llvm_token_ty],
   [llvm_i32_ty, llvm_i32_ty, llvm_i32_ty, llvm_ptr_ty],
diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h b/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h
@@ -135,6 +135,7 @@ class AnyCoroIdInst : public IntrinsicInst {
     auto ID = I->getIntrinsicID();
     return ID == Intrinsic::coro_id || ID == Intrinsic::coro_id_retcon ||
            ID == Intrinsic::coro_id_retcon_once ||
+           ID == Intrinsic::coro_id_retcon_once_dynamic ||
            ID == Intrinsic::coro_id_async;
   }
 
@@ -314,6 +315,72 @@ class CoroIdRetconOnceInst : public AnyCoroIdRetconInst {
   }
 };
 
+/// This represents the llvm.coro.id.retcon.once.dynamic instruction.
+class LLVM_LIBRARY_VISIBILITY CoroIdRetconOnceDynamicInst
+    : public AnyCoroIdInst {
+  enum {
+    SizeArg,
+    AlignArg,
+    CoroFuncPtrArg,
+    AllocatorArg,
+    StorageArg,
+    PrototypeArg,
+    AllocArg,
+    DeallocArg
+  };
+
+public:
+  void checkWellFormed() const;
+
+  uint64_t getStorageSize() const {
+    return cast<ConstantInt>(getArgOperand(SizeArg))->getZExtValue();
+  }
+
+  Align getStorageAlignment() const {
+    return cast<ConstantInt>(getArgOperand(AlignArg))->getAlignValue();
+  }
+
+  Value *getStorage() const { return getArgOperand(StorageArg); }
+
+  /// Return the coro function pointer address. This should be the address of
+  /// a coro function pointer struct for the current coro function.
+  /// struct coro_function_pointer {
+  ///   uint32_t frame size;
+  ///   uint32_t relative_pointer(coro_function);
+  ///  };
+  GlobalVariable *getCoroFunctionPointer() const {
+    return cast<GlobalVariable>(
+        getArgOperand(CoroFuncPtrArg)->stripPointerCasts());
+  }
+
+  /// Return the prototype for the continuation function.  The type,
+  /// attributes, and calling convention of the continuation function(s)
+  /// are taken from this declaration.
+  Function *getPrototype() const {
+    return cast<Function>(getArgOperand(PrototypeArg)->stripPointerCasts());
+  }
+
+  /// Return the function to use for allocating memory.
+  Function *getAllocFunction() const {
+    return cast<Function>(getArgOperand(AllocArg)->stripPointerCasts());
+  }
+
+  /// Return the function to use for deallocating memory.
+  Function *getDeallocFunction() const {
+    return cast<Function>(getArgOperand(DeallocArg)->stripPointerCasts());
+  }
+
+  Value *getAllocator() const { return getArgOperand(AllocatorArg); }
+
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_id_retcon_once_dynamic;
+  }
+  static bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
 /// This represents the llvm.coro.id.async instruction.
 class CoroIdAsyncInst : public AnyCoroIdInst {
   enum { SizeArg, AlignArg, StorageArg, AsyncFuncPtrArg };
diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroShape.h b/llvm/include/llvm/Transforms/Coroutines/CoroShape.h
@@ -45,6 +45,10 @@ enum class ABI {
   /// single continuation function. The continuation function is available as an
   /// intrinsic.
   Async,
+
+  /// The variant of RetconOnce which features a dynamically-sized caller
+  /// allocation.
+  RetconOnceDynamic,
 };
 
 // Holds structural Coroutine Intrinsics for a particular function and other
@@ -127,9 +131,18 @@ struct Shape {
     Function *ResumePrototype;
     Function *Alloc;
     Function *Dealloc;
+    Value *Allocator;
     BasicBlock *ReturnBlock;
     bool IsFrameInlineInStorage;
     ConstantInt* TypeId;
+    GlobalVariable *CoroFuncPointer;
+    Value *Storage;
+    uint64_t StorageSize;
+    Align StorageAlignment;
+    // computed during splitting:
+    uint64_t ContextSize;
+
+    Align getStorageAlignment() const { return Align(StorageAlignment); }
   };
 
   struct AsyncLoweringStorage {
@@ -194,6 +207,7 @@ struct Shape {
                                /*IsVarArg=*/false);
     case coro::ABI::Retcon:
     case coro::ABI::RetconOnce:
+    case coro::ABI::RetconOnceDynamic:
       return RetconLowering.ResumePrototype->getFunctionType();
     case coro::ABI::Async:
       // Not used. The function type depends on the active suspend.
@@ -216,7 +230,8 @@ struct Shape {
   }
 
   ArrayRef<Type *> getRetconResumeTypes() const {
-    assert(ABI == coro::ABI::Retcon || ABI == coro::ABI::RetconOnce);
+    assert(ABI == coro::ABI::Retcon || ABI == coro::ABI::RetconOnce ||
+           ABI == coro::ABI::RetconOnceDynamic);
 
     // The safety of all this is checked by checkWFRetconPrototype.
     auto FTy = RetconLowering.ResumePrototype->getFunctionType();
@@ -230,6 +245,7 @@ struct Shape {
 
     case coro::ABI::Retcon:
     case coro::ABI::RetconOnce:
+    case coro::ABI::RetconOnceDynamic:
       return RetconLowering.ResumePrototype->getCallingConv();
     case coro::ABI::Async:
       return AsyncLowering.AsyncCC;
@@ -262,7 +278,7 @@ struct Shape {
   /// \param CG - if non-null, will be updated for the new call
   void emitDealloc(IRBuilder<> &Builder, Value *Ptr, CallGraph *CG) const;
 
-  Shape() = default;
+  Shape() = delete;
   explicit Shape(Function &F) {
     SmallVector<CoroFrameInst *, 8> CoroFrames;
     SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
diff --git a/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp b/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp
@@ -69,6 +69,7 @@ bool Lowerer::lower(Function &F) {
       case Intrinsic::coro_id:
       case Intrinsic::coro_id_retcon:
       case Intrinsic::coro_id_retcon_once:
+      case Intrinsic::coro_id_retcon_once_dynamic:
       case Intrinsic::coro_id_async:
         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
         break;
diff --git a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
@@ -209,6 +209,7 @@ void Lowerer::lowerEarlyIntrinsics(Function &F) {
         break;
       case Intrinsic::coro_id_retcon:
       case Intrinsic::coro_id_retcon_once:
+      case Intrinsic::coro_id_retcon_once_dynamic:
       case Intrinsic::coro_id_async:
         F.setPresplitCoroutine();
         break;
diff --git a/llvm/lib/Transforms/Coroutines/CoroFrame.cpp b/llvm/lib/Transforms/Coroutines/CoroFrame.cpp
@@ -962,6 +962,19 @@ static StructType *buildFrameType(Function &F, coro::Shape &Shape,
          B.getStructAlign() <= Id->getStorageAlignment());
     break;
   }
+  case coro::ABI::RetconOnceDynamic: {
+    // In the dynamic retcon.once ABI, the frame is always inline in the
+    // storage.
+    Shape.RetconLowering.IsFrameInlineInStorage = true;
+    Shape.RetconLowering.ContextSize =
+        alignTo(Shape.FrameSize, Shape.RetconLowering.StorageAlignment);
+    if (Shape.RetconLowering.StorageAlignment < Shape.FrameAlign) {
+      report_fatal_error(
+          "The alignment requirment of frame variables cannot be higher than "
+          "the alignment of the coro function context");
+    }
+    break;
+  }
   case coro::ABI::Async: {
     Shape.AsyncLowering.FrameOffset =
         alignTo(Shape.AsyncLowering.ContextHeaderSize, Shape.FrameAlign);
@@ -1188,7 +1201,8 @@ static void insertSpills(const FrameDataInfo &FrameData, coro::Shape &Shape) {
 
   // retcon and retcon.once lowering assumes all uses have been sunk.
   if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce ||
-      Shape.ABI == coro::ABI::Async) {
+      Shape.ABI == coro::ABI::Async ||
+      Shape.ABI == coro::ABI::RetconOnceDynamic) {
     // If we found any allocas, replace all of their remaining uses with Geps.
     Builder.SetInsertPoint(SpillBlock, SpillBlock->begin());
     for (const auto &P : FrameData.Allocas) {
@@ -2078,7 +2092,8 @@ void coro::BaseABI::buildCoroutineFrame(bool OptimizeFrame) {
 
   const DominatorTree DT(F);
   if (Shape.ABI != coro::ABI::Async && Shape.ABI != coro::ABI::Retcon &&
-      Shape.ABI != coro::ABI::RetconOnce)
+      Shape.ABI != coro::ABI::RetconOnce &&
+      Shape.ABI != coro::ABI::RetconOnceDynamic)
     sinkLifetimeStartMarkers(F, Shape, Checker, DT);
 
   // All values (that are not allocas) that needs to be spilled to the frame.
@@ -2098,7 +2113,8 @@ void coro::BaseABI::buildCoroutineFrame(bool OptimizeFrame) {
   LLVM_DEBUG(dumpSpills("Spills", Spills));
 
   if (Shape.ABI == coro::ABI::Retcon || Shape.ABI == coro::ABI::RetconOnce ||
-      Shape.ABI == coro::ABI::Async)
+      Shape.ABI == coro::ABI::Async ||
+      Shape.ABI == coro::ABI::RetconOnceDynamic)
     sinkSpillUsesAfterCoroBegin(DT, Shape.CoroBegin, Spills, Allocas);
 
   // Build frame
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
diff --git a/llvm/lib/Transforms/Coroutines/SpillUtils.cpp b/llvm/lib/Transforms/Coroutines/SpillUtils.cpp
