Skip to content

Commit 35c697c

Browse files
aschwaighoferaschwaighofer
authored
Merge pull request #17026 from aschwaighofer/wip_begin_apply_inlining
 SILInliner: Initial support for begin_apply
2 parents ce7de9b + dfec52b commit 35c697c

File tree

3 files changed

+557
-16
lines changed

3 files changed

+557
-16
lines changed

‎lib/SILOptimizer/Utils/PerformanceInlinerUtils.cpp

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -656,12 +656,7 @@ static bool isCallerAndCalleeLayoutConstraintsCompatible(FullApplySite AI) {
656656
// Returns the callee of an apply_inst if it is basically inlinable.
657657
SILFunction *swift::getEligibleFunction(FullApplySite AI,
658658
InlineSelection WhatToInline) {
659-
// For now, we cannot inline begin_apply at all.
660-
if (isa<BeginApplyInst>(AI))
661-
return nullptr;
662-
663659
SILFunction *Callee = AI.getReferencedFunction();
664-
SILFunction *EligibleCallee = nullptr;
665660

666661
if (!Callee) {
667662
return nullptr;
@@ -776,8 +771,7 @@ SILFunction *swift::getEligibleFunction(FullApplySite AI,
776771
return nullptr;
777772
}
778773

779-
EligibleCallee = Callee;
780-
return EligibleCallee;
774+
return Callee;
781775
}
782776

783777
/// Returns true if the instruction \I has any interesting side effects which

‎lib/SILOptimizer/Utils/SILInliner.cpp

Lines changed: 206 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,181 @@
1818
using namespace swift;
1919

2020
bool SILInliner::canInlineFunction(FullApplySite AI) {
21-
// For now, we cannot inline begin_apply at all.
22-
if (isa<BeginApplyInst>(AI))
23-
return false;
24-
2521
return AI.getFunction() != &Original;
2622
}
2723

24+
/// Utility class for rewiring control-flow of inlined begin_apply functions.
25+
class BeginApplySite {
26+
SmallVector<SILBasicBlock *, 4> ExitingBlocks;
27+
SmallVector<AllocStackInst*, 8> YieldedIndirectValues;
28+
SILLocation Loc;
29+
SILBuilder &Builder;
30+
BeginApplyInst *BeginApply;
31+
SILFunction *F;
32+
EndApplyInst *EndApply = nullptr;
33+
SILBasicBlock *EndApplyBB = nullptr;
34+
SILBasicBlock *EndApplyBBMerge = nullptr;
35+
AbortApplyInst *AbortApply = nullptr;
36+
SILBasicBlock *AbortApplyBB = nullptr;
37+
SILBasicBlock *AbortApplyBBMerge = nullptr;
38+
SILArgument *IntToken = nullptr;
39+
40+
unsigned YieldNum = 0;
41+
SmallVector<SILBasicBlock*, 8> YieldResumes;
42+
SmallVector<SILBasicBlock*, 8> YieldUnwinds;
43+
44+
void
45+
getYieldCaseBBs(SmallVectorImpl<std::pair<SILValue, SILBasicBlock *>> &Result,
46+
SmallVectorImpl<SILBasicBlock *> &Dests) {
47+
unsigned Token = 0;
48+
for (auto *Blk : Dests) {
49+
Result.push_back(std::make_pair(
50+
SILValue(Builder.createIntegerLiteral(
51+
Loc,
52+
SILType::getBuiltinIntegerType(
53+
32, Builder.getFunction().getModule().getASTContext()),
54+
Token++)),
55+
Blk));
56+
}
57+
}
58+
59+
public:
60+
BeginApplySite(BeginApplyInst *BeginApply, SILLocation Loc,
61+
SILBuilder &Builder)
62+
: Loc(Loc), Builder(Builder), BeginApply(BeginApply),
63+
F(BeginApply->getFunction()) {}
64+
65+
static Optional<BeginApplySite> isa(FullApplySite AI, SILLocation Loc,
66+
SILBuilder &Builder) {
67+
auto *BeginApply = dyn_cast<BeginApplyInst>(AI);
68+
if (!BeginApply)
69+
return None;
70+
return BeginApplySite(BeginApply, Loc, Builder);
71+
}
72+
73+
void collectCallerExitingBlocks() {
74+
F->findExitingBlocks(ExitingBlocks);
75+
}
76+
77+
void processApply(SILBasicBlock *ReturnToBB) {
78+
// Handle direct and indirect results.
79+
for (auto YieldedValue : BeginApply->getYieldedValues()) {
80+
// Insert an alloc_stack for indirect results.
81+
if (YieldedValue->getType().isAddress()) {
82+
Builder.setInsertionPoint(F->getEntryBlock()->begin());
83+
auto Addr = Builder.createAllocStack(
84+
Loc, YieldedValue->getType().getObjectType());
85+
YieldedValue->replaceAllUsesWith(Addr);
86+
YieldedIndirectValues.push_back(Addr);
87+
for (auto *Exit : ExitingBlocks) {
88+
Builder.setInsertionPoint(Exit->getTerminator());
89+
Builder.createDeallocStack(Loc, Addr);
90+
}
91+
continue;
92+
}
93+
// Insert a phi for direct results.
94+
auto *RetArg = ReturnToBB->createPHIArgument(YieldedValue->getType(),
95+
ValueOwnershipKind::Owned);
96+
// Replace all uses of the ApplyInst with the new argument.
97+
YieldedValue->replaceAllUsesWith(RetArg);
98+
}
99+
100+
// Add a trailing phi argument for the token integer (tells us which yield
101+
// we came from).
102+
IntToken = ReturnToBB->createPHIArgument(
103+
SILType::getBuiltinIntegerType(32, F->getModule().getASTContext()),
104+
ValueOwnershipKind::Owned);
105+
106+
// Get the end_apply, abort_apply instructions.
107+
auto Token = BeginApply->getTokenResult();
108+
for (auto *TokenUse : Token->getUses()) {
109+
EndApply = dyn_cast<EndApplyInst>(TokenUse->getUser());
110+
if (EndApply)
111+
continue;
112+
AbortApply = cast<AbortApplyInst>(TokenUse->getUser());
113+
}
114+
115+
// Split the basic block before the end/abort_apply. We will insert code
116+
// to jump to the resume/unwind blocks depending on the integer token
117+
// later. And the inlined resume/unwind return blocks will jump back to
118+
// the merge blocks.
119+
EndApplyBB = EndApply->getParent();
120+
EndApplyBBMerge = EndApplyBB->split(SILBasicBlock::iterator(EndApply));
121+
if (AbortApply) {
122+
AbortApplyBB = AbortApply->getParent();
123+
AbortApplyBBMerge =
124+
AbortApplyBB->split(SILBasicBlock::iterator(AbortApply));
125+
}
126+
}
127+
128+
void processTerminator(
129+
TermInst *Terminator, SILBasicBlock *ReturnToBB,
130+
llvm::function_ref<SILBasicBlock *(SILBasicBlock *)> remapBlock,
131+
llvm::function_ref<SILValue(SILValue)> remapValue,
132+
llvm::function_ref<void(TermInst *)> mapTerminator) {
133+
// A yield branches to the begin_apply return block passing the yielded
134+
// results as branch arguments. Collect the yields target block for
135+
// resuming later. Pass an integer token to the begin_apply return block
136+
// to mark the yield we came from.
137+
if (auto *Yield = dyn_cast<YieldInst>(Terminator)) {
138+
YieldResumes.push_back(remapBlock(Yield->getResumeBB()));
139+
YieldUnwinds.push_back(remapBlock(Yield->getUnwindBB()));
140+
auto ContextToken = Builder.createIntegerLiteral(
141+
Loc,
142+
SILType::getBuiltinIntegerType(32, F->getModule().getASTContext()),
143+
YieldNum++);
144+
145+
SmallVector<SILValue, 8> BrResults;
146+
unsigned IndirectIdx = 0;
147+
for (auto CalleeYieldedVal : Yield->getYieldedValues()) {
148+
auto YieldedVal = remapValue(CalleeYieldedVal);
149+
if (YieldedVal->getType().isAddress()) {
150+
auto YieldedDestAddr = YieldedIndirectValues[IndirectIdx++];
151+
Builder.createCopyAddr(Loc, YieldedVal, YieldedDestAddr, IsTake,
152+
IsInitialization);
153+
} else
154+
BrResults.push_back(YieldedVal);
155+
}
156+
BrResults.push_back(SILValue(ContextToken));
157+
Builder.createBranch(Loc, ReturnToBB, BrResults);
158+
return;
159+
}
160+
161+
// Return and unwind terminators branch to the end_apply/abort_apply merge
162+
// block respectively.
163+
if (auto *RI = dyn_cast<ReturnInst>(Terminator)) {
164+
Builder.createBranch(Loc, EndApplyBBMerge);
165+
return;
166+
}
167+
if (auto *Unwind = dyn_cast<UnwindInst>(Terminator)) {
168+
Builder.createBranch(Loc, AbortApplyBBMerge);
169+
return;
170+
}
171+
172+
// Otherwise, we just map the branch instruction.
173+
assert(!::isa<ThrowInst>(Terminator) &&
174+
"Unexpected throw instruction in yield_once function");
175+
mapTerminator(Terminator);
176+
}
177+
178+
void dispatchToResumeUnwindBlocks() {
179+
// Resume edge.
180+
Builder.setInsertionPoint(EndApplyBB);
181+
SmallVector<std::pair<SILValue, SILBasicBlock *>, 8> CaseBBs;
182+
getYieldCaseBBs(CaseBBs, YieldResumes);
183+
Builder.createSwitchValue(Loc, IntToken, nullptr, CaseBBs);
184+
EndApply->eraseFromParent();
185+
// Unwind edge.
186+
if (AbortApplyBB) {
187+
Builder.setInsertionPoint(AbortApplyBB);
188+
SmallVector<std::pair<SILValue, SILBasicBlock *>, 8> CaseBBs;
189+
getYieldCaseBBs(CaseBBs, YieldUnwinds);
190+
Builder.createSwitchValue(Loc, IntToken, nullptr, CaseBBs);
191+
AbortApply->eraseFromParent();
192+
}
193+
}
194+
};
195+
28196
/// \brief Inlines the callee of a given ApplyInst (which must be the value of a
29197
/// FunctionRefInst referencing a function with a known body), into the caller
30198
/// containing the ApplyInst, which must be the same function as provided to the
@@ -114,6 +282,12 @@ void SILInliner::inlineFunction(FullApplySite AI, ArrayRef<SILValue> Args) {
114282
ValueMap.insert(std::make_pair(calleeArg, callArg));
115283
}
116284

285+
// Find the existing blocks. We will need them for inlining the co-routine
286+
// call.
287+
auto BeginApply = BeginApplySite::isa(AI, Loc.getValue(), getBuilder());
288+
if (BeginApply)
289+
BeginApply->collectCallerExitingBlocks();
290+
117291
// Recursively visit callee's BB in depth-first preorder, starting with the
118292
// entry block, cloning all instructions other than terminators.
119293
visitSILBasicBlock(CalleeEntryBB);
@@ -131,6 +305,7 @@ void SILInliner::inlineFunction(FullApplySite AI, ArrayRef<SILValue> Args) {
131305

132306
// If we're inlining into a try_apply, we already have a return-to BB.
133307
SILBasicBlock *ReturnToBB;
308+
134309
if (auto tryAI = dyn_cast<TryApplyInst>(AI)) {
135310
ReturnToBB = tryAI->getNormalBB();
136311

@@ -149,17 +324,35 @@ void SILInliner::inlineFunction(FullApplySite AI, ArrayRef<SILValue> Args) {
149324
SILFunction::iterator(ReturnToBB));
150325

151326
// Create an argument on the return-to BB representing the returned value.
152-
auto apply = cast<ApplyInst>(AI.getInstruction());
153-
auto *RetArg = ReturnToBB->createPHIArgument(apply->getType(),
154-
ValueOwnershipKind::Owned);
155-
// Replace all uses of the ApplyInst with the new argument.
156-
apply->replaceAllUsesWith(RetArg);
327+
if (auto apply = dyn_cast<ApplyInst>(AI.getInstruction())) {
328+
auto *RetArg = ReturnToBB->createPHIArgument(apply->getType(),
329+
ValueOwnershipKind::Owned);
330+
// Replace all uses of the ApplyInst with the new argument.
331+
apply->replaceAllUsesWith(RetArg);
332+
} else {
333+
// Handle begin_apply.
334+
BeginApply->processApply(ReturnToBB);
335+
}
157336
}
158337

159338
// Now iterate over the callee BBs and fix up the terminators.
160339
for (auto BI = BBMap.begin(), BE = BBMap.end(); BI != BE; ++BI) {
161340
getBuilder().setInsertionPoint(BI->second);
162341

342+
// Coroutine terminators need special handling.
343+
if (BeginApply) {
344+
BeginApply->processTerminator(
345+
BI->first->getTerminator(), ReturnToBB,
346+
[=](SILBasicBlock *Block) -> SILBasicBlock * {
347+
return this->remapBasicBlock(Block);
348+
},
349+
[=](SILValue Val) -> SILValue {
350+
return this->remapValue(Val);
351+
},
352+
[=](TermInst *Term) { this->visit(Term); });
353+
continue;
354+
}
355+
163356
// Modify return terminators to branch to the return-to BB, rather than
164357
// trying to clone the ReturnInst.
165358
if (auto *RI = dyn_cast<ReturnInst>(BI->first->getTerminator())) {
@@ -190,6 +383,10 @@ void SILInliner::inlineFunction(FullApplySite AI, ArrayRef<SILValue> Args) {
190383
// but remaps basic blocks and values.
191384
visit(BI->first->getTerminator());
192385
}
386+
387+
// Insert dispatch code at end/abort_apply to the resume/unwind target blocks.
388+
if (BeginApply)
389+
BeginApply->dispatchToResumeUnwindBlocks();
193390
}
194391

195392
SILValue SILInliner::borrowFunctionArgument(SILValue callArg,

0 commit comments

Comments
 (0)