lib/Transforms/Scalar/NaryReassociate.cpp

   1 //===- NaryReassociate.cpp - Reassociate n-ary expressions ----------------===//
   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 pass reassociates n-ary add expressions and eliminates the redundancy
  11 // exposed by the reassociation.
  12 //
  13 // A motivating example:
  14 //
  15 //   void foo(int a, int b) {
  16 //     bar(a + b);
  17 //     bar((a + 2) + b);
  18 //   }
  19 //
  20 // An ideal compiler should reassociate (a + 2) + b to (a + b) + 2 and simplify
  21 // the above code to
  22 //
  23 //   int t = a + b;
  24 //   bar(t);
  25 //   bar(t + 2);
  26 //
  27 // However, the Reassociate pass is unable to do that because it processes each
  28 // instruction individually and believes (a + 2) + b is the best form according
  29 // to its rank system.
  30 //
  31 // To address this limitation, NaryReassociate reassociates an expression in a
  32 // form that reuses existing instructions. As a result, NaryReassociate can
  33 // reassociate (a + 2) + b in the example to (a + b) + 2 because it detects that
  34 // (a + b) is computed before.
  35 //
  36 // NaryReassociate works as follows. For every instruction in the form of (a +
  37 // b) + c, it checks whether a + c or b + c is already computed by a dominating
  38 // instruction. If so, it then reassociates (a + b) + c into (a + c) + b or (b +
  39 // c) + a and removes the redundancy accordingly. To efficiently look up whether
  40 // an expression is computed before, we store each instruction seen and its SCEV
  41 // into an SCEV-to-instruction map.
  42 //
  43 // Although the algorithm pattern-matches only ternary additions, it
  44 // automatically handles many >3-ary expressions by walking through the function
  45 // in the depth-first order. For example, given
  46 //
  47 //   (a + c) + d
  48 //   ((a + b) + c) + d
  49 //
  50 // NaryReassociate first rewrites (a + b) + c to (a + c) + b, and then rewrites
  51 // ((a + c) + b) + d into ((a + c) + d) + b.
  52 //
  53 // Finally, the above dominator-based algorithm may need to be run multiple
  54 // iterations before emitting optimal code. One source of this need is that we
  55 // only split an operand when it is used only once. The above algorithm can
  56 // eliminate an instruction and decrease the usage count of its operands. As a
  57 // result, an instruction that previously had multiple uses may become a
  58 // single-use instruction and thus eligible for split consideration. For
  59 // example,
  60 //
  61 //   ac = a + c
  62 //   ab = a + b
  63 //   abc = ab + c
  64 //   ab2 = ab + b
  65 //   ab2c = ab2 + c
  66 //
  67 // In the first iteration, we cannot reassociate abc to ac+b because ab is used
  68 // twice. However, we can reassociate ab2c to abc+b in the first iteration. As a
  69 // result, ab2 becomes dead and ab will be used only once in the second
  70 // iteration.
  71 //
  72 // Limitations and TODO items:
  73 //
  74 // 1) We only considers n-ary adds for now. This should be extended and
  75 // generalized.
  76 //
  77 // 2) Besides arithmetic operations, similar reassociation can be applied to
  78 // GEPs. For example, if
  79 //   X = &arr[a]
  80 // dominates
  81 //   Y = &arr[a + b]
  82 // we may rewrite Y into X + b.
  83 //
  84 //===----------------------------------------------------------------------===//
  85
  86 #include "llvm/Analysis/ScalarEvolution.h"
  87 #include "llvm/Analysis/TargetLibraryInfo.h"
  88 #include "llvm/IR/Dominators.h"
  89 #include "llvm/IR/Module.h"
  90 #include "llvm/IR/PatternMatch.h"
  91 #include "llvm/Transforms/Scalar.h"
  92 #include "llvm/Transforms/Utils/Local.h"
  93 using namespace llvm;
  94 using namespace PatternMatch;
  95
  96 #define DEBUG_TYPE "nary-reassociate"
  97
  98 namespace {
  99 class NaryReassociate : public FunctionPass {
 100 public:
 101   static char ID;
 102
 103   NaryReassociate(): FunctionPass(ID) {
 104     initializeNaryReassociatePass(*PassRegistry::getPassRegistry());
 105   }
 106
 107   bool runOnFunction(Function &F) override;
 108
 109   void getAnalysisUsage(AnalysisUsage &AU) const override {
 110     AU.addPreserved<DominatorTreeWrapperPass>();
 111     AU.addPreserved<ScalarEvolution>();
 112     AU.addPreserved<TargetLibraryInfoWrapperPass>();
 113     AU.addRequired<DominatorTreeWrapperPass>();
 114     AU.addRequired<ScalarEvolution>();
 115     AU.addRequired<TargetLibraryInfoWrapperPass>();
 116     AU.setPreservesCFG();
 117   }
 118
 119 private:
 120   // Runs only one iteration of the dominator-based algorithm. See the header
 121   // comments for why we need multiple iterations.
 122   bool doOneIteration(Function &F);
 123   // Reasssociates I to a better form.
 124   Instruction *tryReassociateAdd(Instruction *I);
 125   // A helper function for tryReassociateAdd. LHS and RHS are explicitly passed.
 126   Instruction *tryReassociateAdd(Value *LHS, Value *RHS, Instruction *I);
 127   // Rewrites I to LHS + RHS if LHS is computed already.
 128   Instruction *tryReassociatedAdd(const SCEV *LHS, Value *RHS, Instruction *I);
 129
 130   DominatorTree *DT;
 131   ScalarEvolution *SE;
 132   TargetLibraryInfo *TLI;
 133   // A lookup table quickly telling which instructions compute the given SCEV.
 134   // Note that there can be multiple instructions at different locations
 135   // computing to the same SCEV, so we map a SCEV to an instruction list.  For
 136   // example,
 137   //
 138   //   if (p1)
 139   //     foo(a + b);
 140   //   if (p2)
 141   //     bar(a + b);
 142   DenseMap<const SCEV *, SmallVector<Instruction *, 2>> SeenExprs;
 143 };
 144 } // anonymous namespace
 145
 146 char NaryReassociate::ID = 0;
 147 INITIALIZE_PASS_BEGIN(NaryReassociate, "nary-reassociate", "Nary reassociation",
 148                       false, false)
 149 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 150 INITIALIZE_PASS_DEPENDENCY(ScalarEvolution)
 151 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 152 INITIALIZE_PASS_END(NaryReassociate, "nary-reassociate", "Nary reassociation",
 153                     false, false)
 154
 155 FunctionPass *llvm::createNaryReassociatePass() {
 156   return new NaryReassociate();
 157 }
 158
 159 bool NaryReassociate::runOnFunction(Function &F) {
 160   if (skipOptnoneFunction(F))
 161     return false;
 162
 163   DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
 164   SE = &getAnalysis<ScalarEvolution>();
 165   TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
 166
 167   bool Changed = false, ChangedInThisIteration;
 168   do {
 169     ChangedInThisIteration = doOneIteration(F);
 170     Changed |= ChangedInThisIteration;
 171   } while (ChangedInThisIteration);
 172   return Changed;
 173 }
 174
 175 bool NaryReassociate::doOneIteration(Function &F) {
 176   bool Changed = false;
 177   SeenExprs.clear();
 178   // Traverse the dominator tree in the depth-first order. This order makes sure
 179   // all bases of a candidate are in Candidates when we process it.
 180   for (auto Node = GraphTraits<DominatorTree *>::nodes_begin(DT);
 181        Node != GraphTraits<DominatorTree *>::nodes_end(DT); ++Node) {
 182     BasicBlock *BB = Node->getBlock();
 183     for (auto I = BB->begin(); I != BB->end(); ++I) {
 184       if (I->getOpcode() == Instruction::Add) {
 185         if (Instruction *NewI = tryReassociateAdd(I)) {
 186           Changed = true;
 187           SE->forgetValue(I);
 188           I->replaceAllUsesWith(NewI);
 189           RecursivelyDeleteTriviallyDeadInstructions(I, TLI);
 190           I = NewI;
 191         }
 192         // We should add the rewritten instruction because tryReassociateAdd may
 193         // have invalidated the original one.
 194         SeenExprs[SE->getSCEV(I)].push_back(I);
 195       }
 196     }
 197   }
 198   return Changed;
 199 }
 200
 201 Instruction *NaryReassociate::tryReassociateAdd(Instruction *I) {
 202   Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
 203   if (auto *NewI = tryReassociateAdd(LHS, RHS, I))
 204     return NewI;
 205   if (auto *NewI = tryReassociateAdd(RHS, LHS, I))
 206     return NewI;
 207   return nullptr;
 208 }
 209
 210 Instruction *NaryReassociate::tryReassociateAdd(Value *LHS, Value *RHS,
 211                                                 Instruction *I) {
 212   Value *A = nullptr, *B = nullptr;
 213   // To be conservative, we reassociate I only when it is the only user of A+B.
 214   if (LHS->hasOneUse() && match(LHS, m_Add(m_Value(A), m_Value(B)))) {
 215     // I = (A + B) + RHS
 216     //   = (A + RHS) + B or (B + RHS) + A
 217     const SCEV *AExpr = SE->getSCEV(A), *BExpr = SE->getSCEV(B);
 218     const SCEV *RHSExpr = SE->getSCEV(RHS);
 219     if (auto *NewI = tryReassociatedAdd(SE->getAddExpr(AExpr, RHSExpr), B, I))
 220       return NewI;
 221     if (auto *NewI = tryReassociatedAdd(SE->getAddExpr(BExpr, RHSExpr), A, I))
 222       return NewI;
 223   }
 224   return nullptr;
 225 }
 226
 227 Instruction *NaryReassociate::tryReassociatedAdd(const SCEV *LHSExpr,
 228                                                  Value *RHS, Instruction *I) {
 229   auto Pos = SeenExprs.find(LHSExpr);
 230   // Bail out if LHSExpr is not previously seen.
 231   if (Pos == SeenExprs.end())
 232     return nullptr;
 233
 234   auto &LHSCandidates = Pos->second;
 235   // Look for the closest dominator LHS of I that computes LHSExpr, and replace
 236   // I with LHS + RHS.
 237   //
 238   // Because we traverse the dominator tree in the pre-order, a
 239   // candidate that doesn't dominate the current instruction won't dominate any
 240   // future instruction either. Therefore, we pop it out of the stack. This
 241   // optimization makes the algorithm O(n).
 242   while (!LHSCandidates.empty()) {
 243     Instruction *LHS = LHSCandidates.back();
 244     if (DT->dominates(LHS, I)) {
 245       Instruction *NewI = BinaryOperator::CreateAdd(LHS, RHS, "", I);
 246       NewI->takeName(I);
 247       return NewI;
 248     }
 249     LHSCandidates.pop_back();
 250   }
 251   return nullptr;
 252 }