Skip to content

Commit 506473d

Browse files
nate-chandlernate-chandler
committed
[Async CC] Supported partial application.
The majority of support comes in the form of emitting partial application forwarders for partial applications of async functions. Such a partial application forwarder must take an async context which has been partially populated at the apply site. It is responsible for populating it "the rest of the way". To do so, like sync partial application forwarders, it takes a second argument, its context, from which it pulls the additional arguments which were capture at partial_apply time. The size of the async context that is passed to the forwarder, however, can't be known at the apply site by simply looking at the signature of the function to be applied (not even by looking at the size associated with the function in the special async function pointer constant which will soon be emitted). The reason is that there are an unknown (at the apply site) number of additional arguments which will be filled by the partial apply forwarder (and in the case of repeated partial applications, further filled in incrementally at each level). To enable this, there will always be a heap object for thick async functions. These heap objects will always store the size of the async context to be allocated as their first element. (Note that it may be possible to apply the same optimization that was applied for thick sync functions where a single refcounted object could be used as the context; doing so, however, must be made to interact properly with the async context size stored in the heap object.) To continue to allow promoting thin async functions to thick async functions without incurring a thunk, at the apply site, a null-check will be performed on the context pointer. If it is null, then the async context size will be determined based on the signature. (When async function pointers become pointers to a constant with a size i32 and a relative address to the underlying function, the size will be read from that constant.) When it is not-null, the size will be pulled from the first field of the context (which will in that case be cast to <{%swift.refcounted, i32}>). To facilitate sharing code and preserving the original structure of emitPartialApplicationForwarder (which weighed in at roughly 700 lines prior to this change), a new small class hierarchy, descending from PartialApplicationForwarderEmission has been added, with subclasses for the sync and async case. The shuffling of arguments into and out of the final explosion that was being performed in the synchronous case has been preserved there, though the arguments are added and removed through a number of methods on the superclass with more descriptive names. That was necessary to enable the async class to handle these different flavors of parameters correctly. To get some initial test coverage, the preexisting IRGen/partial_apply.sil and IRGen/partial_apply_forwarder.sil tests have been duplicated into the async folder. Those tests cases within these files which happened to have been crashing have each been extracted into its own runnable test that both verifies that the compiler does not crash and also that the partial application forwarder behaves correctly. The FileChecks in these tests are extremely minimal, providing only enough information to be sure that arguments are in fact squeezed into an async context.
1 parent 03f4e40 commit 506473d

20 files changed

+2383
-161
lines changed

‎lib/IRGen/GenCall.cpp

