[LoopReroll] Refactor most of reroll() into a helper class

[jmolloy]

Summary

reroll() was slightly monolithic and a pain to modify. Refactor a bunch of its state from local variables to member variables of a helper class, and do some trivial simplification while we're there.

This should set the scene for some changes to the Scale/Root model to make it catch more cases.

[hfinkel]

With modifications suggested below, LGTM. Thanks!

lib/Transforms/Scalar/LoopRerollPass.cpp

143

I'd prefer that we not do this just for access to the analysis variables. Just pass what is needed to the child class's constructor.

698

IIRC, the current version deals with duplicated. I'm not sure this is necessary (one could certainly argue that CSE should take care of anything interesting). Please explain why, however, you're ignoring duplicates in this comment.

[jmolloy]

Hi Hal,

Thanks!

Comments inline.

James

lib/Transforms/Scalar/LoopRerollPass.cpp

143

Sure, I'll make that change.

698

That is true. The main reason is simplification of the code: I'm going to be expanding the model to deal with the case with multiple root sets: consider:

a[i*2+0] = ...; // one root set
a[i*2+1] = ...;
a[i*2+4] = ...; // next root set
a[i*2+5] = ...;
a[(i+45)*2+0] = ...; // and another root set
a[(i+45)*2+1] = ...;

And if one root set can have duplicate indices, it adds another level of indirection I need to look through.

It's doable, but as you say it should be unlikely to be needed given a trivial CSE.

[jmolloy]

r227439.

lib/Transforms/Scalar/LoopRerollPass.cpp

void getAnalysisUsage(AnalysisUsage &AU) const override {
       AU.addPreserved<LoopInfoWrapperPass>();
       AU.addRequired<DominatorTreeWrapperPass>();
       AU.addPreserved<DominatorTreeWrapperPass>();
       AU.addRequired<ScalarEvolution>();
       AU.addRequired<TargetLibraryInfoWrapperPass>();
     }
 
 protected:
+    // DAGRootTracker can access analyses such as SCEV.
+    friend class DAGRootTracker;
+
     AliasAnalysis *AA;
     LoopInfo *LI;
     ScalarEvolution *SE;
     const DataLayout *DL;
     TargetLibraryInfo *TLI;
     DominatorTree *DT;
 
     typedef SmallVector<Instruction *, 16> SmallInstructionVector;
protected:
       // The vector of all possible reductions (for any scale).
       SmallReductionVector PossibleReds;
 
       DenseMap<Instruction *, int> PossibleRedIdx;
       DenseMap<Instruction *, int> PossibleRedIter;
       DenseSet<int> Reds;
     };
 
-    void collectPossibleIVs(Loop *L, SmallInstructionVector &PossibleIVs);
-    void collectPossibleReductions(Loop *L,
-           ReductionTracker &Reductions);
-    void collectInLoopUserSet(Loop *L,
-           const SmallInstructionVector &Roots,
+    // The set of all DAG roots, and state tracking of all roots
+    // for a particular induction variable.
+    struct DAGRootTracker {
+      DAGRootTracker(LoopReroll *Parent, Loop *L, Instruction *IV) :
+        Parent(Parent), SE(Parent->SE), L(L), IV(IV) {
+      }
+
+      /// Stage 1: Find all the DAG roots for the induction variable.
+      bool findRoots();
+      /// Stage 2: Validate if the found roots are valid.
+      bool validate(ReductionTracker &Reductions);
+      /// Stage 3: Assuming validate() returned true, perform the
+      /// replacement.
+      /// @param IterCount The maximum iteration count of L.
+      void replace(const SCEV *IterCount);
+
+    protected:
+      bool findScaleFromMul();
+      bool collectAllRoots();
+
+      void collectInLoopUserSet(const SmallInstructionVector &Roots,
                                 const SmallInstructionSet &Exclude,
                                 const SmallInstructionSet &Final,
                                 DenseSet<Instruction *> &Users);
-    void collectInLoopUserSet(Loop *L,
-           Instruction * Root,
+      void collectInLoopUserSet(Instruction *Root,
                                 const SmallInstructionSet &Exclude,
                                 const SmallInstructionSet &Final,
                                 DenseSet<Instruction *> &Users);
-    bool findScaleFromMul(Instruction *RealIV, uint64_t &Scale,
-                          Instruction *&IV,
-                          SmallInstructionVector &LoopIncs);
-    bool collectAllRoots(Loop *L, uint64_t Inc, uint64_t Scale, Instruction *IV,
-                         SmallVector<SmallInstructionVector, 32> &Roots,
-                         SmallInstructionSet &AllRoots,
-                         SmallInstructionVector &LoopIncs);
+
+      LoopReroll *Parent;
+
+      // Members of Parent, replicated here for brevity.
+      ScalarEvolution *SE;
+      Loop *L;
+
+      // The loop induction variable.
+      Instruction *IV;
+      // Loop step amount.
+      uint64_t Inc;
+      // Loop reroll count; if Inc == 1, this records the scaling applied
+      // to the indvar: a[i*2+0] = ...; a[i*2+1] = ... ;
+      // If Inc is not 1, Scale = Inc.
+      uint64_t Scale;
+      // If Scale != Inc, then RealIV is IV after its multiplication.
+      Instruction *RealIV;
+      // The roots themselves.
+      SmallInstructionVector Roots;
+      // All increment instructions for IV.
+      SmallInstructionVector LoopIncs;
+      // All instructions transitively used by any root.
+      DenseSet<Instruction *> AllRootUses;
+      // All instructions transitively used by the base.
+      DenseSet<Instruction *> BaseUseSet;
+      // All instructions transitively used by the increments.
+      DenseSet<Instruction *> LoopIncUseSet;
+    };
+
+    void collectPossibleIVs(Loop *L, SmallInstructionVector &PossibleIVs);
+    void collectPossibleReductions(Loop *L,
+           ReductionTracker &Reductions);
     bool reroll(Instruction *IV, Loop *L, BasicBlock *Header, const SCEV *IterCount,
                 ReductionTracker &Reductions);
   };
 }
 
 char LoopReroll::ID = 0;
 INITIALIZE_PASS_BEGIN(LoopReroll, "loop-reroll", "Reroll loops", false, false)
 INITIALIZE_AG_DEPENDENCY(AliasAnalysis)

 //   1. Instructions in the set of excluded instructions are never added to the
 //   use set (even if they are users). This is used, for example, to exclude
 //   including root increments in the use set of the primary IV.
 //
 //   2. Instructions in the set of final instructions are added to the use set
 //   if they are users, but their users are not added. This is used, for
 //   example, to prevent a reduction update from forcing all later reduction
 //   updates into the use set.
