[SLSR] simplify ternary adds by reassociation

[jingyue]

Summary

Similar to handling existing candidate forms, this update
pattern-matches

x = a + c
y = (a + b) + c // x dominates y

and rewrites y as

y = x + b.

To be conservative, we only do this when y is the only user of (a + b).

[sanjoy]

Should this rule live inside -reassociate? Optimizing

 define void @reassociate(i32 %a, i32 %b, i32 %c) {
  %1 = add i32 %a, %c
  call void @foo(i32 %1)
  %2 = add i32 %a, %b
  %3 = add i32 %2, %c
  call void @foo(i32 %3)
  ret void
}

to

 define void @reassociate(i32 %a, i32 %b, i32 %c) {
  %1 = add i32 %a, %c
  call void @foo(i32 %1)
  %3 = add i32 %1, %b
  call void @foo(i32 %3)
  ret void
}

looks like a generally good idea to me.

[jingyue]

Hi Sanjoy,

I agree the reassociation performed here is generally good, and I wish I
could put it in -reassociate.

However, I have two major concerns.

1. The algorithm in -reassociate is local in the sense that it ranks nodes

in each individual expression tree and reorders them according to their
ranks. The algorithm implemented here is global: it reassociates an
expression in the way that we can leverage other instructions in the same
function. I also found the workflow of my algorithm nicely fit into SLSR
which scans instructions in the DFS order and simplifies them with respect
to their dominators.

2. -reassociate happens too early. Many opportunities I discovered are

exposed after loop unrolling, but -reassociate happens before loop
unrolling.

Jingyue

[jingyue]

I am thinking of a better alternative to implement this optimization in a separate NaryReassociate pass. I'll see how that goes.

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 folowing forms:
 //
 // Form 1: (B + i) * S
 // Form 2: &B[i * S]
+// Form 3: (B1 + S) + B2
 //
-// where S is an integer variable, and i is a constant integer. If we found two
-// candidates
+// where B, B1, B2, and S are integer variables, and i is a constant integer. If
+// we found two 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
-//
-// or
+// S2: Y = (B + i') * S  => Y = X + (i' - i) * S
 //
 // S1: X = &B[i * S]
-// S2: Y = &B[i' * S]
+// S2: Y = &B[i' * S]    => Y = &X[(i' - i) * S]
 //
-// and S1 dominates S2, we call S1 a basis of S2, and can replace S2 with
+// S1: X = B1 + B2
+// S2: Y = (B1 + S) + B2 => Y = X + S
 //
-// 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.
+// 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
+// - Handle candidates in the form of B + i * S.
+//
+// - Handle integer vectors.
 //
 // - 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/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 
 using namespace llvm;
 using namespace PatternMatch;
 
+#define DEBUG_TYPE "slsr"
+
 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
-      Mul,     // (B + i) * S
-      GEP,     // &B[..][i * S][..]
+      Invalid,    // reserved for the default constructor
+      Mul,        // (B + i) * S
+      GEP,        // &B[..][i * S][..]
+      TernaryAdd, // B + i * S where B is an addition too
     };
 
     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) {}
+    void print(raw_ostream &OS) const;
+
     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;
+    // Stride can be nullptr indicating the wildcard. In that case, Index must
+    // be zero.
     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>
+    // [slsr-mul] base: a, index: 1, stride: b + 2
     //
     // or
     //
-    // <Base: b, Index: 2, Stride: a + 1>
+    // [slsr-mul] 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), DL(nullptr), DT(nullptr), TTI(nullptr) {
     initializeStraightLineStrengthReducePass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addPreserved<DominatorTreeWrapperPass>();
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<ScalarEvolution>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
     // We do not modify the shape of the CFG.
     AU.setPreservesCFG();
   }
 
   bool doInitialization(Module &M) override {
     DL = &M.getDataLayout();
     return false;
   }
 
   bool runOnFunction(Function &F) override;
 
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+  void dumpAllCandidates() const;
+#endif
+
 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);
   // 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:
   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;
+  // TODO: we should split this list into per-candidate-kind lists.
   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
 
+namespace llvm {
+inline raw_ostream &operator<<(raw_ostream &OS,
+                               const StraightLineStrengthReduce::Candidate &C) {
+  C.print(OS);
+  return OS;
+}
+}
+
 char StraightLineStrengthReduce::ID = 0;
 INITIALIZE_PASS_BEGIN(StraightLineStrengthReduce, "slsr",
                       "Straight line strength reduction", false, false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ScalarEvolution)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(StraightLineStrengthReduce, "slsr",
                     "Straight line strength reduction", false, false)
 
 FunctionPass *llvm::createStraightLineStrengthReducePass() {
   return new StraightLineStrengthReduce();
 }
 
