[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.

38–39

extend -> extent

41

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?

49

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

74

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?

92

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?)

99

can corresponds -> can correspond

146

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

193

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?

258

to -> in

276

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.

278

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

289

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

362

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

387

is *s a typo?

496

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

508

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.

38–39

Done.

41

Done.

49

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

74

Done.

92

Done.

99

Done.

146

Done.

193

Done.

258

Done.

276

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?

278

Done.

289

Done.

362

Done.

387

Fixed. I mean ArrayIdx = ArrayIdx *nsw 1.

496

Done. Also added a test on that.

508

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 be able to
+// 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 extend 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.
+// 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 the closest basis or
+// the "immediate" basis.
 //
 // 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.
 #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][..]
   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
     // 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
     // expression. For instance,
     //
     // (a + 1) * (b + 2)
     //
     // can be treated as
     //
     // <Base: a, Index: 1, Stride: b + 2>
     //

 
 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 Add instructions.
+  void allocateCandidateAndFindBasisForAdd(Instruction *I);
+  // Given I = LHS + RHS, factors RHS into i * S and makes (LHS + i * S) a
+  // candidate.
+  void allocateCandidateAndFindBasisForAdd(Value *LHS, Value *RHS,
+                                           Instruction *I);
   // Allocate candidates and find bases for Mul instructions.
   void allocateCandidateAndFindBasisForMul(Instruction *I);
   // Splits LHS into Base + Index and, if succeeds, calls
   // allocateCandidateAndFindBasis.
   void allocateCandidateAndFindBasisForMul(Value *LHS, Value *RHS,
                                            Instruction *I);
   // Allocate candidates and find bases for GetElementPtr instructions.
   void allocateCandidateAndFindBasisForGEP(GetElementPtrInst *GEP);
private:
   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;
 };
 }  // 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)
if (isa<SequentialType>(*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;
   }
 
   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 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.
   Candidates.push_back(C);
 }
 
 void StraightLineStrengthReduce::allocateCandidateAndFindBasis(Instruction *I) {
   switch (I->getOpcode()) {
+  case Instruction::Add:
+    allocateCandidateAndFindBasisForAdd(I);
+    break;
   case Instruction::Mul:
     allocateCandidateAndFindBasisForMul(I);
     break;
   case Instruction::GetElementPtr:
     allocateCandidateAndFindBasisForGEP(cast<GetElementPtrInst>(I));
     break;
   }
 }
 
+void StraightLineStrengthReduce::allocateCandidateAndFindBasisForAdd(
+    Instruction *I) {
+  // Try matching B + i * S.
+  if (!isa<IntegerType>(I->getType()))
+    return;
+
+  Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
+  allocateCandidateAndFindBasisForAdd(LHS, RHS, I);
+  if (LHS != RHS)
+    allocateCandidateAndFindBasisForAdd(RHS, LHS, I);
+}
+
+void StraightLineStrengthReduce::allocateCandidateAndFindBasisForAdd(
+    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
+    allocateCandidateAndFindBasis(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());
+    allocateCandidateAndFindBasis(Candidate::Add, SE->getSCEV(LHS), Idx, S, I);
+  } else {
+    // At least, I = LHS + 1 * RHS
+    ConstantInt *One = ConstantInt::get(cast<IntegerType>(I->getType()), 1);
+    allocateCandidateAndFindBasis(Candidate::Add, SE->getSCEV(LHS), One, RHS, I);
+  }
+}
+
 void StraightLineStrengthReduce::allocateCandidateAndFindBasisForMul(
     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);
   } 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,
                                   I);
   }
 }
 
 void StraightLineStrengthReduce::allocateCandidateAndFindBasisForMul(
     Instruction *I) {
   // Try matching (B + i) * S.
   // TODO: we could extend SLSR to float and vector types.
   if (!isa<IntegerType>(I->getType()))
     return;
 
   Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
   allocateCandidateAndFindBasisForMul(LHS, RHS, I);
   if (LHS != RHS) {
     // Symmetrically, try to split RHS to Base + Index.
     allocateCandidateAndFindBasisForMul(RHS, LHS, I);
   }
 }
 
 void StraightLineStrengthReduce::allocateCandidateAndFindBasisForGEP(
     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);
 }
 
 void StraightLineStrengthReduce::factorArrayIndex(Value *ArrayIdx,
                                                   const SCEV *Base,
                                                   uint64_t ElementSize,
                                                   GetElementPtrInst *GEP) {
   // At least, ArrayIdx = ArrayIdx *s 1.
   allocateCandidateAndFindBasisForGEP(
       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:
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);
+  }
+  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);
 
   // 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:
     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) {

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
+}
+
+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)