Expand loop reroll to handle loop with multiple induction variables and negative increment -part 1/3

[• hulx2000]

Summary

Handle loop with negtive induction variable increment

This patch extends LoopReroll pass to hand the loops which
are similar to the following:

      while (len > 1) {
            sum4 += buf[len];
            sum4 += buf[len-1];
            len -= 2;
        }

Basically, negative loop increment.

The original patch will be split into 3, that's the first one.

[• hulx2000]

handle some of Ana's comments

[hfinkel]

Please enhance your patch description to explain, not just what he patch enables (which you've done), but generically how it does it. There is a bunch of new code here, and it will be good to get a high-level overview of what it does (and I'd expect such a description to appear in the commit message anyway, so this is not "just for review").

Please upload the patch with full context, see:

http://llvm.org/docs/Phabricator.html#requesting-a-review-via-the-web-interface

for instructions.

lib/Transforms/Scalar/LoopRerollPass.cpp

447

We generally abbreviate Instruction to Inst, not Insn. Please rename this function.

Also, please comment what this function does. It is certainly more than just determining if the instruction is a comparison of some kind.

[• hulx2000]

Thanks, Hal, sorry for the delay, just saw your comments this morning

[jmolloy]

Hi,

I'm afraid there are too many different changes bundled together here for me to really review. Please could you split this out into:

1. Adding negative offset support
2. Adding pointer induction variable support
3. Adding multiple induction variable support.

Cheers,

James

[• hulx2000]

Hi, James:

Thank you very much for your comment.

Those patches are bundled together because they are strongly correlated, for an example, if I add support for pointer induction variable, I must support multiple induction variable, otherwise a lot of test will be broken.

I will see if I can split it into three with some extra safety check.

Regards

Lawrence Hu

[• hulx2000]

remove white spaces

[• hulx2000]

upload the full context diff

[jmolloy]

Hi,

Thanks for making these changes - it looks much better now. A few comments still.

Cheers,

James

lib/Transforms/Scalar/LoopRerollPass.cpp

626

I think this line and line 503 can both be merged, and put on line 505. There's no need to ZExt this one and SExt the other - both can be SExted.

629

This can be much simpler. You can just use APInt::abs():

AInt = AInt.abs();
if (AInt.isZero() || AInt.uge(MaxInc))
  continue;
IVToIncMap[I] = AInt.getSExtValue();

812–815

I'm not sure how adding GEP support here relates to negative values. Was this an accidentally added hunk?

867

I think you need to use std::abs() above, on line 721 too. Otherwise the guard on line 721 is useless.

1549

Double semicolon (;;).

Also, you can squash this into just one statement:

bool Negative = IVToIncMap[IV] < 0;

[• hulx2000]

lib/Transforms/Scalar/LoopRerollPass.cpp

812–815

Ｏｈ，　ｓｏｒｒｙ，　ｔｈｉｓ　ｉｓ　ｌｅｆｔ　ｏｖｅｒ　ｗｈｅｎ　Ｉ　ｓｐｌｉｔ　ｍｙ　ｐａｔｃｈ，　ｉｔ　ｉｓ　ｒｅｑｕｉｒｅ　ｆｏｒ　ｐｏｉｎｔｅｒ　ｉｎｄｕｃｔｉｏｎ　ｖａｒｉａｂｌｅ　ｔｈｏｕｇｈ，　ｂｕｔ　Ｉ　ｃａｎ　ｒｅｍｏｖｅ　ｉｔ　ｆｏｒ　ｎｏｗ．

Ｗｉｌｌ　ｆｉｘ　ｔｈｅ　ｒｅｓｔ，　ｔｈａｎｋｓ．``

[• hulx2000]

Thanks, James for your valuable comments.

I have updated my patch.

[jmolloy]

Hi,

This is looking a lot better, thanks!

James

lib/Transforms/Scalar/LoopRerollPass.cpp

373

This should be const, to notify a user that it is not captured or modified.

442

Please add a docstring, like all other variables.

809

I think this is still relating to adding GEP support, is it not?

test/Transforms/LoopReroll/negative.ll

16 ↗