Lines changed: 143 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
#include "swift/AST/ASTContext.h"
2121
#include "swift/AST/GenericEnvironment.h"
2222
#include "swift/Runtime/Config.h"
23+
#include "swift/SIL/SILModule.h"
2324
#include "swift/SIL/SILType.h"
2425
#include "clang/AST/ASTContext.h"
2526
#include "clang/AST/RecordLayout.h"
@@ -135,9 +136,7 @@ AsyncContextLayout irgen::getAsyncContextLayout(
135136
// SelfType self?;
136137
bool hasLocalContextParameter = hasSelfContextParameter(substitutedType);
137138
bool canHaveValidError = substitutedType->hasErrorResult();
138-
bool hasLocalContext = (hasLocalContextParameter || canHaveValidError ||
139-
substitutedType->getRepresentation() ==
140-
SILFunctionTypeRepresentation::Thick);
139+
bool hasLocalContext = (hasLocalContextParameter || canHaveValidError);
141140
SILParameterInfo localContextParameter =
142141
hasLocalContextParameter ? parameters.back() : SILParameterInfo();
143142
if (hasLocalContextParameter) {
@@ -694,6 +693,10 @@ void SignatureExpansion::addAsyncParameters() {
694693
ParamIRTypes.push_back(IGM.SwiftContextPtrTy);
695694
// TODO: Add actor.
696695
// TODO: Add task.
696+
if (FnType->getRepresentation() == SILFunctionTypeRepresentation::Thick) {
697+
IGM.addSwiftSelfAttributes(Attrs, ParamIRTypes.size());
698+
ParamIRTypes.push_back(IGM.RefCountedPtrTy);
699+
}
697700
}
698701

699702
void SignatureExpansion::addCoroutineContextParameter() {
@@ -1740,6 +1743,116 @@ static void externalizeArguments(IRGenFunction &IGF, const Callee &callee,
17401743
Explosion &in, Explosion &out,
17411744
TemporarySet &temporaries, bool isOutlined);
17421745

1746+
llvm::Value *irgen::getDynamicAsyncContextSize(IRGenFunction &IGF,
1747+
AsyncContextLayout layout,
1748+
CanSILFunctionType functionType,
1749+
llvm::Value *thickContext) {
1750+
switch (functionType->getRepresentation()) {
1751+
case SILFunctionTypeRepresentation::Thick: {
1752+
// If the called function is thick, the size of the called function's
1753+
// async context may not be statically knowable.
1754+
//
1755+
// Specifically, if the thick function was produced by a partial_apply,
1756+
// the function which was originally partially applied determines the
1757+
// size of the needed async context. That original function isn't known
1758+
// statically. The dynamic size is available within the context as an
1759+
// i32 at the first index: <{ %swift.refcounted*, /*size*/ i32, ... }>.
1760+
//
1761+
// On the other hand, if the thick function was produced by a
1762+
// thin_to_thick_function, then the context will be nullptr. In that
1763+
// case, the size of the needed async context is known statically to
1764+
// be the size dictated by the function signature.
1765+
//
1766+
// We are currently emitting into some basic block. To handle these two
1767+
// cases, we need to branch based on whether the context is nullptr; each
1768+
// branch must then determine the size in the manner appropriate to it.
1769+
// Finally, both blocks must join back together to make the call:
1770+
//
1771+
// SIL: IR:
1772+
// +-----+ +-------------------------+
1773+
// |.....| |%cond = %ctx == nullptr |
1774+
// |apply| |br %cond, static, dynamic|
1775+
// |.....| +--------/--------------\-+
1776+
// +-----+ / \
1777+
// +-static-------+ +-dynamic----------------------------------------------+
1778+
// |%size = K | |%layout = bitcast %context to <{%swift.context*, i32}>|
1779+
// |br join(%size)| |%size_addr = getelementptr %layout, i32 1, i32 0 |
1780+
// +-----\--------+ |%size = load %size_addr |
1781+
// \ |br join(%size) |
1782+
// \ +------------------------------------------------------+
1783+
// \ /
1784+
// +-join(%size)-----------------------------------------------------------+
1785+
// |%dataAddr = swift_taskAlloc(%task, %size) |
1786+
// |%async_context = bitcast %dataAddr to ASYNC_CONTEXT(static_callee_type)|
1787+
// |... // populate the fields %context with arguments |
1788+
// |call %callee(%async_context, %context) |
1789+
// +-----------------------------------------------------------------------+
1790+
auto *staticSizeBlock = llvm::BasicBlock::Create(IGF.IGM.getLLVMContext());
1791+
auto *dynamicSizeBlock = llvm::BasicBlock::Create(IGF.IGM.getLLVMContext());
1792+
auto *joinBlock = llvm::BasicBlock::Create(IGF.IGM.getLLVMContext());
1793+
1794+
auto hasThickContext =
1795+
IGF.Builder.CreateICmpNE(thickContext, IGF.IGM.RefCountedNull);
1796+
IGF.Builder.CreateCondBr(hasThickContext, dynamicSizeBlock,
1797+
staticSizeBlock);
1798+
1799+
SmallVector<std::pair<llvm::BasicBlock *, llvm::Value *>, 2> phiValues;
1800+
{
1801+
IGF.Builder.emitBlock(staticSizeBlock);
1802+
auto size = getAsyncContextSize(layout);
1803+
auto *sizeValue =
1804+
llvm::ConstantInt::get(IGF.IGM.Int32Ty, size.getValue());
1805+
phiValues.push_back({staticSizeBlock, sizeValue});
1806+
IGF.Builder.CreateBr(joinBlock);
1807+
}
1808+
1809+
{
1810+
IGF.Builder.emitBlock(dynamicSizeBlock);
1811+
SmallVector<const TypeInfo *, 4> argTypeInfos;
1812+
SmallVector<SILType, 4> argValTypes;
1813+
auto int32ASTType =
1814+
BuiltinIntegerType::get(32, IGF.IGM.IRGen.SIL.getASTContext())
1815+
->getCanonicalType();
1816+
auto int32SILType = SILType::getPrimitiveObjectType(int32ASTType);
1817+
const TypeInfo &int32TI = IGF.IGM.getTypeInfo(int32SILType);
1818+
argValTypes.push_back(int32SILType);
1819+
argTypeInfos.push_back(&int32TI);
1820+
HeapLayout layout(IGF.IGM, LayoutStrategy::Optimal, argValTypes,
1821+
argTypeInfos,
1822+
/*typeToFill*/ nullptr, NecessaryBindings());
1823+
auto castThickContext =
1824+
layout.emitCastTo(IGF, thickContext, "context.prefix");
1825+
auto sizeLayout = layout.getElement(0);
1826+
auto sizeAddr = sizeLayout.project(IGF, castThickContext,
1827+
/*NonFixedOffsets*/ llvm::None);
1828+
auto *sizeValue = IGF.Builder.CreateLoad(sizeAddr);
1829+
phiValues.push_back({dynamicSizeBlock, sizeValue});
1830+
IGF.Builder.CreateBr(joinBlock);
1831+
}
1832+
1833+
{
1834+
IGF.Builder.emitBlock(joinBlock);
1835+
auto *phi = IGF.Builder.CreatePHI(IGF.IGM.Int32Ty, phiValues.size());
1836+
for (auto &entry : phiValues) {
1837+
phi->addIncoming(entry.second, entry.first);
1838+
}
1839+
return phi;
1840+
}
1841+
}
1842+
case SILFunctionTypeRepresentation::Thin:
1843+
case SILFunctionTypeRepresentation::CFunctionPointer:
1844+
case SILFunctionTypeRepresentation::Method:
1845+
case SILFunctionTypeRepresentation::ObjCMethod:
1846+
case SILFunctionTypeRepresentation::WitnessMethod:
1847+
case SILFunctionTypeRepresentation::Closure:
1848+
case SILFunctionTypeRepresentation::Block: {
1849+
auto size = getAsyncContextSize(layout);
1850+
auto *sizeValue = llvm::ConstantInt::get(IGF.IGM.Int32Ty, size.getValue());
1851+
return sizeValue;
1852+
}
1853+
}
1854+
}
1855+
17431856
namespace {
17441857

17451858
class SyncCallEmission final : public CallEmission {
@@ -1947,6 +2060,7 @@ class AsyncCallEmission final : public CallEmission {
19472060
Address contextBuffer;
19482061
Size contextSize;
19492062
Address context;
2063+
llvm::Value *thickContext = nullptr;
19502064

19512065
AsyncContextLayout getAsyncContextLayout() {
19522066
return ::getAsyncContextLayout(IGF, getCallee().getOrigFunctionType(),
@@ -1975,9 +2089,14 @@ class AsyncCallEmission final : public CallEmission {
19752089
super::begin();
19762090
assert(!contextBuffer.isValid());
19772091
assert(!context.isValid());
1978-
// Allocate space for the async arguments.
19792092
auto layout = getAsyncContextLayout();
1980-
std::tie(contextBuffer, contextSize) = emitAllocAsyncContext(IGF, layout);
2093+
// Allocate space for the async arguments.
2094+
auto *dynamicContextSize32 = getDynamicAsyncContextSize(
2095+
IGF, layout, CurCallee.getOrigFunctionType(), thickContext);
2096+
auto *dynamicContextSize =
2097+
IGF.Builder.CreateZExt(dynamicContextSize32, IGF.IGM.SizeTy);
2098+
std::tie(contextBuffer, contextSize) = emitAllocAsyncContext(
2099+
IGF, layout, dynamicContextSize, getAsyncContextSize(layout));
19812100
context = layout.emitCastTo(IGF, contextBuffer.getAddress());
19822101
if (layout.canHaveError()) {
19832102
auto fieldLayout = layout.getErrorLayout();
@@ -1993,14 +2112,21 @@ class AsyncCallEmission final : public CallEmission {
19932112
emitDeallocAsyncContext(IGF, contextBuffer, contextSize);
19942113
super::end();
19952114
}
1996-
void setFromCallee() override { super::setFromCallee(); }
2115+
void setFromCallee() override {
2116+
super::setFromCallee();
2117+
thickContext = CurCallee.getSwiftContext();
2118+
}
19972119
SILType getParameterType(unsigned index) override {
19982120
return getAsyncContextLayout().getParameterType(index);
19992121
}
20002122
void setArgs(Explosion &llArgs, bool isOutlined,
20012123
WitnessMetadata *witnessMetadata) override {
20022124
Explosion asyncExplosion;
20032125
asyncExplosion.add(contextBuffer.getAddress());
2126+
if (getCallee().getRepresentation() ==
2127+
SILFunctionTypeRepresentation::Thick) {
2128+
asyncExplosion.add(getCallee().getSwiftContext());
2129+
}
20042130
super::setArgs(asyncExplosion, false, witnessMetadata);
20052131
SILFunctionConventions fnConv(getCallee().getSubstFunctionType(),
20062132
IGF.getSILModule());
@@ -3196,14 +3322,21 @@ void irgen::emitTaskDealloc(IRGenFunction &IGF, Address address,
31963322
llvm::Attribute::ReadNone);
31973323
}
31983324

3325+
std::pair<Address, Size> irgen::emitAllocAsyncContext(IRGenFunction &IGF,
3326+
AsyncContextLayout layout,
3327+
llvm::Value *sizeValue,
3328+
Size sizeLowerBound) {
3329+
auto alignment = getAsyncContextAlignment(IGF.IGM);
3330+
auto address = emitTaskAlloc(IGF, sizeValue, alignment);
3331+
IGF.Builder.CreateLifetimeStart(address, sizeLowerBound);
3332+
return {address, sizeLowerBound};
3333+
}
3334+
31993335
std::pair<Address, Size>
32003336
irgen::emitAllocAsyncContext(IRGenFunction &IGF, AsyncContextLayout layout) {
32013337
auto size = getAsyncContextSize(layout);
32023338
auto *sizeValue = llvm::ConstantInt::get(IGF.IGM.SizeTy, size.getValue());
3203-
auto alignment = getAsyncContextAlignment(IGF.IGM);
3204-
auto address = emitTaskAlloc(IGF, sizeValue, alignment);
3205-
IGF.Builder.CreateLifetimeStart(address, size);
3206-
return {address, size};
3339+
return emitAllocAsyncContext(IGF, layout, sizeValue, size);
32073340
}
32083341

32093342
void irgen::emitDeallocAsyncContext(IRGenFunction &IGF, Address context,

‎lib/IRGen/GenCall.h

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -72,8 +72,6 @@ namespace irgen {
7272
// SwiftPartialFunction * __ptrauth(...) yieldToCaller?;
7373
// SwiftError *errorResult;
7474
// IndirectResultTypes *indirectResults...;
75-
// SelfType self?;
76-
// ArgTypes formalArguments...;
7775
// union {
7876
// struct {
7977
// SwiftPartialFunction * __ptrauth(...) resumeFromYield?;
@@ -173,10 +171,6 @@ namespace irgen {
173171
assert(hasLocalContext());
174172
return getElement(getLocalContextIndex());
175173
}
176-
ParameterConvention getLocalContextConvention() {
177-
assert(hasLocalContext());
178-
return localContextInfo->convention;
179-
}
180174
SILType getLocalContextType() {
181175
assert(hasLocalContext());
182176
return localContextInfo->type;
@@ -230,6 +224,10 @@ namespace irgen {
230224
Optional<ArgumentInfo> localContextInfo);
231225
};
232226

227+
llvm::Value *getDynamicAsyncContextSize(IRGenFunction &IGF,
228+
AsyncContextLayout layout,
229+
CanSILFunctionType functionType,
230+
llvm::Value *thickContext);
233231
AsyncContextLayout getAsyncContextLayout(IRGenFunction &IGF,
234232
SILFunction *function);
235233

@@ -320,6 +318,16 @@ namespace irgen {
320318
Address emitTaskAlloc(IRGenFunction &IGF, llvm::Value *size,
321319
Alignment alignment);
322320
void emitTaskDealloc(IRGenFunction &IGF, Address address, llvm::Value *size);
321+
/// Allocate task local storage for the specified layout but using the
322+
/// provided dynamic size. Allowing the size to be specified dynamically is
323+
/// necessary for applies of thick functions the sizes of whose async contexts
324+
/// are dependent on the underlying, already partially applied, called
325+
/// function. The provided sizeLowerBound will be used to track the lifetime
326+
/// of the allocation that is known statically.
327+
std::pair<Address, Size> emitAllocAsyncContext(IRGenFunction &IGF,
328+
AsyncContextLayout layout,
329+
llvm::Value *sizeValue,
330+
Size sizeLowerBound);
323331
std::pair<Address, Size> emitAllocAsyncContext(IRGenFunction &IGF,
324332
AsyncContextLayout layout);
325333
void emitDeallocAsyncContext(IRGenFunction &IGF, Address context, Size size);

0 commit comments

Comments
 (0)