[SLSR] handle candidate form &B[i * S]

[jingyue]

Summary

This patch enhances SLSR to handle another candidate form &B[i * S]. If
we found two candidates

S1: X = &B[i * S]
S2: Y = &B[i' * S]

and S1 dominates S2, we can replace S2 with

Y = &X[(i' - i) * S]

[hfinkel]

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

242

Shouldn't you test whether or not B isa GlobalVariable? Seems like, if it is, you should set BaseGV to that, and HasBaseReg = false, and otherwise keep the current set of arguments.

317

I don't think this is the right approach, we're inventing ScalarEvolution all over again. SE also already handles the shift case you mention in the TODO comment below.

If I take your test/Transforms/StraightLineStrengthReduce/X86/no-slsr.ll, for example, and run:

opt -analyze -scalar-evolution

We get:

Printing analysis 'Scalar Evolution Analysis' for function 'slsr_gep':
Classifying expressions for: @slsr_gep
  %p0 = getelementptr inbounds i32* %input, i64 0
  -->  %input
  %v0 = load i32* %p0
  -->  %v0
  %p1 = getelementptr inbounds i32* %input, i64 %s
  -->  ((4 * %s)<nsw> + %input)<nsw>
  %v1 = load i32* %p1
  -->  %v1
  %s2 = mul nsw i64 %s, 2
  -->  (2 * %s)
  %p2 = getelementptr inbounds i32* %input, i64 %s2
  -->  ((8 * %s) + %input)<nsw>
  %v2 = load i32* %p2
  -->  %v2
  %1 = add i32 %v0, %v1
  -->  (%v1 + %v0)
  %2 = add i32 %1, %v2
  -->  (%v2 + %v1 + %v0)
Determining loop execution counts for: @slsr_gep

And you can see, it nicely decomposes the GEPs into arithmetic expressions for you.

Andy, please shout if you disagree, and I apologize for not thinking about this before, but I'm inclined to recommend that we rework this pass so that all cases are handled by SE before we get too far. LSR, of course, uses SE, and I think this pass should too for many of the same reasons (even though there is no loop handling here, we need all of the same arithmetic normalization logic, which is fairly significant).

Also, I think you should consider handling the constant-offset case in GEPs because that will let us get cases where the last index in a structure index. like B[i*J].x, which is fairly common. You just need to make sure such an offset gets passed to isLegalAddressingMode (BaseOffset) above.

[jingyue]

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

317

Hi Hal,

Thanks for pointing this out! I original thought SCEV is for loop optimizations and didn't quite think along that line. I'll try to replace BaseExpr or even most of the Candidate class with SCEV. I need to run experiments to see how well this approach works, but in general, I am aligned with you that reusing SCEV is a better approach.

Jingyue

[atrick]

Good observation, Hal. SCEV is there for you to use, so you might as well. I agree that we don't want to reinvent the abstraction over arithmetic expressions in each pass. I will mention two issues:

(1) SCEV may end up recursively evaluating more of the expression than you need, which could be expensive. In the first patch, you had a very simple pattern match, so SCEV didn't seem worthwhile. As you try to cover more patterns, SCEV makes more sense.

(2) SCEVExpander is fraught with peril. If you can use SCEV to help discover candidates, but leave your rewriteCandidateWithBases() code as-is, then you have nothing to worry about. If you want to use SCEVExpander then you need to be very careful about the expressions you generate.

[jingyue]

Take 2 towards handling GEPs in SLSR

• replace BaseExpr with SCEV
• discover more patterns of GEP that can be SLSR'ed

[jingyue]

Hi Andrew and Hal,

Sorry about the long delay in reworking this patch. I hope it is not too cache-cold :)

Two major changes against the previous patch are:

• Replacing BaseExpr with SCEV. As we discussed before, we shouldn't reinvent another abstraction of arithmetic expressions. Reusing SCEV turns out to save lots of code!
• The discovery phase of the current patch tries to match GEPs in the form of &B[..][i * S][..] and make it a candidate (char *)&B[..][0][..] + (i * element_size) * S. The old patch only works when i * S is the last index.

I wish I could use SCEV more extensively in this patch, e.g., to discover candidates and compute bumps. However, I didn't do that for two reasons:

• Given an LLVM value, SCEV recursively evaluates all subexpressions of this value. The current SLSR algorithm does very shallow pattern matching and rewrites candidates based on these simple patterns. Recursively evaluating subexpressions makes SLSR harder to rewrite candidates into desired forms unless we can easily map SCEVs back to LLVM Values. For example, if S in candidate (B + i) * S can be further SCEV-evaluated, we would rewrite the candidate into some form like Y + (i' - i) * expr_of(S) instead of Y + (i' - i) * S.
• ScalarEvolution is designed to be control-flow oblivious, and tends to strip nsw/nuw flags. This makes SLSR harder to factor the multiply of a sext'ed expression GEPs which happens very frequently for array indexing. For example, a GEP &B[i * S] is translated to B + sext(i * S) with the nsw flag on i * S stripped out. Without this flag, SLSR can hardly accept B + sext(i * S) as a valid GEP candidate.

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

242

Done. Please check out function isCompletelyFoldable.

[jingyue]

add some missing comments

[sanjoy]

Few drop-by comments inline.

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

115

Very minor: I'd call this Kind to differentiate from llvm::Type.

323

Unless I'm missing some larger context; for this equivalence to hold, you'll have to prove that Idx * S * ElementSize does not sign-overflow. For instance, consider:

Idx->getType() == i8, pointers are 16 bit, Idx == 4, S == 4 and ElementSize == 24. sext(Idx *s S) *s ElementSize is i16 384 while sext((Idx * ElementSize) *s S) is -128.

399

Have you considered doing a bitcast to i8*, gep on that i8* and then back instead? That will preserve pointer-ness throughout.

406

Maybe use llvm_unreachable here?

[hfinkel]

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

233

Don't need the { } here.

362

Can't you use SCEV to do that here? If you do, then you get the shift case for free. If the problem here is the loss of the information on the NSW, please comment on that.

398

No, please don't do that. Cast the pointer to an i8* (in the appropriate address space), and use a GEP.

[jingyue]

Thanks Hal and Sanjoy for the review. I will fix all other comments.

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

323

Thanks for pointing this out! You are right. I'll fix it.

362

NSW is one problem. The other problem, as I replied earlier, is that this may complicate the rewriting phase because the rewriting would have to translate SCEVs back to IR instructions. For example, suppose S = a + b and the immediate basis of X = S * 4 is Y = S * 3. The current code simply rewrites X as Y + S without tracing into how S is computed. However, if we use SCEV, we would represent X as SCEV (a + b) * 4 and Y as SCEV (a + b) * 3, and rewrite Y as SCEV X + (a + b). After that, we would need to translate SCEV X + (a + b) back into IR instructions X + (a + b) and further into X + S. While I think this is doable, given the relatively simple pattern matching this pass currently has, I don't feel it's worthwhile leveraging SCEVExpander or such to rewrite SCEV into IR instructions.

[hfinkel]

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

362

Okay, please add a detailed comment about this in the code.

[jingyue]

All comments addressed. PTAL

[jingyue]

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

115

Done.

233

Done.

323

Done. We now transform

sext(Idx *s S) *s ElementSize

to

(sext(Idx) * ElementSize) * S

362

Done.

398

Done. Also change BumpWithGEP to BumpWithUglyGEP because we use GEPs in both cases.

399

Done.

406

Done.

[sanjoy]

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

320

Nit: *s is confusing -- there is no difference between signed vs. unsigned integer multiplication. If you mean nsw then *nsw is clearer.

373

You can use m_NSWMul here directly.

[jingyue]

Addresses Sanjoy's follow-up comments

[hfinkel]

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

353

I don't understand this loop, could you please add a comment explaining it. What changes between the two iterations, and why two?

429

I think we always at least have the base implementation available. Just make TTI a hard dependency.

[jingyue]

Adds more comments and makes TTI mandatory

[jingyue]

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

353

Let me know whether this comment helps.

429

Done.

[hfinkel]

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

238

We always have TTI, you don't need to check for it.

259

If Basis.Basis != nullptr, do we want to set C.Basis to that?

353

Yes, that's better.

391

Either remove this comment, or say that we use i8* GEPs, instead of inttoptr/ptrtoint because the latter interferes with pointer-aliasing analysis.

395

This can be just:

unsigned AS = Basis.Ins->getType()->getPointerAddressSpace();

443

Use range-based for.

[jingyue]

some more minor changes

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

238

Thanks. I missed this one.

259

Yes. In that case, C.Ins can still be rewritten with respect to Basis.Ins. For example,

... &a[s]; // Basis. Assume this is the first memory reference to a, so Basis.Basis == nullptr. 
... &a[s * 2]; // C. We still want to rewrite C to Basis + s.

395

Thanks!

443

I may be missing something, but I don't see any API of BasicBlock that returns an iterator_range.
http://llvm.org/docs/doxygen/html/classllvm_1_1BasicBlock.html

[hfinkel]

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

259

Okay, please do. That form should enhance ILP.

443

Don't need one. You can use a range-based for loop with any object that has a begin()/end().

for (auto &i : *B)
  allocateCandidateAndFindBasis(&I);

or something like that.

[jingyue]

use range-based for loops

[meheff]

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

358

I see that this loop enables you to avoid code duplication for the two cases (sext-wrapped and no sext), but it requires a bit of mental effort to see what it is doing. I think this would be clearer if you just broke out the code into a separate function and avoided the loop altogether. Then the code looks like:

foobarfunction(Base, ArrayIdx, ElementSize, GEP)
Value *SextIdx;
if (!match(ArrayIdx, m_SExt(m_Value(SextIdx)))

foobarfunction(Base, SextIdx, ElementSize, GEP);

In this case it is much more obvious that you're considering the two cases (sext and no sext). There is one other location in the file which uses this similar style loop construct and might be clearer expressed as a separate function call.

[jingyue]

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

358

Both places fixed. PTAL.

391

Missed this comment. Fixed in this patch.

[jingyue]

Extract symmetric code into a helper function.

[meheff]

Thanks, Jingyue. LGTM

[atrick]

No further comments from me. I don't see any issues and will defer to the other reviewers.

[jingyue]

Thanks!

Hal and Sanjoy, do you have more comments?

[sanjoy]

In D7459#146494, @jingyue wrote:

Thanks!

Hal and Sanjoy, do you have more comments?

I don't have any more comments.

[hfinkel]

LGTM too.

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 candidate in the form of
+// reduction candidates in two forms:
 //
-// (B + i) * S
+// Form 1: (B + i) * S
+// Form 2: &B[i * S]
 //
-// where B and S are integer constants or variables, and i is a constant
-// integer. If we found two such candidates
+// where S is an integer variable, and i is a constant integer. If we found two
+// candidates
 //
-// S1: X = (B + i) * S S2: Y = (B + i') * S
+// S1: X = (B + i) * S
+// S2: Y = (B + i') * S
+//
+// or
+//
+// 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
+//
+// Y = &X[(i' - i) * S]
+//
 // 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.
 //
 // TODO:
 //
 // - Handle candidates in the form of B + i * S
 //
-// - Handle candidates in the form of pointer arithmetics. e.g., 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.
 #include <vector>
 
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/FoldingSet.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:
+  // Represents the base expression used in a candidate. There are two types of
+  // base expressions so far:
+  // 1. For the candidate form "(B + i) * S", a base expression is simply the
+  //    LLVM value of B.
+  // 2. For the candidate form "&B[i * S]", a base expression is the inner
+  //    pointer B represented by a prefix of a GEP instruction. For example,
+  //    the base expression of "getelementptr a, 1, 2, i * S" is "getelementptr
+  //    a, 1, 2", which is uniquely identified by list [a, 1, 2].
+  class BaseExpr : public FoldingSetNode {
+  private:
+    SmallVector<Value *, 4> Operands;
+
+  public:
+    // for the candidate form "(B + i) * S"
+    BaseExpr(Value *Opnd) : Operands(1, Opnd) {}
+
+    // for the candidate form "&B[i * S]"
+    template <typename ItTy> BaseExpr(ItTy B, ItTy E) : Operands(B, E) {}
+
+    void Profile(FoldingSetNodeID &ID) {
+      for (auto Operand : Operands)
+        ID.AddPointer(Operand);
+    }
+  };
+
   // SLSR candidate. Such a candidate must be in the form of
   //   (Base + Index) * Stride
   struct Candidate : public ilist_node<Candidate> {
-    Candidate(Value *B = nullptr, ConstantInt *Idx = nullptr,
-              Value *S = nullptr, Instruction *I = nullptr)
-        : Base(B), Index(Idx), Stride(S), Ins(I), Basis(nullptr) {}
-    Value *Base;
+    enum Type {
+      Invalid, // reserved for the default constructor
+      Mul,     // (B + i) * S
+      GEP,     // &B[i * S]
+    };
+
+    Candidate()
+        : CandidateType(Invalid), Base(nullptr), Index(nullptr),
+          Stride(nullptr), Ins(nullptr), Basis(nullptr) {}
+    Candidate(Type CT, BaseExpr *B, ConstantInt *Idx, Value *S, Instruction *I)
+        : CandidateType(CT), Base(B), Index(Idx), Stride(S), Ins(I),
+          Basis(nullptr) {}
+    Type CandidateType;
+    BaseExpr *Base;
     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
     // expression. For instance,
     //
     // (a + 1) * (b + 2)
     //
     // can be treated as
     //
     // <Base: a, Index: 1, Stride: b + 2>
     //
     // or
     //
     // <Base: b, Index: 2, Stride: a + 1>
     Instruction *Ins;
     // Points to the immediate basis of this candidate, or nullptr if we cannot
     // find any basis for this candidate.
     Candidate *Basis;
   };
 
   static char ID;
 
-  StraightLineStrengthReduce() : FunctionPass(ID), DT(nullptr) {
+  StraightLineStrengthReduce()
+      : FunctionPass(ID), DL(nullptr), DT(nullptr), TTI(nullptr) {
     initializeStraightLineStrengthReducePass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<DataLayoutPass>();
     AU.addRequired<DominatorTreeWrapperPass>();
+    AU.addRequired<TargetTransformInfoWrapperPass>();
     // We do not modify the shape of the CFG.
     AU.setPreservesCFG();
   }
 
+  bool doInitialization(Module &M) override {
+    DataLayoutPass *DLP = getAnalysisIfAvailable<DataLayoutPass>();
+    if (DLP == nullptr)
+      report_fatal_error("data layout missing");
+    DL = &DLP->getDataLayout();
+    return false;
+  }
+
   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);
   // 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);
+  // Allocate candidates and find bases for Mul instructions.
+  void allocateCandidateAndFindBasisForMul(BinaryOperator *I);
+  // Allocate candidates and find bases for GetElementPtr instructions.
+  void allocateCandidateAndFindBasisForGEP(GetElementPtrInst *GEP);
+  void allocateCandidateAndFindBasis(Candidate::Type CT, Value *B,
+                                     ConstantInt *Idx, Value *S,
+                                     Instruction *I);
   // Given that I is in the form of "(B + Idx) * S", adds this form to
   // Candidates, and finds its immediate basis.
-  void allocateCandidateAndFindBasis(Value *B, ConstantInt *Idx, Value *S,
+  void allocateCandidateAndFindBasis(Candidate::Type CT, BaseExpr *B,
+                                     ConstantInt *Idx, Value *S,
                                      Instruction *I);
   // Rewrites candidate C with respect to Basis.
   void rewriteCandidateWithBasis(const Candidate &C, const Candidate &Basis);
 
+  const DataLayout *DL;
   DominatorTree *DT;
+  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;
+  // To efficiently decide whether two BaseExprs are equivalent,
+  // we use a FoldingSet to keep unique BaseExprs. For example,
+  // the base expression of "getelementptr a, 1, 2, i" and that of
+  // "getelementptr a, 1, 2, i'" share the same folding set node.
+  FoldingSet<BaseExpr> UniqueBaseExpressions;
 };
 }  // anonymous namespace
 
 char StraightLineStrengthReduce::ID = 0;
 INITIALIZE_PASS_BEGIN(StraightLineStrengthReduce, "slsr",
                       "Straight line strength reduction", false, false)
+INITIALIZE_PASS_DEPENDENCY(DataLayoutPass)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(StraightLineStrengthReduce, "slsr",
                     "Straight line strength reduction", false, false)
 
 FunctionPass *llvm::createStraightLineStrengthReducePass() {
   return new StraightLineStrengthReduce();
 }
 
 bool StraightLineStrengthReduce::isBasisFor(const Candidate &Basis,
                                             const Candidate &C) {
   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 and stride.
+          // They share the same base, stride, and candidate type.
           Basis.Base == C.Base &&
-          Basis.Stride == C.Stride);
+          Basis.Stride == C.Stride &&
+          Basis.CandidateType == C.CandidateType);
 }
 
 // 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(Value *B,
