Index: llvm/lib/Transforms/Scalar/LoopPredication.cpp =================================================================== --- llvm/lib/Transforms/Scalar/LoopPredication.cpp +++ llvm/lib/Transforms/Scalar/LoopPredication.cpp @@ -249,6 +249,7 @@ class LoopPredication { AliasAnalysis *AA; ScalarEvolution *SE; + LoopInfo *LI; BranchProbabilityInfo *BPI; Loop *L; @@ -300,10 +301,12 @@ // within the loop. We identify such unprofitable loops through BPI. bool isLoopProfitableToPredicate(); + bool predicateLoopExits(Loop *L, SCEVExpander &Rewriter); + public: - LoopPredication(AliasAnalysis *AA, ScalarEvolution *SE, + LoopPredication(AliasAnalysis *AA, ScalarEvolution *SE, LoopInfo *LI, BranchProbabilityInfo *BPI) - : AA(AA), SE(SE), BPI(BPI){}; + : AA(AA), SE(SE), LI(LI), BPI(BPI) {}; bool runOnLoop(Loop *L); }; @@ -323,10 +326,11 @@ if (skipLoop(L)) return false; auto *SE = &getAnalysis().getSE(); + auto *LI = &getAnalysis().getLoopInfo(); BranchProbabilityInfo &BPI = getAnalysis().getBPI(); auto *AA = &getAnalysis().getAAResults(); - LoopPredication LP(AA, SE, &BPI); + LoopPredication LP(AA, SE, LI, &BPI); return LP.runOnLoop(L); } }; @@ -352,7 +356,7 @@ AM.getResult(L, AR).getManager(); Function *F = L.getHeader()->getParent(); auto *BPI = FAM.getCachedResult(*F); - LoopPredication LP(&AR.AA, &AR.SE, BPI); + LoopPredication LP(&AR.AA, &AR.SE, &AR.LI, BPI); if (!LP.runOnLoop(&L)) return PreservedAnalyses::all(); @@ -953,6 +957,156 @@ return true; } +// Walk back through portions of the CFG which are "likely" to be control +// equivelant. Note that "likely" here is a profitability heuristic, not a +// legality one. We just want to avoid widening over conditions which might +// correlate strongly with the one we skipped. We're fine ignoring exception +// exits (for instance) as we statically predict those to be very rare. +static std::pair +getLikelyControlEqPredecessorForBB(BasicBlock *BB, LoopInfo &LI) { + if (BasicBlock *Pred = BB->getSinglePredecessor()) { + if (BB == Pred->getSingleSuccessor()) + return {Pred, BB}; + } + + // TODO: loop case where all exits are same. + + return {nullptr, nullptr}; +} + + +/// If we can (cheaply) find a widenable branch which controls entry into the +/// loop, return it. +static BranchInst *FindWideableTerminatorAboveLoop(Loop *L, LoopInfo &LI) { + // Skip back through any part of the control flow likely to be control + // equivelent dynamically. + std::pair LastPair; + for (std::pair + Pair(L->getLoopPredecessor(), L->getHeader()); + Pair.first; + Pair = getLikelyControlEqPredecessorForBB(Pair.first, LI)) { + LastPair = Pair; + } + if (auto *BB = LastPair.first) + if (BasicBlock *Pred = BB->getSinglePredecessor()) { + auto *Term = Pred->getTerminator(); + dbgs() << "Considering" << *Term << "\n"; + + Value *Cond, *WC; + BasicBlock *IfTrueBB, *IfFalseBB; + if (parseWidenableBranch(Term, Cond, WC, + IfTrueBB, IfFalseBB) && + IfTrueBB == BB) + return cast(Term); + } + return nullptr; +} + +/// This implements an analogous, but entirely distinct transform from the main +/// loop predication transform. This one is phrased in terms of using a +/// widenable branch *outside* the loop to allow us to simplify loop exits in a +/// following loop. This is closer is spirit to the IndVarSimplify transform +/// of the same name, but is materially different widening loosens legality +/// sharply. +bool LoopPredication::predicateLoopExits(Loop *L, SCEVExpander &Rewriter) { + // The transformation performed here aims to widen a widenable condition + // above the loop such that exits other than the latch are provably dead. It + // assumes that the latch is the dominant exit for profitability, and relies + // on the semantics of widenable expressions for legality. (i.e. being able + // to fall down the widenable path spuriously allows us to ignore exit order, + // nanalyzeable exits, side effects, exceptional exits, and other challenges + // which restrict the applicability of the non-WC based version of this + // transform in IndVarSimplify.) + + SmallVector ExitingBlocks; + L->getExitingBlocks(ExitingBlocks); + + if (ExitingBlocks.size() < 2) + return false; // Nothing to do. + + auto *Latch = L->getLoopLatch(); + if (!Latch) + return false; + + auto *WidenableBR = FindWideableTerminatorAboveLoop(L, *LI); + if (!WidenableBR) + return false; + + const SCEV *LatchEC = SE->getExitCount(L, Latch); + if (isa(LatchEC) || + !SE->isLoopInvariant(LatchEC, L) || + !isSafeToExpandAt(LatchEC, WidenableBR, *SE)) + return false; + + Rewriter.setInsertPoint(WidenableBR); + IRBuilder<> B(WidenableBR); + + bool Changed = false; + Value *LatchECV = nullptr; //lazy generated if needed + for (BasicBlock *ExitingBB : ExitingBlocks) { + if (Latch == ExitingBB) + continue; + + // If our exiting block exits multiple loops, we can only rewrite the + // innermost one. Otherwise, we're changing how many times the innermost + // loop runs before it exits. + if (LI->getLoopFor(ExitingBB) != L) + continue; + + // Can't rewrite non-branch yet. + auto *BI = dyn_cast(ExitingBB->getTerminator()); + if (!BI) + continue; + + // If already constant, nothing to do. + if (isa(BI->getCondition())) + continue; + + const SCEV *ExitCount = SE->getExitCount(L, ExitingBB); + if (isa(ExitCount) || + !isSafeToExpandAt(ExitCount, WidenableBR, *SE)) + continue; + + // If we found a widenable exit condition, do two things: + // 1) fold the widened exit test into the widenable condition + // 2) fold the branch to untaken - avoids infinite looping + + Value *ECV = Rewriter.expandCodeFor(ExitCount); + if (!LatchECV) + LatchECV = Rewriter.expandCodeFor(LatchEC); + Value *RHS = LatchECV; + if (ECV->getType() != RHS->getType()) { + Type *WiderTy = SE->getWiderType(ECV->getType(), RHS->getType()); + ECV = B.CreateZExt(ECV, WiderTy); + RHS = B.CreateZExt(RHS, WiderTy); + } + //assert(DT->dominantes(ExitingBB, Latch)); + Value *NewCond = B.CreateICmp(ICmpInst::ICMP_UGT, ECV, RHS); + + Value *Cond, *WC; + BasicBlock *IfTrueBB, *IfFalseBB; + bool Success = parseWidenableBranch(WidenableBR, Cond, WC, + IfTrueBB, IfFalseBB); + assert(Success && "implied from above"); + NewCond = B.CreateAnd(WC, B.CreateAnd(NewCond, Cond)); + WidenableBR->setCondition(NewCond); + + Value *OldCond = BI->getCondition(); + bool ExitIfTrue = !L->contains(*succ_begin(ExitingBB)); + BI->setCondition(ConstantInt::get(OldCond->getType(), !ExitIfTrue)); + Changed = true; + } + + if (Changed) + // We just mutated a bunch of loop exits changing there exit counts + // widely. We need to force recomputation of the exit counts given these + // changes. Note that all of the inserted exits are never taken, and + // should be removed next time the CFG is modified. + SE->forgetLoop(L); + return Changed; +} + + bool LoopPredication::runOnLoop(Loop *Loop) { L = Loop; @@ -1004,16 +1158,14 @@ cast(BB->getTerminator())); } - if (Guards.empty() && GuardsAsWidenableBranches.empty()) - return false; - SCEVExpander Expander(*SE, *DL, "loop-predication"); - bool Changed = false; - for (auto *Guard : Guards) - Changed |= widenGuardConditions(Guard, Expander); - for (auto *Guard : GuardsAsWidenableBranches) - Changed |= widenWidenableBranchGuardConditions(Guard, Expander); - + if (!Guards.empty() || !GuardsAsWidenableBranches.empty()) { + for (auto *Guard : Guards) + Changed |= widenGuardConditions(Guard, Expander); + for (auto *Guard : GuardsAsWidenableBranches) + Changed |= widenWidenableBranchGuardConditions(Guard, Expander); + } + Changed = predicateLoopExits(L, Expander); return Changed; }