[LoopUnroll] Limit evaluation to 2 iterations for most cases (WIP).

llvm/include/llvm/Analysis/LoopUnrollAnalyzer.h

 class UnrolledInstAnalyzer : private InstVisitor<UnrolledInstAnalyzer, bool> {
   typedef InstVisitor<UnrolledInstAnalyzer, bool> Base;
   friend class InstVisitor<UnrolledInstAnalyzer, bool>;
   struct SimplifiedAddress {
     Value *Base = nullptr;
     ConstantInt *Offset = nullptr;
   };
 
+  bool NeedsResimplify = false;
+
 public:
   UnrolledInstAnalyzer(unsigned Iteration,
                        DenseMap<Value *, Value *> &SimplifiedValues,
                        ScalarEvolution &SE, const Loop *L)
       : SimplifiedValues(SimplifiedValues), SE(SE), L(L) {
       IterationNumber = SE.getConstant(APInt(64, Iteration));
   }
 
   // Allow access to the initial visit method.
   using Base::visit;
 
+  bool needsResimplify() const { return NeedsResimplify; }
+
 private:
   /// A cache of pointer bases and constant-folded offsets corresponding
   /// to GEP (or derived from GEP) instructions.
   ///
   /// In order to find the base pointer one needs to perform non-trivial
   /// traversal of the corresponding SCEV expression, so it's good to have the
   /// results saved.
   DenseMap<Value *, SimplifiedAddress> SimplifiedAddresses;

llvm/lib/Analysis/LoopUnrollAnalyzer.cpp

if (Index >= CDS->getNumElements()) {
     // FIXME: For now we conservatively ignore out of bound accesses, but
     // we're allowed to perform the optimization in this case.
     return false;
   }
 
   Constant *CV = CDS->getElementAsConstant(Index);
   assert(CV && "Constant expected.");
   SimplifiedValues[&I] = CV;
+  NeedsResimplify = true;
 
   return true;
 }
 
 /// Try to simplify cast instruction.
 bool UnrolledInstAnalyzer::visitCastInst(CastInst &I) {
   Value *Op = I.getOperand(0);
   if (Value *Simplified = SimplifiedValues.lookup(Op))

llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp

static Optional<EstimatedUnrollCost> analyzeLoopUnrollCost(
   // A small worklist used to accumulate cost of instructions from each
   // observable and reached root in the loop.
   SmallVector<Instruction *, 16> CostWorklist;
 
   // PHI-used worklist used between iterations while accumulating cost.
   SmallVector<Instruction *, 4> PHIUsedList;
 
   // Helper function to accumulate cost for instructions in the loop.
-  auto AddCostRecursively = [&](Instruction &RootI, int Iteration) {
+  auto AddCostRecursively = [&](Instruction &RootI, int Iteration,
+                                bool SingleIterEvalOnly = false) {
     assert(Iteration >= 0 && "Cannot have a negative iteration!");
     assert(CostWorklist.empty() && "Must start with an empty cost list");
     assert(PHIUsedList.empty() && "Must start with an empty phi used list");
     CostWorklist.push_back(&RootI);
     TargetTransformInfo::TargetCostKind CostKind =
       RootI.getFunction()->hasMinSize() ?
       TargetTransformInfo::TCK_CodeSize :
       TargetTransformInfo::TCK_SizeAndLatency;
for (;; --Iteration) {
         // Mark that we are counting the cost of this instruction now.
         Cost.IsCounted = true;
 
         // If this is a PHI node in the loop header, just add it to the PHI set.
         if (auto *PhiI = dyn_cast<PHINode>(I))
           if (PhiI->getParent() == L->getHeader()) {
             assert(Cost.IsFree && "Loop PHIs shouldn't be evaluated as they "
                                   "inherently simplify during unrolling.");
-            if (Iteration == 0)
+            if (!SingleIterEvalOnly && Iteration == 0)
               continue;
 
             // Push the incoming value from the backedge into the PHI used list
             // if it is an in-loop instruction. We'll use this to populate the
             // cost worklist for the next iteration (as we count backwards).
             if (auto *OpI = dyn_cast<Instruction>(
                     PhiI->getIncomingValueForBlock(L->getLoopLatch())))
-              if (L->contains(OpI))
+              if (L->contains(OpI)) {
+                if (SingleIterEvalOnly)
+                  CostWorklist.push_back(OpI);
+                else
                   PHIUsedList.push_back(OpI);
+              }
             continue;
           }
 
         // First accumulate the cost of this instruction.
         if (!Cost.IsFree) {
           UnrolledCost += TTI.getUserCost(I, CostKind);
           LLVM_DEBUG(dbgs() << "Adding cost of instruction (iteration "
                             << Iteration << "): ");
for (;; --Iteration) {
           CostWorklist.push_back(OpI);
         }
       } while (!CostWorklist.empty());
 
       if (PHIUsedList.empty())
         // We've exhausted the search.
         break;
 
+      if (SingleIterEvalOnly)
+        break;
       assert(Iteration > 0 &&
              "Cannot track PHI-used values past the first iteration!");
       CostWorklist.append(PHIUsedList.begin(), PHIUsedList.end());
       PHIUsedList.clear();
     }
   };
 
   // Ensure that we don't violate the loop structure invariants relied on by
   // this analysis.
   assert(L->isLoopSimplifyForm() && "Must put loop into normal form first.");
   assert(L->isLCSSAForm(DT) &&
          "Must have loops in LCSSA form to track live-out values.");
 
   LLVM_DEBUG(dbgs() << "Starting LoopUnroll profitability analysis...\n");
 
   TargetTransformInfo::TargetCostKind CostKind =
     L->getHeader()->getParent()->hasMinSize() ?
     TargetTransformInfo::TCK_CodeSize : TargetTransformInfo::TCK_SizeAndLatency;
+
+  bool SingleIterEvalOnly = false;
+
+  InstructionCost FirstIterUnrolledCost = 0;
+
   // Simulate execution of each iteration of the loop counting instructions,
   // which would be simplified.
   // Since the same load will take different values on different iterations,
   // we literally have to go through all loop's iterations.
   for (unsigned Iteration = 0; Iteration < TripCount; ++Iteration) {
     LLVM_DEBUG(dbgs() << " Analyzing iteration " << Iteration << "\n");
 
     // Prepare for the iteration by collecting any simplified entry or backedge
for (unsigned Iteration = 0; Iteration < TripCount; ++Iteration) {
 
     // If we found no optimization opportunities on the first iteration, we
     // won't find them on later ones too.
     if (UnrolledCost == RolledDynamicCost) {
       LLVM_DEBUG(dbgs() << "  No opportunities found.. exiting.\n"
                         << "  UnrolledCost: " << UnrolledCost << "\n");
       return None;
     }
+
+    if (Iteration == 0)
+      FirstIterUnrolledCost = UnrolledCost;
+
+    if (Iteration == 1 && !Analyzer.needsResimplify()) {
+      SingleIterEvalOnly = true;
+      UnrolledCost -= FirstIterUnrolledCost;
+      SmallVector<BasicBlock *> ExitingBlocks;
+      L->getExitingBlocks(ExitingBlocks);
+      for (BasicBlock *Exiting : ExitingBlocks) {
+        for (BasicBlock *Succ : successors(Exiting)) {
+          if (L->contains(Succ))
+            continue;
+          ExitWorklist.insert({Exiting, Succ});
+        }
+      }
+      break;
+    }
   }
 
   while (!ExitWorklist.empty()) {
     BasicBlock *ExitingBB, *ExitBB;
     std::tie(ExitingBB, ExitBB) = ExitWorklist.pop_back_val();
 
     for (Instruction &I : *ExitBB) {
       auto *PN = dyn_cast<PHINode>(&I);
       if (!PN)
         break;
 
       Value *Op = PN->getIncomingValueForBlock(ExitingBB);
       if (auto *OpI = dyn_cast<Instruction>(Op))
         if (L->contains(OpI))
-          AddCostRecursively(*OpI, TripCount - 1);
+          AddCostRecursively(*OpI, SingleIterEvalOnly ? 1 : TripCount - 1,
+                             SingleIterEvalOnly);
     }
   }
 
+  if (SingleIterEvalOnly) {
+    UnrolledCost *= TripCount;
+    RolledDynamicCost = (RolledDynamicCost / 2) * TripCount;
+  }
+
   assert(UnrolledCost.isValid() && RolledDynamicCost.isValid() &&
          "All instructions must have a valid cost, whether the "
          "loop is rolled or unrolled.");
 
   LLVM_DEBUG(dbgs() << "Analysis finished:\n"
                     << "UnrolledCost: " << UnrolledCost << ", "
                     << "RolledDynamicCost: " << RolledDynamicCost << "\n");
   return {{unsigned(*UnrolledCost.getValue()),