-void LoopReroll::collectInLoopUserSet(Loop *L,
+void LoopReroll::DAGRootTracker::collectInLoopUserSet(
   Instruction *Root, const SmallInstructionSet &Exclude,
   const SmallInstructionSet &Final,
   DenseSet<Instruction *> &Users) {
   SmallInstructionVector Queue(1, Root);
   while (!Queue.empty()) {
     Instruction *I = Queue.pop_back_val();
     if (!Users.insert(I).second)
       continue;
for (User::op_iterator OI = I->op_begin(),
             !Final.count(Op))
           Queue.push_back(Op);
     }
   }
 }
 
 // Collect all of the users of all of the provided root instructions (combined
 // into a single set).
-void LoopReroll::collectInLoopUserSet(Loop *L,
+void LoopReroll::DAGRootTracker::collectInLoopUserSet(
   const SmallInstructionVector &Roots,
   const SmallInstructionSet &Exclude,
   const SmallInstructionSet &Final,
   DenseSet<Instruction *> &Users) {
   for (SmallInstructionVector::const_iterator I = Roots.begin(),
        IE = Roots.end(); I != IE; ++I)
-    collectInLoopUserSet(L, *I, Exclude, Final, Users);
+    collectInLoopUserSet(*I, Exclude, Final, Users);
 }
 
 static bool isSimpleLoadStore(Instruction *I) {
   if (LoadInst *LI = dyn_cast<LoadInst>(I))
     return LI->isSimple();
   if (StoreInst *SI = dyn_cast<StoreInst>(I))
     return SI->isSimple();
   if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(I))
     return !MI->isVolatile();
   return false;
 }
 
+bool LoopReroll::DAGRootTracker::findRoots() {
+
+  const SCEVAddRecExpr *RealIVSCEV = cast<SCEVAddRecExpr>(SE->getSCEV(IV));
+  Inc = cast<SCEVConstant>(RealIVSCEV->getOperand(1))->
+    getValue()->getZExtValue();
+
+  // The effective induction variable, IV, is normally also the real induction
+  // variable. When we're dealing with a loop like:
+  //   for (int i = 0; i < 500; ++i)
+  //     x[3*i] = ...;
+  //     x[3*i+1] = ...;
+  //     x[3*i+2] = ...;
+  // then the real IV is still i, but the effective IV is (3*i).
+  Scale = Inc;
+  RealIV = IV;
+  if (Inc == 1 && !findScaleFromMul())
+    return false;
+
+  // The set of increment instructions for each increment value.
+  if (!collectAllRoots())
+    return false;
+
+  return true;
+}
+
 // Recognize loops that are setup like this:
 //
 // %iv = phi [ (preheader, ...), (body, %iv.next) ]
 // %scaled.iv = mul %iv, scale
 // f(%scaled.iv)
 // %scaled.iv.1 = add %scaled.iv, 1
 // f(%scaled.iv.1)
 // %scaled.iv.2 = add %scaled.iv, 2
 // f(%scaled.iv.2)
 // %scaled.iv.scale_m_1 = add %scaled.iv, scale-1
 // f(%scaled.iv.scale_m_1)
 // ...
 // %iv.next = add %iv, 1
 // %cmp = icmp(%iv, ...)
 // br %cmp, header, exit
 //
 // and, if found, set IV = %scaled.iv, and add %iv.next to LoopIncs.
