[Coro] RetconOnceDynamic: Add type id.

nate-chandler · nate-chandler · commit c4bbb16a06f8 · 2025-10-29T16:56:18.000-07:00
Add type id to coro function pointer struct, write type id into id
intrinsic, pass type id to allocation functions.
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, llvm_ptr_ty, llvm_ptr_ty, llvm_i64_ty],
     []>;
 def int_coro_alloc : Intrinsic<[llvm_i1_ty], [llvm_token_ty], []>;
 def int_coro_id_async : Intrinsic<[llvm_token_ty],
@@ -1819,7 +1819,7 @@ def int_coro_prepare_retcon : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty],
 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], []>;
+                                      [llvm_anyint_ty, llvm_i32_ty, llvm_i64_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], []>;
diff --git a/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h b/llvm/include/llvm/Transforms/Coroutines/CoroInstr.h
@@ -330,6 +330,7 @@ class LLVM_LIBRARY_VISIBILITY CoroIdRetconOnceDynamicInst
     DeallocArg,
     AllocFrameArg,
     DeallocFrameArg,
+    TypeIdArg,
   };
 
 public:
@@ -385,6 +386,13 @@ class LLVM_LIBRARY_VISIBILITY CoroIdRetconOnceDynamicInst
 
   Value *getAllocator() const { return getArgOperand(AllocatorArg); }
 
+  /// The TypeId to be used for allocating memory.
+  ConstantInt *getTypeId() const {
+    if (arg_size() <= TypeIdArg)
+      return nullptr;
+    return cast<ConstantInt>(getArgOperand(TypeIdArg));
+  }
+
   // 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;
