LoopLoadElimination.cpp [plain text]
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/LoopAccessAnalysis.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/ScalarEvolutionExpander.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Module.h"
#include "llvm/Pass.h"
#include "llvm/Support/Debug.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Utils/LoopVersioning.h"
#include <forward_list>
#define LLE_OPTION "loop-load-elim"
#define DEBUG_TYPE LLE_OPTION
using namespace llvm;
static cl::opt<unsigned> CheckPerElim(
"runtime-check-per-loop-load-elim", cl::Hidden,
cl::desc("Max number of memchecks allowed per eliminated load on average"),
cl::init(1));
static cl::opt<unsigned> LoadElimSCEVCheckThreshold(
"loop-load-elimination-scev-check-threshold", cl::init(8), cl::Hidden,
cl::desc("The maximum number of SCEV checks allowed for Loop "
"Load Elimination"));
STATISTIC(NumLoopLoadEliminted, "Number of loads eliminated by LLE");
namespace {
struct StoreToLoadForwardingCandidate {
LoadInst *Load;
StoreInst *Store;
StoreToLoadForwardingCandidate(LoadInst *Load, StoreInst *Store)
: Load(Load), Store(Store) {}
bool isDependenceDistanceOfOne(PredicatedScalarEvolution &PSE,
Loop *L) const {
Value *LoadPtr = Load->getPointerOperand();
Value *StorePtr = Store->getPointerOperand();
Type *LoadPtrType = LoadPtr->getType();
Type *LoadType = LoadPtrType->getPointerElementType();
assert(LoadPtrType->getPointerAddressSpace() ==
StorePtr->getType()->getPointerAddressSpace() &&
LoadType == StorePtr->getType()->getPointerElementType() &&
"Should be a known dependence");
if (isStridedPtr(PSE, LoadPtr, L) != 1 ||
isStridedPtr(PSE, StorePtr, L) != 1)
return false;
auto &DL = Load->getParent()->getModule()->getDataLayout();
unsigned TypeByteSize = DL.getTypeAllocSize(const_cast<Type *>(LoadType));
auto *LoadPtrSCEV = cast<SCEVAddRecExpr>(PSE.getSCEV(LoadPtr));
auto *StorePtrSCEV = cast<SCEVAddRecExpr>(PSE.getSCEV(StorePtr));
auto *Dist = cast<SCEVConstant>(
PSE.getSE()->getMinusSCEV(StorePtrSCEV, LoadPtrSCEV));
const APInt &Val = Dist->getAPInt();
return Val == TypeByteSize;
}
Value *getLoadPtr() const { return Load->getPointerOperand(); }
#ifndef NDEBUG
friend raw_ostream &operator<<(raw_ostream &OS,
const StoreToLoadForwardingCandidate &Cand) {
OS << *Cand.Store << " -->\n";
OS.indent(2) << *Cand.Load << "\n";
return OS;
}
#endif
};
bool doesStoreDominatesAllLatches(BasicBlock *StoreBlock, Loop *L,
DominatorTree *DT) {
SmallVector<BasicBlock *, 8> Latches;
L->getLoopLatches(Latches);
return std::all_of(Latches.begin(), Latches.end(),
[&](const BasicBlock *Latch) {
return DT->dominates(StoreBlock, Latch);
});
}
static bool isLoadConditional(LoadInst *Load, Loop *L) {
return Load->getParent() != L->getHeader();
}
class LoadEliminationForLoop {
public:
LoadEliminationForLoop(Loop *L, LoopInfo *LI, const LoopAccessInfo &LAI,
DominatorTree *DT)
: L(L), LI(LI), LAI(LAI), DT(DT), PSE(LAI.PSE) {}
std::forward_list<StoreToLoadForwardingCandidate>
findStoreToLoadDependences(const LoopAccessInfo &LAI) {
std::forward_list<StoreToLoadForwardingCandidate> Candidates;
const auto *Deps = LAI.getDepChecker().getDependences();
if (!Deps)
return Candidates;
SmallSet<Instruction *, 4> LoadsWithUnknownDepedence;
for (const auto &Dep : *Deps) {
Instruction *Source = Dep.getSource(LAI);
Instruction *Destination = Dep.getDestination(LAI);
if (Dep.Type == MemoryDepChecker::Dependence::Unknown) {
if (isa<LoadInst>(Source))
LoadsWithUnknownDepedence.insert(Source);
if (isa<LoadInst>(Destination))
LoadsWithUnknownDepedence.insert(Destination);
continue;
}
if (Dep.isBackward())
std::swap(Source, Destination);
else
assert(Dep.isForward() && "Needs to be a forward dependence");
auto *Store = dyn_cast<StoreInst>(Source);
if (!Store)
continue;
auto *Load = dyn_cast<LoadInst>(Destination);
if (!Load)
continue;
if (Store->getPointerOperand()->getType() !=
Load->getPointerOperand()->getType())
continue;
Candidates.emplace_front(Load, Store);
}
if (!LoadsWithUnknownDepedence.empty())
Candidates.remove_if([&](const StoreToLoadForwardingCandidate &C) {
return LoadsWithUnknownDepedence.count(C.Load);
});
return Candidates;
}
unsigned getInstrIndex(Instruction *Inst) {
auto I = InstOrder.find(Inst);
assert(I != InstOrder.end() && "No index for instruction");
return I->second;
}
void removeDependencesFromMultipleStores(
std::forward_list<StoreToLoadForwardingCandidate> &Candidates) {
typedef DenseMap<LoadInst *, const StoreToLoadForwardingCandidate *>
LoadToSingleCandT;
LoadToSingleCandT LoadToSingleCand;
for (const auto &Cand : Candidates) {
bool NewElt;
LoadToSingleCandT::iterator Iter;
std::tie(Iter, NewElt) =
LoadToSingleCand.insert(std::make_pair(Cand.Load, &Cand));
if (!NewElt) {
const StoreToLoadForwardingCandidate *&OtherCand = Iter->second;
if (OtherCand == nullptr)
continue;
if (Cand.Store->getParent() == OtherCand->Store->getParent() &&
Cand.isDependenceDistanceOfOne(PSE, L) &&
OtherCand->isDependenceDistanceOfOne(PSE, L)) {
if (getInstrIndex(OtherCand->Store) < getInstrIndex(Cand.Store))
OtherCand = &Cand;
} else
OtherCand = nullptr;
}
}
Candidates.remove_if([&](const StoreToLoadForwardingCandidate &Cand) {
if (LoadToSingleCand[Cand.Load] != &Cand) {
DEBUG(dbgs() << "Removing from candidates: \n" << Cand
<< " The load may have multiple stores forwarding to "
<< "it\n");
return true;
}
return false;
});
}
bool needsChecking(unsigned PtrIdx1, unsigned PtrIdx2,
const SmallSet<Value *, 4> &PtrsWrittenOnFwdingPath,
const std::set<Value *> &CandLoadPtrs) {
Value *Ptr1 =
LAI.getRuntimePointerChecking()->getPointerInfo(PtrIdx1).PointerValue;
Value *Ptr2 =
LAI.getRuntimePointerChecking()->getPointerInfo(PtrIdx2).PointerValue;
return ((PtrsWrittenOnFwdingPath.count(Ptr1) && CandLoadPtrs.count(Ptr2)) ||
(PtrsWrittenOnFwdingPath.count(Ptr2) && CandLoadPtrs.count(Ptr1)));
}
SmallSet<Value *, 4> findPointersWrittenOnForwardingPath(
const SmallVectorImpl<StoreToLoadForwardingCandidate> &Candidates) {
LoadInst *LastLoad =
std::max_element(Candidates.begin(), Candidates.end(),
[&](const StoreToLoadForwardingCandidate &A,
const StoreToLoadForwardingCandidate &B) {
return getInstrIndex(A.Load) < getInstrIndex(B.Load);
})
->Load;
StoreInst *FirstStore =
std::min_element(Candidates.begin(), Candidates.end(),
[&](const StoreToLoadForwardingCandidate &A,
const StoreToLoadForwardingCandidate &B) {
return getInstrIndex(A.Store) <
getInstrIndex(B.Store);
})
->Store;
SmallSet<Value *, 4> PtrsWrittenOnFwdingPath;
auto InsertStorePtr = [&](Instruction *I) {
if (auto *S = dyn_cast<StoreInst>(I))
PtrsWrittenOnFwdingPath.insert(S->getPointerOperand());
};
const auto &MemInstrs = LAI.getDepChecker().getMemoryInstructions();
std::for_each(MemInstrs.begin() + getInstrIndex(FirstStore) + 1,
MemInstrs.end(), InsertStorePtr);
std::for_each(MemInstrs.begin(), &MemInstrs[getInstrIndex(LastLoad)],
InsertStorePtr);
return PtrsWrittenOnFwdingPath;
}
SmallVector<RuntimePointerChecking::PointerCheck, 4> collectMemchecks(
const SmallVectorImpl<StoreToLoadForwardingCandidate> &Candidates) {
SmallSet<Value *, 4> PtrsWrittenOnFwdingPath =
findPointersWrittenOnForwardingPath(Candidates);
std::set<Value *> CandLoadPtrs;
std::transform(Candidates.begin(), Candidates.end(),
std::inserter(CandLoadPtrs, CandLoadPtrs.begin()),
std::mem_fn(&StoreToLoadForwardingCandidate::getLoadPtr));
const auto &AllChecks = LAI.getRuntimePointerChecking()->getChecks();
SmallVector<RuntimePointerChecking::PointerCheck, 4> Checks;
std::copy_if(AllChecks.begin(), AllChecks.end(), std::back_inserter(Checks),
[&](const RuntimePointerChecking::PointerCheck &Check) {
for (auto PtrIdx1 : Check.first->Members)
for (auto PtrIdx2 : Check.second->Members)
if (needsChecking(PtrIdx1, PtrIdx2,
PtrsWrittenOnFwdingPath, CandLoadPtrs))
return true;
return false;
});
DEBUG(dbgs() << "\nPointer Checks (count: " << Checks.size() << "):\n");
DEBUG(LAI.getRuntimePointerChecking()->printChecks(dbgs(), Checks));
return Checks;
}
void
propagateStoredValueToLoadUsers(const StoreToLoadForwardingCandidate &Cand,
SCEVExpander &SEE) {
Value *Ptr = Cand.Load->getPointerOperand();
auto *PtrSCEV = cast<SCEVAddRecExpr>(PSE.getSCEV(Ptr));
auto *PH = L->getLoopPreheader();
Value *InitialPtr = SEE.expandCodeFor(PtrSCEV->getStart(), Ptr->getType(),
PH->getTerminator());
Value *Initial =
new LoadInst(InitialPtr, "load_initial", PH->getTerminator());
PHINode *PHI = PHINode::Create(Initial->getType(), 2, "store_forwarded",
&L->getHeader()->front());
PHI->addIncoming(Initial, PH);
PHI->addIncoming(Cand.Store->getOperand(0), L->getLoopLatch());
Cand.Load->replaceAllUsesWith(PHI);
}
bool processLoop() {
DEBUG(dbgs() << "\nIn \"" << L->getHeader()->getParent()->getName()
<< "\" checking " << *L << "\n");
auto StoreToLoadDependences = findStoreToLoadDependences(LAI);
if (StoreToLoadDependences.empty())
return false;
InstOrder = LAI.getDepChecker().generateInstructionOrderMap();
removeDependencesFromMultipleStores(StoreToLoadDependences);
if (StoreToLoadDependences.empty())
return false;
SmallVector<StoreToLoadForwardingCandidate, 4> Candidates;
unsigned NumForwarding = 0;
for (const StoreToLoadForwardingCandidate Cand : StoreToLoadDependences) {
DEBUG(dbgs() << "Candidate " << Cand);
if (!doesStoreDominatesAllLatches(Cand.Store->getParent(), L, DT))
continue;
if (isLoadConditional(Cand.Load, L))
continue;
if (!Cand.isDependenceDistanceOfOne(PSE, L))
continue;
++NumForwarding;
DEBUG(dbgs()
<< NumForwarding
<< ". Valid store-to-load forwarding across the loop backedge\n");
Candidates.push_back(Cand);
}
if (Candidates.empty())
return false;
SmallVector<RuntimePointerChecking::PointerCheck, 4> Checks =
collectMemchecks(Candidates);
if (Checks.size() > Candidates.size() * CheckPerElim) {
DEBUG(dbgs() << "Too many run-time checks needed.\n");
return false;
}
if (LAI.PSE.getUnionPredicate().getComplexity() >
LoadElimSCEVCheckThreshold) {
DEBUG(dbgs() << "Too many SCEV run-time checks needed.\n");
return false;
}
if (!Checks.empty() || !LAI.PSE.getUnionPredicate().isAlwaysTrue()) {
if (L->getHeader()->getParent()->optForSize()) {
DEBUG(dbgs() << "Versioning is needed but not allowed when optimizing "
"for size.\n");
return false;
}
LoopVersioning LV(LAI, L, LI, DT, PSE.getSE(), false);
LV.setAliasChecks(std::move(Checks));
LV.setSCEVChecks(LAI.PSE.getUnionPredicate());
LV.versionLoop();
}
SCEVExpander SEE(*PSE.getSE(), L->getHeader()->getModule()->getDataLayout(),
"storeforward");
for (const auto &Cand : Candidates)
propagateStoredValueToLoadUsers(Cand, SEE);
NumLoopLoadEliminted += NumForwarding;
return true;
}
private:
Loop *L;
DenseMap<Instruction *, unsigned> InstOrder;
LoopInfo *LI;
const LoopAccessInfo &LAI;
DominatorTree *DT;
PredicatedScalarEvolution PSE;
};
class LoopLoadElimination : public FunctionPass {
public:
LoopLoadElimination() : FunctionPass(ID) {
initializeLoopLoadEliminationPass(*PassRegistry::getPassRegistry());
}
bool runOnFunction(Function &F) override {
auto *LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
auto *LAA = &getAnalysis<LoopAccessAnalysis>();
auto *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
SmallVector<Loop *, 8> Worklist;
for (Loop *TopLevelLoop : *LI)
for (Loop *L : depth_first(TopLevelLoop))
if (L->empty())
Worklist.push_back(L);
bool Changed = false;
for (Loop *L : Worklist) {
const LoopAccessInfo &LAI = LAA->getInfo(L, ValueToValueMap());
LoadEliminationForLoop LEL(L, LI, LAI, DT);
Changed |= LEL.processLoop();
}
return Changed;
}
void getAnalysisUsage(AnalysisUsage &AU) const override {
AU.addRequiredID(LoopSimplifyID);
AU.addRequired<LoopInfoWrapperPass>();
AU.addPreserved<LoopInfoWrapperPass>();
AU.addRequired<LoopAccessAnalysis>();
AU.addRequired<ScalarEvolutionWrapperPass>();
AU.addRequired<DominatorTreeWrapperPass>();
AU.addPreserved<DominatorTreeWrapperPass>();
}
static char ID;
};
}
char LoopLoadElimination::ID;
static const char LLE_name[] = "Loop Load Elimination";
INITIALIZE_PASS_BEGIN(LoopLoadElimination, LLE_OPTION, LLE_name, false, false)
INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
INITIALIZE_PASS_DEPENDENCY(LoopAccessAnalysis)
INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
INITIALIZE_PASS_END(LoopLoadElimination, LLE_OPTION, LLE_name, false, false)
namespace llvm {
FunctionPass *createLoopLoadEliminationPass() {
return new LoopLoadElimination();
}
}