[SLSR] handle candidate form (B + i * S)
[oota-llvm.git] / lib / Transforms / Scalar / StraightLineStrengthReduce.cpp
1 //===-- StraightLineStrengthReduce.cpp - ------------------------*- C++ -*-===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file implements straight-line strength reduction (SLSR). Unlike loop
11 // strength reduction, this algorithm is designed to reduce arithmetic
12 // redundancy in straight-line code instead of loops. It has proven to be
13 // effective in simplifying arithmetic statements derived from an unrolled loop.
14 // It can also simplify the logic of SeparateConstOffsetFromGEP.
15 //
16 // There are many optimizations we can perform in the domain of SLSR. This file
17 // for now contains only an initial step. Specifically, we look for strength
18 // reduction candidates in the following forms:
19 //
20 // Form 1: B + i * S
21 // Form 2: (B + i) * S
22 // Form 3: &B[i * S]
23 //
24 // where S is an integer variable, and i is a constant integer. If we found two
25 // candidates S1 and S2 in the same form and S1 dominates S2, we may rewrite S2
26 // in a simpler way with respect to S1. For example,
27 //
28 // S1: X = B + i * S
29 // S2: Y = B + i' * S   => X + (i' - i) * S
30 //
31 // S1: X = (B + i) * S
32 // S2: Y = (B + i') * S => X + (i' - i) * S
33 //
34 // S1: X = &B[i * S]
35 // S2: Y = &B[i' * S]   => &X[(i' - i) * S]
36 //
37 // Note: (i' - i) * S is folded to the extent possible.
38 //
39 // This rewriting is in general a good idea. The code patterns we focus on
40 // usually come from loop unrolling, so (i' - i) * S is likely the same
41 // across iterations and can be reused. When that happens, the optimized form
42 // takes only one add starting from the second iteration.
43 //
44 // When such rewriting is possible, we call S1 a "basis" of S2. When S2 has
45 // multiple bases, we choose to rewrite S2 with respect to its "immediate"
46 // basis, the basis that is the closest ancestor in the dominator tree.
47 //
48 // TODO:
49 //
50 // - Floating point arithmetics when fast math is enabled.
51 //
52 // - SLSR may decrease ILP at the architecture level. Targets that are very
53 //   sensitive to ILP may want to disable it. Having SLSR to consider ILP is
54 //   left as future work.
55 //
56 // - When (i' - i) is constant but i and i' are not, we could still perform
57 //   SLSR.
58 #include <vector>
59
60 #include "llvm/ADT/DenseSet.h"
61 #include "llvm/ADT/FoldingSet.h"
62 #include "llvm/Analysis/ScalarEvolution.h"
63 #include "llvm/Analysis/TargetTransformInfo.h"
64 #include "llvm/IR/DataLayout.h"
65 #include "llvm/IR/Dominators.h"
66 #include "llvm/IR/IRBuilder.h"
67 #include "llvm/IR/Module.h"
68 #include "llvm/IR/PatternMatch.h"
69 #include "llvm/Support/raw_ostream.h"
70 #include "llvm/Transforms/Scalar.h"
71
72 using namespace llvm;
73 using namespace PatternMatch;
74
75 namespace {
76
77 class StraightLineStrengthReduce : public FunctionPass {
78 public:
79   // SLSR candidate. Such a candidate must be in one of the forms described in
80   // the header comments.
81   struct Candidate : public ilist_node<Candidate> {
82     enum Kind {
83       Invalid, // reserved for the default constructor
84       Add,     // B + i * S
85       Mul,     // (B + i) * S
86       GEP,     // &B[..][i * S][..]
87     };
88
89     Candidate()
90         : CandidateKind(Invalid), Base(nullptr), Index(nullptr),
91           Stride(nullptr), Ins(nullptr), Basis(nullptr) {}
92     Candidate(Kind CT, const SCEV *B, ConstantInt *Idx, Value *S,
93               Instruction *I)
94         : CandidateKind(CT), Base(B), Index(Idx), Stride(S), Ins(I),
95           Basis(nullptr) {}
96     Kind CandidateKind;
97     const SCEV *Base;
98     // Note that Index and Stride of a GEP candidate do not necessarily have the
99     // same integer type. In that case, during rewriting, Stride will be
100     // sign-extended or truncated to Index's type.
101     ConstantInt *Index;
102     Value *Stride;
103     // The instruction this candidate corresponds to. It helps us to rewrite a
104     // candidate with respect to its immediate basis. Note that one instruction
105     // can correspond to multiple candidates depending on how you associate the
106     // expression. For instance,
107     //
108     // (a + 1) * (b + 2)
109     //
110     // can be treated as
111     //
112     // <Base: a, Index: 1, Stride: b + 2>
113     //
114     // or
115     //
116     // <Base: b, Index: 2, Stride: a + 1>
117     Instruction *Ins;
118     // Points to the immediate basis of this candidate, or nullptr if we cannot
119     // find any basis for this candidate.
120     Candidate *Basis;
121   };
122
123   static char ID;
124
125   StraightLineStrengthReduce()
126       : FunctionPass(ID), DL(nullptr), DT(nullptr), TTI(nullptr) {
127     initializeStraightLineStrengthReducePass(*PassRegistry::getPassRegistry());
128   }
129
130   void getAnalysisUsage(AnalysisUsage &AU) const override {
131     AU.addRequired<DominatorTreeWrapperPass>();
132     AU.addRequired<ScalarEvolution>();
133     AU.addRequired<TargetTransformInfoWrapperPass>();
134     // We do not modify the shape of the CFG.
135     AU.setPreservesCFG();
136   }
137
138   bool doInitialization(Module &M) override {
139     DL = &M.getDataLayout();
140     return false;
141   }
142
143   bool runOnFunction(Function &F) override;
144
145 private:
146   // Returns true if Basis is a basis for C, i.e., Basis dominates C and they
147   // share the same base and stride.
148   bool isBasisFor(const Candidate &Basis, const Candidate &C);
149   // Returns whether the candidate can be folded into an addressing mode.
150   bool isFoldable(const Candidate &C, TargetTransformInfo *TTI,
151                   const DataLayout *DL);
152   // Returns true if C is already in a simplest form and not worth being
153   // rewritten.
154   bool isSimplestForm(const Candidate &C);
155   // Checks whether I is in a candidate form. If so, adds all the matching forms
156   // to Candidates, and tries to find the immediate basis for each of them.
157   void allocateCandidatesAndFindBasis(Instruction *I);
158   // Allocate candidates and find bases for Add instructions.
159   void allocateCandidatesAndFindBasisForAdd(Instruction *I);
160   // Given I = LHS + RHS, factors RHS into i * S and makes (LHS + i * S) a
161   // candidate.
162   void allocateCandidatesAndFindBasisForAdd(Value *LHS, Value *RHS,
163                                             Instruction *I);
164   // Allocate candidates and find bases for Mul instructions.
165   void allocateCandidatesAndFindBasisForMul(Instruction *I);
166   // Splits LHS into Base + Index and, if succeeds, calls
167   // allocateCandidatesAndFindBasis.
168   void allocateCandidatesAndFindBasisForMul(Value *LHS, Value *RHS,
169                                             Instruction *I);
170   // Allocate candidates and find bases for GetElementPtr instructions.
171   void allocateCandidatesAndFindBasisForGEP(GetElementPtrInst *GEP);
172   // A helper function that scales Idx with ElementSize before invoking
173   // allocateCandidatesAndFindBasis.
174   void allocateCandidatesAndFindBasisForGEP(const SCEV *B, ConstantInt *Idx,
175                                             Value *S, uint64_t ElementSize,
176                                             Instruction *I);
177   // Adds the given form <CT, B, Idx, S> to Candidates, and finds its immediate
178   // basis.
179   void allocateCandidatesAndFindBasis(Candidate::Kind CT, const SCEV *B,
180                                       ConstantInt *Idx, Value *S,
181                                       Instruction *I);
182   // Rewrites candidate C with respect to Basis.
183   void rewriteCandidateWithBasis(const Candidate &C, const Candidate &Basis);
184   // A helper function that factors ArrayIdx to a product of a stride and a
185   // constant index, and invokes allocateCandidatesAndFindBasis with the
186   // factorings.
187   void factorArrayIndex(Value *ArrayIdx, const SCEV *Base, uint64_t ElementSize,
188                         GetElementPtrInst *GEP);
189   // Emit code that computes the "bump" from Basis to C. If the candidate is a
190   // GEP and the bump is not divisible by the element size of the GEP, this
191   // function sets the BumpWithUglyGEP flag to notify its caller to bump the
192   // basis using an ugly GEP.
193   static Value *emitBump(const Candidate &Basis, const Candidate &C,
194                          IRBuilder<> &Builder, const DataLayout *DL,
195                          bool &BumpWithUglyGEP);
196
197   const DataLayout *DL;
198   DominatorTree *DT;
199   ScalarEvolution *SE;
200   TargetTransformInfo *TTI;
201   ilist<Candidate> Candidates;
202   // Temporarily holds all instructions that are unlinked (but not deleted) by
203   // rewriteCandidateWithBasis. These instructions will be actually removed
204   // after all rewriting finishes.
205   std::vector<Instruction *> UnlinkedInstructions;
206 };
207 }  // anonymous namespace
208
209 char StraightLineStrengthReduce::ID = 0;
210 INITIALIZE_PASS_BEGIN(StraightLineStrengthReduce, "slsr",
211                       "Straight line strength reduction", false, false)
212 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
213 INITIALIZE_PASS_DEPENDENCY(ScalarEvolution)
214 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
215 INITIALIZE_PASS_END(StraightLineStrengthReduce, "slsr",
216                     "Straight line strength reduction", false, false)
217
218 FunctionPass *llvm::createStraightLineStrengthReducePass() {
219   return new StraightLineStrengthReduce();
220 }
221
222 bool StraightLineStrengthReduce::isBasisFor(const Candidate &Basis,
223                                             const Candidate &C) {
224   return (Basis.Ins != C.Ins && // skip the same instruction
225           // Basis must dominate C in order to rewrite C with respect to Basis.
226           DT->dominates(Basis.Ins->getParent(), C.Ins->getParent()) &&
227           // They share the same base, stride, and candidate kind.
228           Basis.Base == C.Base &&
229           Basis.Stride == C.Stride &&
230           Basis.CandidateKind == C.CandidateKind);
231 }
232
233 static bool isGEPFoldable(GetElementPtrInst *GEP,
234                           const TargetTransformInfo *TTI,
235                           const DataLayout *DL) {
236   GlobalVariable *BaseGV = nullptr;
237   int64_t BaseOffset = 0;
238   bool HasBaseReg = false;
239   int64_t Scale = 0;
240
241   if (GlobalVariable *GV = dyn_cast<GlobalVariable>(GEP->getPointerOperand()))
242     BaseGV = GV;
243   else
244     HasBaseReg = true;
245
246   gep_type_iterator GTI = gep_type_begin(GEP);
247   for (auto I = GEP->idx_begin(); I != GEP->idx_end(); ++I, ++GTI) {
248     if (isa<SequentialType>(*GTI)) {
249       int64_t ElementSize = DL->getTypeAllocSize(GTI.getIndexedType());
250       if (ConstantInt *ConstIdx = dyn_cast<ConstantInt>(*I)) {
251         BaseOffset += ConstIdx->getSExtValue() * ElementSize;
252       } else {
253         // Needs scale register.
254         if (Scale != 0) {
255           // No addressing mode takes two scale registers.
256           return false;
257         }
258         Scale = ElementSize;
259       }
260     } else {
261       StructType *STy = cast<StructType>(*GTI);
262       uint64_t Field = cast<ConstantInt>(*I)->getZExtValue();
263       BaseOffset += DL->getStructLayout(STy)->getElementOffset(Field);
264     }
265   }
266   return TTI->isLegalAddressingMode(GEP->getType()->getElementType(), BaseGV,
267                                     BaseOffset, HasBaseReg, Scale);
268 }
269
270 // Returns whether (Base + Index * Stride) can be folded to an addressing mode.
271 static bool isAddFoldable(const SCEV *Base, ConstantInt *Index, Value *Stride,
272                           TargetTransformInfo *TTI) {
273   return TTI->isLegalAddressingMode(Base->getType(), nullptr, 0, true,
274                                     Index->getSExtValue());
275 }
276
277 bool StraightLineStrengthReduce::isFoldable(const Candidate &C,
278                                             TargetTransformInfo *TTI,
279                                             const DataLayout *DL) {
280   if (C.CandidateKind == Candidate::Add)
281     return isAddFoldable(C.Base, C.Index, C.Stride, TTI);
282   if (C.CandidateKind == Candidate::GEP)
283     return isGEPFoldable(cast<GetElementPtrInst>(C.Ins), TTI, DL);
284   return false;
285 }
286
287 // Returns true if GEP has zero or one non-zero index.
288 static bool hasOnlyOneNonZeroIndex(GetElementPtrInst *GEP) {
289   unsigned NumNonZeroIndices = 0;
290   for (auto I = GEP->idx_begin(); I != GEP->idx_end(); ++I) {
291     ConstantInt *ConstIdx = dyn_cast<ConstantInt>(*I);
292     if (ConstIdx == nullptr || !ConstIdx->isZero())
293       ++NumNonZeroIndices;
294   }
295   return NumNonZeroIndices <= 1;
296 }
297
298 bool StraightLineStrengthReduce::isSimplestForm(const Candidate &C) {
299   if (C.CandidateKind == Candidate::Add) {
300     // B + 1 * S or B + (-1) * S
301     return C.Index->isOne() || C.Index->isMinusOne();
302   }
303   if (C.CandidateKind == Candidate::Mul) {
304     // (B + 0) * S
305     return C.Index->isZero();
306   }
307   if (C.CandidateKind == Candidate::GEP) {
308     // (char*)B + S or (char*)B - S
309     return ((C.Index->isOne() || C.Index->isMinusOne()) &&
310             hasOnlyOneNonZeroIndex(cast<GetElementPtrInst>(C.Ins)));
311   }
312   return false;
313 }
314
315 // TODO: We currently implement an algorithm whose time complexity is linear in
316 // the number of existing candidates. However, we could do better by using
317 // ScopedHashTable. Specifically, while traversing the dominator tree, we could
318 // maintain all the candidates that dominate the basic block being traversed in
319 // a ScopedHashTable. This hash table is indexed by the base and the stride of
320 // a candidate. Therefore, finding the immediate basis of a candidate boils down
321 // to one hash-table look up.
322 void StraightLineStrengthReduce::allocateCandidatesAndFindBasis(
323     Candidate::Kind CT, const SCEV *B, ConstantInt *Idx, Value *S,
324     Instruction *I) {
325   Candidate C(CT, B, Idx, S, I);
326   // SLSR can complicate an instruction in two cases:
327   //
328   // 1. If we can fold I into an addressing mode, computing I is likely free or
329   // takes only one instruction.
330   //
331   // 2. I is already in a simplest form. For example, when
332   //      X = B + 8 * S
333   //      Y = B + S,
334   //    rewriting Y to X - 7 * S is probably a bad idea.
335   //
336   // In the above cases, we still add I to the candidate list so that I can be
337   // the basis of other candidates, but we leave I's basis blank so that I
338   // won't be rewritten.
339   if (!isFoldable(C, TTI, DL) && !isSimplestForm(C)) {
340     // Try to compute the immediate basis of C.
341     unsigned NumIterations = 0;
342     // Limit the scan radius to avoid running in quadratice time.
343     static const unsigned MaxNumIterations = 50;
344     for (auto Basis = Candidates.rbegin();
345          Basis != Candidates.rend() && NumIterations < MaxNumIterations;
346          ++Basis, ++NumIterations) {
347       if (isBasisFor(*Basis, C)) {
348         C.Basis = &(*Basis);
349         break;
350       }
351     }
352   }
353   // Regardless of whether we find a basis for C, we need to push C to the
354   // candidate list so that it can be the basis of other candidates.
355   Candidates.push_back(C);
356 }
357
358 void StraightLineStrengthReduce::allocateCandidatesAndFindBasis(
359     Instruction *I) {
360   switch (I->getOpcode()) {
361   case Instruction::Add:
362     allocateCandidatesAndFindBasisForAdd(I);
363     break;
364   case Instruction::Mul:
365     allocateCandidatesAndFindBasisForMul(I);
366     break;
367   case Instruction::GetElementPtr:
368     allocateCandidatesAndFindBasisForGEP(cast<GetElementPtrInst>(I));
369     break;
370   }
371 }
372
373 void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForAdd(
374     Instruction *I) {
375   // Try matching B + i * S.
376   if (!isa<IntegerType>(I->getType()))
377     return;
378
379   assert(I->getNumOperands() == 2 && "isn't I an add?");
380   Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
381   allocateCandidatesAndFindBasisForAdd(LHS, RHS, I);
382   if (LHS != RHS)
383     allocateCandidatesAndFindBasisForAdd(RHS, LHS, I);
384 }
385
386 void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForAdd(
387     Value *LHS, Value *RHS, Instruction *I) {
388   Value *S = nullptr;
389   ConstantInt *Idx = nullptr;
390   if (match(RHS, m_Mul(m_Value(S), m_ConstantInt(Idx)))) {
391     // I = LHS + RHS = LHS + Idx * S
392     allocateCandidatesAndFindBasis(Candidate::Add, SE->getSCEV(LHS), Idx, S, I);
393   } else if (match(RHS, m_Shl(m_Value(S), m_ConstantInt(Idx)))) {
394     // I = LHS + RHS = LHS + (S << Idx) = LHS + S * (1 << Idx)
395     APInt One(Idx->getBitWidth(), 1);
396     Idx = ConstantInt::get(Idx->getContext(), One << Idx->getValue());
397     allocateCandidatesAndFindBasis(Candidate::Add, SE->getSCEV(LHS), Idx, S, I);
398   } else {
399     // At least, I = LHS + 1 * RHS
400     ConstantInt *One = ConstantInt::get(cast<IntegerType>(I->getType()), 1);
401     allocateCandidatesAndFindBasis(Candidate::Add, SE->getSCEV(LHS), One, RHS,
402                                    I);
403   }
404 }
405
406 void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForMul(
407     Value *LHS, Value *RHS, Instruction *I) {
408   Value *B = nullptr;
409   ConstantInt *Idx = nullptr;
410   // Only handle the canonical operand ordering.
411   if (match(LHS, m_Add(m_Value(B), m_ConstantInt(Idx)))) {
412     // If LHS is in the form of "Base + Index", then I is in the form of
413     // "(Base + Index) * RHS".
414     allocateCandidatesAndFindBasis(Candidate::Mul, SE->getSCEV(B), Idx, RHS, I);
415   } else {
416     // Otherwise, at least try the form (LHS + 0) * RHS.
417     ConstantInt *Zero = ConstantInt::get(cast<IntegerType>(I->getType()), 0);
418     allocateCandidatesAndFindBasis(Candidate::Mul, SE->getSCEV(LHS), Zero, RHS,
419                                   I);
420   }
421 }
422
423 void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForMul(
424     Instruction *I) {
425   // Try matching (B + i) * S.
426   // TODO: we could extend SLSR to float and vector types.
427   if (!isa<IntegerType>(I->getType()))
428     return;
429
430   assert(I->getNumOperands() == 2 && "isn't I a mul?");
431   Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
432   allocateCandidatesAndFindBasisForMul(LHS, RHS, I);
433   if (LHS != RHS) {
434     // Symmetrically, try to split RHS to Base + Index.
435     allocateCandidatesAndFindBasisForMul(RHS, LHS, I);
436   }
437 }
438
439 void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForGEP(
440     const SCEV *B, ConstantInt *Idx, Value *S, uint64_t ElementSize,
441     Instruction *I) {
442   // I = B + sext(Idx *nsw S) * ElementSize
443   //   = B + (sext(Idx) * sext(S)) * ElementSize
444   //   = B + (sext(Idx) * ElementSize) * sext(S)
445   // Casting to IntegerType is safe because we skipped vector GEPs.
446   IntegerType *IntPtrTy = cast<IntegerType>(DL->getIntPtrType(I->getType()));
447   ConstantInt *ScaledIdx = ConstantInt::get(
448       IntPtrTy, Idx->getSExtValue() * (int64_t)ElementSize, true);
449   allocateCandidatesAndFindBasis(Candidate::GEP, B, ScaledIdx, S, I);
450 }
451
452 void StraightLineStrengthReduce::factorArrayIndex(Value *ArrayIdx,
453                                                   const SCEV *Base,
454                                                   uint64_t ElementSize,
455                                                   GetElementPtrInst *GEP) {
456   // At least, ArrayIdx = ArrayIdx *nsw 1.
457   allocateCandidatesAndFindBasisForGEP(
458       Base, ConstantInt::get(cast<IntegerType>(ArrayIdx->getType()), 1),
459       ArrayIdx, ElementSize, GEP);
460   Value *LHS = nullptr;
461   ConstantInt *RHS = nullptr;
462   // One alternative is matching the SCEV of ArrayIdx instead of ArrayIdx
463   // itself. This would allow us to handle the shl case for free. However,
464   // matching SCEVs has two issues:
465   //
466   // 1. this would complicate rewriting because the rewriting procedure
467   // would have to translate SCEVs back to IR instructions. This translation
468   // is difficult when LHS is further evaluated to a composite SCEV.
469   //
470   // 2. ScalarEvolution is designed to be control-flow oblivious. It tends
471   // to strip nsw/nuw flags which are critical for SLSR to trace into
472   // sext'ed multiplication.
473   if (match(ArrayIdx, m_NSWMul(m_Value(LHS), m_ConstantInt(RHS)))) {
474     // SLSR is currently unsafe if i * S may overflow.
475     // GEP = Base + sext(LHS *nsw RHS) * ElementSize
476     allocateCandidatesAndFindBasisForGEP(Base, RHS, LHS, ElementSize, GEP);
477   } else if (match(ArrayIdx, m_NSWShl(m_Value(LHS), m_ConstantInt(RHS)))) {
478     // GEP = Base + sext(LHS <<nsw RHS) * ElementSize
479     //     = Base + sext(LHS *nsw (1 << RHS)) * ElementSize
480     APInt One(RHS->getBitWidth(), 1);
481     ConstantInt *PowerOf2 =
482         ConstantInt::get(RHS->getContext(), One << RHS->getValue());
483     allocateCandidatesAndFindBasisForGEP(Base, PowerOf2, LHS, ElementSize, GEP);
484   }
485 }
486
487 void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForGEP(
488     GetElementPtrInst *GEP) {
489   // TODO: handle vector GEPs
490   if (GEP->getType()->isVectorTy())
491     return;
492
493   const SCEV *GEPExpr = SE->getSCEV(GEP);
494   Type *IntPtrTy = DL->getIntPtrType(GEP->getType());
495
496   gep_type_iterator GTI = gep_type_begin(GEP);
497   for (auto I = GEP->idx_begin(); I != GEP->idx_end(); ++I) {
498     if (!isa<SequentialType>(*GTI++))
499       continue;
500     Value *ArrayIdx = *I;
501     // Compute the byte offset of this index.
502     uint64_t ElementSize = DL->getTypeAllocSize(*GTI);
503     const SCEV *ElementSizeExpr = SE->getSizeOfExpr(IntPtrTy, *GTI);
504     const SCEV *ArrayIdxExpr = SE->getSCEV(ArrayIdx);
505     ArrayIdxExpr = SE->getTruncateOrSignExtend(ArrayIdxExpr, IntPtrTy);
506     const SCEV *LocalOffset =
507         SE->getMulExpr(ArrayIdxExpr, ElementSizeExpr, SCEV::FlagNSW);
508     // The base of this candidate equals GEPExpr less the byte offset of this
509     // index.
510     const SCEV *Base = SE->getMinusSCEV(GEPExpr, LocalOffset);
511     factorArrayIndex(ArrayIdx, Base, ElementSize, GEP);
512     // When ArrayIdx is the sext of a value, we try to factor that value as
513     // well.  Handling this case is important because array indices are
514     // typically sign-extended to the pointer size.
515     Value *TruncatedArrayIdx = nullptr;
516     if (match(ArrayIdx, m_SExt(m_Value(TruncatedArrayIdx))))
517       factorArrayIndex(TruncatedArrayIdx, Base, ElementSize, GEP);
518   }
519 }
520
521 // A helper function that unifies the bitwidth of A and B.
522 static void unifyBitWidth(APInt &A, APInt &B) {
523   if (A.getBitWidth() < B.getBitWidth())
524     A = A.sext(B.getBitWidth());
525   else if (A.getBitWidth() > B.getBitWidth())
526     B = B.sext(A.getBitWidth());
527 }
528
529 Value *StraightLineStrengthReduce::emitBump(const Candidate &Basis,
530                                             const Candidate &C,
531                                             IRBuilder<> &Builder,
532                                             const DataLayout *DL,
533                                             bool &BumpWithUglyGEP) {
534   APInt Idx = C.Index->getValue(), BasisIdx = Basis.Index->getValue();
535   unifyBitWidth(Idx, BasisIdx);
536   APInt IndexOffset = Idx - BasisIdx;
537
538   BumpWithUglyGEP = false;
539   if (Basis.CandidateKind == Candidate::GEP) {
540     APInt ElementSize(
541         IndexOffset.getBitWidth(),
542         DL->getTypeAllocSize(
543             cast<GetElementPtrInst>(Basis.Ins)->getType()->getElementType()));
544     APInt Q, R;
545     APInt::sdivrem(IndexOffset, ElementSize, Q, R);
546     if (R.getSExtValue() == 0)
547       IndexOffset = Q;
548     else
549       BumpWithUglyGEP = true;
550   }
551
552   // Compute Bump = C - Basis = (i' - i) * S.
553   // Common case 1: if (i' - i) is 1, Bump = S.
554   if (IndexOffset.getSExtValue() == 1)
555     return C.Stride;
556   // Common case 2: if (i' - i) is -1, Bump = -S.
557   if (IndexOffset.getSExtValue() == -1)
558     return Builder.CreateNeg(C.Stride);
559
560   // Otherwise, Bump = (i' - i) * sext/trunc(S). Note that (i' - i) and S may
561   // have different bit widths.
562   IntegerType *DeltaType =
563       IntegerType::get(Basis.Ins->getContext(), IndexOffset.getBitWidth());
564   Value *ExtendedStride = Builder.CreateSExtOrTrunc(C.Stride, DeltaType);
565   if (IndexOffset.isPowerOf2()) {
566     // If (i' - i) is a power of 2, Bump = sext/trunc(S) << log(i' - i).
567     ConstantInt *Exponent = ConstantInt::get(DeltaType, IndexOffset.logBase2());
568     return Builder.CreateShl(ExtendedStride, Exponent);
569   }
570   if ((-IndexOffset).isPowerOf2()) {
571     // If (i - i') is a power of 2, Bump = -sext/trunc(S) << log(i' - i).
572     ConstantInt *Exponent =
573         ConstantInt::get(DeltaType, (-IndexOffset).logBase2());
574     return Builder.CreateNeg(Builder.CreateShl(ExtendedStride, Exponent));
575   }
576   Constant *Delta = ConstantInt::get(DeltaType, IndexOffset);
577   return Builder.CreateMul(ExtendedStride, Delta);
578 }
579
580 void StraightLineStrengthReduce::rewriteCandidateWithBasis(
581     const Candidate &C, const Candidate &Basis) {
582   assert(C.CandidateKind == Basis.CandidateKind && C.Base == Basis.Base &&
583          C.Stride == Basis.Stride);
584   // We run rewriteCandidateWithBasis on all candidates in a post-order, so the
585   // basis of a candidate cannot be unlinked before the candidate.
586   assert(Basis.Ins->getParent() != nullptr && "the basis is unlinked");
587
588   // An instruction can correspond to multiple candidates. Therefore, instead of
589   // simply deleting an instruction when we rewrite it, we mark its parent as
590   // nullptr (i.e. unlink it) so that we can skip the candidates whose
591   // instruction is already rewritten.
592   if (!C.Ins->getParent())
593     return;
594
595   IRBuilder<> Builder(C.Ins);
596   bool BumpWithUglyGEP;
597   Value *Bump = emitBump(Basis, C, Builder, DL, BumpWithUglyGEP);
598   Value *Reduced = nullptr; // equivalent to but weaker than C.Ins
599   switch (C.CandidateKind) {
600   case Candidate::Add:
601   case Candidate::Mul:
602     if (BinaryOperator::isNeg(Bump)) {
603       Reduced =
604           Builder.CreateSub(Basis.Ins, BinaryOperator::getNegArgument(Bump));
605     } else {
606       Reduced = Builder.CreateAdd(Basis.Ins, Bump);
607     }
608     break;
609   case Candidate::GEP:
610     {
611       Type *IntPtrTy = DL->getIntPtrType(C.Ins->getType());
612       bool InBounds = cast<GetElementPtrInst>(C.Ins)->isInBounds();
613       if (BumpWithUglyGEP) {
614         // C = (char *)Basis + Bump
615         unsigned AS = Basis.Ins->getType()->getPointerAddressSpace();
616         Type *CharTy = Type::getInt8PtrTy(Basis.Ins->getContext(), AS);
617         Reduced = Builder.CreateBitCast(Basis.Ins, CharTy);
618         if (InBounds)
619           Reduced =
620               Builder.CreateInBoundsGEP(Builder.getInt8Ty(), Reduced, Bump);
621         else
622           Reduced = Builder.CreateGEP(Builder.getInt8Ty(), Reduced, Bump);
623         Reduced = Builder.CreateBitCast(Reduced, C.Ins->getType());
624       } else {
625         // C = gep Basis, Bump
626         // Canonicalize bump to pointer size.
627         Bump = Builder.CreateSExtOrTrunc(Bump, IntPtrTy);
628         if (InBounds)
629           Reduced = Builder.CreateInBoundsGEP(nullptr, Basis.Ins, Bump);
630         else
631           Reduced = Builder.CreateGEP(nullptr, Basis.Ins, Bump);
632       }
633     }
634     break;
635   default:
636     llvm_unreachable("C.CandidateKind is invalid");
637   };
638   Reduced->takeName(C.Ins);
639   C.Ins->replaceAllUsesWith(Reduced);
640   C.Ins->dropAllReferences();
641   // Unlink C.Ins so that we can skip other candidates also corresponding to
642   // C.Ins. The actual deletion is postponed to the end of runOnFunction.
643   C.Ins->removeFromParent();
644   UnlinkedInstructions.push_back(C.Ins);
645 }
646
647 bool StraightLineStrengthReduce::runOnFunction(Function &F) {
648   if (skipOptnoneFunction(F))
649     return false;
650
651   TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
652   DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
653   SE = &getAnalysis<ScalarEvolution>();
654   // Traverse the dominator tree in the depth-first order. This order makes sure
655   // all bases of a candidate are in Candidates when we process it.
656   for (auto node = GraphTraits<DominatorTree *>::nodes_begin(DT);
657        node != GraphTraits<DominatorTree *>::nodes_end(DT); ++node) {
658     for (auto &I : *node->getBlock())
659       allocateCandidatesAndFindBasis(&I);
660   }
661
662   // Rewrite candidates in the reverse depth-first order. This order makes sure
663   // a candidate being rewritten is not a basis for any other candidate.
664   while (!Candidates.empty()) {
665     const Candidate &C = Candidates.back();
666     if (C.Basis != nullptr) {
667       rewriteCandidateWithBasis(C, *C.Basis);
668     }
669     Candidates.pop_back();
670   }
671
672   // Delete all unlink instructions.
673   for (auto I : UnlinkedInstructions) {
674     delete I;
675   }
676   bool Ret = !UnlinkedInstructions.empty();
677   UnlinkedInstructions.clear();
678   return Ret;
679 }