Skip to content

Commit b69976b

Browse files
Merge pull request #37040 from aschwaighofer/fix_asyn_dynamic_replacement_5.5
[5.5] Runtime/IRGen: Fix async dynamic replacements
2 parents 8990114 + 34196d5 commit b69976b

File tree

12 files changed

+400
-175
lines changed

12 files changed

+400
-175
lines changed

include/swift/ABI/Metadata.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4897,6 +4897,9 @@ struct DynamicReplacementKey {
48974897
uint16_t getExtraDiscriminator() const {
48984898
return flags & 0x0000FFFF;
48994899
}
4900+
bool isAsync() const {
4901+
return ((flags >> 16 ) & 0x1);
4902+
}
49004903
};
49014904

49024905
/// A record describing a dynamic function replacement.

lib/IRGen/Callee.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,13 @@ namespace irgen {
135135
assert(isSigned());
136136
return Discriminator;
137137
}
138+
PointerAuthInfo getCorrespondingCodeAuthInfo() const {
139+
if (auto authInfo = *this) {
140+
return PointerAuthInfo(authInfo.getCorrespondingCodeKey(),
141+
authInfo.getDiscriminator());
142+
}
143+
return *this;
144+
}
138145

139146
/// Are the auth infos obviously the same?
140147
friend bool operator==(const PointerAuthInfo &lhs,

lib/IRGen/GenCall.cpp

Lines changed: 78 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -219,19 +219,26 @@ llvm::Value *IRGenFunction::getAsyncContext() {
219219
return Builder.CreateLoad(asyncContextLocation);
220220
}
221221

222-
llvm::CallInst *
223-
IRGenFunction::emitSuspendAsyncCall(unsigned asyncContextIndex,
224-
llvm::StructType *resultTy,
225-
ArrayRef<llvm::Value *> args) {
226-
auto *id = Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_suspend_async,
227-
{resultTy}, args);
228-
llvm::Value *calleeContext =
229-
Builder.CreateExtractValue(id, asyncContextIndex);
230-
calleeContext = Builder.CreateBitOrPointerCast(calleeContext, IGM.Int8PtrTy);
231-
llvm::Constant *projectFn = cast<llvm::Constant>(args[2])->stripPointerCasts();
232-
llvm::Value *context = Builder.CreateCall(projectFn, {calleeContext});
222+
void IRGenFunction::storeCurrentAsyncContext(llvm::Value *context) {
233223
context = Builder.CreateBitCast(context, IGM.SwiftContextPtrTy);
234224
Builder.CreateStore(context, asyncContextLocation);
225+
}
226+
227+
llvm::CallInst *IRGenFunction::emitSuspendAsyncCall(
228+
unsigned asyncContextIndex, llvm::StructType *resultTy,
229+
ArrayRef<llvm::Value *> args, bool restoreCurrentContext) {
230+
auto *id = Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_suspend_async,
231+
{resultTy}, args);
232+
if (restoreCurrentContext) {
233+
llvm::Value *calleeContext =
234+
Builder.CreateExtractValue(id, asyncContextIndex);
235+
calleeContext =
236+
Builder.CreateBitOrPointerCast(calleeContext, IGM.Int8PtrTy);
237+
llvm::Constant *projectFn =
238+
cast<llvm::Constant>(args[2])->stripPointerCasts();
239+
llvm::Value *context = Builder.CreateCall(projectFn, {calleeContext});
240+
storeCurrentAsyncContext(context);
241+
}
235242

236243
return id;
237244
}
@@ -2023,10 +2030,9 @@ std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
20232030
Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
20242031
/*expectedType*/ functionPointer.getFunctionType()->getPointerTo());
20252032
}
2026-
if (auto authInfo = functionPointer.getAuthInfo()) {
2027-
auto newAuthInfo = PointerAuthInfo(authInfo.getCorrespondingCodeKey(),
2028-
authInfo.getDiscriminator());
2029-
fn = emitPointerAuthSign(IGF, fn, newAuthInfo);
2033+
if (auto authInfo =
2034+
functionPointer.getAuthInfo().getCorrespondingCodeAuthInfo()) {
2035+
fn = emitPointerAuthSign(IGF, fn, authInfo);
20302036
}
20312037
}
20322038
llvm::Value *size = nullptr;
@@ -2354,13 +2360,11 @@ class AsyncCallEmission final : public CallEmission {
23542360
}
23552361

