[SLSR] handle candidate form (B + i * S)

[jingyue]

Summary

With this patch, SLSR may rewrite

S1: X = B + i * S
S2: Y = B + i' * S

to

S2: Y = X + (i' - i) * S

A secondary improvement: if (i' - i) is a power of 2, emit Y as X + (S << log(i' - i)). (S << log(i' -i)) is in a canonical form and thus more likely GVN'ed than (i' - i) * S.

[jingyue]

Do not SLSR if (B + i * S) can be folded to an addressing mode.

[meheff]

LGTM.

Thanks.
Mark

[broune]

I added some comments not directly related to the patch, since you're the author of the whole pass, but feel free to ignore those parts for this patch.

Would it make sense to avoid transformations like this one?

X = B + 8 * S
Y = B + S

to

X = B + 8 * S
Y = X - 7 * S

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

29

Might be good to explain here why this is an improvement and where the strength is reduced. The rewritten expression still has the form B + i * S, just with a different B and i. Right? Perhaps we're hoping that (i' - i) * S will be easier to optimize, e.g. if it appears more than once or if i' - i is now -1, 0, 1 or a power of 2?

37

Same comment as above.

39–42

extend -> extent

46

I think this is a definition of what an immediate basis is, though it reads to me as a statement that we either use the closest basis or the immediate basis, which suggests that they are not necessarily the same. Perhaps: with respect to the immediate basis, which is the basis that is the closest ancestor in the dominator tree." Is that what "closest" means here?

54–57

Would it make sense to add a TODO for the case where i - i' is constant but i and i' are not?

79–80

Should there be another case added above for Add, or perhaps just refer to the comments on the entries in Kind instead of duplicating it here?

97–98

This is ambiguous as "may not" can mean "must not" or "might not". I'd do "may not" -> "might not" or rewrite as "does not necessarily have". (I know you didn't change this in this patch, but you're the author of the pass, right?)

105

can corresponds -> can correspond

158

The function name says singular candidate and the comment says plural candidates. So one of them is not right.

205

We also keep track of unlinked instructions by their parent being null, and therefore never insert the same element twice nor query whether an element is in UnlinkedInstructions. So might this better be a lower-overhead type like a vector?

269–270

to -> in

274

Maybe add: in case this expression is later used as an address in a way that is foldable as an addressing mode.

Or am I misunderstanding this? The significant difference seems to me to be that above, we know it's an address and thus will be folded if possible as it's included directly in a GEP and here this expression may or may not feed directly into a GEP later on.

325–326

running forever -> taking too much time. It would terminate.

354

Could add: so that it can be the basis of other candidates.

431

Would it make sense to assert that I has exactly 2 operands?

455–456

is *s a typo?

565

would it make sense to catch the negative of a power of 2?

583–586

Would it make sense to assert that Basis.Ins->getParent() != nullptr ?

[jingyue]

Broune's comments

[jingyue]

That's a very good point. I think isSimplestForm in the updated patch handles that case. I also added a test. PTAL

Thanks for the detailed review! These comments improved the patch a lot.

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

29

Done.

37

Done.

39–42

Done.

46

Done.

54–57

Done. I don't think it's a common case, unless you encountered that already? Anyway, it doesn't hurt adding a TODO.

79–80

Done.

97–98

Done.

105

Done.

158

Done.

205

Done.

269–270

Done.

274

Even if an add instruction is not used as an address by a load/store, it can be likely computed in one machine instruction (e.g., leal in x86). I agree this heuristic is inaccurate, but I don't have a better way of characterizing it. Any better idea?

325–326

Done.

354

Done.

431

Done.

455–456

Fixed. I mean ArrayIdx = ArrayIdx *nsw 1.

565

Done. Also added a test on that.

583–586

Done.

[eliben]

LGTM

[jingyue]

Minor changes in comments

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

 // This file implements straight-line strength reduction (SLSR). Unlike loop
 // strength reduction, this algorithm is designed to reduce arithmetic
 // redundancy in straight-line code instead of loops. It has proven to be
 // effective in simplifying arithmetic statements derived from an unrolled loop.
 // It can also simplify the logic of SeparateConstOffsetFromGEP.
 //
 // There are many optimizations we can perform in the domain of SLSR. This file
 // for now contains only an initial step. Specifically, we look for strength
-// reduction candidates in two forms:
+// reduction candidates in the following forms:
 //
-// Form 1: (B + i) * S
-// Form 2: &B[i * S]
+// Form 1: B + i * S
+// Form 2: (B + i) * S
+// Form 3: &B[i * S]
 //
 // where S is an integer variable, and i is a constant integer. If we found two
-// candidates
+// candidates S1 and S2 in the same form and S1 dominates S2, we may rewrite S2
+// in a simpler way with respect to S1. For example,
 //
-// S1: X = (B + i) * S
-// S2: Y = (B + i') * S
+// S1: X = B + i * S
+// S2: Y = B + i' * S   => X + (i' - i) * S
 //
-// or
+// S1: X = (B + i) * S
+// S2: Y = (B + i') * S => X + (i' - i) * S
 //
 // S1: X = &B[i * S]
-// S2: Y = &B[i' * S]
-//
-// and S1 dominates S2, we call S1 a basis of S2, and can replace S2 with
-//
-// Y = X + (i' - i) * S
-//
-// or
+// S2: Y = &B[i' * S]   => &X[(i' - i) * S]
 //
-// Y = &X[(i' - i) * S]
+// Note: (i' - i) * S is folded to the extent possible.
 //
-// where (i' - i) * S is folded to the extent possible. When S2 has multiple
-// bases, we pick the one that is closest to S2, or S2's "immediate" basis.
+// This rewriting is in general a good idea. The code patterns we focus on
+// usually come from loop unrolling, so (i' - i) * S is likely the same
+// across iterations and can be reused. When that happens, the optimized form
+// takes only one add starting from the second iteration.
+//
+// When such rewriting is possible, we call S1 a "basis" of S2. When S2 has
+// multiple bases, we choose to rewrite S2 with respect to its "immediate"
+// basis, the basis that is the closest ancestor in the dominator tree.
 //
 // TODO:
 //
-// - Handle candidates in the form of B + i * S
-//
 // - Floating point arithmetics when fast math is enabled.
 //
 // - SLSR may decrease ILP at the architecture level. Targets that are very
 //   sensitive to ILP may want to disable it. Having SLSR to consider ILP is
 //   left as future work.
+//
+// - When (i' - i) is constant but i and i' are not, we could still perform
+//   SLSR.
 #include <vector>
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/FoldingSet.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 
 using namespace llvm;
 using namespace PatternMatch;
 
 namespace {
 
 class StraightLineStrengthReduce : public FunctionPass {
 public:
-  // SLSR candidate. Such a candidate must be in the form of
-  //   (Base + Index) * Stride
-  // or
-  //   Base[..][Index * Stride][..]
+  // SLSR candidate. Such a candidate must be in one of the forms described in
+  // the header comments.
   struct Candidate : public ilist_node<Candidate> {
     enum Kind {
       Invalid, // reserved for the default constructor
+      Add,     // B + i * S
       Mul,     // (B + i) * S
       GEP,     // &B[..][i * S][..]
     };
 
     Candidate()
         : CandidateKind(Invalid), Base(nullptr), Index(nullptr),
           Stride(nullptr), Ins(nullptr), Basis(nullptr) {}
     Candidate(Kind CT, const SCEV *B, ConstantInt *Idx, Value *S,
               Instruction *I)
         : CandidateKind(CT), Base(B), Index(Idx), Stride(S), Ins(I),
           Basis(nullptr) {}
     Kind CandidateKind;
     const SCEV *Base;
-    // Note that Index and Stride of a GEP candidate may not have the same
-    // integer type. In that case, during rewriting, Stride will be
+    // Note that Index and Stride of a GEP candidate do not necessarily have the
+    // same integer type. In that case, during rewriting, Stride will be
     // sign-extended or truncated to Index's type.
     ConstantInt *Index;
     Value *Stride;
     // The instruction this candidate corresponds to. It helps us to rewrite a
     // candidate with respect to its immediate basis. Note that one instruction
-    // can corresponds to multiple candidates depending on how you associate the
+    // can correspond to multiple candidates depending on how you associate the
     // expression. For instance,
     //
     // (a + 1) * (b + 2)
     //
     // can be treated as
     //
     // <Base: a, Index: 1, Stride: b + 2>
     //
public:
   }
 
   bool runOnFunction(Function &F) override;
 
 private:
   // Returns true if Basis is a basis for C, i.e., Basis dominates C and they
   // share the same base and stride.
   bool isBasisFor(const Candidate &Basis, const Candidate &C);
+  // Returns whether the candidate can be folded into an addressing mode.
+  bool isFoldable(const Candidate &C, TargetTransformInfo *TTI,
+                  const DataLayout *DL);
+  // Returns true if C is already in a simplest form and not worth being
+  // rewritten.
+  bool isSimplestForm(const Candidate &C);
   // Checks whether I is in a candidate form. If so, adds all the matching forms
   // to Candidates, and tries to find the immediate basis for each of them.
-  void allocateCandidateAndFindBasis(Instruction *I);
+  void allocateCandidatesAndFindBasis(Instruction *I);
+  // Allocate candidates and find bases for Add instructions.
+  void allocateCandidatesAndFindBasisForAdd(Instruction *I);
+  // Given I = LHS + RHS, factors RHS into i * S and makes (LHS + i * S) a
+  // candidate.
+  void allocateCandidatesAndFindBasisForAdd(Value *LHS, Value *RHS,
+                                            Instruction *I);
   // Allocate candidates and find bases for Mul instructions.
-  void allocateCandidateAndFindBasisForMul(Instruction *I);
+  void allocateCandidatesAndFindBasisForMul(Instruction *I);
   // Splits LHS into Base + Index and, if succeeds, calls
-  // allocateCandidateAndFindBasis.
-  void allocateCandidateAndFindBasisForMul(Value *LHS, Value *RHS,
+  // allocateCandidatesAndFindBasis.
+  void allocateCandidatesAndFindBasisForMul(Value *LHS, Value *RHS,
                                             Instruction *I);
   // Allocate candidates and find bases for GetElementPtr instructions.
-  void allocateCandidateAndFindBasisForGEP(GetElementPtrInst *GEP);
+  void allocateCandidatesAndFindBasisForGEP(GetElementPtrInst *GEP);
   // A helper function that scales Idx with ElementSize before invoking
-  // allocateCandidateAndFindBasis.
-  void allocateCandidateAndFindBasisForGEP(const SCEV *B, ConstantInt *Idx,
+  // allocateCandidatesAndFindBasis.
+  void allocateCandidatesAndFindBasisForGEP(const SCEV *B, ConstantInt *Idx,
                                             Value *S, uint64_t ElementSize,
                                             Instruction *I);
   // Adds the given form <CT, B, Idx, S> to Candidates, and finds its immediate
   // basis.
-  void allocateCandidateAndFindBasis(Candidate::Kind CT, const SCEV *B,
+  void allocateCandidatesAndFindBasis(Candidate::Kind CT, const SCEV *B,
                                       ConstantInt *Idx, Value *S,
                                       Instruction *I);
   // Rewrites candidate C with respect to Basis.
   void rewriteCandidateWithBasis(const Candidate &C, const Candidate &Basis);
   // A helper function that factors ArrayIdx to a product of a stride and a
-  // constant index, and invokes allocateCandidateAndFindBasis with the
+  // constant index, and invokes allocateCandidatesAndFindBasis with the
   // factorings.
   void factorArrayIndex(Value *ArrayIdx, const SCEV *Base, uint64_t ElementSize,
                         GetElementPtrInst *GEP);
   // Emit code that computes the "bump" from Basis to C. If the candidate is a
   // GEP and the bump is not divisible by the element size of the GEP, this
   // function sets the BumpWithUglyGEP flag to notify its caller to bump the
   // basis using an ugly GEP.
   static Value *emitBump(const Candidate &Basis, const Candidate &C,
                          IRBuilder<> &Builder, const DataLayout *DL,
                          bool &BumpWithUglyGEP);
 
   const DataLayout *DL;
   DominatorTree *DT;
   ScalarEvolution *SE;
   TargetTransformInfo *TTI;
   ilist<Candidate> Candidates;
   // Temporarily holds all instructions that are unlinked (but not deleted) by
   // rewriteCandidateWithBasis. These instructions will be actually removed
   // after all rewriting finishes.
-  DenseSet<Instruction *> UnlinkedInstructions;
+  std::vector<Instruction *> UnlinkedInstructions;
 };
 }  // anonymous namespace
 
 char StraightLineStrengthReduce::ID = 0;
 INITIALIZE_PASS_BEGIN(StraightLineStrengthReduce, "slsr",
                       "Straight line strength reduction", false, false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ScalarEvolution)
return (Basis.Ins != C.Ins && // skip the same instruction
           // Basis must dominate C in order to rewrite C with respect to Basis.
           DT->dominates(Basis.Ins->getParent(), C.Ins->getParent()) &&
           // They share the same base, stride, and candidate kind.
           Basis.Base == C.Base &&
           Basis.Stride == C.Stride &&
           Basis.CandidateKind == C.CandidateKind);
 }
 
-static bool isCompletelyFoldable(GetElementPtrInst *GEP,
-                                 const TargetTransformInfo *TTI,
-                                 const DataLayout *DL) {
+static bool isGEPFoldable(GetElementPtrInst *GEP,
   GlobalVariable *BaseGV = nullptr;
   int64_t BaseOffset = 0;
   bool HasBaseReg = false;
   int64_t Scale = 0;
 
   if (GlobalVariable *GV = dyn_cast<GlobalVariable>(GEP->getPointerOperand()))
     BaseGV = GV;
   else
if (isa<SequentialType>(*GTI)) {
       StructType *STy = cast<StructType>(*GTI);
       uint64_t Field = cast<ConstantInt>(*I)->getZExtValue();
       BaseOffset += DL->getStructLayout(STy)->getElementOffset(Field);
     }
   }
   return TTI->isLegalAddressingMode(GEP->getType()->getElementType(), BaseGV,
                                     BaseOffset, HasBaseReg, Scale);
 }
 
-// TODO: We currently implement an algorithm whose time complexity is linear to
-// the number of existing candidates. However, a better algorithm exists. We
-// could depth-first search the dominator tree, and maintain a hash table that
-// contains all candidates that dominate the node being traversed.  This hash
-// table is indexed by the base and the stride of a candidate.  Therefore,
-// finding the immediate basis of a candidate boils down to one hash-table look
-// up.
-void StraightLineStrengthReduce::allocateCandidateAndFindBasis(
-    Candidate::Kind CT, const SCEV *B, ConstantInt *Idx, Value *S,
-    Instruction *I) {
-  if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(I)) {
-    // If &B[Idx * S] fits into an addressing mode, do not turn it into
-    // non-free computation.
-    if (isCompletelyFoldable(GEP, TTI, DL))
-      return;
-  }
+// Returns whether (Base + Index * Stride) can be folded to an addressing mode.
+static bool isAddFoldable(const SCEV *Base, ConstantInt *Index, Value *Stride,
+                          TargetTransformInfo *TTI) {
+  return TTI->isLegalAddressingMode(Base->getType(), nullptr, 0, true,
+                                    Index->getSExtValue());
+}
+
+bool StraightLineStrengthReduce::isFoldable(const Candidate &C,
+                                            TargetTransformInfo *TTI,
+                                            const DataLayout *DL) {
+  if (C.CandidateKind == Candidate::Add)
+    return isAddFoldable(C.Base, C.Index, C.Stride, TTI);
+  if (C.CandidateKind == Candidate::GEP)
+    return isGEPFoldable(cast<GetElementPtrInst>(C.Ins), TTI, DL);
+  return false;
+}
+
+// Returns true if GEP has zero or one non-zero index.
+static bool hasOnlyOneNonZeroIndex(GetElementPtrInst *GEP) {
+  unsigned NumNonZeroIndices = 0;
+  for (auto I = GEP->idx_begin(); I != GEP->idx_end(); ++I) {
+    ConstantInt *ConstIdx = dyn_cast<ConstantInt>(*I);
+    if (ConstIdx == nullptr || !ConstIdx->isZero())
+      ++NumNonZeroIndices;
+  }
+  return NumNonZeroIndices <= 1;
+}
+
+bool StraightLineStrengthReduce::isSimplestForm(const Candidate &C) {
+  if (C.CandidateKind == Candidate::Add) {
+    // B + 1 * S or B + (-1) * S
+    return C.Index->isOne() || C.Index->isMinusOne();
+  }
+  if (C.CandidateKind == Candidate::Mul) {
+    // (B + 0) * S
+    return C.Index->isZero();
+  }
+  if (C.CandidateKind == Candidate::GEP) {
+    // (char*)B + S or (char*)B - S
+    return ((C.Index->isOne() || C.Index->isMinusOne()) &&
+            hasOnlyOneNonZeroIndex(cast<GetElementPtrInst>(C.Ins)));
+  }
+  return false;
 
+// TODO: We currently implement an algorithm whose time complexity is linear in
+// the number of existing candidates. However, we could do better by using
+// ScopedHashTable. Specifically, while traversing the dominator tree, we could
+// maintain all the candidates that dominate the basic block being traversed in
+// a ScopedHashTable. This hash table is indexed by the base and the stride of
+// a candidate. Therefore, finding the immediate basis of a candidate boils down
+// to one hash-table look up.
+void StraightLineStrengthReduce::allocateCandidatesAndFindBasis(
+    Candidate::Kind CT, const SCEV *B, ConstantInt *Idx, Value *S,
+    Instruction *I) {
   Candidate C(CT, B, Idx, S, I);
+  // SLSR can complicate an instruction in two cases:
+  //
+  // 1. If we can fold I into an addressing mode, computing I is likely free or
+  // takes only one instruction.
+  //
+  // 2. I is already in a simplest form. For example, when
+  //      X = B + 8 * S
+  //      Y = B + S,
+  //    rewriting Y to X - 7 * S is probably a bad idea.
+  //
+  // In the above cases, we still add I to the candidate list so that I can be
+  // the basis of other candidates, but we leave I's basis blank so that I
+  // won't be rewritten.
+  if (!isFoldable(C, TTI, DL) && !isSimplestForm(C)) {
     // Try to compute the immediate basis of C.
     unsigned NumIterations = 0;
-  // Limit the scan radius to avoid running forever.
+    // Limit the scan radius to avoid running in quadratice time.
     static const unsigned MaxNumIterations = 50;
     for (auto Basis = Candidates.rbegin();
          Basis != Candidates.rend() && NumIterations < MaxNumIterations;
          ++Basis, ++NumIterations) {
       if (isBasisFor(*Basis, C)) {
         C.Basis = &(*Basis);
         break;
       }
     }
+  }
   // Regardless of whether we find a basis for C, we need to push C to the
-  // candidate list.
+  // candidate list so that it can be the basis of other candidates.
   Candidates.push_back(C);
 }
 
-void StraightLineStrengthReduce::allocateCandidateAndFindBasis(Instruction *I) {
+void StraightLineStrengthReduce::allocateCandidatesAndFindBasis(
+    Instruction *I) {
   switch (I->getOpcode()) {
+  case Instruction::Add:
+    allocateCandidatesAndFindBasisForAdd(I);
+    break;
   case Instruction::Mul:
-    allocateCandidateAndFindBasisForMul(I);
+    allocateCandidatesAndFindBasisForMul(I);
     break;
   case Instruction::GetElementPtr:
-    allocateCandidateAndFindBasisForGEP(cast<GetElementPtrInst>(I));
+    allocateCandidatesAndFindBasisForGEP(cast<GetElementPtrInst>(I));
     break;
   }
 }
 
-void StraightLineStrengthReduce::allocateCandidateAndFindBasisForMul(
+void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForAdd(
+    Instruction *I) {
+  // Try matching B + i * S.
+  if (!isa<IntegerType>(I->getType()))
+    return;
+
+  assert(I->getNumOperands() == 2 && "isn't I an add?");
+  Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
+  allocateCandidatesAndFindBasisForAdd(LHS, RHS, I);
+  if (LHS != RHS)
+    allocateCandidatesAndFindBasisForAdd(RHS, LHS, I);
+}
+
+void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForAdd(
+    Value *LHS, Value *RHS, Instruction *I) {
+  Value *S = nullptr;
+  ConstantInt *Idx = nullptr;
+  if (match(RHS, m_Mul(m_Value(S), m_ConstantInt(Idx)))) {
+    // I = LHS + RHS = LHS + Idx * S
+    allocateCandidatesAndFindBasis(Candidate::Add, SE->getSCEV(LHS), Idx, S, I);
+  } else if (match(RHS, m_Shl(m_Value(S), m_ConstantInt(Idx)))) {
+    // I = LHS + RHS = LHS + (S << Idx) = LHS + S * (1 << Idx)
+    APInt One(Idx->getBitWidth(), 1);
+    Idx = ConstantInt::get(Idx->getContext(), One << Idx->getValue());
+    allocateCandidatesAndFindBasis(Candidate::Add, SE->getSCEV(LHS), Idx, S, I);
+  } else {
+    // At least, I = LHS + 1 * RHS
+    ConstantInt *One = ConstantInt::get(cast<IntegerType>(I->getType()), 1);
+    allocateCandidatesAndFindBasis(Candidate::Add, SE->getSCEV(LHS), One, RHS,
+                                   I);
+  }
+}
+
+void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForMul(
     Value *LHS, Value *RHS, Instruction *I) {
   Value *B = nullptr;
   ConstantInt *Idx = nullptr;
   // Only handle the canonical operand ordering.
   if (match(LHS, m_Add(m_Value(B), m_ConstantInt(Idx)))) {
     // If LHS is in the form of "Base + Index", then I is in the form of
     // "(Base + Index) * RHS".
-    allocateCandidateAndFindBasis(Candidate::Mul, SE->getSCEV(B), Idx, RHS, I);
+    allocateCandidatesAndFindBasis(Candidate::Mul, SE->getSCEV(B), Idx, RHS, I);
   } else {
     // Otherwise, at least try the form (LHS + 0) * RHS.
     ConstantInt *Zero = ConstantInt::get(cast<IntegerType>(I->getType()), 0);
-    allocateCandidateAndFindBasis(Candidate::Mul, SE->getSCEV(LHS), Zero, RHS,
+    allocateCandidatesAndFindBasis(Candidate::Mul, SE->getSCEV(LHS), Zero, RHS,
                                   I);
   }
 }
 
-void StraightLineStrengthReduce::allocateCandidateAndFindBasisForMul(
+void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForMul(
     Instruction *I) {
   // Try matching (B + i) * S.
   // TODO: we could extend SLSR to float and vector types.
   if (!isa<IntegerType>(I->getType()))
     return;
 
+  assert(I->getNumOperands() == 2 && "isn't I a mul?");
   Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
-  allocateCandidateAndFindBasisForMul(LHS, RHS, I);
+  allocateCandidatesAndFindBasisForMul(LHS, RHS, I);
   if (LHS != RHS) {
     // Symmetrically, try to split RHS to Base + Index.
-    allocateCandidateAndFindBasisForMul(RHS, LHS, I);
+    allocateCandidatesAndFindBasisForMul(RHS, LHS, I);
   }
 }
 
-void StraightLineStrengthReduce::allocateCandidateAndFindBasisForGEP(
+void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForGEP(
     const SCEV *B, ConstantInt *Idx, Value *S, uint64_t ElementSize,
     Instruction *I) {
   // I = B + sext(Idx *nsw S) * ElementSize
   //   = B + (sext(Idx) * sext(S)) * ElementSize
   //   = B + (sext(Idx) * ElementSize) * sext(S)
   // Casting to IntegerType is safe because we skipped vector GEPs.
   IntegerType *IntPtrTy = cast<IntegerType>(DL->getIntPtrType(I->getType()));
   ConstantInt *ScaledIdx = ConstantInt::get(
       IntPtrTy, Idx->getSExtValue() * (int64_t)ElementSize, true);
-  allocateCandidateAndFindBasis(Candidate::GEP, B, ScaledIdx, S, I);
+  allocateCandidatesAndFindBasis(Candidate::GEP, B, ScaledIdx, S, I);
 }
 
 void StraightLineStrengthReduce::factorArrayIndex(Value *ArrayIdx,
                                                   const SCEV *Base,
                                                   uint64_t ElementSize,
                                                   GetElementPtrInst *GEP) {
-  // At least, ArrayIdx = ArrayIdx *s 1.
-  allocateCandidateAndFindBasisForGEP(
+  // At least, ArrayIdx = ArrayIdx *nsw 1.
+  allocateCandidatesAndFindBasisForGEP(
       Base, ConstantInt::get(cast<IntegerType>(ArrayIdx->getType()), 1),
       ArrayIdx, ElementSize, GEP);
   Value *LHS = nullptr;
   ConstantInt *RHS = nullptr;
   // One alternative is matching the SCEV of ArrayIdx instead of ArrayIdx
   // itself. This would allow us to handle the shl case for free. However,
   // matching SCEVs has two issues:
   //
   // 1. this would complicate rewriting because the rewriting procedure
   // would have to translate SCEVs back to IR instructions. This translation
   // is difficult when LHS is further evaluated to a composite SCEV.
   //
   // 2. ScalarEvolution is designed to be control-flow oblivious. It tends
   // to strip nsw/nuw flags which are critical for SLSR to trace into
   // sext'ed multiplication.
   if (match(ArrayIdx, m_NSWMul(m_Value(LHS), m_ConstantInt(RHS)))) {
     // SLSR is currently unsafe if i * S may overflow.
     // GEP = Base + sext(LHS *nsw RHS) * ElementSize
-    allocateCandidateAndFindBasisForGEP(Base, RHS, LHS, ElementSize, GEP);
+    allocateCandidatesAndFindBasisForGEP(Base, RHS, LHS, ElementSize, GEP);
   } else if (match(ArrayIdx, m_NSWShl(m_Value(LHS), m_ConstantInt(RHS)))) {
     // GEP = Base + sext(LHS <<nsw RHS) * ElementSize
     //     = Base + sext(LHS *nsw (1 << RHS)) * ElementSize
     APInt One(RHS->getBitWidth(), 1);
     ConstantInt *PowerOf2 =
         ConstantInt::get(RHS->getContext(), One << RHS->getValue());
-    allocateCandidateAndFindBasisForGEP(Base, PowerOf2, LHS, ElementSize, GEP);
+    allocateCandidatesAndFindBasisForGEP(Base, PowerOf2, LHS, ElementSize, GEP);
   }
 }
 
-void StraightLineStrengthReduce::allocateCandidateAndFindBasisForGEP(
+void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForGEP(
     GetElementPtrInst *GEP) {
   // TODO: handle vector GEPs
   if (GEP->getType()->isVectorTy())
     return;
 
   const SCEV *GEPExpr = SE->getSCEV(GEP);
   Type *IntPtrTy = DL->getIntPtrType(GEP->getType());
 
APInt ElementSize(
             cast<GetElementPtrInst>(Basis.Ins)->getType()->getElementType()));
     APInt Q, R;
     APInt::sdivrem(IndexOffset, ElementSize, Q, R);
     if (R.getSExtValue() == 0)
       IndexOffset = Q;
     else
       BumpWithUglyGEP = true;
   }
+
   // Compute Bump = C - Basis = (i' - i) * S.
   // Common case 1: if (i' - i) is 1, Bump = S.
   if (IndexOffset.getSExtValue() == 1)
     return C.Stride;
   // Common case 2: if (i' - i) is -1, Bump = -S.
   if (IndexOffset.getSExtValue() == -1)
     return Builder.CreateNeg(C.Stride);
-  // Otherwise, Bump = (i' - i) * sext/trunc(S).
-  ConstantInt *Delta = ConstantInt::get(Basis.Ins->getContext(), IndexOffset);
-  Value *ExtendedStride = Builder.CreateSExtOrTrunc(C.Stride, Delta->getType());
+
+  // Otherwise, Bump = (i' - i) * sext/trunc(S). Note that (i' - i) and S may
+  // have different bit widths.
+  IntegerType *DeltaType =
+      IntegerType::get(Basis.Ins->getContext(), IndexOffset.getBitWidth());
+  Value *ExtendedStride = Builder.CreateSExtOrTrunc(C.Stride, DeltaType);
+  if (IndexOffset.isPowerOf2()) {
+    // If (i' - i) is a power of 2, Bump = sext/trunc(S) << log(i' - i).
+    ConstantInt *Exponent = ConstantInt::get(DeltaType, IndexOffset.logBase2());
+    return Builder.CreateShl(ExtendedStride, Exponent);
+  }
+  if ((-IndexOffset).isPowerOf2()) {
+    // If (i - i') is a power of 2, Bump = -sext/trunc(S) << log(i' - i).
+    ConstantInt *Exponent =
+        ConstantInt::get(DeltaType, (-IndexOffset).logBase2());
+    return Builder.CreateNeg(Builder.CreateShl(ExtendedStride, Exponent));
+  }
+  Constant *Delta = ConstantInt::get(DeltaType, IndexOffset);
   return Builder.CreateMul(ExtendedStride, Delta);
 }
 
 void StraightLineStrengthReduce::rewriteCandidateWithBasis(
     const Candidate &C, const Candidate &Basis) {
   assert(C.CandidateKind == Basis.CandidateKind && C.Base == Basis.Base &&
          C.Stride == Basis.Stride);
+  // We run rewriteCandidateWithBasis on all candidates in a post-order, so the
+  // basis of a candidate cannot be unlinked before the candidate.
+  assert(Basis.Ins->getParent() != nullptr && "the basis is unlinked");
 
   // An instruction can correspond to multiple candidates. Therefore, instead of
   // simply deleting an instruction when we rewrite it, we mark its parent as
   // nullptr (i.e. unlink it) so that we can skip the candidates whose
   // instruction is already rewritten.
   if (!C.Ins->getParent())
     return;
 
   IRBuilder<> Builder(C.Ins);
   bool BumpWithUglyGEP;
   Value *Bump = emitBump(Basis, C, Builder, DL, BumpWithUglyGEP);
   Value *Reduced = nullptr; // equivalent to but weaker than C.Ins
   switch (C.CandidateKind) {
+  case Candidate::Add:
   case Candidate::Mul:
+    if (BinaryOperator::isNeg(Bump)) {
+      Reduced =
+          Builder.CreateSub(Basis.Ins, BinaryOperator::getNegArgument(Bump));
+    } else {
       Reduced = Builder.CreateAdd(Basis.Ins, Bump);
+    }
     break;
   case Candidate::GEP:
     {
       Type *IntPtrTy = DL->getIntPtrType(C.Ins->getType());
       bool InBounds = cast<GetElementPtrInst>(C.Ins)->isInBounds();
       if (BumpWithUglyGEP) {
         // C = (char *)Basis + Bump
         unsigned AS = Basis.Ins->getType()->getPointerAddressSpace();
default:
     llvm_unreachable("C.CandidateKind is invalid");
   };
   Reduced->takeName(C.Ins);
   C.Ins->replaceAllUsesWith(Reduced);
   C.Ins->dropAllReferences();
   // Unlink C.Ins so that we can skip other candidates also corresponding to
   // C.Ins. The actual deletion is postponed to the end of runOnFunction.
   C.Ins->removeFromParent();
-  UnlinkedInstructions.insert(C.Ins);
+  UnlinkedInstructions.push_back(C.Ins);
 }
 
 bool StraightLineStrengthReduce::runOnFunction(Function &F) {
   if (skipOptnoneFunction(F))
     return false;
 
   TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
   DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   SE = &getAnalysis<ScalarEvolution>();
   // Traverse the dominator tree in the depth-first order. This order makes sure
   // all bases of a candidate are in Candidates when we process it.
   for (auto node = GraphTraits<DominatorTree *>::nodes_begin(DT);
        node != GraphTraits<DominatorTree *>::nodes_end(DT); ++node) {
     for (auto &I : *node->getBlock())
-      allocateCandidateAndFindBasis(&I);
+      allocateCandidatesAndFindBasis(&I);
   }
 
   // Rewrite candidates in the reverse depth-first order. This order makes sure
   // a candidate being rewritten is not a basis for any other candidate.
   while (!Candidates.empty()) {
     const Candidate &C = Candidates.back();
     if (C.Basis != nullptr) {
       rewriteCandidateWithBasis(C, *C.Basis);

test/Transforms/StraightLineStrengthReduce/X86/no-slsr.ll

 ; RUN: opt < %s -slsr -gvn -dce -S | FileCheck %s
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 ; Do not perform SLSR on &input[s] and &input[s * 2] which fit into addressing
 ; modes of X86.
-define i32 @slsr_gep(i32* %input, i64 %s) {
-; CHECK-LABEL: @slsr_gep(
+define i32 @no_slsr_gep(i32* %input, i64 %s) {
+; CHECK-LABEL: @no_slsr_gep(
   ; v0 = input[0];
   %p0 = getelementptr inbounds i32, i32* %input, i64 0
   %v0 = load i32, i32* %p0
 
   ; v1 = input[s];
   %p1 = getelementptr inbounds i32, i32* %input, i64 %s
 ; CHECK: %p1 = getelementptr inbounds i32, i32* %input, i64 %s
   %v1 = load i32, i32* %p1
 
   ; v2 = input[s * 2];
   %s2 = mul nsw i64 %s, 2
   %p2 = getelementptr inbounds i32, i32* %input, i64 %s2
 ; CHECK: %p2 = getelementptr inbounds i32, i32* %input, i64 %s2
   %v2 = load i32, i32* %p2
 
   ; return v0 + v1 + v2;
   %1 = add i32 %v0, %v1
   %2 = add i32 %1, %v2
   ret i32 %2
 }
 
+define void @no_slsr_add(i32 %b, i32 %s) {
+; CHECK-LABEL: @no_slsr_add(
+  %1 = add i32 %b, %s
+; CHECK: add i32 %b, %s
+  call void @foo(i32 %1)
+  %s2 = mul i32 %s, 2
+; CHECK: %s2 = mul i32 %s, 2
+  %2 = add i32 %b, %s2
+; CHECK: add i32 %b, %s2
+  call void @foo(i32 %2)
+  ret void
+}
+
+declare void @foo(i32 %a)

test/Transforms/StraightLineStrengthReduce/slsr-add.ll

+; RUN: opt < %s -slsr -gvn -dce -S | FileCheck %s
+
+target datalayout = "e-i64:64-v16:16-v32:32-n16:32:64"
+
+define void @shl(i32 %b, i32 %s) {
+; CHECK-LABEL: @shl(
+  %1 = add i32 %b, %s
+; [[BASIS:%[a-zA-Z0-9]+]] = add i32 %b, %s
+  call void @foo(i32 %1)
+  %s2 = shl i32 %s, 1
+  %2 = add i32 %b, %s2
+; add i32 [[BASIS]], %s
+  call void @foo(i32 %2)
+  ret void
+}
+
+define void @stride_is_2s(i32 %b, i32 %s) {
+; CHECK-LABEL: @stride_is_2s(
+  %s2 = shl i32 %s, 1
+; CHECK: %s2 = shl i32 %s, 1
+  %1 = add i32 %b, %s2
+; CHECK: [[t1:%[a-zA-Z0-9]+]] = add i32 %b, %s2
+  call void @foo(i32 %1)
+  %s4 = shl i32 %s, 2
+  %2 = add i32 %b, %s4
+; CHECK: [[t2:%[a-zA-Z0-9]+]] = add i32 [[t1]], %s2
+  call void @foo(i32 %2)
+  %s6 = mul i32 %s, 6
+  %3 = add i32 %b, %s6
+; CHECK: add i32 [[t2]], %s2
+  call void @foo(i32 %3)
+  ret void
+}
+
+define void @stride_is_3s(i32 %b, i32 %s) {
+; CHECK-LABEL: @stride_is_3s(
+  %1 = add i32 %s, %b
+; CHECK: [[t1:%[a-zA-Z0-9]+]] = add i32 %s, %b
+  call void @foo(i32 %1)
+  %s4 = shl i32 %s, 2
+  %2 = add i32 %s4, %b
+; CHECK: [[bump:%[a-zA-Z0-9]+]] = mul i32 %s, 3
+; CHECK: [[t2:%[a-zA-Z0-9]+]] = add i32 [[t1]], [[bump]]
+  call void @foo(i32 %2)
+  %s7 = mul i32 %s, 7
+  %3 = add i32 %s7, %b
+; CHECK: add i32 [[t2]], [[bump]]
+  call void @foo(i32 %3)
+  ret void
+}
+
+; foo(b + 6 * s);
+; foo(b + 4 * s);
+; foo(b + 2 * s);
+;   =>
+; t1 = b + 6 * s;
+; foo(t1);
+; s2 = 2 * s;
+; t2 = t1 - s2;
+; foo(t2);
+; t3 = t2 - s2;
+; foo(t3);
+define void @stride_is_minus_2s(i32 %b, i32 %s) {
+; CHECK-LABEL: @stride_is_minus_2s(
+  %s6 = mul i32 %s, 6
+  %1 = add i32 %b, %s6
+; CHECK: [[t1:%[a-zA-Z0-9]+]] = add i32 %b, %s6
+; CHECK: call void @foo(i32 [[t1]])
+  call void @foo(i32 %1)
+  %s4 = shl i32 %s, 2
+  %2 = add i32 %b, %s4
+; CHECK: [[bump:%[a-zA-Z0-9]+]] = shl i32 %s, 1
+; CHECK: [[t2:%[a-zA-Z0-9]+]] = sub i32 [[t1]], [[bump]]
+  call void @foo(i32 %2)
+; CHECK: call void @foo(i32 [[t2]])
+  %s2 = shl i32 %s, 1
+  %3 = add i32 %b, %s2
+; CHECK: [[t3:%[a-zA-Z0-9]+]] = sub i32 [[t2]], [[bump]]
+  call void @foo(i32 %3)
+; CHECK: call void @foo(i32 [[t3]])
+  ret void
+}
+
+; t = b + (s << 3);
+; foo(t);
+; foo(b + s);
+;
+; do not rewrite b + s to t - 7 * s because the latter is more complicated.
+define void @simple_enough(i32 %b, i32 %s) {
+; CHECK-LABEL: @simple_enough(
+  %s8 = shl i32 %s, 3
+  %1 = add i32 %b, %s8
+  call void @foo(i32 %1)
+  %2 = add i32 %b, %s
+; CHECK: [[t:%[a-zA-Z0-9]+]] = add i32 %b, %s{{$}}
+  call void @foo(i32 %2)
+; CHECK: call void @foo(i32 [[t]])
+  ret void
+}
+
+declare void @foo(i32 %a)

test/Transforms/StraightLineStrengthReduce/slsr-mul.ll

 define i32 @slsr3(i32 %b, i32 %s) {
 ; CHECK-LABEL: @slsr3(
   %mul0 = mul i32 %b, %s
 ; CHECK: mul i32
   %v0 = call i32 @foo(i32 %mul0)
 
   %b1 = add i32 %b, 2
   %mul1 = mul i32 %b1, %s
-; CHECK: [[BUMP:%[a-zA-Z0-9]+]] = mul i32 %s, 2
+; CHECK: [[BUMP:%[a-zA-Z0-9]+]] = shl i32 %s, 1
 ; CHECK: %mul1 = add i32 %mul0, [[BUMP]]
   %v1 = call i32 @foo(i32 %mul1)
 
   %b2 = add i32 %b, 4
   %mul2 = mul i32 %b2, %s
 ; CHECK: %mul2 = add i32 %mul1, [[BUMP]]
   %v2 = call i32 @foo(i32 %mul2)