Merge pull request #10120 from nate-chandler/coro-retcon-dynamic

nate-chandler · web-flow · commit a3e32cdaae72 · 2025-02-26T21:26:14.000-08:00
[Coro] Add variant of retcon.once ABI.
diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
@@ -1680,6 +1680,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/lib/Transforms/Coroutines/CoroCleanup.cpp b/llvm/lib/Transforms/Coroutines/CoroCleanup.cpp
@@ -68,6 +68,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
@@ -213,6 +213,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
@@ -1394,6 +1394,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);
@@ -1969,7 +1982,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) {
@@ -2857,7 +2871,8 @@ static void collectFrameAlloca(AllocaInst *AI, coro::Shape &Shape,
   // code.
   bool ShouldUseLifetimeStartInfo =
       (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);
   AllocaUseVisitor Visitor{AI->getDataLayout(), DT, Shape, Checker,
                            ShouldUseLifetimeStartInfo};
   Visitor.visitPtr(*AI);
@@ -3160,7 +3175,8 @@ void coro::buildCoroutineFrame(
   SmallVector<CoroAllocaAllocInst*, 4> LocalAllocas;
   SmallVector<Instruction*, 4> DeadInstructions;
   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);
 
   // Collect the spills for arguments and other not-materializable values.
@@ -3238,7 +3254,8 @@ void coro::buildCoroutineFrame(
 
   LLVM_DEBUG(dumpSpills("Spills", FrameData.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(F, FrameData, Shape.CoroBegin);
   Shape.FrameTy = buildFrameType(F, Shape, FrameData);
   Shape.FramePtr = Shape.CoroBegin;
diff --git a/llvm/lib/Transforms/Coroutines/CoroInstr.h b/llvm/lib/Transforms/Coroutines/CoroInstr.h
@@ -134,6 +134,7 @@ class LLVM_LIBRARY_VISIBILITY 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;
   }
 
@@ -317,6 +318,72 @@ class LLVM_LIBRARY_VISIBILITY CoroIdRetconOnceInst
   }
 };
 
+/// 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 LLVM_LIBRARY_VISIBILITY CoroIdAsyncInst : public AnyCoroIdInst {
   enum { SizeArg, AlignArg, StorageArg, AsyncFuncPtrArg };
diff --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h
@@ -73,6 +73,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
@@ -128,9 +132,18 @@ struct LLVM_LIBRARY_VISIBILITY 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 {
@@ -196,6 +209,7 @@ struct LLVM_LIBRARY_VISIBILITY 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.
@@ -206,8 +220,8 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
   }
 
   ArrayRef<Type*> getRetconResultTypes() const {
-    assert(ABI == coro::ABI::Retcon ||
-           ABI == coro::ABI::RetconOnce);
+    assert(ABI == coro::ABI::Retcon || ABI == coro::ABI::RetconOnce ||
+           ABI == coro::ABI::RetconOnceDynamic);
     auto FTy = CoroBegin->getFunction()->getFunctionType();
 
     // The safety of all this is checked by checkWFRetconPrototype.
@@ -219,8 +233,8 @@ struct LLVM_LIBRARY_VISIBILITY 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();
@@ -234,6 +248,7 @@ struct LLVM_LIBRARY_VISIBILITY Shape {
 
     case coro::ABI::Retcon:
     case coro::ABI::RetconOnce:
+    case coro::ABI::RetconOnceDynamic:
       return RetconLowering.ResumePrototype->getCallingConv();
     case coro::ABI::Async:
       return AsyncLowering.AsyncCC;
@@ -266,12 +281,15 @@ struct LLVM_LIBRARY_VISIBILITY 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, bool OptimizeFrame = false)
       : OptimizeFrame(OptimizeFrame) {
     buildFrom(F);
   }
   void buildFrom(Function &F);
+
+private:
+  void buildShapeFromRetconInst(AnyCoroIdInst *Id, Function *Prototype);
 };
 
 bool defaultMaterializable(Instruction &V);
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