(On Diff #29943)

This test will fail in release builds because registers will lose their names. Please either make it not rely on named variables or add a

; REQUIRES: asserts

to the top (please also check the line above is actually correct, I've just pasted it from memory!)

[• hulx2000]

Thanks Jame.

Sorry that missed a few thing last time, this time should be better.

Best, I didn't change the const thing, compiler complains, and I didn't see testcase asset in release build, I tried pass, it is fine, are yours a special build?

Regards

Lawrence Hu

[• hulx2000]

lib/Transforms/Scalar/LoopRerollPass.cpp

373

unfortunately, compiler always contains if I do that.

[jmolloy]

LGTM, thanks for making those changes.

If you don't add REQUIRES: asserts, you'll get complaints on the list when this goes in. I forget exactly which cmake variable you change to remove value names completely, but rest assured there are people who use it - I've found that out myself :)

Cheers,

James

[• hulx2000]

Great, Thanks, James.

I will add Assert from the next test and on.

Would you mind to commit it for you so I can start the next one? I don't have commit right yet.

[• hulx2000]

typo: for me.

[mzolotukhin]

lib/Transforms/Scalar/LoopRerollPass.cpp

169–172

s/it's/its/

172

I'd rather commit this typo fix immediately when I found it (no pre-commit review is required for this kind of changes).

442

s/it's/its/

808

No need to add a blank line here.

1548

Redundant whitespaces.

test/Transforms/LoopReroll/negative.ll

19–27 ↗

(On Diff #30013)

Do we really need to check that we generate *exactly* these statements? Would it be possible to look for, say, a specific branch (or absence of the branch) - this would make the test much more stable.

29–38 ↗

(On Diff #30013)

Do we really need to have this body, or would the test work with almost empty loop?
Like this:

while (len > 1) {
            len -= 2;
        }

instead of

while (len > 1) {
            sum4 += buf[len];
            sum4 += buf[len-1];
            len -= 2;
        }

Since you only run loop-reroll, it shouldn't be collapsed to nop.

53–54 ↗

(On Diff #30013)

Maybe just return i32 %sum4.0.lcssa to get rid of the goo function?

[• hulx2000]

test/Transforms/LoopReroll/negative.ll

29–38 ↗

(On Diff #30013)

The simple loop won't do enough test for the logic in LoopReroll

[• hulx2000]

Addressed Michael's comment

[• hulx2000]

Hi, Michael

Could you please accept it if you have now further comment? It will be easier for me to prepare next two patches.

Thanks.

Lawrence Hu

[• hulx2000]

Hi, Hal:

Do you have any other comments about this patch, could you please accept it if not?

Thanks

Lawrence Hu

[hfinkel]

lib/Transforms/Scalar/LoopRerollPass.cpp

1547–1548

Line is too long.

1550–1551

Line is too long.

test/Transforms/LoopReroll/negative.ll

2 ↗

(On Diff #30307)

Remove this. There is no valid reason why a test like this needs to be asserts-mode only. And it is really important that we have as much testing as possible for the NDEBUG builds (even if that were to require using more regexes in the test). That having been said, what James said was not quite right. The NDEBUG opt does not remove all variable names (certain passes did that, or used to, but that's another matter). What is true is that NDEBUG Clang does not produce variable names for locals.

[• hulx2000]

Thanks, Hal,

Updated.

Regards

[hfinkel]

LGTM now too, thanks!

[• hulx2000]

committed as git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@243171 91177308-0d34-0410-b5e6-96231b3b80d8

lib/Transforms/Scalar/LoopRerollPass.cpp

protected:
     LoopInfo *LI;
     ScalarEvolution *SE;
     TargetLibraryInfo *TLI;
     DominatorTree *DT;
 
     typedef SmallVector<Instruction *, 16> SmallInstructionVector;
     typedef SmallSet<Instruction *, 16>   SmallInstructionSet;
 
-    // A chain of isomorphic instructions, indentified by a single-use PHI,
+    Instruction *LoopCtrlOnlyIV;
+    DenseMap<Instruction *, int64_t> IVToIncMap;
+
+    // A chain of isomorphic instructions, identified by a single-use PHI,
     // representing a reduction. Only the last value may be used outside the
     // loop.
     struct SimpleLoopReduction {
       SimpleLoopReduction(Instruction *P, Loop *L)
         : Valid(false), Instructions(1, P) {
         assert(isa<PHINode>(P) && "First reduction instruction must be a PHI");
         add(L);
       }
protected:
 
     // A DAGRootSet models an induction variable being used in a rerollable
     // loop. For example,
     //
     //   x[i*3+0] = y1
     //   x[i*3+1] = y2
     //   x[i*3+2] = y3
     //
-    //   Base instruction -> i*3               
+    //   Base instruction -> i*3
     //                    +---+----+
     //                   /    |     \
     //               ST[y1]  +1     +2  <-- Roots
     //                        |      |
     //                      ST[y2] ST[y3]
     //
     // There may be multiple DAGRoots, for example:
     //
struct DAGRootSet {
       SmallInstructionSet SubsumedInsts;
     };
 
     // 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,
                      ScalarEvolution *SE, AliasAnalysis *AA,
-                     TargetLibraryInfo *TLI)
-          : Parent(Parent), L(L), SE(SE), AA(AA), TLI(TLI), IV(IV) {}
+                     TargetLibraryInfo *TLI, ReductionTracker &Reducs,
+                     DenseMap<Instruction *, int64_t> &IncrMap,
+                     Instruction *LoopCtrlIV)
+          : Parent(Parent), L(L), SE(SE), AA(AA), TLI(TLI), IV(IV),
+            Reductions(Reducs), IVToIncMap(IncrMap),
+            LoopCtrlOnlyIV(LoopCtrlIV) {}
 
       /// 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);
+      void replaceIV(Instruction *Inst, Instruction *IV, const SCEV *IterCount);
+      void replaceNonLoopCtrlIV();
 
     protected:
       typedef MapVector<Instruction*, SmallBitVector> UsesTy;
 
       bool findRootsRecursive(Instruction *IVU,
                               SmallInstructionSet SubsumedInsts);
       bool findRootsBase(Instruction *IVU, SmallInstructionSet SubsumedInsts);
       bool collectPossibleRoots(Instruction *Base,
protected:
       Loop *L;
       ScalarEvolution *SE;
       AliasAnalysis *AA;
       TargetLibraryInfo *TLI;
 
       // The loop induction variable.
       Instruction *IV;
       // Loop step amount.
-      uint64_t Inc;
+      int64_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;
       // The roots themselves.
       SmallVector<DAGRootSet,16> RootSets;
       // All increment instructions for IV.
       SmallInstructionVector LoopIncs;
       // Map of all instructions in the loop (in order) to the iterations
       // they are used in (or specially, IL_All for instructions
       // used in the loop increment mechanism).
       UsesTy Uses;
+      ReductionTracker &Reductions;
+      DenseMap<Instruction *, int64_t> &IVToIncMap;
+      Instruction *LoopCtrlOnlyIV;
     };
 
+    // Check if it is a compare-like instruction whose user is a branch
+    bool isCompareInst(Instruction *I, BasicBlock *Header) {
+      Instruction *Branch = dyn_cast<BranchInst>(Header->getTerminator());
+      if (I->hasOneUse() && Branch &&
+          (dyn_cast<BranchInst>(*(I->user_begin())) == Branch))
+        return true;
+      return false;
+    };
+
+    bool isLoopCtrlOnlyIV(Loop *L, Instruction *IV);
     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);
   };
 }
 

 static bool hasUsesOutsideLoop(Instruction *I, Loop *L) {
   for (User *U : I->users()) {
     if (!L->contains(cast<Instruction>(U)))
       return true;
   }
   return false;
 }
 
+static const SCEVConstant *getConstantCoefFromStep(ScalarEvolution *SE,
+                                                   const SCEV *SCEVExpr,
+                                                   Instruction *IV) {
+  const SCEVConstant *CIncSCEV = nullptr;
+  const SCEVMulExpr *MulSCEV = dyn_cast<SCEVMulExpr>(SCEVExpr);
+
+  // If StepRecurrence of a SCEVExpr is a constant (c1 * c2, c2 = sizeof(x)),
+  // Return c1.
+  if (!MulSCEV && IV->getType()->isPointerTy())
+    if (const SCEVConstant *IncSCEV = dyn_cast<SCEVConstant>(SCEVExpr)) {
+      const PointerType *PTy = dyn_cast<PointerType>(IV->getType());
+      Type *ElTy = PTy->getElementType();
+      const SCEV *SizeOfExpr =
+          SE->getSizeOfExpr(SE->getEffectiveSCEVType(IV->getType()), ElTy);
+      const SCEV *NewSCEV = nullptr;
+      if (IncSCEV->getValue()->getValue().isNegative()) {
+        NewSCEV = SE->getUDivExpr(SE->getNegativeSCEV(SCEVExpr), SizeOfExpr);
+        CIncSCEV = dyn_cast<SCEVConstant>(SE->getNegativeSCEV(NewSCEV));
+      } else
+        CIncSCEV =
+            dyn_cast<SCEVConstant>(SE->getUDivExpr(SCEVExpr, SizeOfExpr));
+      return CIncSCEV;
+    }
+
+  if (!MulSCEV)
+    return nullptr;
+
+  // If StepRecurrence of a SCEVExpr is a c * sizeof(x), where c is constant,
+  // Return c.
+  for (const SCEV *Operand : MulSCEV->operands()) {
+    if (const SCEVConstant *Constant = dyn_cast<SCEVConstant>(Operand))
+      CIncSCEV = Constant;
+    else if (const SCEVUnknown *Unknown = dyn_cast<SCEVUnknown>(Operand)) {
+      Type *AllocTy;
+      if (Unknown->isSizeOf(AllocTy)) {
+        DEBUG(dbgs() << "LRR: SizeOf SCEVExpr: " << *Unknown << "\n");
+      } else
+        break;
+    }
+  }
+  return CIncSCEV;
+}
+
+// Check if an IV is only for loop control purpose:
+// 1. It only has one use which is loop increment, then it is only used
+//      by compare and PHI, compare is only used by branch.
+// 2. It is used by loop increment and compare, loop increment is only
+//    used by PHI, compare is only used by branch.
+bool LoopReroll::isLoopCtrlOnlyIV(Loop *L, Instruction *IV) {
+
+  int32_t IVUses = IV->getNumUses();
+  if (!(IVUses == 2 || IVUses == 1))
+    return false;
+
+  BasicBlock *Header = L->getHeader();
+
+  for (auto *User : IV->users()) {
+    int32_t IncOrCmpUses = User->getNumUses();
+    bool isCompInst = isCompareInst(dyn_cast<Instruction>(User), Header);
+
+    // User can only have one or two uses.
+    if (!(IncOrCmpUses == 2 || IncOrCmpUses == 1))
+      return false;
+
+    // Case 1?
+    if (IVUses == 1) {
+      // The only user must be loop increment.
+      if (isCompInst)
+         return false;
+
+      // Loop increment must have two uses.
+      if (IncOrCmpUses != 2)
+        return false;
+    }
+
+    // Case 2?
+    if (IVUses == 2 && IncOrCmpUses != 1)
+      return false;
+
+    // All users of IV must be binary operation or cmp
+    if (auto *BO = dyn_cast<BinaryOperator>(User)) {
+      if (BO->getOpcode() == Instruction::Add) {
+        // Loop Increment
+        // User of Loop Increment should be either PHI or CMP
+        for (auto *UU : User->users()) {
+          if (PHINode *PN = dyn_cast<PHINode>(UU)) {
+            if (PN != IV)
+              return false;
+          }
+          // Must be a CMP
+          else if (!isCompareInst(dyn_cast<Instruction>(UU), Header))
+            return false;
+        }
+      } else
+        return false;
+    // Compare : can only have one use, and must be branch
+    } else if (!isCompInst)
+      return false;
+  }
+  return true;
+}
+
 // Collect the list of loop induction variables with respect to which it might
 // be possible to reroll the loop.
 void LoopReroll::collectPossibleIVs(Loop *L,
                                     SmallInstructionVector &PossibleIVs) {
   BasicBlock *Header = L->getHeader();
   for (BasicBlock::iterator I = Header->begin(),
        IE = Header->getFirstInsertionPt(); I != IE; ++I) {
     if (!isa<PHINode>(I))
       continue;
-    if (!I->getType()->isIntegerTy())
+    if (!(I->getType()->isIntegerTy() || I->getType()->isPointerTy()))
       continue;
 
     if (const SCEVAddRecExpr *PHISCEV =
         dyn_cast<SCEVAddRecExpr>(SE->getSCEV(I))) {
       if (PHISCEV->getLoop() != L)
         continue;
       if (!PHISCEV->isAffine())
         continue;
-      if (const SCEVConstant *IncSCEV =
-          dyn_cast<SCEVConstant>(PHISCEV->getStepRecurrence(*SE))) {
-        if (!IncSCEV->getValue()->getValue().isStrictlyPositive())
-          continue;
+      const SCEVConstant *IncSCEV = nullptr;
+      if (I->getType()->isPointerTy()) {
+        IncSCEV =
+            getConstantCoefFromStep(SE, PHISCEV->getStepRecurrence(*SE), I);
+      } else
+        IncSCEV = dyn_cast<SCEVConstant>(PHISCEV->getStepRecurrence(*SE));
+
+      if (!IncSCEV)
+        return;
+
+      const APInt &AInt = IncSCEV->getValue()->getValue();
+      if (AInt.isStrictlyPositive()) {
         if (IncSCEV->getValue()->uge(MaxInc))
           continue;
+        DEBUG(dbgs() << "LRR: Possible IV: " << *I << " = " << *PHISCEV
+                     << "\n");
+        IVToIncMap[I] = IncSCEV->getValue()->getZExtValue();
+      }
 
-        DEBUG(dbgs() << "LRR: Possible IV: " << *I << " = " <<
-              *PHISCEV << "\n");
-        PossibleIVs.push_back(I);
+      if (AInt.isNegative()) {
+        const SCEVConstant *PIncSCEV =
+            dyn_cast<SCEVConstant>(SE->getNegativeSCEV(IncSCEV));
+        if (!PIncSCEV || PIncSCEV->getValue()->uge(MaxInc))
+          continue;
+        DEBUG(dbgs() << "LRR: Possible negative IV: " << *I << " = " << *PHISCEV
+                     << "\n");
+        IVToIncMap[I] = IncSCEV->getValue()->getSExtValue();
       }
+
+      if (isLoopCtrlOnlyIV(L, I)) {
+        assert(!LoopCtrlOnlyIV && "Found two loop control only IV");
+        LoopCtrlOnlyIV = I;
+        DEBUG(dbgs() << "LRR: Possible loop control only IV: " << *I << " = "
+                     << *PHISCEV << "\n");
+      } else
+        PossibleIVs.push_back(I);
     }
   }
 }
 
 // Add the remainder of the reduction-variable chain to the instruction vector
 // (the initial PHINode has already been added). If successful, the object is
 // marked as valid.
 void LoopReroll::SimpleLoopReduction::add(Loop *L) {
case Instruction::SExt:
       return true;
     }
   }
   return false;
 }
 
 static bool isLoopIncrement(User *U, Instruction *IV) {
   BinaryOperator *BO = dyn_cast<BinaryOperator>(U);
-  if (!BO || BO->getOpcode() != Instruction::Add)
+
+  if (BO && BO->getOpcode() != Instruction::Add)
     return false;
 
-  for (auto *UU : BO->users()) {
+  if (!BO && !isa<GetElementPtrInst>(U))
+    return false;
+
+  for (auto *UU : U->users()) {
     PHINode *PN = dyn_cast<PHINode>(UU);
     if (PN && PN == IV)
       return true;
   }
+
   return false;
 }
 
 bool LoopReroll::DAGRootTracker::
 collectPossibleRoots(Instruction *Base, std::map<int64_t,Instruction*> &Roots) {
   SmallInstructionVector BaseUsers;
 
   for (auto *I : Base->users()) {
if (Roots.find(V) != Roots.end())
       return false;
 
     // FIXME: Add support for negative values.
     if (V < 0) {
       DEBUG(dbgs() << "LRR: Aborting due to negative value: " << V << "\n");
       return false;
     }
 
     Roots[V] = cast<Instruction>(I);
   }
 
   if (Roots.empty())
     return false;
 
   // If we found non-loop-inc, non-root users of Base, assume they are
   // for the zeroth root index. This is because "add %a, 0" gets optimized
   // away.
   if (BaseUsers.size()) {
     if (Roots.find(0) != Roots.end()) {
       DEBUG(dbgs() << "LRR: Multiple roots found for base - aborting!\n");
       return false;
     }
     Roots[0] = Base;
   }
 
   // Calculate the number of users of the base, or lowest indexed, iteration.
   unsigned NumBaseUses = BaseUsers.size();
   if (NumBaseUses == 0)
     NumBaseUses = Roots.begin()->second->getNumUses();
-  
   // Check that every node has the same number of users.
   for (auto &KV : Roots) {
     if (KV.first == 0)
       continue;
     if (KV.second->getNumUses() != NumBaseUses) {
       DEBUG(dbgs() << "LRR: Aborting - Root and Base #users not the same: "
             << "#Base=" << NumBaseUses << ", #Root=" <<
             KV.second->getNumUses() << "\n");
       return false;
     }
   }
 
-  return true; 
+  return true;
 }
 
 bool LoopReroll::DAGRootTracker::
 findRootsRecursive(Instruction *I, SmallInstructionSet SubsumedInsts) {
   // Does the user look like it could be part of a root set?
   // All its users must be simple arithmetic ops.
   if (I->getNumUses() > IL_MaxRerollIterations)
     return false;
findRootsBase(Instruction *IVU, SmallInstructionSet SubsumedInsts) {
   if (IVU->getOpcode() != Instruction::Mul &&
       IVU->getOpcode() != Instruction::PHI)
     return false;
 
   std::map<int64_t, Instruction*> V;
   if (!collectPossibleRoots(IVU, V))
     return false;
 
-  // If we didn't get a root for index zero, then IVU must be 
+  // If we didn't get a root for index zero, then IVU must be
   // subsumed.
   if (V.find(0) == V.end())
     SubsumedInsts.insert(IVU);
 
   // Partition the vector into monotonically increasing indexes.
   DAGRootSet DRS;
   DRS.BaseInst = nullptr;
 
for (auto &KV : V) {
     }
   }
   RootSets.push_back(DRS);
 
   return true;
 }
 
 bool LoopReroll::DAGRootTracker::findRoots() {
-
-  const SCEVAddRecExpr *RealIVSCEV = cast<SCEVAddRecExpr>(SE->getSCEV(IV));
-  Inc = cast<SCEVConstant>(RealIVSCEV->getOperand(1))->
-    getValue()->getZExtValue();
+  Inc = IVToIncMap[IV];
 
   assert(RootSets.empty() && "Unclean state!");
-  if (Inc == 1) {
+  if (std::abs(Inc) == 1) {
     for (auto *IVU : IV->users()) {
       if (isLoopIncrement(IVU, IV))
         LoopIncs.push_back(cast<Instruction>(IVU));
     }
     if (!findRootsRecursive(IV, SmallInstructionSet()))
       return false;
     LoopIncs.push_back(IV);
   } else {
bool LoopReroll::DAGRootTracker::validate(ReductionTracker &Reductions) {
   if (!collectUsedInstructions(PossibleRedSet))
     return false;
 
   // Make sure we mark the reduction PHIs as used in all iterations.
   for (auto *I : PossibleRedPHISet) {
     Uses[I].set(IL_All);
   }
 
+  // Make sure we mark the loop ctrl only PHIs as used in all iterations.
+  // Including loop increment, compares, branches.
+  BasicBlock *Header = L->getHeader();
+  if (LoopCtrlOnlyIV && LoopCtrlOnlyIV != IV) {
+    for (auto *U : LoopCtrlOnlyIV->users()) {
+      Instruction *UI = dyn_cast<Instruction>(U);
+      Uses[UI].set(IL_All);
+      for (auto *UU : UI->users()) {
+        Instruction *UUI = dyn_cast<Instruction>(UU);
+        Uses[UUI].set(IL_All);
+        if (UUI->hasOneUse()) {
+          Instruction *UUUI = dyn_cast<BranchInst>(*(UUI->user_begin()));
+          if (UUUI == dyn_cast<BranchInst>(Header->getTerminator()))
+            Uses[UUUI].set(IL_All);
+        }
+      }
+    }
+  }
+
   // Make sure all instructions in the loop are in one and only one
   // set.
   for (auto &KV : Uses) {
     if (KV.second.count() != 1) {
       DEBUG(dbgs() << "LRR: Aborting - instruction is not used in 1 iteration: "
             << *KV.first << " (#uses=" << KV.second.count() << ")\n");
       return false;
     }
while (BaseIt != Uses.end() && RootIt != Uses.end()) {
         }
 
         if (TryIt == Uses.end() || TryIt == RootIt ||
             instrDependsOn(TryIt->first, RootIt, TryIt)) {
           DEBUG(dbgs() << "LRR: iteration root match failed at " << *BaseInst <<
                 " vs. " << *RootInst << "\n");
           return false;
         }
-        
         RootIt = TryIt;
         RootInst = TryIt->first;
       }
 
       // All instructions between the last root and this root
-      // may belong to some other iteration. If they belong to a 
+      // may belong to some other iteration. If they belong to a
       // future iteration, then they're dangerous to alias with.
-      // 
+      //
       // Note that because we allow a limited amount of flexibility in the order
       // that we visit nodes, LastRootIt might be *before* RootIt, in which
       // case we've already checked this set of instructions so we shouldn't
       // do anything.
       for (; LastRootIt < RootIt; ++LastRootIt) {
         Instruction *I = LastRootIt->first;
         if (LastRootIt->second.find_first() < (int)Iter)
           continue;
bool LoopReroll::DAGRootTracker::validate(ReductionTracker &Reductions) {
   }
 
   DEBUG(dbgs() << "LRR: Matched all iteration increments for " <<
                   *IV << "\n");
 
   return true;
 }
 
-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();) {
-    unsigned I = Uses[&*J].find_first();
-    if (I > 0 && I < IL_All) {
-      Instruction *D = &*J;
-      DEBUG(dbgs() << "LRR: removing: " << *D << "\n");
-      D->eraseFromParent();
-      continue;
+// For non loop control IVs, we only need to update the last increment
+// with right amount.
+void LoopReroll::DAGRootTracker::replaceNonLoopCtrlIV() {
+  const SCEV *NewIncr = nullptr;
+  for (auto *LoopInc : LoopIncs) {
+    GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(LoopInc);
+    const SCEVConstant *COp = nullptr;
+    if (GEP && LoopInc->getOperand(0)->getType()->isPointerTy()) {
+      COp = dyn_cast<SCEVConstant>(SE->getSCEV(LoopInc->getOperand(1)));
+    } else {
+      COp = dyn_cast<SCEVConstant>(SE->getSCEV(LoopInc->getOperand(0)));
+      if (!COp)
+        COp = dyn_cast<SCEVConstant>(SE->getSCEV(LoopInc->getOperand(1)));
     }
 
-    ++J;
-  }
-  const DataLayout &DL = Header->getModule()->getDataLayout();
+    if (!COp)
+      assert(false && "Didn't find constant operand of LoopInc!\n");
+
+    const APInt &AInt = COp->getValue()->getValue();
+    const SCEV *ScaleSCEV = SE->getConstant(COp->getType(), Scale);
+    if (AInt.isNegative()) {
+      NewIncr = SE->getNegativeSCEV(COp);
+      NewIncr = SE->getUDivExpr(NewIncr, ScaleSCEV);
+      NewIncr = SE->getNegativeSCEV(NewIncr);
+    } else
+      NewIncr = SE->getUDivExpr(COp, ScaleSCEV);
+
+    LoopInc->setOperand(1, dyn_cast<SCEVConstant>(NewIncr)->getValue());
+  }
 
-  // We need to create a new induction variable for each different BaseInst.
-  for (auto &DRS : RootSets) {
-    // Insert the new induction variable.
-    const SCEVAddRecExpr *RealIVSCEV =
-      cast<SCEVAddRecExpr>(SE->getSCEV(DRS.BaseInst));
-    const SCEV *Start = RealIVSCEV->getStart();
-    const SCEVAddRecExpr *H = cast<SCEVAddRecExpr>
-      (SE->getAddRecExpr(Start,
-                         SE->getConstant(RealIVSCEV->getType(), 1),
-                         L, SCEV::FlagAnyWrap));
-    { // Limit the lifetime of SCEVExpander.
-      SCEVExpander Expander(*SE, DL, "reroll");
-      Value *NewIV = Expander.expandCodeFor(H, IV->getType(), Header->begin());
+void LoopReroll::DAGRootTracker::replaceIV(Instruction *Inst,
+                                           Instruction *InstIV,
+                                           const SCEV *IterCount) {
+  BasicBlock *Header = L->getHeader();
+  int64_t Inc = IVToIncMap[InstIV];
+  bool NeedNewIV = InstIV == LoopCtrlOnlyIV;
+  bool Negative = !NeedNewIV && Inc < 0;
+
+  const SCEVAddRecExpr *RealIVSCEV = cast<SCEVAddRecExpr>(SE->getSCEV(Inst));
+
+  if (NeedNewIV)
+    Start = SE->getConstant(Start->getType(), 0);
+
+  const SCEVAddRecExpr *H = nullptr;
+  const SCEV *SizeOfExpr = nullptr;
+  const SCEV *OneExpr =
+      SE->getConstant(RealIVSCEV->getType(), Negative ? -1 : 1);
+  if (Inst->getType()->isPointerTy()) {
+    const PointerType *PTy = dyn_cast<PointerType>(Inst->getType());
+    Type *ElTy = PTy->getElementType();
+    SizeOfExpr =
+        SE->getSizeOfExpr(SE->getEffectiveSCEVType(Inst->getType()), ElTy);
+
+    OneExpr = SE->getMulExpr(OneExpr, SizeOfExpr);
+  }
+  H = cast<SCEVAddRecExpr>(
+      SE->getAddRecExpr(Start, OneExpr, L, SCEV::FlagAnyWrap));
+  // Limit the lifetime of SCEVExpander.
+  const DataLayout &DL = Header->getModule()->getDataLayout();
+  Value *NewIV = Expander.expandCodeFor(H, InstIV->getType(), Header->begin());
 
-      for (auto &KV : Uses) {
-        if (KV.second.find_first() == 0)
-          KV.first->replaceUsesOfWith(DRS.BaseInst, NewIV);
-      }
+  for (auto &KV : Uses)
+      KV.first->replaceUsesOfWith(Inst, NewIV);
 
-      if (BranchInst *BI = dyn_cast<BranchInst>(Header->getTerminator())) {
-        // FIXME: Why do we need this check?
-        if (Uses[BI].find_first() == IL_All) {
-          const SCEV *ICSCEV = RealIVSCEV->evaluateAtIteration(IterCount, *SE);
+      const SCEV *ICSCEV = nullptr;
 
-          // Iteration count SCEV minus 1
-          const SCEV *ICMinus1SCEV =
-            SE->getMinusSCEV(ICSCEV, SE->getConstant(ICSCEV->getType(), 1));
+      if (NeedNewIV)
+        ICSCEV = SE->getMulExpr(IterCount,
+                                SE->getConstant(IterCount->getType(), Scale));
+      else
+        ICSCEV = RealIVSCEV->evaluateAtIteration(IterCount, *SE);
 
-          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, Parent);
-
-            ICMinus1 = Expander.expandCodeFor(ICMinus1SCEV, NewIV->getType(),
-                                              Preheader->getTerminator());
-          }
+      // Iteration count SCEV minus 1
+      const SCEV *Minus1SCEV =
+          SE->getConstant(ICSCEV->getType(), Negative ? -1 : 1);
+      if (Inst->getType()->isPointerTy())
+        Minus1SCEV = SE->getMulExpr(Minus1SCEV, SizeOfExpr);
+      const SCEV *ICMinus1SCEV = SE->getMinusSCEV(ICSCEV, Minus1SCEV);
+
+      // Iteration count minus 1
+      Value *ICMinus1 = nullptr;
+      if (isa<SCEVConstant>(ICMinus1SCEV))
+      else {
 
-          Value *Cond =
-            new ICmpInst(BI, CmpInst::ICMP_EQ, NewIV, ICMinus1, "exitcond");
-          BI->setCondition(Cond);
 
-          if (BI->getSuccessor(1) != Header)
-            BI->swapSuccessors();
-        }
-      }
-    }
+
+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();) {
+    unsigned I = Uses[&*J].find_first();
+    if (I > 0 && I < IL_All) {
+      Instruction *D = &*J;
+      DEBUG(dbgs() << "LRR: removing: " << *D << "\n");
+      D->eraseFromParent();
+      continue;
+    }
+
+    ++J;
   }
 
+  bool HasTwoIVs = LoopCtrlOnlyIV && LoopCtrlOnlyIV != IV;
+
+  // We need to handle loop control only IVs
+  if (HasTwoIVs) {
+    replaceNonLoopCtrlIV();
+    replaceIV(LoopCtrlOnlyIV, LoopCtrlOnlyIV, IterCount);
+  } else
+    // We need to create a new induction variable for each different BaseInst.
+    for (auto &DRS : RootSets)
+      // Insert the new induction variable.
+      replaceIV(DRS.BaseInst, IV, IterCount);
+
   SimplifyInstructionsInBlock(Header, TLI);
   DeleteDeadPHIs(Header, 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() {

 //
 // 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, SE, AA, TLI);
+  DAGRootTracker DAGRoots(this, L, IV, SE, AA, TLI, Reductions, IVToIncMap,
+                          LoopCtrlOnlyIV);
 
   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();
bool LoopReroll::runOnLoop(Loop *L, LPPassManager &LPM) {
   const SCEV *LIBETC = SE->getBackedgeTakenCount(L);
   const SCEV *IterCount =
     SE->getAddExpr(LIBETC, SE->getConstant(LIBETC->getType(), 1));
   DEBUG(dbgs() << "LRR: iteration count = " << *IterCount << "\n");
 
   // First, we need to find the induction variable with respect to which we can
   // reroll (there may be several possible options).
   SmallInstructionVector PossibleIVs;
+  IVToIncMap.clear();
+  LoopCtrlOnlyIV = nullptr;
   collectPossibleIVs(L, PossibleIVs);
 
   if (PossibleIVs.empty()) {
     DEBUG(dbgs() << "LRR: No possible IVs found\n");
     return Changed;
   }
 
   ReductionTracker Reductions;

test/Transforms/LoopReroll/complex_reroll.ll

+; RUN: opt -S  -loop-reroll   %s | FileCheck %s
+declare i32 @goo(i32, i32)
+
+@buf = external global i8*
+@aaa = global [16 x i8] c"\01\02\03\04\05\06\07\08\09\0A\0B\0C\0D\0E\0F\10", align 1
+
+define i32 @test1(i32 %len) {
+entry:
+  br label %while.body
+
+while.body:
+;CHECK-LABEL: while.body:
+;CHECK-NEXT:    %indvar = phi i32 [ %indvar.next, %while.body ], [ 0, %entry ]
+;CHECK-NEXT:    %buf.021 = phi i8* [ getelementptr inbounds ([16 x i8], [16 x i8]* @aaa, i64 0, i64 0), %entry ], [ %add.ptr, %while.body ]
+;CHECK-NEXT:    %sum44.020 = phi i64 [ 0, %entry ], [ %add, %while.body ]
+;CHECK-NEXT:    %0 = load i8, i8* %buf.021, align 1
+;CHECK-NEXT:    %conv = zext i8 %0 to i64
+;CHECK-NEXT:    %add = add i64 %conv, %sum44.020
+;CHECK-NEXT:    %add.ptr = getelementptr inbounds i8, i8* %buf.021, i64 1
+;CHECK-NEXT:    %indvar.next = add i32 %indvar, 1
+;CHECK-NEXT:    %exitcond = icmp eq i32 %indvar, 15
+;CHECK-NEXT:    br i1 %exitcond, label %while.end, label %while.body
+
+  %dec22 = phi i32 [ 4, %entry ], [ %dec, %while.body ]
+  %buf.021 = phi i8* [ getelementptr inbounds ([16 x i8], [16 x i8]* @aaa, i64 0, i64 0), %entry ], [ %add.ptr, %while.body ]
+  %sum44.020 = phi i64 [ 0, %entry ], [ %add9, %while.body ]
+  %0 = load i8, i8* %buf.021, align 1
+  %conv = zext i8 %0 to i64
+  %add = add i64 %conv, %sum44.020
+  %arrayidx1 = getelementptr inbounds i8, i8* %buf.021, i64 1
+  %1 = load i8, i8* %arrayidx1, align 1
+  %conv2 = zext i8 %1 to i64
+  %add3 = add i64 %add, %conv2
+  %arrayidx4 = getelementptr inbounds i8, i8* %buf.021, i64 2
+  %2 = load i8, i8* %arrayidx4, align 1
+  %conv5 = zext i8 %2 to i64
+  %add6 = add i64 %add3, %conv5
+  %arrayidx7 = getelementptr inbounds i8, i8* %buf.021, i64 3
+  %3 = load i8, i8* %arrayidx7, align 1
+  %conv8 = zext i8 %3 to i64
+  %add9 = add i64 %add6, %conv8
+  %add.ptr = getelementptr inbounds i8, i8* %buf.021, i64 4
+  %dec = add nsw i32 %dec22, -1
+  %tobool = icmp eq i32 %dec, 0
+  br i1 %tobool, label %while.end, label %while.body
+
+while.end:                                        ; preds = %while.body
+  %conv11 = trunc i64 %add9 to i32
+  %call = tail call i32 @goo(i32 0, i32 %conv11)
+  unreachable
+}
+
+define i32 @test2(i32 %N, i32* nocapture readonly %a, i32 %S) {
+entry:
+  %cmp.9 = icmp sgt i32 %N, 0
+  br i1 %cmp.9, label %for.body.lr.ph, label %for.cond.cleanup
+
+for.body.lr.ph:
+  br label %for.body
+
+for.cond.for.cond.cleanup_crit_edge:
+  br label %for.cond.cleanup
+
+for.cond.cleanup:
+  %S.addr.0.lcssa = phi i32 [ %add2, %for.cond.for.cond.cleanup_crit_edge ], [ %S, %entry ]
+  ret i32 %S.addr.0.lcssa
+
+for.body:
+;CHECK-LABEL: for.body:
+;CHECK-NEXT:    %indvar = phi i32 [ %indvar.next, %for.body ], [ 0, %for.body.lr.ph ]
+;CHECK-NEXT:    %S.addr.011 = phi i32 [ %S, %for.body.lr.ph ], [ %add, %for.body ]
+;CHECK-NEXT:    %a.addr.010 = phi i32* [ %a, %for.body.lr.ph ], [ %incdec.ptr1, %for.body ]
+;CHECK-NEXT:    %4 = load i32, i32* %a.addr.010, align 4
+;CHECK-NEXT:    %add = add nsw i32 %4, %S.addr.011
+;CHECK-NEXT:    %incdec.ptr1 = getelementptr inbounds i32, i32* %a.addr.010, i64 1
+;CHECK-NEXT:    %indvar.next = add i32 %indvar, 1
+;CHECK-NEXT:    %exitcond = icmp eq i32 %indvar, %3
+;CHECK-NEXT:    br i1 %exitcond, label %for.cond.for.cond.cleanup_crit_edge, label %for.body
+
+  %i.012 = phi i32 [ 0, %for.body.lr.ph ], [ %add3, %for.body ]
+  %S.addr.011 = phi i32 [ %S, %for.body.lr.ph ], [ %add2, %for.body ]
+  %a.addr.010 = phi i32* [ %a, %for.body.lr.ph ], [ %incdec.ptr1, %for.body ]
+  %incdec.ptr = getelementptr inbounds i32, i32* %a.addr.010, i64 1
+  %0 = load i32, i32* %a.addr.010, align 4
+  %add = add nsw i32 %0, %S.addr.011
+  %incdec.ptr1 = getelementptr inbounds i32, i32* %a.addr.010, i64 2
+  %1 = load i32, i32* %incdec.ptr, align 4
+  %add2 = add nsw i32 %add, %1
+  %add3 = add nsw i32 %i.012, 2
+  %cmp = icmp slt i32 %add3, %N
+  br i1 %cmp, label %for.body, label %for.cond.for.cond.cleanup_crit_edge
+}
+