[closure-spec] Refactor closure spec so that we can in a subsequent commit handle both partial_apply and thin_to_thick_function.

Since this is just a refactoring, it should be NFC.

This is to fix a dumb bug where we will propagate closures if the
closure captures variables (i.e. partial_apply), but we will not
propagate closures that do not capture any values (i.e.
thin_to_thick_function).

  a.map {
    return $0 + v
  }

But this will not be:

  a.map {
    return $0 + 1
  }

This will also give us the flexibility in the future to support the
propagation of function pointers (at which point we should probably
rename this pass).

Swift SVN r24759
This commit is contained in:
Michael Gottesman
2015-01-27 21:34:30 +00:00
parent bf3c92c743
commit e06dc22976

View File

@@ -32,11 +32,13 @@ STATISTIC(NumClosureSpecialized,
"Number of functions with closures specialized"); "Number of functions with closures specialized");
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// Closure Spec Cloner // Closure Spec Cloner Interface
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
namespace { namespace {
class CallSiteDescriptor;
/// \brief A SILCloner subclass which clones a function that takes a closure /// \brief A SILCloner subclass which clones a function that takes a closure
/// argument. We update the parameter list to remove the parameter for the /// argument. We update the parameter list to remove the parameter for the
/// closure argument and to append the variables captured in the closure. /// closure argument and to append the variables captured in the closure.
@@ -49,33 +51,158 @@ public:
friend class SILVisitor<ClosureSpecCloner>; friend class SILVisitor<ClosureSpecCloner>;
friend class SILCloner<ClosureSpecCloner>; friend class SILCloner<ClosureSpecCloner>;
ClosureSpecCloner(SILFunction *PAIUser, PartialApplyInst *PAI, ClosureSpecCloner(const CallSiteDescriptor &CallSiteDesc,
unsigned ClosureIndex, StringRef ClonedName) StringRef ClonedName)
: SuperTy(*initCloned(PAIUser, PAI, ClosureIndex, ClonedName)), : SuperTy(*initCloned(CallSiteDesc, ClonedName)),
PAIUser(PAIUser), ClosureIndex(ClosureIndex), PAI(PAI) { CallSiteDesc(CallSiteDesc) {}
}
void populateCloned(); void populateCloned();
SILFunction *getCloned() { return &getBuilder().getFunction(); } SILFunction *getCloned() { return &getBuilder().getFunction(); }
static SILFunction *cloneFunction(SILFunction *F, PartialApplyInst *PAI, static SILFunction *cloneFunction(const CallSiteDescriptor &CallSiteDesc,
unsigned ClosureIndex, StringRef NewName) { StringRef NewName) {
ClosureSpecCloner C(F, PAI, ClosureIndex, NewName); ClosureSpecCloner C(CallSiteDesc, NewName);
C.populateCloned(); C.populateCloned();
++NumClosureSpecialized; ++NumClosureSpecialized;
return C.getCloned(); return C.getCloned();
}; };
private: private:
static SILFunction *initCloned(SILFunction *PAIUser, PartialApplyInst *PAI, static SILFunction *initCloned(const CallSiteDescriptor &CallSiteDesc,
unsigned ClosureIndex, StringRef ClonedName); StringRef ClonedName);
SILFunction *PAIUser; const CallSiteDescriptor &CallSiteDesc;
unsigned ClosureIndex;
PartialApplyInst *PAI;
}; };
} // end anonymous namespace } // end anonymous namespace
//===----------------------------------------------------------------------===//
// Arg Spec Descriptor
//===----------------------------------------------------------------------===//
namespace {
class CallSiteDescriptor {
SILInstruction *Closure;
ApplyInst *AI;
unsigned ClosureIndex;
public:
CallSiteDescriptor(SILInstruction *ClosureInst, ApplyInst *AI,
unsigned ClosureIndex)
: Closure(ClosureInst), AI(AI), ClosureIndex(ClosureIndex) {}
SILFunction *getApplyCallee() const {
return cast<FunctionRefInst>(AI->getCallee())->getReferencedFunction();
}
SILFunction *getClosureCallee() const {
auto *PAI = cast<PartialApplyInst>(Closure);
return cast<FunctionRefInst>(PAI->getCallee())->getReferencedFunction();
}
unsigned getClosureIndex() const { return ClosureIndex; }
SILInstruction *
createNewClosure(SILBuilder &B, SILValue V,
llvm::SmallVectorImpl<SILValue> &Args) const {
return B.createPartialApply(Closure->getLoc(), V, V.getType(), {}, Args,
Closure->getType(0));
}
ApplyInst *getApplyInst() const { return AI; }
void specializeClosure() const;
void createName(llvm::SmallString<64> &NewName) const;
OperandValueArrayRef getArguments() const {
if (auto *PAI = dyn_cast<PartialApplyInst>(Closure))
return PAI->getArguments();
return OperandValueArrayRef(ArrayRef<Operand>());
}
SILInstruction *getClosure() const { return Closure; }
unsigned getNumArguments() const {
auto *PAI = cast<PartialApplyInst>(Closure);
return PAI->getNumArguments();
}
SILLocation getLoc() const { return Closure->getLoc(); }
SILModule &getModule() const { return AI->getModule(); }
};
} // end anonymous namespace
/// Update the callsite to pass in the correct arguments.
static void rewriteApplyInst(const CallSiteDescriptor &CSDesc,
SILFunction *NewF) {
ApplyInst *AI = CSDesc.getApplyInst();
SILBuilderWithScope<2> Builder(AI);
FunctionRefInst *FRI = Builder.createFunctionRef(AI->getLoc(), NewF);
// Create the args for the new apply by removing the closure argument and
// appending the captured argument.
llvm::SmallVector<SILValue, 8> NewArgs;
unsigned Index = 0;
for (auto Arg : AI->getArguments()) {
if (Index != CSDesc.getClosureIndex())
NewArgs.push_back(Arg);
Index++;
}
for (auto Arg : CSDesc.getArguments())
NewArgs.push_back(Arg);
SILType LoweredType = NewF->getLoweredType();
SILType ResultType = LoweredType.castTo<SILFunctionType>()->getSILResult();
ApplyInst *NewAI = Builder.createApply(AI->getLoc(), FRI, LoweredType,
ResultType, ArrayRef<Substitution>(),
NewArgs, NewF->isTransparent());
// Replace all uses of the old apply with the new apply.
AI->replaceAllUsesWith(NewAI);
// Erase the old apply.
AI->eraseFromParent();
SILInstruction *Closure = CSDesc.getClosure();
if (Closure->use_empty())
Closure->eraseFromParent();
}
void CallSiteDescriptor::createName(llvm::SmallString<64> &NewName) const {
llvm::raw_svector_ostream buffer(NewName);
Mangle::Mangler M(buffer);
auto P = Mangle::SpecializationPass::ClosureSpecializer;
Mangle::FunctionSignatureSpecializationMangler FSSM(P, M, getApplyCallee());
auto *PAI = cast<PartialApplyInst>(Closure);
FSSM.setArgumentClosureProp(getClosureIndex(), PAI);
FSSM.mangle();
}
void CallSiteDescriptor::specializeClosure() const {
llvm::SmallString<64> NewFName;
createName(NewFName);
DEBUG(llvm::dbgs() << " Perform optimizations with new name " << NewFName
<< '\n');
// Then see if we already have a specialized version of this function in our
// module.
SILFunction *NewF = getModule().lookUpFunction(NewFName);
// If not, create a specialized version of ApplyCallee calling the closure
// directly.
if (!NewF)
NewF = ClosureSpecCloner::cloneFunction(*this, NewFName);
// Rewrite the call
rewriteApplyInst(*this, NewF);
}
//===----------------------------------------------------------------------===//
// Closure Spec Cloner Implementation
//===----------------------------------------------------------------------===//
/// In this function we create the actual cloned function and its proper cloned /// In this function we create the actual cloned function and its proper cloned
/// type. But we do not create any body. This implies that the creation of the /// type. But we do not create any body. This implies that the creation of the
/// actual arguments in the function is in populateCloned. /// actual arguments in the function is in populateCloned.
@@ -85,19 +212,20 @@ private:
/// \arg ClosureIndex The index of the partial apply in PAUser's function /// \arg ClosureIndex The index of the partial apply in PAUser's function
/// signature. /// signature.
/// \arg ClonedName The name of the cloned function that we will create. /// \arg ClonedName The name of the cloned function that we will create.
SILFunction *ClosureSpecCloner::initCloned(SILFunction *PAIUser, SILFunction *
PartialApplyInst *PAI, ClosureSpecCloner::initCloned(const CallSiteDescriptor &CallSiteDesc,
unsigned ClosureIndex,
StringRef ClonedName) { StringRef ClonedName) {
SILFunction *ClosureUser = CallSiteDesc.getApplyCallee();
// This is the list of new interface parameters of the cloned function. // This is the list of new interface parameters of the cloned function.
llvm::SmallVector<SILParameterInfo, 4> NewParameterInfoList; llvm::SmallVector<SILParameterInfo, 4> NewParameterInfoList;
// First add to NewParameterInfoList all of the SILParameterInfo in the // First add to NewParameterInfoList all of the SILParameterInfo in the
// original function except for the closure. // original function except for the closure.
CanSILFunctionType PAIUserFunTy = PAIUser->getLoweredFunctionType(); CanSILFunctionType ClosureUserFunTy = ClosureUser->getLoweredFunctionType();
unsigned Index = 0; unsigned Index = 0;
for (auto &param : PAIUserFunTy->getParameters()) { for (auto &param : ClosureUserFunTy->getParameters()) {
if (Index != ClosureIndex) if (Index != CallSiteDesc.getClosureIndex())
NewParameterInfoList.push_back(param); NewParameterInfoList.push_back(param);
++Index; ++Index;
} }
@@ -105,33 +233,30 @@ SILFunction *ClosureSpecCloner::initCloned(SILFunction *PAIUser,
// Then add any arguments that are captured in the closure to the function's // Then add any arguments that are captured in the closure to the function's
// argument type. Since they are captured, we need to pass them directly into // argument type. Since they are captured, we need to pass them directly into
// the new specialized function. // the new specialized function.
auto *ClosedOverFunFRI = cast<FunctionRefInst>(PAI->getCallee()); SILFunction *ClosedOverFun = CallSiteDesc.getClosureCallee();
auto ClosedOverFunTy = ClosedOverFunFRI->getFunctionType(); CanSILFunctionType ClosedOverFunTy = ClosedOverFun->getLoweredFunctionType();
// Captured parameters are always appended to the function signature. So grab the // Captured parameters are always appended to the function signature. So grab the
unsigned NumTotalParams = ClosedOverFunTy->getParameters().size(); unsigned NumTotalParams = ClosedOverFunTy->getParameters().size();
unsigned NumNotCaptured = NumTotalParams - PAI->getNumArguments(); unsigned NumNotCaptured = NumTotalParams - CallSiteDesc.getNumArguments();
for (auto &PInfo : ClosedOverFunTy->getParameters().slice(NumNotCaptured)) for (auto &PInfo : ClosedOverFunTy->getParameters().slice(NumNotCaptured))
NewParameterInfoList.push_back(PInfo); NewParameterInfoList.push_back(PInfo);
SILModule &M = PAIUser->getModule(); SILModule &M = ClosureUser->getModule();
auto ClonedTy = auto ClonedTy = SILFunctionType::get(
SILFunctionType::get(PAIUserFunTy->getGenericSignature(), ClosureUserFunTy->getGenericSignature(), ClosureUserFunTy->getExtInfo(),
PAIUserFunTy->getExtInfo(), ClosureUserFunTy->getCalleeConvention(), NewParameterInfoList,
PAIUserFunTy->getCalleeConvention(), ClosureUserFunTy->getResult(), M.getASTContext());
NewParameterInfoList,
PAIUserFunTy->getResult(),
M.getASTContext());
// We make this function bare so we don't have to worry about decls in the // We make this function bare so we don't have to worry about decls in the
// SILArgument. // SILArgument.
auto Fn = SILFunction::create( auto Fn = SILFunction::create(
M, PAIUser->getLinkage(), ClonedName, ClonedTy, M, ClosureUser->getLinkage(), ClonedName, ClonedTy,
PAIUser->getContextGenericParams(), PAIUser->getLocation(), IsBare, ClosureUser->getContextGenericParams(), ClosureUser->getLocation(),
PAIUser->isTransparent(), PAIUser->isFragile(), IsBare, ClosureUser->isTransparent(), ClosureUser->isFragile(),
PAIUser->getClassVisibility(), PAIUser->getInlineStrategy(), ClosureUser->getClassVisibility(), ClosureUser->getInlineStrategy(),
PAIUser->getEffectsKind(), PAIUser, PAIUser->getDebugScope()); ClosureUser->getEffectsKind(), ClosureUser, ClosureUser->getDebugScope());
Fn->setSemanticsAttr(PAIUser->getSemanticsAttr()); Fn->setSemanticsAttr(ClosureUser->getSemanticsAttr());
return Fn; return Fn;
} }
@@ -141,15 +266,17 @@ void ClosureSpecCloner::populateCloned() {
SILFunction *Cloned = getCloned(); SILFunction *Cloned = getCloned();
SILModule &M = Cloned->getModule(); SILModule &M = Cloned->getModule();
SILFunction *ClosureUser = CallSiteDesc.getApplyCallee();
// Create arguments for the entry block. // Create arguments for the entry block.
SILBasicBlock *PAIUserEntryBB = PAIUser->begin(); SILBasicBlock *ClosureUserEntryBB = ClosureUser->begin();
SILBasicBlock *ClonedEntryBB = new (M) SILBasicBlock(Cloned); SILBasicBlock *ClonedEntryBB = new (M) SILBasicBlock(Cloned);
// Remove the closure argument. // Remove the closure argument.
SILArgument *ClosureArg = nullptr; SILArgument *ClosureArg = nullptr;
for (size_t i = 0, e = PAIUserEntryBB->bbarg_size(); i != e; ++i) { for (size_t i = 0, e = ClosureUserEntryBB->bbarg_size(); i != e; ++i) {
SILArgument *Arg = PAIUserEntryBB->getBBArg(i); SILArgument *Arg = ClosureUserEntryBB->getBBArg(i);
if (i == ClosureIndex) { if (i == CallSiteDesc.getClosureIndex()) {
ClosureArg = Arg; ClosureArg = Arg;
continue; continue;
} }
@@ -168,12 +295,10 @@ void ClosureSpecCloner::populateCloned() {
// definition is nothing in the partial apply user function that references // definition is nothing in the partial apply user function that references
// such arguments. After this pass is done the only thing that will reference // such arguments. After this pass is done the only thing that will reference
// the arguments is the partial apply that we will create. // the arguments is the partial apply that we will create.
SILFunction *ClosedOverFun = CallSiteDesc.getClosureCallee();
auto *ClosedOverFunFRI = cast<FunctionRefInst>(PAI->getCallee()); CanSILFunctionType ClosedOverFunTy = ClosedOverFun->getLoweredFunctionType();
auto *ClosedOverFun = ClosedOverFunFRI->getReferencedFunction();
auto ClosedOverFunTy = ClosedOverFunFRI->getFunctionType();
unsigned NumTotalParams = ClosedOverFunTy->getParameters().size(); unsigned NumTotalParams = ClosedOverFunTy->getParameters().size();
unsigned NumNotCaptured = NumTotalParams - PAI->getNumArguments(); unsigned NumNotCaptured = NumTotalParams - CallSiteDesc.getNumArguments();
llvm::SmallVector<SILValue, 4> NewPAIArgs; llvm::SmallVector<SILValue, 4> NewPAIArgs;
for (auto &PInfo : ClosedOverFunTy->getParameters().slice(NumNotCaptured)) { for (auto &PInfo : ClosedOverFunTy->getParameters().slice(NumNotCaptured)) {
SILValue MappedValue = SILValue MappedValue =
@@ -184,17 +309,16 @@ void ClosureSpecCloner::populateCloned() {
getBuilder().setInsertionPoint(ClonedEntryBB); getBuilder().setInsertionPoint(ClonedEntryBB);
// Clone FRI and PAI, and replace usage of the removed closure argument // Clone FRI and PAI, and replace usage of the removed closure argument
// with result of cloned PAI. // with result of cloned PAI.
SILValue FnVal = getBuilder().createFunctionRef(ClosedOverFunFRI->getLoc(), SILValue FnVal =
ClosedOverFun); getBuilder().createFunctionRef(CallSiteDesc.getLoc(), ClosedOverFun);
auto *NewPAI = getBuilder().createPartialApply(PAI->getLoc(), FnVal, auto *NewClosure =
FnVal.getType(), {}, CallSiteDesc.createNewClosure(getBuilder(), FnVal, NewPAIArgs);
NewPAIArgs, PAI->getType()); ValueMap.insert(std::make_pair(ClosureArg, SILValue(NewClosure)));
ValueMap.insert(std::make_pair(ClosureArg, SILValue(NewPAI, 0)));
BBMap.insert(std::make_pair(PAIUserEntryBB, ClonedEntryBB)); BBMap.insert(std::make_pair(ClosureUserEntryBB, ClonedEntryBB));
// Recursively visit original BBs in depth-first preorder, starting with the // Recursively visit original BBs in depth-first preorder, starting with the
// entry block, cloning all instructions other than terminators. // entry block, cloning all instructions other than terminators.
visitSILBasicBlock(PAIUserEntryBB); visitSILBasicBlock(ClosureUserEntryBB);
// Now iterate over the BBs and fix up the terminators. // Now iterate over the BBs and fix up the terminators.
for (auto BI = BBMap.begin(), BE = BBMap.end(); BI != BE; ++BI) { for (auto BI = BBMap.begin(), BE = BBMap.end(); BI != BE; ++BI) {
@@ -203,56 +327,6 @@ void ClosureSpecCloner::populateCloned() {
} }
} }
//===----------------------------------------------------------------------===//
// Arg Spec Descriptor
//===----------------------------------------------------------------------===//
namespace {
struct ArgDescriptor {
PartialApplyInst *PAI;
ApplyInst *AI;
unsigned ClosureIndex;
ArgDescriptor(PartialApplyInst *PAI, ApplyInst *AI,
unsigned ClosureIndex) :
PAI(PAI), AI(AI), ClosureIndex(ClosureIndex) {
}
};
} // end anonymous namespace
/// Update the callsite to pass in the correct arguments.
static void rewriteApplyInst(ArgDescriptor &AD, SILFunction *NewF) {
SILBuilderWithScope<2> Builder(AD.AI);
FunctionRefInst *FRI = Builder.createFunctionRef(AD.AI->getLoc(), NewF);
// Create the args for the new apply by removing the closure argument and
// appending the captured argument.
llvm::SmallVector<SILValue, 8> NewArgs;
unsigned Index = 0;
for (auto Arg : AD.AI->getArguments()) {
if (Index != AD.ClosureIndex)
NewArgs.push_back(Arg);
Index++;
}
for (auto Arg : AD.PAI->getArguments())
NewArgs.push_back(Arg);
SILType LoweredType = NewF->getLoweredType();
SILType ResultType = LoweredType.castTo<SILFunctionType>()->getSILResult();
ApplyInst *NewAI = Builder.createApply(AD.AI->getLoc(), FRI, LoweredType,
ResultType, ArrayRef<Substitution>(),
NewArgs, NewF->isTransparent());
// Replace all uses of the old apply with the new apply.
AD.AI->replaceAllUsesWith(NewAI);
// Erase the old apply.
AD.AI->eraseFromParent();
if (AD.PAI->use_empty())
AD.PAI->eraseFromParent();
}
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// Closure Specializer // Closure Specializer
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
@@ -267,28 +341,15 @@ struct ClosureSpecializer {
} }
void gatherCallSites(SILFunction *Caller, void gatherCallSites(SILFunction *Caller,
llvm::SmallVectorImpl<ArgDescriptor> &CallSites, llvm::SmallVectorImpl<CallSiteDescriptor> &CallSites,
llvm::SmallPtrSet<ApplyInst *, 4> &MultipleClosureAI); llvm::SmallPtrSet<ApplyInst *, 4> &MultipleClosureAI);
bool specialize(SILFunction *Caller); bool specialize(SILFunction *Caller);
}; };
} // end anonymous namespace } // end anonymous namespace
static void createName(SILFunction *Callee, SILArgument *Arg, void ClosureSpecializer::gatherCallSites(
PartialApplyInst *PAI, SILFunction *Caller, llvm::SmallVectorImpl<CallSiteDescriptor> &CallSites,
llvm::SmallString<64> &Name) {
llvm::raw_svector_ostream buffer(Name);
Mangle::Mangler M(buffer);
auto P = Mangle::SpecializationPass::ClosureSpecializer;
Mangle::FunctionSignatureSpecializationMangler FSSM(P, M, Callee);
FSSM.setArgumentClosureProp(Arg->getIndex(), PAI);
FSSM.mangle();
}
void
ClosureSpecializer::
gatherCallSites(SILFunction *Caller,
llvm::SmallVectorImpl<ArgDescriptor> &CallSites,
llvm::SmallPtrSet<ApplyInst *, 4> &MultipleClosureAI) { llvm::SmallPtrSet<ApplyInst *, 4> &MultipleClosureAI) {
// A set of apply inst that we have associated with a closure. We use this to // A set of apply inst that we have associated with a closure. We use this to
@@ -360,9 +421,9 @@ gatherCallSites(SILFunction *Caller,
if (!PAIIndex.hasValue()) if (!PAIIndex.hasValue())
continue; continue;
// Now we know that AD is profitable to specialize. Add it to our call // Now we know that CSDesc is profitable to specialize. Add it to our call
// site list. // site list.
CallSites.push_back(ArgDescriptor(PAI, AI, PAIIndex.getValue())); CallSites.push_back(CallSiteDescriptor(PAI, AI, PAIIndex.getValue()));
} }
} }
} }
@@ -373,32 +434,18 @@ bool ClosureSpecializer::specialize(SILFunction *Caller) {
// Collect all of the PartialApplyInsts that are used as arguments to // Collect all of the PartialApplyInsts that are used as arguments to
// ApplyInsts. Check the profitability of specializing the closure argument. // ApplyInsts. Check the profitability of specializing the closure argument.
llvm::SmallVector<ArgDescriptor, 8> CallSites; llvm::SmallVector<CallSiteDescriptor, 8> CallSites;
llvm::SmallPtrSet<ApplyInst *, 4> MultipleClosureAI; llvm::SmallPtrSet<ApplyInst *, 4> MultipleClosureAI;
gatherCallSites(Caller, CallSites, MultipleClosureAI); gatherCallSites(Caller, CallSites, MultipleClosureAI);
bool Changed = false; bool Changed = false;
for (auto &AD : CallSites) { for (auto &CSDesc : CallSites) {
// Do not specialize apply insts that take in multiple closures. This pass // Do not specialize apply insts that take in multiple closures. This pass
// does not know how to do this yet. // does not know how to do this yet.
if (MultipleClosureAI.count(AD.AI)) if (MultipleClosureAI.count(CSDesc.getApplyInst()))
continue; continue;
auto *CalleeFRI = cast<FunctionRefInst>(AD.AI->getCallee()); CSDesc.specializeClosure();
auto *Callee = CalleeFRI->getReferencedFunction();
llvm::SmallString<64> NewFName;
createName(Callee, Callee->getArgument(AD.ClosureIndex), AD.PAI,
NewFName);
DEBUG(llvm::dbgs() << " Perform optimizations with new name "
<< NewFName << '\n');
SILFunction *NewF = Callee->getModule().lookUpFunction(NewFName);
if (!NewF)
NewF = ClosureSpecCloner::cloneFunction(Callee, AD.PAI, AD.ClosureIndex,
NewFName);
rewriteApplyInst(AD, NewF);
Changed = true; Changed = true;
} }