diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td index e5c6f59541e4e..24b2636eba1be 100644 --- 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 index e2f32bcc5119c..ef9fe3e31b050 100644 --- 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(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 index be6a69abfdb41..8f9b9234402d3 100644 --- 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 index 8761ba4e5b2fc..2678fae903251 100644 --- 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,8 @@ Value *coro::Shape::emitAlloc(IRBuilder<> &Builder, Value *Size, case coro::ABI::RetconOnceDynamic: { unsigned sizeParamIndex = 0; Function *Alloc = nullptr; - if (isa_and_nonnull(AI)) { - assert(ABI == coro::ABI::RetconOnceDynamic); + if ((ABI == coro::ABI::RetconOnceDynamic) && + isa(AI)) { Alloc = RetconLowering.AllocFrame; } else { Alloc = RetconLowering.Alloc; @@ -553,10 +554,18 @@ 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) + if (ABI == coro::ABI::RetconOnce || ABI == coro::ABI::RetconOnceDynamic) { + Value *TypeId = nullptr; + if (RetconLowering.TypeId) { + if (auto *CAAFI = dyn_cast_or_null(AI)) { + // If malloc type ids are being used, use the callee's type id when + // allocating the callee's frame. + TypeId = CAAFI->getTypeId(); + } else { + TypeId = RetconLowering.TypeId; + } Args.push_back(TypeId); + } } auto *Call = Builder.CreateCall(Alloc, Args); if (ABI == coro::ABI::RetconOnceDynamic) { @@ -583,7 +592,6 @@ void coro::Shape::emitDealloc(IRBuilder<> &Builder, Value *Ptr, case coro::ABI::RetconOnceDynamic: { Function *Dealloc = nullptr; if (isa_and_nonnull(AI)) { - assert(ABI == coro::ABI::RetconOnceDynamic); Dealloc = RetconLowering.DeallocFrame; } else { Dealloc = RetconLowering.Dealloc; @@ -737,6 +745,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 index bcd270b8653b2..6b9977d2d2279 100644 --- 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 index 5d1c3420de16a..a8ff770fc8dde 100644 --- 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)