-                                                               ConstantInt *Idx,
-                                                               Value *S,
-                                                               Instruction *I) {
-  Candidate C(B, Idx, S, I);
+void StraightLineStrengthReduce::allocateCandidateAndFindBasis(
+    Candidate::Type CT, BaseExpr *B, ConstantInt *Idx, Value *S,
+  assert(Idx->getType() == S->getType());
+
+  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 (TTI &&
+        TTI->isLegalAddressingMode(
+            GEP->getType(), /*BaseGV=*/nullptr, /*BaseOffset=*/0,
+            /*HasBaseReg=*/true,
+            /*Scale=*/DL->getTypeAllocSize(GEP->getType()->getElementType()) *
+                Idx->getSExtValue())) {
+      return;
+    }
+  }
+
+  Candidate C(CT, B, Idx, S, I);
   // Try to compute the immediate basis of C.
   unsigned NumIterations = 0;
   // Limit the scan radius to avoid running forever.
   static const int 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.
   Candidates.push_back(C);
 }
 
+void StraightLineStrengthReduce::allocateCandidateAndFindBasis(
+    Candidate::Type CT, Value *B, ConstantInt *Idx, Value *S, Instruction *I) {
+  BaseExpr *Base = new BaseExpr(B);
+  BaseExpr *UniqueBase = UniqueBaseExpressions.GetOrInsertNode(Base);
+  if (UniqueBase != Base)
+    delete Base;
+  allocateCandidateAndFindBasis(CT, UniqueBase, Idx, S, I);
+}
+
 void StraightLineStrengthReduce::allocateCandidateAndFindBasis(Instruction *I) {
-  Value *B = nullptr;
-  ConstantInt *Idx = nullptr;
-  // "(Base + Index) * Stride" must be a Mul instruction at the first hand.
-  if (I->getOpcode() == Instruction::Mul) {
-    if (IntegerType *ITy = dyn_cast<IntegerType>(I->getType())) {
-      Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
-      for (unsigned Swapped = 0; Swapped < 2; ++Swapped) {
-        // 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(B, Idx, RHS, I);
-        } else {
-          // Otherwise, at least try the form (LHS + 0) * RHS.
-          allocateCandidateAndFindBasis(LHS, ConstantInt::get(ITy, 0), RHS, I);
-        }
-        // Swap LHS and RHS so that we also cover the cases where LHS is the
-        // stride.
-        if (LHS == RHS)
-          break;
-        std::swap(LHS, RHS);
-      }
-    }
+  switch (I->getOpcode()) {
+  case Instruction::Mul:
+    allocateCandidateAndFindBasisForMul(cast<BinaryOperator>(I));
+    break;
+  case Instruction::GetElementPtr:
+    allocateCandidateAndFindBasisForGEP(cast<GetElementPtrInst>(I));
+    break;
+  }
+}
+
+void StraightLineStrengthReduce::allocateCandidateAndFindBasisForMul(
+    BinaryOperator *I) {
+  // Try matching (B + i) * S.
+  // TODO: we could extend SLSR to float and vector types.
+  IntegerType *ITy = dyn_cast<IntegerType>(I->getType());
+  if (!ITy)
+    return;
+
+    Value *B = nullptr;
+    ConstantInt *Idx = nullptr;
+      allocateCandidateAndFindBasis(Candidate::Mul, B, Idx, RHS, I);
+      allocateCandidateAndFindBasis(Candidate::Mul, LHS,
+                                    ConstantInt::get(ITy, 0), RHS, I);
+    // Swap LHS and RHS so that we also cover the cases where LHS is the stride.
+
+void StraightLineStrengthReduce::allocateCandidateAndFindBasisForGEP(
+    GetElementPtrInst *I) {
+  // Try matching B[...][i * S].
+  // TODO: we exploit the SLSR opportunities only in the last index, because
+  // this case is most common.
+  assert(I->getNumOperands() > 1 && "A GEP should have at least one index");
+
+  // LastType is the type of the last dimension, i.e., typeof(B[...]). SLSR
+  // needs to make sure this type is an array type (i.e., SequentialType).
+  gep_type_iterator LastType = gep_type_begin(*I);
+  for (unsigned Index = 0; Index < I->getNumIndices() - 1; ++Index, ++LastType);
+  if (!isa<SequentialType>(*LastType)) {
+    return;
+  }
+
+  Value *LastIdx = *(I->value_op_end() - 1);
+  if (!isa<IntegerType>(LastIdx->getType()))
+    return;
+
+  BaseExpr *Base = new BaseExpr(I->value_op_begin(), I->value_op_end() - 1);
+  BaseExpr *UniqueBase = UniqueBaseExpressions.GetOrInsertNode(Base);
+  if (UniqueBase != Base) // Base already exists.
+    delete Base;
+
+  for (unsigned Stripped = 0; Stripped < 2; ++Stripped) {
+    // Factorize LastIdx to the form of i * S.
+    // At least LastIdx = 1 * LastIdx.
+    allocateCandidateAndFindBasis(
+        Candidate::GEP, UniqueBase,
+        ConstantInt::get(cast<IntegerType>(LastIdx->getType()), 1), LastIdx, I);
+    Value *LHS = nullptr;
+    ConstantInt *RHS = nullptr;
+    // TODO: handle shl. e.g., we could treat (S << 2) as (S * 4).
+    if (match(LastIdx, m_Mul(m_Value(LHS), m_ConstantInt(RHS)))) {
+      // This transformation is unsafe if i * S may overflow.
+      if (cast<OverflowingBinaryOperator>(LastIdx)->hasNoSignedWrap())
+        allocateCandidateAndFindBasis(Candidate::GEP, UniqueBase, RHS, LHS, I);
+    }
+    // Strips the sext from LastIdx and try factoring again.
+    if (!match(LastIdx, m_SExt(m_Value(LastIdx))))
+      break;
   }
 }
 
 void StraightLineStrengthReduce::rewriteCandidateWithBasis(
     const Candidate &C, const Candidate &Basis) {
+  assert(C.CandidateType == Basis.CandidateType && C.Base == Basis.Base &&
+         C.Stride == Basis.Stride);
+
   // 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;
-  assert(C.Base == Basis.Base && C.Stride == Basis.Stride);
-  // Basis = (B + i) * S
-  // C     = (B + i') * S
-  //   ==>
-  // C     = Basis + (i' - i) * S
+
   IRBuilder<> Builder(C.Ins);
   ConstantInt *IndexOffset = ConstantInt::get(
       C.Ins->getContext(), C.Index->getValue() - Basis.Index->getValue());
-  Value *Reduced;
-  // TODO: preserve nsw/nuw in some cases.
+  // Compute Bump = C - Basis = (i' - i) * S.
+  Value *Bump = nullptr;
   if (IndexOffset->isOne()) {
-    // If (i' - i) is 1, fold C into Basis + S.
-    Reduced = Builder.CreateAdd(Basis.Ins, C.Stride);
+    // If (i' - i) is 1, Bump = S.
+    Bump = C.Stride;
   } else if (IndexOffset->isMinusOne()) {
-    // If (i' - i) is -1, fold C into Basis - S.
-    Reduced = Builder.CreateSub(Basis.Ins, C.Stride);
+    // If (i' - i) is -1, Bump = -S.
+    Bump = Builder.CreateNeg(C.Stride);
   } else {
-    Value *Bump = Builder.CreateMul(C.Stride, IndexOffset);
+    // Otherwise, Bump = (i' - i) * S.
+    Bump = Builder.CreateMul(C.Stride, IndexOffset);
+  }
+
+  Value *Reduced = nullptr;
+  switch (C.CandidateType) {
+  case Candidate::Mul:
+    // Basis = (B + i) * S
+    // C     = (B + i') * S
+    //   ==>
+    // C     = Basis + (i' - i) * S
     Reduced = Builder.CreateAdd(Basis.Ins, Bump);
+    break;
+  case Candidate::GEP:
+    // Basis = &B[..][i * S]
+    // C     = &B[..][i' * S]
+    //   ==>
+    // C     = &Basis[(i' - i) * S]
+    {
+      Type *IntPtrTy = DL->getIntPtrType(C.Ins->getType());
+      // Canonicalize bump to pointer size.
+      Bump = Builder.CreateSExtOrTrunc(Bump, IntPtrTy);
+      if (cast<GetElementPtrInst>(C.Ins)->isInBounds())
+        Reduced = Builder.CreateInBoundsGEP(Basis.Ins, Bump);
+      else
+        Reduced = Builder.CreateGEP(Basis.Ins, Bump);
     }
+    break;
+  default:
+    assert(false && "CandidateType 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);
 }
 
 bool StraightLineStrengthReduce::runOnFunction(Function &F) {
   if (skipOptnoneFunction(F))
     return false;
 
+  if (TargetTransformInfoWrapperPass *TTIP =
+          getAnalysisIfAvailable<TargetTransformInfoWrapperPass>()) {
+    TTI = &TTIP->getTTI(F);
+  } else {
+    TTI = nullptr;
+  }
+
+  UniqueBaseExpressions.clear();
   DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   // 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) {
     BasicBlock *B = node->getBlock();
     for (auto I = B->begin(); I != B->end(); ++I) {
       allocateCandidateAndFindBasis(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();

test/Transforms/StraightLineStrengthReduce/X86/lit.local.cfg

+if not 'X86' in config.root.targets:
+    config.unsupported = True

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(
+  ; v0 = input[0];
+  %p0 = getelementptr inbounds i32* %input, i64 0
+  %v0 = load i32* %p0
+
+  ; v1 = input[s];
+  %p1 = getelementptr inbounds i32* %input, i64 %s
+; CHECK: %p1 = getelementptr inbounds i32* %input, i64 %s
+  %v1 = load i32* %p1
+
+  ; v2 = input[s * 2];
+  %s2 = mul nsw i64 %s, 2
+  %p2 = getelementptr inbounds i32* %input, i64 %s2
+; CHECK: %p2 = getelementptr inbounds i32* %input, i64 %s2
+  %v2 = load i32* %p2
+
+  ; return v0 + v1 + v2;
+  %1 = add i32 %v0, %v1
+  %2 = add i32 %1, %v2
+  ret i32 %2
+}
+

test/Transforms/StraightLineStrengthReduce/slsr-gep.ll

+; RUN: opt < %s -slsr -gvn -dce -S | FileCheck %s
+
+target datalayout = "e-i64:64-v16:16-v32:32-n16:32:64"
+
+define i32 @slsr_gep(i32* %input, i64 %s) {
+; CHECK-LABEL: @slsr_gep(
+  ; v0 = input[0];
+  %p0 = getelementptr inbounds i32* %input, i64 0
+  %v0 = load i32* %p0
+
+  ; v1 = input[s];
+  %p1 = getelementptr inbounds i32* %input, i64 %s
+; CHECK: %p1 = getelementptr inbounds i32* %input, i64 %s
+  %v1 = load i32* %p1
+
+  ; v2 = input[s * 2];
+  %s2 = mul nsw i64 %s, 2
+  %p2 = getelementptr inbounds i32* %input, i64 %s2
+; CHECK: %p2 = getelementptr inbounds i32* %p1, i64 %s
+  %v2 = load i32* %p2
+
+  ; return v0 + v1 + v2;
+  %1 = add i32 %v0, %v1
+  %2 = add i32 %1, %v2
+  ret i32 %2
+}
+
+define i32 @slsr_gep_sext(i32* %input, i32 %s) {
+; CHECK-LABEL: @slsr_gep_sext(
+  ; v0 = input[0];
+  %p0 = getelementptr inbounds i32* %input, i64 0
+  %v0 = load i32* %p0
+
+  ; v1 = input[(long)s];
+  %t = sext i32 %s to i64
+  %p1 = getelementptr inbounds i32* %input, i64 %t
+; CHECK: %p1 = getelementptr inbounds i32* %input, i64 %t
+  %v1 = load i32* %p1
+
+  ; v2 = input[(long)(s * 2)];
+  %s2 = mul nsw i32 %s, 2
+  %t2 = sext i32 %s2 to i64
+  %p2 = getelementptr inbounds i32* %input, i64 %t2
+; CHECK: %p2 = getelementptr inbounds i32* %p1, i64 %t
+  %v2 = load i32* %p2
+
+  ; return v0 + v1 + v2;
+  %1 = add i32 %v0, %v1
+  %2 = add i32 %1, %v2
+  ret i32 %2
+}

test/Transforms/StraightLineStrengthReduce/slsr-mul.ll

 ; RUN: opt < %s -slsr -gvn -dce -S | FileCheck %s
 
+target datalayout = "e-i64:64-v16:16-v32:32-n16:32:64"
+
 declare i32 @foo(i32 %a)
 
 define i32 @slsr1(i32 %b, i32 %s) {
 ; CHECK-LABEL: @slsr1(
   ; v0 = foo(b * s);
   %mul0 = mul i32 %b, %s
 ; CHECK: mul i32
 ; CHECK-NOT: mul i32