LoopUnswitch.cpp

//===-- LoopUnswitch.cpp - Hoist loop-invariant conditionals in loop ------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file was developed by the LLVM research group and is distributed under
// the University of Illinois Open Source License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This pass transforms loops that contain branches on loop-invariant conditions
// to have multiple loops.  For example, it turns the left into the right code:
//
//  for (...)                  if (lic)
//    A                          for (...)
//    if (lic)                     A; B; C
//      B                      else
//    C                          for (...)
//                                 A; C
//
// This can increase the size of the code exponentially (doubling it every time
// a loop is unswitched) so we only unswitch if the resultant code will be
// smaller than a threshold.
//
// This pass expects LICM to be run before it to hoist invariant conditions out
// of the loop, to make the unswitching opportunity obvious.
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "loop-unswitch"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Constants.h"
#include "llvm/Function.h"
#include "llvm/Instructions.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/CommandLine.h"
#include <algorithm>
#include <iostream>
#include <set>
using namespace llvm;

namespace {
  Statistic<> NumBranches("loop-unswitch", "Number of branches unswitched");
  Statistic<> NumSwitches("loop-unswitch", "Number of switches unswitched");
  Statistic<> NumSelects ("loop-unswitch", "Number of selects unswitched");
  Statistic<> NumTrivial ("loop-unswitch",
                          "Number of unswitches that are trivial");
  Statistic<> NumSimplify("loop-unswitch", 
                          "Number of simplifications of unswitched code");
  cl::opt<unsigned>
  Threshold("loop-unswitch-threshold", cl::desc("Max loop size to unswitch"),
            cl::init(10), cl::Hidden);
  
  class LoopUnswitch : public FunctionPass {
    LoopInfo *LI;  // Loop information

    // LoopProcessWorklist - List of loops we need to process.
    std::vector<Loop*> LoopProcessWorklist;
  public:
    virtual bool runOnFunction(Function &F);
    bool visitLoop(Loop *L);

    /// This transformation requires natural loop information & requires that
    /// loop preheaders be inserted into the CFG...
    ///
    virtual void getAnalysisUsage(AnalysisUsage &AU) const {
      AU.addRequiredID(LoopSimplifyID);
      AU.addPreservedID(LoopSimplifyID);
      AU.addRequired<LoopInfo>();
      AU.addPreserved<LoopInfo>();
    }

  private:
    /// RemoveLoopFromWorklist - If the specified loop is on the loop worklist,
    /// remove it.
    void RemoveLoopFromWorklist(Loop *L) {
      std::vector<Loop*>::iterator I = std::find(LoopProcessWorklist.begin(),
                                                 LoopProcessWorklist.end(), L);
      if (I != LoopProcessWorklist.end())
        LoopProcessWorklist.erase(I);
    }
      
    bool UnswitchIfProfitable(Value *LoopCond, Constant *Val,Loop *L);
    unsigned getLoopUnswitchCost(Loop *L, Value *LIC);
    void UnswitchTrivialCondition(Loop *L, Value *Cond, Constant *Val,
                                  BasicBlock *ExitBlock);
    void UnswitchNontrivialCondition(Value *LIC, Constant *OnVal, Loop *L);
    BasicBlock *SplitEdge(BasicBlock *From, BasicBlock *To);
    BasicBlock *SplitBlock(BasicBlock *Old, Instruction *SplitPt);

    void RewriteLoopBodyWithConditionConstant(Loop *L, Value *LIC,
                                              Constant *Val, bool isEqual);
    
    void SimplifyCode(std::vector<Instruction*> &Worklist);
    void RemoveBlockIfDead(BasicBlock *BB,
                           std::vector<Instruction*> &Worklist);
    void RemoveLoopFromHierarchy(Loop *L);
  };
  RegisterOpt<LoopUnswitch> X("loop-unswitch", "Unswitch loops");
}

FunctionPass *llvm::createLoopUnswitchPass() { return new LoopUnswitch(); }

