lib/Transforms/IPO/PartialSpecialization.cpp

   1 //===-- PartialSpecialization.cpp - Specialize for common constants--------===//
   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 finds function arguments that are often a common constant and
  11 // specializes a version of the called function for that constant.
  12 //
  13 // This pass simply does the cloning for functions it specializes.  It depends
  14 // on IPSCCP and DAE to clean up the results.
  15 //
  16 // The initial heuristic favors constant arguments that are used in control
  17 // flow.
  18 //
  19 //===----------------------------------------------------------------------===//
  20
  21 #define DEBUG_TYPE "partialspecialization"
  22 #include "llvm/Transforms/IPO.h"
  23 #include "llvm/Constant.h"
  24 #include "llvm/Instructions.h"
  25 #include "llvm/Module.h"
  26 #include "llvm/Pass.h"
  27 #include "llvm/ADT/Statistic.h"
  28 #include "llvm/Analysis/InlineCost.h"
  29 #include "llvm/Transforms/Utils/Cloning.h"
  30 #include "llvm/Support/CallSite.h"
  31 #include "llvm/ADT/DenseSet.h"
  32 #include <map>
  33 using namespace llvm;
  34
  35 STATISTIC(numSpecialized, "Number of specialized functions created");
  36 STATISTIC(numReplaced, "Number of callers replaced by specialization");
  37
  38 // Maximum number of arguments markable interested
  39 static const int MaxInterests = 6;
  40
  41 namespace {
  42   typedef SmallVector<int, MaxInterests> InterestingArgVector;
  43   class PartSpec : public ModulePass {
  44     void scanForInterest(Function&, InterestingArgVector&);
  45     int scanDistribution(Function&, int, std::map<Constant*, int>&);
  46     InlineCostAnalyzer CA;
  47   public :
  48     static char ID; // Pass identification, replacement for typeid
  49     PartSpec() : ModulePass(ID) {
  50       initializePartSpecPass(*PassRegistry::getPassRegistry());
  51     }
  52     bool runOnModule(Module &M);
  53   };
  54 }
  55
  56 char PartSpec::ID = 0;
  57 INITIALIZE_PASS(PartSpec, "partialspecialization",
  58                 "Partial Specialization", false, false)
  59
  60 // Specialize F by replacing the arguments (keys) in replacements with the
  61 // constants (values).  Replace all calls to F with those constants with
  62 // a call to the specialized function.  Returns the specialized function
  63 static Function*
  64 SpecializeFunction(Function* F,
  65                    ValueToValueMapTy& replacements) {
  66   // arg numbers of deleted arguments
  67   DenseMap<unsigned, const Argument*> deleted;
  68   for (ValueToValueMapTy::iterator
  69          repb = replacements.begin(), repe = replacements.end();
  70        repb != repe; ++repb) {
  71     Argument const *arg = cast<const Argument>(repb->first);
  72     deleted[arg->getArgNo()] = arg;
  73   }
  74
  75   Function* NF = CloneFunction(F, replacements,
  76                                /*ModuleLevelChanges=*/false);
  77   NF->setLinkage(GlobalValue::InternalLinkage);
  78   F->getParent()->getFunctionList().push_back(NF);
  79
  80   // FIXME: Specialized versions getting the same constants should also get
  81   // the same name.  That way, specializations for public functions can be
  82   // marked linkonce_odr and reused across modules.
  83
  84   for (Value::use_iterator ii = F->use_begin(), ee = F->use_end();
  85        ii != ee; ) {
  86     Value::use_iterator i = ii;
  87     ++ii;
  88     User *U = *i;
  89     CallSite CS(U);
  90     if (CS) {
  91       if (CS.getCalledFunction() == F) {
  92         SmallVector<Value*, 6> args;
  93         // Assemble the non-specialized arguments for the updated callsite.
  94         // In the process, make sure that the specialized arguments are
  95         // constant and match the specialization.  If that's not the case,
  96         // this callsite needs to call the original or some other
  97         // specialization; don't change it here.
  98         CallSite::arg_iterator as = CS.arg_begin(), ae = CS.arg_end();
  99         for (CallSite::arg_iterator ai = as; ai != ae; ++ai) {
 100           DenseMap<unsigned, const Argument*>::iterator delit = deleted.find(
 101             std::distance(as, ai));
 102           if (delit == deleted.end())
 103             args.push_back(cast<Value>(ai));
 104           else {
 105             Constant *ci = dyn_cast<Constant>(ai);
 106             if (!(ci && ci == replacements[delit->second]))
 107               goto next_use;
 108           }
 109         }
 110         Value* NCall;
 111         if (CallInst *CI = dyn_cast<CallInst>(U)) {
 112           NCall = CallInst::Create(NF, args.begin(), args.end(),
 113                                    CI->getName(), CI);
 114           cast<CallInst>(NCall)->setTailCall(CI->isTailCall());
 115           cast<CallInst>(NCall)->setCallingConv(CI->getCallingConv());
 116         } else {
 117           InvokeInst *II = cast<InvokeInst>(U);
 118           NCall = InvokeInst::Create(NF, II->getNormalDest(),
 119                                      II->getUnwindDest(),
 120                                      args.begin(), args.end(),
 121                                      II->getName(), II);
 122           cast<InvokeInst>(NCall)->setCallingConv(II->getCallingConv());
 123         }
 124         CS.getInstruction()->replaceAllUsesWith(NCall);
 125         CS.getInstruction()->eraseFromParent();
 126         ++numReplaced;
 127       }
 128     }
 129     next_use:;
 130   }
 131   return NF;
 132 }
 133
 134
 135 bool PartSpec::runOnModule(Module &M) {
 136   bool Changed = false;
 137   for (Module::iterator I = M.begin(); I != M.end(); ++I) {
 138     Function &F = *I;
 139     if (F.isDeclaration() || F.mayBeOverridden()) continue;
 140     InterestingArgVector interestingArgs;
 141     scanForInterest(F, interestingArgs);
 142
 143     // Find the first interesting Argument that we can specialize on
 144     // If there are multiple interesting Arguments, then those will be found
 145     // when processing the cloned function.
 146     bool breakOuter = false;
 147     for (unsigned int x = 0; !breakOuter && x < interestingArgs.size(); ++x) {
 148       std::map<Constant*, int> distribution;
 149       scanDistribution(F, interestingArgs[x], distribution);
 150       for (std::map<Constant*, int>::iterator ii = distribution.begin(),
 151              ee = distribution.end(); ii != ee; ++ii) {
 152         // The distribution map might have an entry for NULL (i.e., one or more
 153         // callsites were passing a non-constant there).  We allow that to
 154         // happen so that we can see whether any callsites pass a non-constant;
 155         // if none do and the function is internal, we might have an opportunity
 156         // to kill the original function.
 157         if (!ii->first) continue;
 158         int bonus = ii->second;
 159         SmallVector<unsigned, 1> argnos;
 160         argnos.push_back(interestingArgs[x]);
 161         InlineCost cost = CA.getSpecializationCost(&F, argnos);
 162         // FIXME: If this is the last constant entry, and no non-constant
 163         // entries exist, and the target function is internal, the cost should
 164         // be reduced by the original size of the target function, almost
 165         // certainly making it negative and causing a specialization that will
 166         // leave the original function dead and removable.
 167         if (cost.isAlways() ||
 168            (cost.isVariable() && cost.getValue() < bonus)) {
 169           ValueToValueMapTy m;
 170           Function::arg_iterator arg = F.arg_begin();
 171           for (int y = 0; y < interestingArgs[x]; ++y)
 172             ++arg;
 173           m[&*arg] = ii->first;
 174           SpecializeFunction(&F, m);
 175           ++numSpecialized;
 176           breakOuter = true;
 177           Changed = true;
 178         }
 179       }
 180     }
 181   }
 182   return Changed;
 183 }
 184
 185 /// scanForInterest - This function decides which arguments would be worth
 186 /// specializing on.
 187 void PartSpec::scanForInterest(Function& F, InterestingArgVector& args) {
 188   for(Function::arg_iterator ii = F.arg_begin(), ee = F.arg_end();
 189       ii != ee; ++ii) {
 190     int argno = std::distance(F.arg_begin(), ii);
 191     SmallVector<unsigned, 1> argnos;
 192     argnos.push_back(argno);
 193     int bonus = CA.getSpecializationBonus(&F, argnos);
 194     if (bonus > 0) {
 195       args.push_back(argno);
 196     }
 197   }
 198 }
 199
 200 /// scanDistribution - Construct a histogram of constants for arg of F at arg.
 201 /// For each distinct constant, we'll compute the total of the specialization
 202 /// bonus across all callsites passing that constant; if that total exceeds
 203 /// the specialization cost, we will create the specialization.
 204 int PartSpec::scanDistribution(Function& F, int arg,
 205                                std::map<Constant*, int>& dist) {
 206   bool hasIndirect = false;
 207   int total = 0;
 208   for (Value::use_iterator ii = F.use_begin(), ee = F.use_end();
 209       ii != ee; ++ii) {
 210     User *U = *ii;
 211     CallSite CS(U);
 212     if (CS && CS.getCalledFunction() == &F) {
 213       SmallVector<unsigned, 1> argnos;
 214       argnos.push_back(arg);
 215       dist[dyn_cast<Constant>(CS.getArgument(arg))] +=
 216            CA.getSpecializationBonus(&F, argnos);
 217       ++total;
 218     } else
 219       hasIndirect = true;
 220   }
 221
 222   // Preserve the original address taken function even if all other uses
 223   // will be specialized.
 224   if (hasIndirect) ++total;
 225   return total;
 226 }
 227
 228 ModulePass* llvm::createPartialSpecializationPass() { return new PartSpec(); }