23562362
FunctionPointer getCalleeFunctionPointer() override {
2357-
PointerAuthInfo newAuthInfo;
2358-
if (auto authInfo = CurCallee.getFunctionPointer().getAuthInfo()) {
2359-
newAuthInfo = PointerAuthInfo(authInfo.getCorrespondingCodeKey(),
2360-
authInfo.getDiscriminator());
2361-
}
2363+
PointerAuthInfo codeAuthInfo = CurCallee.getFunctionPointer()
2364+
.getAuthInfo()
2365+
.getCorrespondingCodeAuthInfo();
23622366
return FunctionPointer(
2363-
FunctionPointer::Kind::Function, calleeFunction, newAuthInfo,
2367+
FunctionPointer::Kind::Function, calleeFunction, codeAuthInfo,
23642368
Signature::forAsyncAwait(IGF.IGM, getCallee().getOrigFunctionType()));
23652369
}
23662370

@@ -2571,20 +2575,6 @@ class AsyncCallEmission final : public CallEmission {
25712575
return IGF.getCalleeErrorResultSlot(errorType);
25722576
}
25732577

2574-
FunctionPointer getFunctionPointerForDispatchCall(const FunctionPointer &fn) {
2575-
auto &IGM = IGF.IGM;
2576-
// Strip off the return type. The original function pointer signature
2577-
// captured both the entry point type and the resume function type.
2578-
auto *fnTy = llvm::FunctionType::get(
2579-
IGM.VoidTy, fn.getSignature().getType()->params(), false /*vaargs*/);
2580-
auto signature =
2581-
Signature(fnTy, fn.getSignature().getAttributes(), IGM.SwiftAsyncCC);
2582-
auto fnPtr =
2583-
FunctionPointer(FunctionPointer::Kind::Function, fn.getRawPointer(),
2584-
fn.getAuthInfo(), signature);
2585-
return fnPtr;
2586-
}
2587-
25882578
llvm::CallInst *createCall(const FunctionPointer &fn,
25892579
ArrayRef<llvm::Value *> args) override {
25902580
auto &IGM = IGF.IGM;
@@ -2604,8 +2594,8 @@ class AsyncCallEmission final : public CallEmission {
26042594
auto resumeProjFn = IGF.getOrCreateResumePrjFn();
26052595
arguments.push_back(
26062596
Builder.CreateBitOrPointerCast(resumeProjFn, IGM.Int8PtrTy));
2607-
auto dispatchFn =
2608-
IGF.createAsyncDispatchFn(getFunctionPointerForDispatchCall(fn), args);
2597+
auto dispatchFn = IGF.createAsyncDispatchFn(
2598+
getFunctionPointerForDispatchCall(IGM, fn), args);
26092599
arguments.push_back(
26102600
Builder.CreateBitOrPointerCast(dispatchFn, IGM.Int8PtrTy));
26112601
arguments.push_back(
@@ -4731,10 +4721,8 @@ llvm::Value *FunctionPointer::getPointer(IRGenFunction &IGF) const {
47314721
auto *result = IGF.emitLoadOfRelativePointer(
47324722
Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
47334723
/*expectedType*/ getFunctionType()->getPointerTo());
4734-
if (auto authInfo = AuthInfo) {
4735-
auto newAuthInfo = PointerAuthInfo(authInfo.getCorrespondingCodeKey(),
4736-
authInfo.getDiscriminator());
4737-
result = emitPointerAuthSign(IGF, result, newAuthInfo);
4724+
if (auto codeAuthInfo = AuthInfo.getCorrespondingCodeAuthInfo()) {
4725+
result = emitPointerAuthSign(IGF, result, codeAuthInfo);
47384726
}
47394727
return result;
47404728
}
@@ -4775,11 +4763,7 @@ FunctionPointer FunctionPointer::getAsFunction(IRGenFunction &IGF) const {
47754763
case FunctionPointer::BasicKind::Function:
47764764
return *this;
47774765
case FunctionPointer::BasicKind::AsyncFunctionPointer: {
4778-
auto authInfo = AuthInfo;
4779-
if (authInfo) {
4780-
authInfo = PointerAuthInfo(AuthInfo.getCorrespondingCodeKey(),
4781-
AuthInfo.getDiscriminator());
4782-
}
4766+
auto authInfo = AuthInfo.getCorrespondingCodeAuthInfo();
47834767
return FunctionPointer(Kind::Function, getPointer(IGF), authInfo, Sig);
47844768
}
47854769
}
@@ -4924,3 +4908,51 @@ Address irgen::emitAutoDiffAllocateSubcontext(
49244908
call->setCallingConv(IGF.IGM.SwiftCC);
49254909
return Address(call, IGF.IGM.getPointerAlignment());
49264910
}
4911+
4912+
FunctionPointer
4913+
irgen::getFunctionPointerForDispatchCall(IRGenModule &IGM,
4914+
const FunctionPointer &fn) {
4915+
// Strip off the return type. The original function pointer signature
4916+
// captured both the entry point type and the resume function type.
4917+
auto *fnTy = llvm::FunctionType::get(
4918+
IGM.VoidTy, fn.getSignature().getType()->params(), false /*vaargs*/);
4919+
auto signature =
4920+
Signature(fnTy, fn.getSignature().getAttributes(), IGM.SwiftAsyncCC);
4921+
auto fnPtr = FunctionPointer(FunctionPointer::Kind::Function,
4922+
fn.getRawPointer(), fn.getAuthInfo(), signature);
4923+
return fnPtr;
4924+
}
4925+
4926+
void irgen::forwardAsyncCallResult(IRGenFunction &IGF,
4927+
CanSILFunctionType fnType,
4928+
AsyncContextLayout &layout,
4929+
llvm::CallInst *call) {
4930+
auto &IGM = IGF.IGM;
4931+
auto numAsyncContextParams =
4932+
Signature::forAsyncReturn(IGM, fnType).getAsyncContextIndex() + 1;
4933+
llvm::Value *result = call;
4934+
auto *suspendResultTy = cast<llvm::StructType>(result->getType());
4935+
Explosion resultExplosion;
4936+
Explosion errorExplosion;
4937+
auto hasError = fnType->hasErrorResult();
4938+
Optional<ArrayRef<llvm::Value *>> nativeResults = llvm::None;
4939+
SmallVector<llvm::Value *, 16> nativeResultsStorage;
4940+
4941+
if (suspendResultTy->getNumElements() == numAsyncContextParams) {
4942+
// no result to forward.
4943+
assert(!hasError);
4944+
} else {
4945+
auto &Builder = IGF.Builder;
4946+
auto resultTys =
4947+
makeArrayRef(suspendResultTy->element_begin() + numAsyncContextParams,
4948+
suspendResultTy->element_end());
4949+
4950+
for (unsigned i = 0, e = resultTys.size(); i != e; ++i) {
4951+
llvm::Value *elt =
4952+
Builder.CreateExtractValue(result, numAsyncContextParams + i);
4953+
nativeResultsStorage.push_back(elt);
4954+
}
4955+
nativeResults = nativeResultsStorage;
4956+
}
4957+
emitAsyncReturn(IGF, layout, fnType, nativeResults);
4958+
}

lib/IRGen/GenCall.h

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,12 @@ namespace irgen {
272272
IRGenFunction &IGF, Address context);
273273
Address emitAutoDiffAllocateSubcontext(
274274
IRGenFunction &IGF, Address context, llvm::Value *size);
275+
276+
FunctionPointer getFunctionPointerForDispatchCall(IRGenModule &IGM,
277+
const FunctionPointer &fn);
278+
void forwardAsyncCallResult(IRGenFunction &IGF, CanSILFunctionType fnType,
279+
AsyncContextLayout &layout, llvm::CallInst *call);
280+
275281
} // end namespace irgen
276282
} // end namespace swift
277283

0 commit comments

Comments
 (0)