bool LoopUnswitch::runOnFunction(Function &F) {
  bool Changed = false;
  LI = &getAnalysis<LoopInfo>();

  // Populate the worklist of loops to process in post-order.
  for (LoopInfo::iterator I = LI->begin(), E = LI->end(); I != E; ++I)
    for (po_iterator<Loop*> LI = po_begin(*I), E = po_end(*I); LI != E; ++LI)
      LoopProcessWorklist.push_back(*LI);

  // Process the loops in worklist order, this is a post-order visitation of
  // the loops.  We use a worklist of loops so that loops can be removed at any
  // time if they are deleted (e.g. the backedge of a loop is removed).
  while (!LoopProcessWorklist.empty()) {
    Loop *L = LoopProcessWorklist.back();
    LoopProcessWorklist.pop_back();    
    Changed |= visitLoop(L);
  }

  return Changed;
}

/// FindLIVLoopCondition - Cond is a condition that occurs in L.  If it is
/// invariant in the loop, or has an invariant piece, return the invariant.
/// Otherwise, return null.
static Value *FindLIVLoopCondition(Value *Cond, Loop *L, bool &Changed) {
  // Constants should be folded, not unswitched on!
  if (isa<Constant>(Cond)) return false;
  
  // TODO: Handle: br (VARIANT|INVARIANT).
  // TODO: Hoist simple expressions out of loops.
  if (L->isLoopInvariant(Cond)) return Cond;
  
  if (BinaryOperator *BO = dyn_cast<BinaryOperator>(Cond))
    if (BO->getOpcode() == Instruction::And ||
        BO->getOpcode() == Instruction::Or) {
      // If either the left or right side is invariant, we can unswitch on this,
      // which will cause the branch to go away in one loop and the condition to
      // simplify in the other one.
      if (Value *LHS = FindLIVLoopCondition(BO->getOperand(0), L, Changed))
        return LHS;
      if (Value *RHS = FindLIVLoopCondition(BO->getOperand(1), L, Changed))
        return RHS;
    }
      
      return 0;
}

bool LoopUnswitch::visitLoop(Loop *L) {
  bool Changed = false;
  
  // Loop over all of the basic blocks in the loop.  If we find an interior
  // block that is branching on a loop-invariant condition, we can unswitch this
  // loop.
  for (Loop::block_iterator I = L->block_begin(), E = L->block_end();
       I != E; ++I) {
    TerminatorInst *TI = (*I)->getTerminator();
    if (BranchInst *BI = dyn_cast<BranchInst>(TI)) {
      // If this isn't branching on an invariant condition, we can't unswitch
      // it.
      if (BI->isConditional()) {
        // See if this, or some part of it, is loop invariant.  If so, we can
        // unswitch on it if we desire.
        Value *LoopCond = FindLIVLoopCondition(BI->getCondition(), L, Changed);
        if (LoopCond && UnswitchIfProfitable(LoopCond, ConstantBool::True, L)) {
          ++NumBranches;
          return true;
        }
      }      
    } else if (SwitchInst *SI = dyn_cast<SwitchInst>(TI)) {
      Value *LoopCond = FindLIVLoopCondition(SI->getCondition(), L, Changed);
      if (LoopCond && SI->getNumCases() > 1) {
        // Find a value to unswitch on:
        // FIXME: this should chose the most expensive case!
        Constant *UnswitchVal = SI->getCaseValue(1);
        if (UnswitchIfProfitable(LoopCond, UnswitchVal, L)) {
          ++NumSwitches;
          return true;
        }
      }
    }
    
    // Scan the instructions to check for unswitchable values.
    for (BasicBlock::iterator BBI = (*I)->begin(), E = (*I)->end(); 
         BBI != E; ++BBI)
      if (SelectInst *SI = dyn_cast<SelectInst>(BBI)) {
        Value *LoopCond = FindLIVLoopCondition(SI->getCondition(), L, Changed);
        if (LoopCond && UnswitchIfProfitable(LoopCond, ConstantBool::True, L)) {
          ++NumSelects;
          return true;
        }
      }
  }