+void StraightLineStrengthReduce::Candidate::print(raw_ostream &OS) const {
+  switch (CandidateKind) {
+    case Mul: OS << "[slsr-mul]"; break;
+    case GEP: OS << "[slsr-gep]"; break;
+    case TernaryAdd: OS << "[slsr-ternary-add]"; break;
+    default: llvm_unreachable("CandidateKind is invalid");
+  }
+  OS << " base: " << *Base << ";";
+  OS << " index: " << *Index << ";";
+  OS << " stride: ";
+  if (Stride == nullptr)
+    OS << "<wildcard>";
+  else
+    OS << *Stride;
+}
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+void StraightLineStrengthReduce::dumpAllCandidates() const {
+  for (auto &C : Candidates)
+    dbgs() << C << "\n";
+}
+#endif
+
 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, stride, and candidate kind.
-          Basis.Base == C.Base &&
-          Basis.Stride == C.Stride &&
-          Basis.CandidateKind == C.CandidateKind);
+  // Skip the same instruction.
+  if (Basis.Ins == C.Ins)
+    return false;
+
+  if (!DT->dominates(Basis.Ins->getParent(), C.Ins->getParent()))
+    return false;
+
+  // They must share the same candidate kind.
+  if (Basis.CandidateKind != C.CandidateKind)
+    return false;
+
+  // They share the same base.
+  if (Basis.Base != C.Base)
+    return false;
+
+  // They share the same stride considering wildcards.
+  if (Basis.Stride != nullptr && C.Stride != nullptr &&
+      Basis.Stride != C.Stride)
+    return false;
+
+  // For ternary-add candidates, we expect Basis has index zero and C has index
+  // one; otherwise, the rewriting won't be beneficial.
+  if (C.CandidateKind == Candidate::TernaryAdd) {
+    if (!(Basis.Index->isZero() && C.Index->isOne()))
+      return false;
+  }
+
+  return true;
 }
 
 static bool isCompletelyFoldable(GetElementPtrInst *GEP,
                                  const TargetTransformInfo *TTI,
                                  const DataLayout *DL) {
   GlobalVariable *BaseGV = nullptr;
   int64_t BaseOffset = 0;
   bool HasBaseReg = false;

 // 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::Mul:
     allocateCandidateAndFindBasisForMul(I);
     break;
   case Instruction::GetElementPtr:
+    // If &B[Idx * S] fits into an addressing mode, do not turn it into
+    // non-free computation.
+    if (!isCompletelyFoldable(cast<GetElementPtrInst>(I), TTI, DL))
       allocateCandidateAndFindBasisForGEP(cast<GetElementPtrInst>(I));
     break;
-  }
+  case Instruction::Add:
+    allocateCandidateAndFindBasisForAdd(I);
+    break;
+  }
+
+void StraightLineStrengthReduce::allocateCandidateAndFindBasisForAdd(
+    Instruction *I) {
+  // TODO: skip vector types for now.
+  if (I->getType()->isVectorTy())
+    return;
+
+  Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
+  const SCEV *LHSExpr = SE->getSCEV(LHS), *RHSExpr = SE->getSCEV(RHS);
+  ConstantInt *One = ConstantInt::get(cast<IntegerType>(I->getType()), 1);
+  ConstantInt *Zero = ConstantInt::get(cast<IntegerType>(I->getType()), 0);
+  Value *A = nullptr, *B = nullptr;
+  // To be conservative, we reassociate I only when it is the only user of
+  // (A + B).
+  if (LHS->hasOneUse() && match(LHS, m_Add(m_Value(A), m_Value(B)))) {
+    // I = (A + B) + RHS
+    //   = (A + RHS) + B or (B + RHS) + A
+    const SCEV *AExpr = SE->getSCEV(A), *BExpr = SE->getSCEV(B);
+    allocateCandidateAndFindBasis(
+        Candidate::TernaryAdd, SE->getAddExpr(AExpr, RHSExpr), One, B, I);
+    allocateCandidateAndFindBasis(
+        Candidate::TernaryAdd, SE->getAddExpr(BExpr, RHSExpr), One, A, I);
+  } else if (RHS->hasOneUse() && match(RHS, m_Add(m_Value(A), m_Value(B)))) {
+    // I = LHS + (A + B)
+    //   = (A + LHS) + B or (B + LHS) + A
+    const SCEV *AExpr = SE->getSCEV(A), *BExpr = SE->getSCEV(B);
+    allocateCandidateAndFindBasis(
+        Candidate::TernaryAdd, SE->getAddExpr(AExpr, LHSExpr), One, B, I);
+    allocateCandidateAndFindBasis(
+        Candidate::TernaryAdd, SE->getAddExpr(BExpr, LHSExpr), One, A, I);
+  }
+  // Even I is not a ternary add, it can serve as the basis of other candidates.
+  // Therefore, we add the form
+  //   I = (LHS + RHS) + 0 * <wildcard>
+  // to the candidate list.
+  allocateCandidateAndFindBasis(Candidate::TernaryAdd, SE->getSCEV(I), Zero,
+                                nullptr, 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)))) {
Value *StraightLineStrengthReduce::emitBump(const Candidate &Basis,
                                             const Candidate &C,
                                             IRBuilder<> &Builder,
                                             const DataLayout *DL,
                                             bool &BumpWithUglyGEP) {
   APInt Idx = C.Index->getValue(), BasisIdx = Basis.Index->getValue();
   unifyBitWidth(Idx, BasisIdx);
   APInt IndexOffset = Idx - BasisIdx;
 
+  Value *Stride = C.Stride;
+  if (Stride == nullptr)
+    Stride = Basis.Stride;
+  assert(Stride != nullptr &&
+         "at least one of the two strides should be non-wildcard");
+
   BumpWithUglyGEP = false;
   if (Basis.CandidateKind == Candidate::GEP) {
     APInt ElementSize(
         IndexOffset.getBitWidth(),
         DL->getTypeAllocSize(
             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;
+    return Stride;
   // Common case 2: if (i' - i) is -1, Bump = -S.
   if (IndexOffset.getSExtValue() == -1)
-    return Builder.CreateNeg(C.Stride);
+    return Builder.CreateNeg(Stride);
   // Otherwise, Bump = (i' - i) * sext/trunc(S).
   ConstantInt *Delta = ConstantInt::get(Basis.Ins->getContext(), IndexOffset);
-  Value *ExtendedStride = Builder.CreateSExtOrTrunc(C.Stride, Delta->getType());
+  Value *ExtendedStride = Builder.CreateSExtOrTrunc(Stride, Delta->getType());
   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::TernaryAdd:
   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) {
bool StraightLineStrengthReduce::runOnFunction(Function &F) {
   // 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);
   }
 
+  DEBUG(dumpAllCandidates());
+
   // 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);
     }
     Candidates.pop_back();

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"
+
+declare void @foo(i32 %a)
+
+; foo(a + c);
+; foo((a + b) + c);
+; foo(a + (b + c));
+;   =>
+; t1 = a + c;
+; foo(t);
+; t2 = t1 + b;
+; foo(t2);
+; foo(t2);
+define void @reassociate(i32 %a, i32 %b, i32 %c) {
+; CHECK-LABEL: @reassociate(
+  %1 = add i32 %a, %c
+; CHECK: [[BASE:%[a-zA-Z0-9]+]] = add i32 %a, %c
+  call void @foo(i32 %1)
+  %2 = add i32 %a, %b
+  %3 = add i32 %2, %c
+; CHECK: [[COMMON:%[a-zA-Z0-9]+]] = add i32 [[BASE]], %b
+  call void @foo(i32 %3)
+; CHECK: call void @foo(i32 [[COMMON]])
+  %4 = add i32 %b, %c
+  %5 = add i32 %a, %4
+  call void @foo(i32 %5)
+; CHECK: call void @foo(i32 [[COMMON]])
+  ret void
+}
+
+; t1 = a + c;
+; foo(t1);
+; t2 = a + b;
+; foo(t2);
+; t3 = t2 + c;
+; foo(t3);
+;
+; Do not rewrite t3 into t1 + b because t2 is used elsewhere and is likely free.
+define void @no_reassociate(i32 %a, i32 %b, i32 %c) {
+; CHECK-LABEL: @no_reassociate(
+  %1 = add i32 %a, %c
+; CHECK: add i32 %a, %c
+  call void @foo(i32 %1)
+  %2 = add i32 %a, %b
+; CHECK: add i32 %a, %b
+  call void @foo(i32 %2)
+  %3 = add i32 %2, %c
+; CHECK: add i32 %2, %c
+  call void @foo(i32 %3)
+  ret void
+}