LoopUnrollRuntime.cpp [plain text]
#include "llvm/Transforms/Utils/UnrollLoop.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/LoopIterator.h"
#include "llvm/Analysis/LoopPass.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/ScalarEvolutionExpander.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Metadata.h"
#include "llvm/IR/Module.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include <algorithm>
using namespace llvm;
#define DEBUG_TYPE "loop-unroll"
STATISTIC(NumRuntimeUnrolled,
"Number of loops unrolled with run-time trip counts");
static void ConnectProlog(Loop *L, Value *BECount, unsigned Count,
BasicBlock *LastPrologBB, BasicBlock *PrologEnd,
BasicBlock *OrigPH, BasicBlock *NewPH,
ValueToValueMapTy &VMap, DominatorTree *DT,
LoopInfo *LI, Pass *P) {
BasicBlock *Latch = L->getLoopLatch();
assert(Latch && "Loop must have a latch");
for (succ_iterator SBI = succ_begin(Latch), SBE = succ_end(Latch);
SBI != SBE; ++SBI) {
for (BasicBlock::iterator BBI = (*SBI)->begin();
PHINode *PN = dyn_cast<PHINode>(BBI); ++BBI) {
PHINode *NewPN = PHINode::Create(PN->getType(), 2, PN->getName()+".unr",
PrologEnd->getTerminator());
if (L->contains(PN)) {
NewPN->addIncoming(PN->getIncomingValueForBlock(NewPH), OrigPH);
} else {
NewPN->addIncoming(UndefValue::get(PN->getType()), OrigPH);
}
Value *V = PN->getIncomingValueForBlock(Latch);
if (Instruction *I = dyn_cast<Instruction>(V)) {
if (L->contains(I)) {
V = VMap[I];
}
}
NewPN->addIncoming(V, LastPrologBB);
if (L->contains(PN)) {
PN->setIncomingValue(PN->getBasicBlockIndex(NewPH), NewPN);
} else {
PN->addIncoming(NewPN, PrologEnd);
}
}
}
Instruction *InsertPt = PrologEnd->getTerminator();
IRBuilder<> B(InsertPt);
assert(Count != 0 && "nonsensical Count!");
Value *BrLoopExit =
B.CreateICmpULT(BECount, ConstantInt::get(BECount->getType(), Count - 1));
BasicBlock *Exit = L->getUniqueExitBlock();
assert(Exit && "Loop must have a single exit block only");
SmallVector<BasicBlock*, 4> Preds(pred_begin(Exit), pred_end(Exit));
SplitBlockPredecessors(Exit, Preds, ".unr-lcssa", DT, LI,
P->mustPreserveAnalysisID(LCSSAID));
B.CreateCondBr(BrLoopExit, Exit, NewPH);
InsertPt->eraseFromParent();
}
static void CloneLoopBlocks(Loop *L, Value *NewIter, const bool UnrollProlog,
BasicBlock *InsertTop, BasicBlock *InsertBot,
std::vector<BasicBlock *> &NewBlocks,
LoopBlocksDFS &LoopBlocks, ValueToValueMapTy &VMap,
LoopInfo *LI) {
BasicBlock *Preheader = L->getLoopPreheader();
BasicBlock *Header = L->getHeader();
BasicBlock *Latch = L->getLoopLatch();
Function *F = Header->getParent();
LoopBlocksDFS::RPOIterator BlockBegin = LoopBlocks.beginRPO();
LoopBlocksDFS::RPOIterator BlockEnd = LoopBlocks.endRPO();
Loop *NewLoop = nullptr;
Loop *ParentLoop = L->getParentLoop();
if (!UnrollProlog) {
NewLoop = new Loop();
if (ParentLoop)
ParentLoop->addChildLoop(NewLoop);
else
LI->addTopLevelLoop(NewLoop);
}
for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
BasicBlock *NewBB = CloneBasicBlock(*BB, VMap, ".prol", F);
NewBlocks.push_back(NewBB);
if (NewLoop)
NewLoop->addBasicBlockToLoop(NewBB, *LI);
else if (ParentLoop)
ParentLoop->addBasicBlockToLoop(NewBB, *LI);
VMap[*BB] = NewBB;
if (Header == *BB) {
InsertTop->getTerminator()->setSuccessor(0, NewBB);
}
if (Latch == *BB) {
VMap.erase((*BB)->getTerminator());
BasicBlock *FirstLoopBB = cast<BasicBlock>(VMap[Header]);
BranchInst *LatchBR = cast<BranchInst>(NewBB->getTerminator());
IRBuilder<> Builder(LatchBR);
if (UnrollProlog) {
Builder.CreateBr(InsertBot);
} else {
PHINode *NewIdx = PHINode::Create(NewIter->getType(), 2, "prol.iter",
FirstLoopBB->getFirstNonPHI());
Value *IdxSub =
Builder.CreateSub(NewIdx, ConstantInt::get(NewIdx->getType(), 1),
NewIdx->getName() + ".sub");
Value *IdxCmp =
Builder.CreateIsNotNull(IdxSub, NewIdx->getName() + ".cmp");
Builder.CreateCondBr(IdxCmp, FirstLoopBB, InsertBot);
NewIdx->addIncoming(NewIter, InsertTop);
NewIdx->addIncoming(IdxSub, NewBB);
}
LatchBR->eraseFromParent();
}
}
for (BasicBlock::iterator I = Header->begin(); isa<PHINode>(I); ++I) {
PHINode *NewPHI = cast<PHINode>(VMap[&*I]);
if (UnrollProlog) {
VMap[&*I] = NewPHI->getIncomingValueForBlock(Preheader);
cast<BasicBlock>(VMap[Header])->getInstList().erase(NewPHI);
} else {
unsigned idx = NewPHI->getBasicBlockIndex(Preheader);
NewPHI->setIncomingBlock(idx, InsertTop);
BasicBlock *NewLatch = cast<BasicBlock>(VMap[Latch]);
idx = NewPHI->getBasicBlockIndex(Latch);
Value *InVal = NewPHI->getIncomingValue(idx);
NewPHI->setIncomingBlock(idx, NewLatch);
if (VMap[InVal])
NewPHI->setIncomingValue(idx, VMap[InVal]);
}
}
if (NewLoop) {
SmallVector<Metadata *, 4> MDs;
MDs.push_back(nullptr);
MDNode *LoopID = NewLoop->getLoopID();
if (LoopID) {
for (unsigned i = 1, ie = LoopID->getNumOperands(); i < ie; ++i) {
bool IsUnrollMetadata = false;
MDNode *MD = dyn_cast<MDNode>(LoopID->getOperand(i));
if (MD) {
const MDString *S = dyn_cast<MDString>(MD->getOperand(0));
IsUnrollMetadata = S && S->getString().startswith("llvm.loop.unroll.");
}
if (!IsUnrollMetadata)
MDs.push_back(LoopID->getOperand(i));
}
}
LLVMContext &Context = NewLoop->getHeader()->getContext();
SmallVector<Metadata *, 1> DisableOperands;
DisableOperands.push_back(MDString::get(Context, "llvm.loop.unroll.disable"));
MDNode *DisableNode = MDNode::get(Context, DisableOperands);
MDs.push_back(DisableNode);
MDNode *NewLoopID = MDNode::get(Context, MDs);
NewLoopID->replaceOperandWith(0, NewLoopID);
NewLoop->setLoopID(NewLoopID);
}
}
bool llvm::UnrollRuntimeLoopProlog(Loop *L, unsigned Count,
bool AllowExpensiveTripCount, LoopInfo *LI,
LPPassManager *LPM) {
if (!L->getExitingBlock())
return false;
if (!L->isLoopSimplifyForm() || !L->getUniqueExitBlock())
return false;
if (!LPM)
return false;
auto *SEWP = LPM->getAnalysisIfAvailable<ScalarEvolutionWrapperPass>();
if (!SEWP)
return false;
ScalarEvolution &SE = SEWP->getSE();
const SCEV *BECountSC = SE.getBackedgeTakenCount(L);
if (isa<SCEVCouldNotCompute>(BECountSC) ||
!BECountSC->getType()->isIntegerTy())
return false;
unsigned BEWidth = cast<IntegerType>(BECountSC->getType())->getBitWidth();
const SCEV *TripCountSC =
SE.getAddExpr(BECountSC, SE.getConstant(BECountSC->getType(), 1));
if (isa<SCEVCouldNotCompute>(TripCountSC))
return false;
BasicBlock *Header = L->getHeader();
const DataLayout &DL = Header->getModule()->getDataLayout();
SCEVExpander Expander(SE, DL, "loop-unroll");
if (!AllowExpensiveTripCount && Expander.isHighCostExpansion(TripCountSC, L))
return false;
if (!isPowerOf2_32(Count))
return false;
if (Log2_32(Count) > BEWidth)
return false;
if (Loop *ParentLoop = L->getParentLoop())
SE.forgetLoop(ParentLoop);
auto *DTWP = LPM->getAnalysisIfAvailable<DominatorTreeWrapperPass>();
auto *DT = DTWP ? &DTWP->getDomTree() : nullptr;
BasicBlock *PH = L->getLoopPreheader();
BasicBlock *Latch = L->getLoopLatch();
BasicBlock *PEnd = SplitEdge(PH, Header, DT, LI);
BasicBlock *NewPH = SplitBlock(PEnd, PEnd->getTerminator(), DT, LI);
BranchInst *PreHeaderBR = cast<BranchInst>(PH->getTerminator());
Value *TripCount = Expander.expandCodeFor(TripCountSC, TripCountSC->getType(),
PreHeaderBR);
Value *BECount = Expander.expandCodeFor(BECountSC, BECountSC->getType(),
PreHeaderBR);
IRBuilder<> B(PreHeaderBR);
Value *ModVal = B.CreateAnd(TripCount, Count - 1, "xtraiter");
Value *BranchVal = B.CreateIsNotNull(ModVal, "lcmp.mod");
B.CreateCondBr(BranchVal, PEnd, PEnd);
assert(PreHeaderBR->isUnconditional() &&
PreHeaderBR->getSuccessor(0) == PEnd &&
"CFG edges in Preheader are not correct");
PreHeaderBR->eraseFromParent();
Function *F = Header->getParent();
LoopBlocksDFS LoopBlocks(L);
LoopBlocks.perform(LI);
std::vector<BasicBlock *> NewBlocks;
ValueToValueMapTy VMap;
bool UnrollPrologue = Count == 2;
CloneLoopBlocks(L, ModVal, UnrollPrologue, PH, PEnd, NewBlocks, LoopBlocks,
VMap, LI);
F->getBasicBlockList().splice(PEnd->getIterator(), F->getBasicBlockList(),
NewBlocks[0]->getIterator(), F->end());
for (unsigned i = 0, e = NewBlocks.size(); i != e; ++i) {
for (BasicBlock::iterator I = NewBlocks[i]->begin(),
E = NewBlocks[i]->end();
I != E; ++I) {
RemapInstruction(&*I, VMap,
RF_NoModuleLevelChanges | RF_IgnoreMissingEntries);
}
}
BasicBlock *LastLoopBB = cast<BasicBlock>(VMap[Latch]);
ConnectProlog(L, BECount, Count, LastLoopBB, PEnd, PH, NewPH, VMap, DT, LI,
LPM->getAsPass());
NumRuntimeUnrolled++;
return true;
}