@@ -869,7 +877,16 @@ class CoroAllocaAllocInst : public AnyCoroAllocaAllocInst {
 
 /// This represents the llvm.coro.alloca.alloc.frame instruction.
 class CoroAllocaAllocFrameInst : public AnyCoroAllocaAllocInst {
+  enum { TypeIdArg = 2 };
+
 public:
+  /// Return the TypeId to be used for allocating typed memory
+  Value *getTypeId() const {
+    if (arg_size() <= TypeIdArg)
+      return nullptr;
+    return getArgOperand(TypeIdArg);
+  }
+
   // 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;
diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -1186,8 +1186,10 @@ static void updateCoroFuncPointerContextSize(coro::Shape &Shape) {
   auto *OrigContextSize = FuncPtrStruct->getOperand(1);
   auto *NewContextSize = ConstantInt::get(OrigContextSize->getType(),
                                           Shape.RetconLowering.ContextSize);
-  auto *NewFuncPtrStruct = ConstantStruct::get(
-      FuncPtrStruct->getType(), OrigRelativeFunOffset, NewContextSize);
+  auto *OrigTypeID = FuncPtrStruct->getOperand(2);
+  auto *NewFuncPtrStruct =
+      ConstantStruct::get(FuncPtrStruct->getType(), OrigRelativeFunOffset,
+                          NewContextSize, OrigTypeID);
 
   Shape.RetconLowering.CoroFuncPointer->setInitializer(NewFuncPtrStruct);
 }
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -336,6 +336,7 @@ void coro::Shape::analyze(Function &F,
     RetconLowering.StorageSize = ContinuationId->getStorageSize();
     RetconLowering.StorageAlignment = ContinuationId->getStorageAlignment();
     RetconLowering.CoroFuncPointer = ContinuationId->getCoroFunctionPointer();
+    RetconLowering.TypeId = ContinuationId->getTypeId();
     break;
   }
   case Intrinsic::coro_id_retcon:
@@ -535,8 +536,7 @@ Value *coro::Shape::emitAlloc(IRBuilder<> &Builder, Value *Size,
   case coro::ABI::RetconOnceDynamic: {
     unsigned sizeParamIndex = 0;
     Function *Alloc = nullptr;
-    if (isa_and_nonnull<CoroAllocaAllocFrameInst>(AI)) {
-      assert(ABI == coro::ABI::RetconOnceDynamic);
+    if ((ABI == coro::ABI::RetconOnceDynamic) && isa<CoroAllocaAllocFrameInst>(AI)) {
       Alloc = RetconLowering.AllocFrame;
     } else {
       Alloc = RetconLowering.Alloc;
@@ -553,11 +553,14 @@ Value *coro::Shape::emitAlloc(IRBuilder<> &Builder, Value *Size,
         Size, Alloc->getFunctionType()->getParamType(sizeParamIndex),
         /*is signed*/ false);
     Args.push_back(Size);
-    if (ABI == coro::ABI::RetconOnce) {
-      ConstantInt *TypeId = RetconLowering.TypeId;
-      if (TypeId != nullptr)
-        Args.push_back(TypeId);
+    Value *TypeId = nullptr;
+    if (auto *CAAFI = dyn_cast_or_null<CoroAllocaAllocFrameInst>(AI)) {
+      TypeId = CAAFI->getTypeId();
+    } else if (ABI == coro::ABI::RetconOnce || ABI == coro::ABI::RetconOnceDynamic) {
+      TypeId = RetconLowering.TypeId;
     }
+    if (TypeId != nullptr)
+      Args.push_back(TypeId);
     auto *Call = Builder.CreateCall(Alloc, Args);
     if (ABI == coro::ABI::RetconOnceDynamic) {
       Call->addParamAttr(1, Attribute::SwiftCoro);
@@ -583,7 +586,6 @@ void coro::Shape::emitDealloc(IRBuilder<> &Builder, Value *Ptr,
   case coro::ABI::RetconOnceDynamic: {
     Function *Dealloc = nullptr;
     if (isa_and_nonnull<CoroAllocaAllocFrameInst>(AI)) {
-      assert(ABI == coro::ABI::RetconOnceDynamic);
       Dealloc = RetconLowering.DeallocFrame;
     } else {
       Dealloc = RetconLowering.Dealloc;
@@ -737,6 +739,8 @@ void CoroIdRetconOnceDynamicInst::checkWellFormed() const {
   checkWFRetconPrototype(this, getArgOperand(PrototypeArg));
   checkWFAlloc(this, getArgOperand(AllocArg));
   checkWFDealloc(this, getArgOperand(DeallocArg));
+  checkWFAlloc(this, getArgOperand(AllocFrameArg));
+  checkWFDealloc(this, getArgOperand(DeallocFrameArg));
 }
 
 static void checkAsyncFuncPointer(const Instruction *I, Value *V) {
diff --git a/llvm/test/Transforms/Coroutines/coro-retcon-once-dynamic-nocleanup.ll b/llvm/test/Transforms/Coroutines/coro-retcon-once-dynamic-nocleanup.ll
@@ -4,14 +4,15 @@ target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "arm64-apple-macos99.99"
 
 
-@func_cfp = constant <{ i32, i32 }>
+@func_cfp = constant <{ i32, i32, i64 }>
   <{ i32 trunc (
        i64 sub (
          i64 ptrtoint (ptr @func to i64),
          i64 ptrtoint (ptr getelementptr inbounds (<{ i32, i32 }>, ptr @func_cfp, i32 0, i32 1) to i64)
        )
      to i32),
-     i32 64
+     i32 64, ; frame size
+     i64 1010101010 ; type_id
 }>
 
 
@@ -45,7 +46,8 @@ entry:
     ptr nonnull @allocate,
     ptr nonnull @deallocate,
     ptr nonnull @allocate_frame,
-    ptr nonnull @deallocate_frame
+    ptr nonnull @deallocate_frame,
+    i64 1010101010
   )
   %handle = call ptr @llvm.coro.begin(token %3, ptr null)
   %yielded = getelementptr inbounds %func_self, ptr %2, i32 0, i32 0
diff --git a/llvm/test/Transforms/Coroutines/coro-retcon-once-dynamic.ll b/llvm/test/Transforms/Coroutines/coro-retcon-once-dynamic.ll
@@ -3,34 +3,36 @@
 target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "arm64-apple-macos99.99"
 
-%coro_func_pointer = type <{ i32, i32 }>
+%coro_func_pointer = type <{ i32, i32, i64 }>
 
 ; CHECK-LABEL: %func.Frame = type { ptr }
 ; CHECK-LABEL: %big_types.Frame = type { <32 x i8>, [16 x i8], i64, ptr, %Integer8 }
 
-; CHECK-LABEL: @func_cfp = constant <{ i32, i32 }>
+; CHECK-LABEL: @func_cfp = constant <{ i32, i32, i64 }>
 ; CHECK-SAME:  <{
 ; CHECK-SAME:    i32 trunc
 ; CHECK-SAME:    i32 16
 ; CHECK-SAME:  }>
-@func_cfp = constant <{ i32, i32 }>
+@func_cfp = constant <{ i32, i32, i64 }>
   <{ i32 trunc ( ; offset to @func from @func_cfp
        i64 sub (
          i64 ptrtoint (ptr @func to i64),
          i64 ptrtoint (ptr getelementptr inbounds (<{ i32, i32 }>, ptr @func_cfp, i32 0, i32 1) to i64)
        )
      to i32),
-     i32 64 ; frame size
+     i32 64, ; frame size
+     i64 1010101010 ; type_id
 }>
 
-@big_types_cfp = constant <{ i32, i32 }>
+@big_types_cfp = constant <{ i32, i32, i64 }>
   <{ i32 trunc ( ; offset to @func from @big_types_cfp
        i64 sub (
          i64 ptrtoint (ptr @big_types to i64),
          i64 ptrtoint (ptr getelementptr inbounds (<{ i32, i32 }>, ptr @big_types_cfp, i32 0, i32 1) to i64)
        )
      to i32),
-     i32 64 ; frame size
+     i32 64, ; frame size
+     i64 1010101010 ; type_id
 }>
 
 
@@ -90,7 +92,8 @@ entry:
     ptr nonnull @allocate,
     ptr nonnull @deallocate,
     ptr nonnull @allocate_frame,
-    ptr nonnull @deallocate_frame
+    ptr nonnull @deallocate_frame,
+    i64 1010101010
   )
   %handle = call ptr @llvm.coro.begin(token %id, ptr null)
   %load = load i32, ptr %array
@@ -114,9 +117,9 @@ cleanup:
 
 declare void @continuation_prototype(ptr, ptr)
 
-declare swiftcorocc noalias ptr @allocate(ptr %frame, ptr swiftcoro %cator, i32 %size)
+declare swiftcorocc noalias ptr @allocate(ptr %frame, ptr swiftcoro %cator, i32 %size, i64 %typeid)
 declare void @deallocate(ptr %frame, ptr swiftcoro %cator, ptr %ptr)
-declare swiftcorocc noalias ptr @allocate_frame(ptr %frame, ptr swiftcoro %cator, i32 %size)
+declare swiftcorocc noalias ptr @allocate_frame(ptr %frame, ptr swiftcoro %cator, i32 %size, i64 %typeid)
 declare void @deallocate_frame(ptr %frame, ptr swiftcoro %cator, ptr %ptr)
 
 declare void @use(ptr %ptr)
@@ -146,7 +149,8 @@ define swiftcorocc { ptr, ptr } @big_types(ptr noalias %frame, ptr swiftcoro %al
     ptr nonnull @allocate,
     ptr nonnull @deallocate,
     ptr nonnull @allocate_frame,
-    ptr nonnull @deallocate_frame
+    ptr nonnull @deallocate_frame,
+    i64 1010101010
   )
   %handle = tail call ptr @llvm.coro.begin(token %id, ptr null)
   call void @llvm.lifetime.start.p0(i64 1, ptr nonnull %element_addr)
@@ -168,14 +172,15 @@ define swiftcorocc { ptr, ptr } @big_types(ptr noalias %frame, ptr swiftcoro %al
 }
 
 declare i32 @getSize()
-@allocating_something_else_cfp = constant <{ i32, i32 }>
+@allocating_something_else_cfp = constant <{ i32, i32, i64 }>
   <{ i32 trunc ( ; offset to @func from @allocating_something_else_cfp
        i64 sub (
          i64 ptrtoint (ptr @allocating_something_else to i64),
          i64 ptrtoint (ptr getelementptr inbounds (<{ i32, i32 }>, ptr @allocating_something_else_cfp, i32 0, i32 1) to i64)
        )
      to i32),
-     i32 64 ; frame size
+     i32 64, ; frame size
+     i64 1010101010 ; type_id
 }>
 declare { ptr, ptr } @allocating_something_else(ptr noalias %frame, ptr swiftcoro %allocator)
 
@@ -210,14 +215,15 @@ declare { ptr, ptr } @allocating_something_else(ptr noalias %frame, ptr swiftcor
 ; CHECK-SAME:      ptr swiftcoro %1,
 ; CHECK-SAME:      ptr
 ; CHECK-SAME:    )
-@allocating_something_cfp = constant <{ i32, i32 }>
+@allocating_something_cfp = constant <{ i32, i32, i64 }>
   <{ i32 trunc ( ; offset to @func from @allocating_something_cfp
        i64 sub (
          i64 ptrtoint (ptr @allocating_something to i64),
          i64 ptrtoint (ptr getelementptr inbounds (<{ i32, i32 }>, ptr @allocating_something_cfp, i32 0, i32 1) to i64)
        )
      to i32),
-     i32 64 ; frame size
+     i32 64, ; frame size
+     i64 1010101010 ; type_id
 }>
 define { ptr, ptr } @allocating_something(ptr noalias %frame, ptr swiftcoro %allocator) {
 entry:
@@ -231,7 +237,8 @@ entry:
     ptr nonnull @allocate,
     ptr nonnull @deallocate,
     ptr nonnull @allocate_frame,
-    ptr nonnull @deallocate_frame
+    ptr nonnull @deallocate_frame,
+    i64 1010101010
   )
   %hdl = call ptr @llvm.coro.begin(token %id, ptr null)
   %size = call i32 @getSize()
@@ -240,7 +247,7 @@ entry:
 
   %callee_frame_size = load i32, ptr getelementptr inbounds nuw (%coro_func_pointer, ptr @allocating_something_else_cfp, i32 0, i32 1), align 8
   %callee_frame_size_64 = zext i32 %callee_frame_size to i64
-  %callee_frame_token = call token @llvm.coro.alloca.alloc.frame.i64(i64 %callee_frame_size_64, i32 16)
+  %callee_frame_token = call token @llvm.coro.alloca.alloc.frame.i64(i64 %callee_frame_size_64, i32 16, i64 1010101010)
   %callee_frame = call ptr @llvm.coro.alloca.get(token %callee_frame_token)
   call void @llvm.lifetime.start.p0(i64 -1, ptr %callee_frame)
   %allocating_something_else = call ptr @llvm.coro.prepare.retcon(ptr @allocating_something_else)
