[WIP] An alternate approach to widening exit conditions using widenable conditions

[reames]

Summary

This is an alternative approach to D69454. Despite the lack of code sharing, I like this version a lot better. The interweaving of ideas in the IndVars version is really hard to follow.

[reames]

Recording some thoughts.

There's a potential serious problem here with both approaches. We're hoisting logic (mostly arithmetic) out of the loop and adding a new use outside the loop. As such, there's a potentially serious problem with undef and poison as the existing flags may no longer be correct for the new user.

This is exactly the same problem which unswitch has. Interestingly, this proposed transform basically implements unswitch, but without the duplication. As such, any solution (say, the proposed freeze intrinsic) which works for unswitch, should work here as well.

While waiting for freeze, some practical notes:

• propagatesFullPoison does not include and. And unswitch appears to essentially rely on this. As such, from a poison perspective, anding in a new condition to the widenable branch is probably okay in practice. (Caveat: unless poison is exploited by some transform above the loop following use simplification? But unswitch would have this problem too.)
• Undef is harder. I haven't found a chain of logic to convince myself it's harmless in practice just yet.

[reames]

Subsumed by D69830

llvm/lib/Transforms/Scalar/LoopPredication.cpp

void dump() {
     dbgs() << "LoopICmp Pred = " << Pred << ", IV = " << *IV
            << ", Limit = " << *Limit << "\n";
   }
 };
 
 class LoopPredication {
   AliasAnalysis *AA;
   ScalarEvolution *SE;
+  LoopInfo *LI;
   BranchProbabilityInfo *BPI;
 
   Loop *L;
   const DataLayout *DL;
   BasicBlock *Preheader;
   LoopICmp LatchCheck;
 
   bool isSupportedStep(const SCEV* Step);
class LoopPredication {
   bool widenGuardConditions(IntrinsicInst *II, SCEVExpander &Expander);
   bool widenWidenableBranchGuardConditions(BranchInst *Guard, SCEVExpander &Expander);
   // If the loop always exits through another block in the loop, we should not
   // predicate based on the latch check. For example, the latch check can be a
   // very coarse grained check and there can be more fine grained exit checks
   // 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);
 };
 
 class LoopPredicationLegacyPass : public LoopPass {
 public:
   static char ID;
   LoopPredicationLegacyPass() : LoopPass(ID) {
     initializeLoopPredicationLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<BranchProbabilityInfoWrapperPass>();
     getLoopAnalysisUsage(AU);
   }
 
   bool runOnLoop(Loop *L, LPPassManager &LPM) override {
     if (skipLoop(L))
       return false;
     auto *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
+    auto *LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     BranchProbabilityInfo &BPI =
         getAnalysis<BranchProbabilityInfoWrapperPass>().getBPI();
     auto *AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
-    LoopPredication LP(AA, SE, &BPI);
+    LoopPredication LP(AA, SE, LI, &BPI);
     return LP.runOnLoop(L);
   }
 };
 
 char LoopPredicationLegacyPass::ID = 0;
 } // end namespace llvm
 
 INITIALIZE_PASS_BEGIN(LoopPredicationLegacyPass, "loop-predication",

 
 PreservedAnalyses LoopPredicationPass::run(Loop &L, LoopAnalysisManager &AM,
                                            LoopStandardAnalysisResults &AR,
                                            LPMUpdater &U) {
   const auto &FAM =
       AM.getResult<FunctionAnalysisManagerLoopProxy>(L, AR).getManager();
   Function *F = L.getHeader()->getParent();
   auto *BPI = FAM.getCachedResult<BranchProbabilityAnalysis>(*F);
-  LoopPredication LP(&AR.AA, &AR.SE, BPI);
+  LoopPredication LP(&AR.AA, &AR.SE, &AR.LI, BPI);
   if (!LP.runOnLoop(&L))
     return PreservedAnalyses::all();
 
   return getLoopPassPreservedAnalyses();
 }
 
 Optional<LoopICmp>
 LoopPredication::parseLoopICmp(ICmpInst *ICI) {
if (ExitingBlockProbability > LatchProbabilityThreshold)
       return false;
   }
   // Using BPI, we have concluded that the most probable way to exit from the
   // loop is through the latch (or there's no profile information and all
   // exits are equally likely).
   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<BasicBlock *, BasicBlock *>
+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<BasicBlock *, BasicBlock *> LastPair;
+  for (std::pair<BasicBlock *, BasicBlock *>
+         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<BranchInst>(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<BasicBlock*, 16> 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<SCEVCouldNotCompute>(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<BranchInst>(ExitingBB->getTerminator());
+    if (!BI)
+      continue;
+
+    // If already constant, nothing to do.
+    if (isa<Constant>(BI->getCondition()))
+      continue;
+
+    const SCEV *ExitCount = SE->getExitCount(L, ExitingBB);
+    if (isa<SCEVCouldNotCompute>(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;
 
   LLVM_DEBUG(dbgs() << "Analyzing ");
   LLVM_DEBUG(L->dump());
 
   Module *M = L->getHeader()->getModule();
 
for (auto &I : *BB)
       if (isGuard(&I))
         Guards.push_back(cast<IntrinsicInst>(&I));
     if (PredicateWidenableBranchGuards &&
         isGuardAsWidenableBranch(BB->getTerminator()))
       GuardsAsWidenableBranches.push_back(
           cast<BranchInst>(BB->getTerminator()));
   }
 
-  if (Guards.empty() && GuardsAsWidenableBranches.empty())
-    return false;
-
   SCEVExpander Expander(*SE, *DL, "loop-predication");
-
   bool Changed = false;
+  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;
 }