-bool LoopReroll::findScaleFromMul(Instruction *RealIV, uint64_t &Scale,
-                                  Instruction *&IV,
-                                  SmallInstructionVector &LoopIncs) {
+bool LoopReroll::DAGRootTracker::findScaleFromMul() {
+
   // This is a special case: here we're looking for all uses (except for
   // the increment) to be multiplied by a common factor. The increment must
   // be by one. This is to capture loops like:
   //   for (int i = 0; i < 500; ++i) {
   //     foo(3*i); foo(3*i+1); foo(3*i+2);
   //   }
   if (RealIV->getNumUses() != 2)
     return false;
if (const SCEVConstant *MulScale =
     if (!MulScaleCI->uge(2) || MulScaleCI->uge(MaxInc))
       return false;
     Scale = MulScaleCI->getZExtValue();
     IV = User1;
   } else
     return false;
 
   DEBUG(dbgs() << "LRR: Found possible scaling " << *User1 << "\n");
+
+  assert(Scale <= MaxInc && "Scale is too large");
+  assert(Scale > 1 && "Scale must be at least 2");
+
   return true;
 }
 
 // Collect all root increments with respect to the provided induction variable
 // (normally the PHI, but sometimes a multiply). A root increment is an
 // instruction, normally an add, with a positive constant less than Scale. In a
 // rerollable loop, each of these increments is the root of an instruction
 // graph isomorphic to the others. Also, we collect the final induction
 // increment (the increment equal to the Scale), and its users in LoopIncs.
-bool LoopReroll::collectAllRoots(Loop *L, uint64_t Inc, uint64_t Scale,
-                                 Instruction *IV,
-                                 SmallVector<SmallInstructionVector, 32> &Roots,
-                                 SmallInstructionSet &AllRoots,
-                                 SmallInstructionVector &LoopIncs) {
+bool LoopReroll::DAGRootTracker::collectAllRoots() {
+  Roots.resize(Scale-1);
+  
   for (User *U : IV->users()) {
     Instruction *UI = cast<Instruction>(U);
     if (!SE->isSCEVable(UI->getType()))
       continue;
     if (UI->getType() != IV->getType())
       continue;
     if (!L->contains(UI))
       continue;
     if (hasUsesOutsideLoop(UI, L))
       continue;
 
     if (const SCEVConstant *Diff = dyn_cast<SCEVConstant>(SE->getMinusSCEV(
           SE->getSCEV(UI), SE->getSCEV(IV)))) {
       uint64_t Idx = Diff->getValue()->getValue().getZExtValue();
       if (Idx > 0 && Idx < Scale) {
-        Roots[Idx-1].push_back(UI);
-        AllRoots.insert(UI);
+        if (Roots[Idx-1])
+          // No duplicates allowed.
+          return false;
+        Roots[Idx-1] = UI;
       } else if (Idx == Scale && Inc > 1) {
         LoopIncs.push_back(UI);
       }
     }
   }
 
-  if (Roots[0].empty())
-    return false;
-  bool AllSame = true;
-  for (unsigned i = 1; i < Scale-1; ++i)
-    if (Roots[i].size() != Roots[0].size()) {
-      AllSame = false;
-      break;
-    }
-
-  if (!AllSame)
-    return false;
-
-  return true;
-}
-
-// Validate the selected reductions. All iterations must have an isomorphic
-// part of the reduction chain and, for non-associative reductions, the chain
-// entries must appear in order.
-bool LoopReroll::ReductionTracker::validateSelected() {
-  // For a non-associative reduction, the chain entries must appear in order.
-  for (DenseSet<int>::iterator RI = Reds.begin(), RIE = Reds.end();
-       RI != RIE; ++RI) {
-    int i = *RI;
-    int PrevIter = 0, BaseCount = 0, Count = 0;
-    for (Instruction *J : PossibleReds[i]) {
-      // Note that all instructions in the chain must have been found because
-      // all instructions in the function must have been assigned to some
-      // iteration.
-      int Iter = PossibleRedIter[J];
-      if (Iter != PrevIter && Iter != PrevIter + 1 &&
-          !PossibleReds[i].getReducedValue()->isAssociative()) {
-        DEBUG(dbgs() << "LRR: Out-of-order non-associative reduction: " <<
-                        J << "\n");
-        return false;
-      }
-
-      if (Iter != PrevIter) {
-        if (Count != BaseCount) {
-          DEBUG(dbgs() << "LRR: Iteration " << PrevIter <<
-                " reduction use count " << Count <<
-                " is not equal to the base use count " <<
-                BaseCount << "\n");
-          return false;
-        }
+  for (unsigned i = 0; i < Scale-1; ++i) {
+    if (!Roots[i])
 
-        Count = 0;
-      }
-
-      ++Count;
-      if (Iter == 0)
-        ++BaseCount;
-
-      PrevIter = Iter;
-    }
-  }
-
   return true;
 }
 
-// For all selected reductions, remove all parts except those in the first
-// iteration (and the PHI). Replace outside uses of the reduced value with uses
-// of the first-iteration reduced value (in other words, reroll the selected
-// reductions).
-void LoopReroll::ReductionTracker::replaceSelected() {
-  // Fixup reductions to refer to the last instruction associated with the
-  // first iteration (not the last).
-  for (DenseSet<int>::iterator RI = Reds.begin(), RIE = Reds.end();
-       RI != RIE; ++RI) {
-    int i = *RI;
-    int j = 0;
-    for (int e = PossibleReds[i].size(); j != e; ++j)
-      if (PossibleRedIter[PossibleReds[i][j]] != 0) {
-        --j;
-        break;
-      }
-
-    // Replace users with the new end-of-chain value.
-    SmallInstructionVector Users;
-    for (User *U : PossibleReds[i].getReducedValue()->users())
-      Users.push_back(cast<Instruction>(U));
-
-    for (SmallInstructionVector::iterator J = Users.begin(),
-         JE = Users.end(); J != JE; ++J)
-      (*J)->replaceUsesOfWith(PossibleReds[i].getReducedValue(),
-                              PossibleReds[i][j]);
-  }
-}
-
-// Reroll the provided loop with respect to the provided induction variable.
-// Generally, we're looking for a loop like this:
-//
-// %iv = phi [ (preheader, ...), (body, %iv.next) ]
-// f(%iv)
-// %iv.1 = add %iv, 1                <-- a root increment
-// f(%iv.1)
-// %iv.2 = add %iv, 2                <-- a root increment
-// f(%iv.2)
-// %iv.scale_m_1 = add %iv, scale-1  <-- a root increment
-// f(%iv.scale_m_1)
-// ...
-// %iv.next = add %iv, scale
-// %cmp = icmp(%iv, ...)
-// br %cmp, header, exit
-//
-// Notably, we do not require that f(%iv), f(%iv.1), etc. be isolated groups of
-// instructions. In other words, the instructions in f(%iv), f(%iv.1), etc. can
-// be intermixed with eachother. The restriction imposed by this algorithm is
-// that the relative order of the isomorphic instructions in f(%iv), f(%iv.1),
-// etc. be the same.
-//
-// First, we collect the use set of %iv, excluding the other increment roots.
-// This gives us f(%iv). Then we iterate over the loop instructions (scale-1)
-// times, having collected the use set of f(%iv.(i+1)), during which we:
-//   - Ensure that the next unmatched instruction in f(%iv) is isomorphic to
-//     the next unmatched instruction in f(%iv.(i+1)).
-//   - Ensure that both matched instructions don't have any external users
-//     (with the exception of last-in-chain reduction instructions).
-//   - Track the (aliasing) write set, and other side effects, of all
-//     instructions that belong to future iterations that come before the matched
-//     instructions. If the matched instructions read from that write set, then
-//     f(%iv) or f(%iv.(i+1)) has some dependency on instructions in
-//     f(%iv.(j+1)) for some j > i, and we cannot reroll the loop. Similarly,
-//     if any of these future instructions had side effects (could not be
-//     speculatively executed), and so do the matched instructions, when we
-//     cannot reorder those side-effect-producing instructions, and rerolling
-//     fails.
-//
-// Finally, we make sure that all loop instructions are either loop increment
-// roots, belong to simple latch code, parts of validated reductions, part of
-// f(%iv) or part of some f(%iv.i). If all of that is true (and all reductions
-// have been validated), then we reroll the loop.
-bool LoopReroll::reroll(Instruction *IV, Loop *L, BasicBlock *Header,
-                        const SCEV *IterCount,
-                        ReductionTracker &Reductions) {
-  const SCEVAddRecExpr *RealIVSCEV = cast<SCEVAddRecExpr>(SE->getSCEV(IV));
-  uint64_t Inc = cast<SCEVConstant>(RealIVSCEV->getOperand(1))->
-                   getValue()->getZExtValue();
-  // The collection of loop increment instructions.
-  SmallInstructionVector LoopIncs;
-  uint64_t Scale = Inc;
-
-  // The effective induction variable, IV, is normally also the real induction
-  // variable. When we're dealing with a loop like:
-  //   for (int i = 0; i < 500; ++i)
-  //     x[3*i] = ...;
-  //     x[3*i+1] = ...;
-  //     x[3*i+2] = ...;
-  // then the real IV is still i, but the effective IV is (3*i).
-  Instruction *RealIV = IV;
-  if (Inc == 1 && !findScaleFromMul(RealIV, Scale, IV, LoopIncs))
-    return false;
-
-  assert(Scale <= MaxInc && "Scale is too large");
-  assert(Scale > 1 && "Scale must be at least 2");
-
-  // The set of increment instructions for each increment value.
-  SmallVector<SmallInstructionVector, 32> Roots(Scale-1);
-  SmallInstructionSet AllRoots;
-  if (!collectAllRoots(L, Inc, Scale, IV, Roots, AllRoots, LoopIncs))
-    return false;
-
-  DEBUG(dbgs() << "LRR: Found all root induction increments for: " <<
-                  *RealIV << "\n");
-
-  // An array of just the possible reductions for this scale factor. When we
-  // collect the set of all users of some root instructions, these reduction
-  // instructions are treated as 'final' (their uses are not considered).
-  // This is important because we don't want the root use set to search down
-  // the reduction chain.
-  SmallInstructionSet PossibleRedSet;
-  SmallInstructionSet PossibleRedLastSet, PossibleRedPHISet;
-  Reductions.restrictToScale(Scale, PossibleRedSet, PossibleRedPHISet,
-                             PossibleRedLastSet);
+bool LoopReroll::DAGRootTracker::validate(ReductionTracker &Reductions) {
+  BasicBlock *Header = L->getHeader();
 
   // We now need to check for equivalence of the use graph of each root with
   // that of the primary induction variable (excluding the roots). Our goal
   // here is not to solve the full graph isomorphism problem, but rather to
   // catch common cases without a lot of work. As a result, we will assume
   // that the relative order of the instructions in each unrolled iteration
   // is the same (although we will not make an assumption about how the
   // different iterations are intermixed). Note that while the order must be
   // the same, the instructions may not be in the same basic block.
-  SmallInstructionSet Exclude(AllRoots);
+  SmallInstructionSet Exclude;
+  Exclude.insert(Roots.begin(), Roots.end());
   Exclude.insert(LoopIncs.begin(), LoopIncs.end());
 
-  DenseSet<Instruction *> BaseUseSet;
-  collectInLoopUserSet(L, IV, Exclude, PossibleRedSet, BaseUseSet);
+  // An array of just the possible reductions for this scale factor. When we
+  // collect the set of all users of some root instructions, these reduction
+  // instructions are treated as 'final' (their uses are not considered).
+  // This is important because we don't want the root use set to search down
+  // the reduction chain.
+  SmallInstructionSet PossibleRedSet;
+  SmallInstructionSet PossibleRedLastSet;
+  SmallInstructionSet PossibleRedPHISet;
+  Reductions.restrictToScale(Scale, PossibleRedSet,
+                             PossibleRedPHISet, PossibleRedLastSet);
+                                          
+
+  collectInLoopUserSet(IV, Exclude, PossibleRedSet, BaseUseSet);
 
-  DenseSet<Instruction *> AllRootUses;
   std::vector<DenseSet<Instruction *> > RootUseSets(Scale-1);
 
   bool MatchFailed = false;
   for (unsigned i = 0; i < Scale-1 && !MatchFailed; ++i) {
     DenseSet<Instruction *> &RootUseSet = RootUseSets[i];
-    collectInLoopUserSet(L, Roots[i], SmallInstructionSet(),
+    collectInLoopUserSet(Roots[i], SmallInstructionSet(),
                          PossibleRedSet, RootUseSet);
 
     DEBUG(dbgs() << "LRR: base use set size: " << BaseUseSet.size() <<
                     " vs. iteration increment " << (i+1) <<
                     " use set size: " << RootUseSet.size() << "\n");
 
     if (BaseUseSet.size() != RootUseSet.size()) {
       MatchFailed = true;
       break;
     }
 
     // In addition to regular aliasing information, we need to look for
     // instructions from later (future) iterations that have side effects
     // preventing us from reordering them past other instructions with side
     // effects.
     bool FutureSideEffects = false;
-    AliasSetTracker AST(*AA);
+    AliasSetTracker AST(*Parent->AA);
 
     // The map between instructions in f(%iv.(i+1)) and f(%iv).
     DenseMap<Value *, Value *> BaseMap;
 
     assert(L->getNumBlocks() == 1 && "Cannot handle multi-block loops");
     for (BasicBlock::iterator J1 = Header->begin(), J2 = Header->begin(),
          JE = Header->end(); J1 != JE && !MatchFailed; ++J1) {
       if (cast<Instruction>(J1) == RealIV)
         continue;
       if (cast<Instruction>(J1) == IV)
         continue;
       if (!BaseUseSet.count(J1))
         continue;
       if (PossibleRedPHISet.count(J1)) // Skip reduction PHIs.
         continue;
 
-      while (J2 != JE && (!RootUseSet.count(J2) ||
-             std::find(Roots[i].begin(), Roots[i].end(), J2) !=
-               Roots[i].end())) {
+      while (J2 != JE && (!RootUseSet.count(J2) || Roots[i] ==  J2)) {
         // As we iterate through the instructions, instructions that don't
         // belong to previous iterations (or the base case), must belong to
         // future iterations. We want to track the alias set of writes from
         // previous iterations.
         if (!isa<PHINode>(J2) && !BaseUseSet.count(J2) &&
             !AllRootUses.count(J2)) {
           if (J2->mayWriteToMemory())
             AST.add(J2);
 
           // Note: This is specifically guarded by a check on isa<PHINode>,
           // which while a valid (somewhat arbitrary) micro-optimization, is
           // needed because otherwise isSafeToSpeculativelyExecute returns
           // false on PHI nodes.
-          if (!isSimpleLoadStore(J2) && !isSafeToSpeculativelyExecute(J2, DL))
+          if (!isSimpleLoadStore(J2) && !isSafeToSpeculativelyExecute(J2, Parent->DL))
             FutureSideEffects = true;
         }
 
         ++J2;
       }
 
       if (!J1->isSameOperationAs(J2)) {
         DEBUG(dbgs() << "LRR: iteration root match failed at " << *J1 <<
for (BasicBlock::iterator J1 = Header->begin(), J2 = Header->begin(),
       }
 
       // Make sure that we don't alias with any instruction in the alias set
       // tracker. If we do, then we depend on a future iteration, and we
       // can't reroll.
       if (J2->mayReadFromMemory()) {
         for (AliasSetTracker::iterator K = AST.begin(), KE = AST.end();
              K != KE && !MatchFailed; ++K) {
-          if (K->aliasesUnknownInst(J2, *AA)) {
+          if (K->aliasesUnknownInst(J2, *Parent->AA)) {
             DEBUG(dbgs() << "LRR: iteration root match failed at " << *J1 <<
                             " vs. " << *J2 << " (depends on future store)\n");
             MatchFailed = true;
             break;
           }
         }
       }
 
       // If we've past an instruction from a future iteration that may have
       // side effects, and this instruction might also, then we can't reorder
       // them, and this matching fails. As an exception, we allow the alias
       // set tracker to handle regular (simple) load/store dependencies.
       if (FutureSideEffects &&
             ((!isSimpleLoadStore(J1) &&
-              !isSafeToSpeculativelyExecute(J1, DL)) ||
+              !isSafeToSpeculativelyExecute(J1, Parent->DL)) ||
              (!isSimpleLoadStore(J2) &&
-              !isSafeToSpeculativelyExecute(J2, DL)))) {
+              !isSafeToSpeculativelyExecute(J2, Parent->DL)))) {
         DEBUG(dbgs() << "LRR: iteration root match failed at " << *J1 <<
                         " vs. " << *J2 <<
                         " (side effects prevent reordering)\n");
         MatchFailed = true;
         break;
       }
 
       // For instructions that are part of a reduction, if the operation is
for (BasicBlock::iterator J1 = Header->begin(), J2 = Header->begin(),
           if (InReduction)
             if (Instruction *Op2I = dyn_cast<Instruction>(Op2))
               if (Reductions.isPairInSame(J2, Op2I))
                 continue;
 
           DenseMap<Value *, Value *>::iterator BMI = BaseMap.find(Op2);
           if (BMI != BaseMap.end())
             Op2 = BMI->second;
-          else if (std::find(Roots[i].begin(), Roots[i].end(),
-                             (Instruction*) Op2) != Roots[i].end())
+          else if (Roots[i] == (Instruction*) Op2)
             Op2 = IV;
 
           if (J1->getOperand(Swapped ? unsigned(!j) : j) != Op2) {
             // If we've not already decided to swap the matched operands, and
             // we've not already matched our first operand (note that we could
             // have skipped matching the first operand because it is part of a
             // reduction above), and the instruction is commutative, then try
             // the swapped match.
bool LoopReroll::DAGRootTracker::validate(ReductionTracker &Reductions) {
   }
 
   if (MatchFailed)
     return false;
 
   DEBUG(dbgs() << "LRR: Matched all iteration increments for " <<
                   *RealIV << "\n");
 
-  DenseSet<Instruction *> LoopIncUseSet;
-  collectInLoopUserSet(L, LoopIncs, SmallInstructionSet(),
+  collectInLoopUserSet(LoopIncs, SmallInstructionSet(),
                        SmallInstructionSet(), LoopIncUseSet);
   DEBUG(dbgs() << "LRR: Loop increment set size: " <<
                   LoopIncUseSet.size() << "\n");
 
   // Make sure that all instructions in the loop have been included in some
   // use set.
   for (BasicBlock::iterator J = Header->begin(), JE = Header->end();
        J != JE; ++J) {
     if (isa<DbgInfoIntrinsic>(J))
       continue;
     if (cast<Instruction>(J) == RealIV)
       continue;
     if (cast<Instruction>(J) == IV)
       continue;
     if (BaseUseSet.count(J) || AllRootUses.count(J) ||
         (LoopIncUseSet.count(J) && (J->isTerminator() ||
-                                    isSafeToSpeculativelyExecute(J, DL))))
+                                    isSafeToSpeculativelyExecute(J, Parent->DL))))
       continue;
 
-    if (AllRoots.count(J))
+    if (std::find(Roots.begin(), Roots.end(), J) != Roots.end())
       continue;
 
     if (Reductions.isSelectedPHI(J))
       continue;
 
     DEBUG(dbgs() << "LRR: aborting reroll based on " << *RealIV <<
                     " unprocessed instruction found: " << *J << "\n");
     MatchFailed = true;
     break;
   }
 
   if (MatchFailed)
     return false;
 
   DEBUG(dbgs() << "LRR: all instructions processed from " <<
                   *RealIV << "\n");
+  return true;
+}
 
-  if (!Reductions.validateSelected())
-    return false;
-
-  // At this point, we've validated the rerolling, and we're committed to
-  // making changes!
-
-  Reductions.replaceSelected();
-
+void LoopReroll::DAGRootTracker::replace(const SCEV *IterCount) {
+  BasicBlock *Header = L->getHeader();
   // Remove instructions associated with non-base iterations.
   for (BasicBlock::reverse_iterator J = Header->rbegin();
        J != Header->rend();) {
     if (AllRootUses.count(&*J)) {
       Instruction *D = &*J;
       DEBUG(dbgs() << "LRR: removing: " << *D << "\n");
       D->eraseFromParent();
       continue;
     }
 
     ++J;
   }
 
   // Insert the new induction variable.
+  const SCEVAddRecExpr *RealIVSCEV = cast<SCEVAddRecExpr>(SE->getSCEV(RealIV));
   const SCEV *Start = RealIVSCEV->getStart();
   if (Inc == 1)
     Start = SE->getMulExpr(Start,
                            SE->getConstant(Start->getType(), Scale));
   const SCEVAddRecExpr *H =
     cast<SCEVAddRecExpr>(SE->getAddRecExpr(Start,
                            SE->getConstant(RealIVSCEV->getType(), 1),
                            L, SCEV::FlagAnyWrap));
if (BranchInst *BI = dyn_cast<BranchInst>(Header->getTerminator())) {
           SE->getMinusSCEV(ICSCEV, SE->getConstant(ICSCEV->getType(), 1));
 
         Value *ICMinus1; // Iteration count minus 1
         if (isa<SCEVConstant>(ICMinus1SCEV)) {
           ICMinus1 = Expander.expandCodeFor(ICMinus1SCEV, NewIV->getType(), BI);
         } else {
           BasicBlock *Preheader = L->getLoopPreheader();
           if (!Preheader)
-            Preheader = InsertPreheaderForLoop(L, this);
+            Preheader = InsertPreheaderForLoop(L, Parent);
 
           ICMinus1 = Expander.expandCodeFor(ICMinus1SCEV, NewIV->getType(),
                                             Preheader->getTerminator());
         }
 
         Value *Cond =
             new ICmpInst(BI, CmpInst::ICMP_EQ, NewIV, ICMinus1, "exitcond");
         BI->setCondition(Cond);
 
         if (BI->getSuccessor(1) != Header)
           BI->swapSuccessors();
       }
     }
   }
 
-  SimplifyInstructionsInBlock(Header, DL, TLI);
-  DeleteDeadPHIs(Header, TLI);
+  SimplifyInstructionsInBlock(Header, Parent->DL, Parent->TLI);
+  DeleteDeadPHIs(Header, Parent->TLI);
+}
+
+// Validate the selected reductions. All iterations must have an isomorphic
+// part of the reduction chain and, for non-associative reductions, the chain
+// entries must appear in order.
+bool LoopReroll::ReductionTracker::validateSelected() {
+  // For a non-associative reduction, the chain entries must appear in order.
+  for (DenseSet<int>::iterator RI = Reds.begin(), RIE = Reds.end();
+       RI != RIE; ++RI) {
+    int i = *RI;
+    int PrevIter = 0, BaseCount = 0, Count = 0;
+    for (Instruction *J : PossibleReds[i]) {
+      // Note that all instructions in the chain must have been found because
+      // all instructions in the function must have been assigned to some
+      // iteration.
+      int Iter = PossibleRedIter[J];
+      if (Iter != PrevIter && Iter != PrevIter + 1 &&
+          !PossibleReds[i].getReducedValue()->isAssociative()) {
+        DEBUG(dbgs() << "LRR: Out-of-order non-associative reduction: " <<
+                        J << "\n");
+        return false;
+      }
+
+      if (Iter != PrevIter) {
+        if (Count != BaseCount) {
+          DEBUG(dbgs() << "LRR: Iteration " << PrevIter <<
+                " reduction use count " << Count <<
+                " is not equal to the base use count " <<
+                BaseCount << "\n");
+          return false;
+        }
+
+        Count = 0;
+      }
+
+      ++Count;
+      if (Iter == 0)
+        ++BaseCount;
+
+      PrevIter = Iter;
+    }
+  }
+
+  return true;
+}
+
+// For all selected reductions, remove all parts except those in the first
+// iteration (and the PHI). Replace outside uses of the reduced value with uses
+// of the first-iteration reduced value (in other words, reroll the selected
+// reductions).
+void LoopReroll::ReductionTracker::replaceSelected() {
+  // Fixup reductions to refer to the last instruction associated with the
+  // first iteration (not the last).
+  for (DenseSet<int>::iterator RI = Reds.begin(), RIE = Reds.end();
+       RI != RIE; ++RI) {
+    int i = *RI;
+    int j = 0;
+    for (int e = PossibleReds[i].size(); j != e; ++j)
+      if (PossibleRedIter[PossibleReds[i][j]] != 0) {
+        --j;
+        break;
+      }
+
+    // Replace users with the new end-of-chain value.
+    SmallInstructionVector Users;
+    for (User *U : PossibleReds[i].getReducedValue()->users())
+      Users.push_back(cast<Instruction>(U));
+
+    for (SmallInstructionVector::iterator J = Users.begin(),
+         JE = Users.end(); J != JE; ++J)
+      (*J)->replaceUsesOfWith(PossibleReds[i].getReducedValue(),
+                              PossibleReds[i][j]);
+  }
+}
+
+// Reroll the provided loop with respect to the provided induction variable.
+// Generally, we're looking for a loop like this:
+//
+// %iv = phi [ (preheader, ...), (body, %iv.next) ]
+// f(%iv)
+// %iv.1 = add %iv, 1                <-- a root increment
+// f(%iv.1)
+// %iv.2 = add %iv, 2                <-- a root increment
+// f(%iv.2)
+// %iv.scale_m_1 = add %iv, scale-1  <-- a root increment
+// f(%iv.scale_m_1)
+// ...
+// %iv.next = add %iv, scale
+// %cmp = icmp(%iv, ...)
+// br %cmp, header, exit
+//
+// Notably, we do not require that f(%iv), f(%iv.1), etc. be isolated groups of
+// instructions. In other words, the instructions in f(%iv), f(%iv.1), etc. can
+// be intermixed with eachother. The restriction imposed by this algorithm is
+// that the relative order of the isomorphic instructions in f(%iv), f(%iv.1),
+// etc. be the same.
+//
+// First, we collect the use set of %iv, excluding the other increment roots.
+// This gives us f(%iv). Then we iterate over the loop instructions (scale-1)
+// times, having collected the use set of f(%iv.(i+1)), during which we:
+//   - Ensure that the next unmatched instruction in f(%iv) is isomorphic to
+//     the next unmatched instruction in f(%iv.(i+1)).
+//   - Ensure that both matched instructions don't have any external users
+//     (with the exception of last-in-chain reduction instructions).
+//   - Track the (aliasing) write set, and other side effects, of all
+//     instructions that belong to future iterations that come before the matched
+//     instructions. If the matched instructions read from that write set, then
+//     f(%iv) or f(%iv.(i+1)) has some dependency on instructions in
+//     f(%iv.(j+1)) for some j > i, and we cannot reroll the loop. Similarly,
+//     if any of these future instructions had side effects (could not be
+//     speculatively executed), and so do the matched instructions, when we
+//     cannot reorder those side-effect-producing instructions, and rerolling
+//     fails.
+//
+// Finally, we make sure that all loop instructions are either loop increment
+// roots, belong to simple latch code, parts of validated reductions, part of
+// f(%iv) or part of some f(%iv.i). If all of that is true (and all reductions
+// have been validated), then we reroll the loop.
+bool LoopReroll::reroll(Instruction *IV, Loop *L, BasicBlock *Header,
+                        const SCEV *IterCount,
+                        ReductionTracker &Reductions) {
+  DAGRootTracker DAGRoots(this, L, IV);
+
+  if (!DAGRoots.findRoots())
+    return false;
+  DEBUG(dbgs() << "LRR: Found all root induction increments for: " <<
+                  *IV << "\n");
+  
+  if (!DAGRoots.validate(Reductions))
+    return false;
+  if (!Reductions.validateSelected())
+    return false;
+  // At this point, we've validated the rerolling, and we're committed to
+  // making changes!
+
+  Reductions.replaceSelected();
+  DAGRoots.replace(IterCount);
+
   ++NumRerolledLoops;
   return true;
 }
 
 bool LoopReroll::runOnLoop(Loop *L, LPPassManager &LPM) {
   if (skipOptnoneFunction(L))
     return false;