Reassociate.cpp

//===- Reassociate.cpp - Reassociate binary expressions -------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This pass reassociates commutative expressions in an order that is designed
// to promote better constant propagation, GCSE, LICM, PRE, etc.
//
// For example: 4 + (x + 5) -> x + (4 + 5)
//
// In the implementation of this algorithm, constants are assigned rank = 0,
// function arguments are rank = 1, and other values are assigned ranks
// corresponding to the reverse post order traversal of current function
// (starting at 2), which effectively gives values in deep loops higher rank
// than values not in loops.
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "reassociate"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Constants.h"
#include "llvm/DerivedTypes.h"
#include "llvm/Function.h"
#include "llvm/Instructions.h"
#include "llvm/IntrinsicInst.h"
#include "llvm/Pass.h"
#include "llvm/Assembly/Writer.h"
#include "llvm/Support/CFG.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/ValueHandle.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/ADT/DenseMap.h"
#include <algorithm>
using namespace llvm;

STATISTIC(NumLinear , "Number of insts linearized");
STATISTIC(NumChanged, "Number of insts reassociated");
STATISTIC(NumAnnihil, "Number of expr tree annihilated");
STATISTIC(NumFactor , "Number of multiplies factored");

namespace {
  struct ValueEntry {
    unsigned Rank;
    Value *Op;
    ValueEntry(unsigned R, Value *O) : Rank(R), Op(O) {}
  };
  inline bool operator<(const ValueEntry &LHS, const ValueEntry &RHS) {
    return LHS.Rank > RHS.Rank;   // Sort so that highest rank goes to start.
  }
}

#ifndef NDEBUG
/// PrintOps - Print out the expression identified in the Ops list.
///
static void PrintOps(Instruction *I, const SmallVectorImpl<ValueEntry> &Ops) {
  Module *M = I->getParent()->getParent()->getParent();
  dbgs() << Instruction::getOpcodeName(I->getOpcode()) << " "
       << *Ops[0].Op->getType() << '\t';
  for (unsigned i = 0, e = Ops.size(); i != e; ++i) {
    dbgs() << "[ ";
    WriteAsOperand(dbgs(), Ops[i].Op, false, M);
    dbgs() << ", #" << Ops[i].Rank << "] ";
  }
}
#endif
  
namespace {
  class Reassociate : public FunctionPass {
    DenseMap<BasicBlock*, unsigned> RankMap;
    DenseMap<AssertingVH<>, unsigned> ValueRankMap;
    SmallVector<WeakVH, 8> RedoInsts;
    SmallVector<WeakVH, 8> DeadInsts;
    bool MadeChange;
  public:
    static char ID; // Pass identification, replacement for typeid
    Reassociate() : FunctionPass(ID) {
      initializeReassociatePass(*PassRegistry::getPassRegistry());
    }

    bool runOnFunction(Function &F);

    virtual void getAnalysisUsage(AnalysisUsage &AU) const {
      AU.setPreservesCFG();
    }
  private:
    void BuildRankMap(Function &F);
    unsigned getRank(Value *V);
    Value *ReassociateExpression(BinaryOperator *I);
    void RewriteExprTree(BinaryOperator *I, SmallVectorImpl<ValueEntry> &Ops,
                         unsigned Idx = 0);
    Value *OptimizeExpression(BinaryOperator *I,
                              SmallVectorImpl<ValueEntry> &Ops);
    Value *OptimizeAdd(Instruction *I, SmallVectorImpl<ValueEntry> &Ops);
    void LinearizeExprTree(BinaryOperator *I, SmallVectorImpl<ValueEntry> &Ops);
    void LinearizeExpr(BinaryOperator *I);
    Value *RemoveFactorFromExpression(Value *V, Value *Factor);
    void ReassociateInst(BasicBlock::iterator &BBI);
    
    void RemoveDeadBinaryOp(Value *V);
  };
}

char Reassociate::ID = 0;
INITIALIZE_PASS(Reassociate, "reassociate",
                "Reassociate expressions", false, false)

// Public interface to the Reassociate pass
FunctionPass *llvm::createReassociatePass() { return new Reassociate(); }

void Reassociate::RemoveDeadBinaryOp(Value *V) {
  Instruction *Op = dyn_cast<Instruction>(V);
  if (!Op || !isa<BinaryOperator>(Op))
    return;
  
  Value *LHS = Op->getOperand(0), *RHS = Op->getOperand(1);
  
  ValueRankMap.erase(Op);
  DeadInsts.push_back(Op);
  RemoveDeadBinaryOp(LHS);
  RemoveDeadBinaryOp(RHS);
}


static bool isUnmovableInstruction(Instruction *I) {
  if (I->getOpcode() == Instruction::PHI ||
      I->getOpcode() == Instruction::Alloca ||
      I->getOpcode() == Instruction::Load ||
      I->getOpcode() == Instruction::Invoke ||
      (I->getOpcode() == Instruction::Call &&
       !isa<DbgInfoIntrinsic>(I)) ||
      I->getOpcode() == Instruction::UDiv || 
      I->getOpcode() == Instruction::SDiv ||
      I->getOpcode() == Instruction::FDiv ||
      I->getOpcode() == Instruction::URem ||
      I->getOpcode() == Instruction::SRem ||
      I->getOpcode() == Instruction::FRem)
    return true;
  return false;
}

void Reassociate::BuildRankMap(Function &F) {
  unsigned i = 2;

  // Assign distinct ranks to function arguments
  for (Function::arg_iterator I = F.arg_begin(), E = F.arg_end(); I != E; ++I)
    ValueRankMap[&*I] = ++i;

  ReversePostOrderTraversal<Function*> RPOT(&F);
  for (ReversePostOrderTraversal<Function*>::rpo_iterator I = RPOT.begin(),
         E = RPOT.end(); I != E; ++I) {
    BasicBlock *BB = *I;
    unsigned BBRank = RankMap[BB] = ++i << 16;

    // Walk the basic block, adding precomputed ranks for any instructions that
    // we cannot move.  This ensures that the ranks for these instructions are
    // all different in the block.
    for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E; ++I)
      if (isUnmovableInstruction(I))
        ValueRankMap[&*I] = ++BBRank;
  }
}

unsigned Reassociate::getRank(Value *V) {
  Instruction *I = dyn_cast<Instruction>(V);
  if (I == 0) {
    if (isa<Argument>(V)) return ValueRankMap[V];   // Function argument.
    return 0;  // Otherwise it's a global or constant, rank 0.
  }

  if (unsigned Rank = ValueRankMap[I])
    return Rank;    // Rank already known?

  // If this is an expression, return the 1+MAX(rank(LHS), rank(RHS)) so that
  // we can reassociate expressions for code motion!  Since we do not recurse
  // for PHI nodes, we cannot have infinite recursion here, because there
  // cannot be loops in the value graph that do not go through PHI nodes.
  unsigned Rank = 0, MaxRank = RankMap[I->getParent()];
  for (unsigned i = 0, e = I->getNumOperands();
       i != e && Rank != MaxRank; ++i)
    Rank = std::max(Rank, getRank(I->getOperand(i)));

  // If this is a not or neg instruction, do not count it for rank.  This
  // assures us that X and ~X will have the same rank.
  if (!I->getType()->isIntegerTy() ||
      (!BinaryOperator::isNot(I) && !BinaryOperator::isNeg(I)))
    ++Rank;

  //DEBUG(dbgs() << "Calculated Rank[" << V->getName() << "] = "
  //     << Rank << "\n");

  return ValueRankMap[I] = Rank;
}