lib/Transforms/IPO/PoolAllocate.cpp

   1 //===-- PoolAllocate.cpp - Pool Allocation Pass ---------------------------===//
   2 //
   3 // This transform changes programs so that disjoint data structures are
   4 // allocated out of different pools of memory, increasing locality.
   5 //
   6 //===----------------------------------------------------------------------===//
   7
   8 #include "llvm/Transforms/PoolAllocate.h"
   9 #include "llvm/Transforms/Utils/Cloning.h"
  10 #include "llvm/Analysis/DataStructure.h"
  11 #include "llvm/Analysis/DSGraph.h"
  12 #include "llvm/Module.h"
  13 #include "llvm/DerivedTypes.h"
  14 #include "llvm/Constants.h"
  15 #include "llvm/Instructions.h"
  16 #include "llvm/Target/TargetData.h"
  17 #include "llvm/Support/InstVisitor.h"
  18 #include "Support/Statistic.h"
  19 #include "Support/VectorExtras.h"
  20
  21 using namespace PA;
  22
  23 namespace {
  24   const Type *VoidPtrTy = PointerType::get(Type::SByteTy);
  25   // The type to allocate for a pool descriptor: { sbyte*, uint }
  26   const Type *PoolDescType =
  27     StructType::get(make_vector<const Type*>(VoidPtrTy, Type::UIntTy, 0));
  28   const PointerType *PoolDescPtr = PointerType::get(PoolDescType);
  29
  30   RegisterOpt<PoolAllocate>
  31   X("poolalloc", "Pool allocate disjoint data structures");
  32 }
  33
  34 void PoolAllocate::getAnalysisUsage(AnalysisUsage &AU) const {
  35   AU.addRequired<BUDataStructures>();
  36   AU.addRequired<TargetData>();
  37 }
  38
  39 bool PoolAllocate::run(Module &M) {
  40   if (M.begin() == M.end()) return false;
  41   CurModule = &M;
  42
  43   AddPoolPrototypes();
  44   BU = &getAnalysis<BUDataStructures>();
  45
  46   std::map<Function*, Function*> FuncMap;
  47
  48   // Loop over only the function initially in the program, don't traverse newly
  49   // added ones.  If the function uses memory, make it's clone.
  50   Module::iterator LastOrigFunction = --M.end();
  51   for (Module::iterator I = M.begin(); ; ++I) {
  52     if (!I->isExternal())
  53       if (Function *R = MakeFunctionClone(*I))
  54         FuncMap[I] = R;
  55     if (I == LastOrigFunction) break;
  56   }
  57
  58   ++LastOrigFunction;
  59
  60   // Now that all call targets are available, rewrite the function bodies of the
  61   // clones.
  62   for (Module::iterator I = M.begin(); I != LastOrigFunction; ++I)
  63     if (!I->isExternal()) {
  64       std::map<Function*, Function*>::iterator FI = FuncMap.find(I);
  65       ProcessFunctionBody(*I, FI != FuncMap.end() ? *FI->second : *I);
  66     }
  67
  68   FunctionInfo.clear();
  69   return true;
  70 }
  71
  72
  73 // AddPoolPrototypes - Add prototypes for the pool functions to the specified
  74 // module and update the Pool* instance variables to point to them.
  75 //
  76 void PoolAllocate::AddPoolPrototypes() {
  77   CurModule->addTypeName("PoolDescriptor", PoolDescType);
  78
  79   // Get poolinit function...
  80   FunctionType *PoolInitTy =
  81     FunctionType::get(Type::VoidTy,
  82                       make_vector<const Type*>(PoolDescPtr, Type::UIntTy, 0),
  83                       false);
  84   PoolInit = CurModule->getOrInsertFunction("poolinit", PoolInitTy);
  85
  86   // Get pooldestroy function...
  87   std::vector<const Type*> PDArgs(1, PoolDescPtr);
  88   FunctionType *PoolDestroyTy =
  89     FunctionType::get(Type::VoidTy, PDArgs, false);
  90   PoolDestroy = CurModule->getOrInsertFunction("pooldestroy", PoolDestroyTy);
  91
  92   // Get the poolalloc function...
  93   FunctionType *PoolAllocTy = FunctionType::get(VoidPtrTy, PDArgs, false);
  94   PoolAlloc = CurModule->getOrInsertFunction("poolalloc", PoolAllocTy);
  95
  96   // Get the poolfree function...
  97   PDArgs.push_back(VoidPtrTy);       // Pointer to free
  98   FunctionType *PoolFreeTy = FunctionType::get(Type::VoidTy, PDArgs, false);
  99   PoolFree = CurModule->getOrInsertFunction("poolfree", PoolFreeTy);
 100
 101 #if 0
 102   Args[0] = Type::UIntTy;            // Number of slots to allocate
 103   FunctionType *PoolAllocArrayTy = FunctionType::get(VoidPtrTy, Args, true);
 104   PoolAllocArray = CurModule->getOrInsertFunction("poolallocarray",
 105                                                   PoolAllocArrayTy);
 106 #endif
 107 }
 108
 109
 110 // MakeFunctionClone - If the specified function needs to be modified for pool
 111 // allocation support, make a clone of it, adding additional arguments as
 112 // neccesary, and return it.  If not, just return null.
 113 //
 114 Function *PoolAllocate::MakeFunctionClone(Function &F) {
 115   DSGraph &G = BU->getDSGraph(F);
 116   std::vector<DSNode*> &Nodes = G.getNodes();
 117   if (Nodes.empty()) return 0;  // No memory activity, nothing is required
 118
 119   FuncInfo &FI = FunctionInfo[&F];   // Create a new entry for F
 120   FI.Clone = 0;
 121
 122   // Find DataStructure nodes which are allocated in pools non-local to the
 123   // current function.  This set will contain all of the DSNodes which require
 124   // pools to be passed in from outside of the function.
 125   hash_set<DSNode*> &MarkedNodes = FI.MarkedNodes;
 126
 127   // Mark globals and incomplete nodes as live... (this handles arguments)
 128   if (F.getName() != "main")
 129     for (unsigned i = 0, e = Nodes.size(); i != e; ++i)
 130       if (Nodes[i]->NodeType & (DSNode::GlobalNode | DSNode::Incomplete) &&
 131           Nodes[i]->NodeType & (DSNode::HeapNode))
 132         Nodes[i]->markReachableNodes(MarkedNodes);
 133
 134   // Marked the returned node as alive...
 135   G.getRetNode().getNode()->markReachableNodes(MarkedNodes);
 136
 137   if (MarkedNodes.empty())   // We don't need to clone the function if there
 138     return 0;                // are no incoming arguments to be added.
 139
 140   // Figure out what the arguments are to be for the new version of the function
 141   const FunctionType *OldFuncTy = F.getFunctionType();
 142   std::vector<const Type*> ArgTys;
 143   ArgTys.reserve(OldFuncTy->getParamTypes().size() + MarkedNodes.size());
 144
 145   FI.ArgNodes.reserve(MarkedNodes.size());
 146   for (hash_set<DSNode*>::iterator I = MarkedNodes.begin(),
 147          E = MarkedNodes.end(); I != E; ++I)
 148     if ((*I)->NodeType & DSNode::Incomplete) {
 149       ArgTys.push_back(PoolDescPtr);      // Add the appropriate # of pool descs
 150       FI.ArgNodes.push_back(*I);
 151     }
 152   if (FI.ArgNodes.empty()) return 0;      // No nodes to be pool allocated!
 153
 154   ArgTys.insert(ArgTys.end(), OldFuncTy->getParamTypes().begin(),
 155                 OldFuncTy->getParamTypes().end());
 156
 157
 158   // Create the new function prototype
 159   FunctionType *FuncTy = FunctionType::get(OldFuncTy->getReturnType(), ArgTys,
 160                                            OldFuncTy->isVarArg());
 161   // Create the new function...
 162   Function *New = new Function(FuncTy, true, F.getName(), F.getParent());
 163
 164   // Set the rest of the new arguments names to be PDa<n> and add entries to the
 165   // pool descriptors map
 166   std::map<DSNode*, Value*> &PoolDescriptors = FI.PoolDescriptors;
 167   Function::aiterator NI = New->abegin();
 168   for (unsigned i = 0, e = FI.ArgNodes.size(); i != e; ++i, ++NI) {
 169     NI->setName("PDa");  // Add pd entry
 170     PoolDescriptors.insert(std::make_pair(FI.ArgNodes[i], NI));
 171   }
 172
 173   // Map the existing arguments of the old function to the corresponding
 174   // arguments of the new function.
 175   std::map<const Value*, Value*> ValueMap;
 176   for (Function::aiterator I = F.abegin(), E = F.aend(); I != E; ++I, ++NI) {
 177     ValueMap[I] = NI;
 178     NI->setName(I->getName());
 179   }
 180
 181   // Populate the value map with all of the globals in the program.
 182   // FIXME: This should be unneccesary!
 183   Module &M = *F.getParent();
 184   for (Module::iterator I = M.begin(), E=M.end(); I!=E; ++I)    ValueMap[I] = I;
 185   for (Module::giterator I = M.gbegin(), E=M.gend(); I!=E; ++I) ValueMap[I] = I;
 186
 187   // Perform the cloning.
 188   std::vector<ReturnInst*> Returns;
 189   CloneFunctionInto(New, &F, ValueMap, Returns);
 190
 191   // Invert the ValueMap into the NewToOldValueMap
 192   std::map<Value*, const Value*> &NewToOldValueMap = FI.NewToOldValueMap;
 193   for (std::map<const Value*, Value*>::iterator I = ValueMap.begin(),
 194          E = ValueMap.end(); I != E; ++I)
 195     NewToOldValueMap.insert(std::make_pair(I->second, I->first));
 196
 197   return FI.Clone = New;
 198 }
 199
 200
 201 // processFunction - Pool allocate any data structures which are contained in
 202 // the specified function...
 203 //
 204 void PoolAllocate::ProcessFunctionBody(Function &F, Function &NewF) {
 205   DSGraph &G = BU->getDSGraph(F);
 206   std::vector<DSNode*> &Nodes = G.getNodes();
 207   if (Nodes.empty()) return;     // Quick exit if nothing to do...
 208
 209   FuncInfo &FI = FunctionInfo[&F];   // Get FuncInfo for F
 210   hash_set<DSNode*> &MarkedNodes = FI.MarkedNodes;
 211
 212   DEBUG(std::cerr << "[" << F.getName() << "] Pool Allocate: ");
 213
 214   // Loop over all of the nodes which are non-escaping, adding pool-allocatable
 215   // ones to the NodesToPA vector.
 216   std::vector<DSNode*> NodesToPA;
 217   for (unsigned i = 0, e = Nodes.size(); i != e; ++i)
 218     if (Nodes[i]->NodeType & DSNode::HeapNode &&   // Pick nodes with heap elems
 219         !(Nodes[i]->NodeType & DSNode::Array) &&   // Doesn't handle arrays yet.
 220         !MarkedNodes.count(Nodes[i]))              // Can't be marked
 221       NodesToPA.push_back(Nodes[i]);
 222
 223   DEBUG(std::cerr << NodesToPA.size() << " nodes to pool allocate\n");
 224   if (!NodesToPA.empty()) {
 225     // Create pool construction/destruction code
 226     std::map<DSNode*, Value*> &PoolDescriptors = FI.PoolDescriptors;
 227     CreatePools(NewF, NodesToPA, PoolDescriptors);
 228   }
 229
 230   // Transform the body of the function now...
 231   TransformFunctionBody(NewF, G, FI);
 232 }
 233
 234
 235 // CreatePools - This creates the pool initialization and destruction code for
 236 // the DSNodes specified by the NodesToPA list.  This adds an entry to the
 237 // PoolDescriptors map for each DSNode.
 238 //
 239 void PoolAllocate::CreatePools(Function &F,
 240                                const std::vector<DSNode*> &NodesToPA,
 241                                std::map<DSNode*, Value*> &PoolDescriptors) {
 242   // Find all of the return nodes in the CFG...
 243   std::vector<BasicBlock*> ReturnNodes;
 244   for (Function::iterator I = F.begin(), E = F.end(); I != E; ++I)
 245     if (isa<ReturnInst>(I->getTerminator()))
 246       ReturnNodes.push_back(I);
 247
 248   TargetData &TD = getAnalysis<TargetData>();
 249
 250   // Loop over all of the pools, inserting code into the entry block of the
 251   // function for the initialization and code in the exit blocks for
 252   // destruction.
 253   //
 254   Instruction *InsertPoint = F.front().begin();
 255   for (unsigned i = 0, e = NodesToPA.size(); i != e; ++i) {
 256     DSNode *Node = NodesToPA[i];
 257
 258     // Create a new alloca instruction for the pool...
 259     Value *AI = new AllocaInst(PoolDescType, 0, "PD", InsertPoint);
 260
 261     Value *ElSize =
 262       ConstantUInt::get(Type::UIntTy, TD.getTypeSize(Node->getType()));
 263
 264     // Insert the call to initialize the pool...
 265     new CallInst(PoolInit, make_vector(AI, ElSize, 0), "", InsertPoint);
 266
 267     // Update the PoolDescriptors map
 268     PoolDescriptors.insert(std::make_pair(Node, AI));
 269
 270     // Insert a call to pool destroy before each return inst in the function
 271     for (unsigned r = 0, e = ReturnNodes.size(); r != e; ++r)
 272       new CallInst(PoolDestroy, make_vector(AI, 0), "",
 273                    ReturnNodes[r]->getTerminator());
 274   }
 275 }
 276
 277
 278 namespace {
 279   /// FuncTransform - This class implements transformation required of pool
 280   /// allocated functions.
 281   struct FuncTransform : public InstVisitor<FuncTransform> {
 282     PoolAllocate &PAInfo;
 283     DSGraph &G;
 284     FuncInfo &FI;
 285
 286     FuncTransform(PoolAllocate &P, DSGraph &g, FuncInfo &fi)
 287       : PAInfo(P), G(g), FI(fi) {}
 288
 289     void visitMallocInst(MallocInst &MI);
 290     void visitFreeInst(FreeInst &FI);
 291     void visitCallInst(CallInst &CI);
 292
 293   private:
 294     DSNode *getDSNodeFor(Value *V) {
 295       if (!FI.NewToOldValueMap.empty()) {
 296         // If the NewToOldValueMap is in effect, use it.
 297         std::map<Value*,const Value*>::iterator I = FI.NewToOldValueMap.find(V);
 298         if (I != FI.NewToOldValueMap.end())
 299           V = (Value*)I->second;
 300       }
 301
 302       return G.getScalarMap()[V].getNode();
 303     }
 304     Value *getPoolHandle(Value *V) {
 305       DSNode *Node = getDSNodeFor(V);
 306       // Get the pool handle for this DSNode...
 307       std::map<DSNode*, Value*>::iterator I = FI.PoolDescriptors.find(Node);
 308       return I != FI.PoolDescriptors.end() ? I->second : 0;
 309     }
 310   };
 311 }
 312
 313 void PoolAllocate::TransformFunctionBody(Function &F, DSGraph &G, FuncInfo &FI){
 314   FuncTransform(*this, G, FI).visit(F);
 315 }
 316
 317
 318 void FuncTransform::visitMallocInst(MallocInst &MI) {
 319   // Get the pool handle for the node that this contributes to...
 320   Value *PH = getPoolHandle(&MI);
 321   if (PH == 0) return;
 322
 323   // Insert a call to poolalloc
 324   Value *V = new CallInst(PAInfo.PoolAlloc, make_vector(PH, 0),
 325                           MI.getName(), &MI);
 326   MI.setName("");  // Nuke MIs name
 327
 328   // Cast to the appropriate type...
 329   Value *Casted = new CastInst(V, MI.getType(), V->getName(), &MI);
 330
 331   // Update def-use info
 332   MI.replaceAllUsesWith(Casted);
 333
 334   // Remove old malloc instruction
 335   MI.getParent()->getInstList().erase(&MI);
 336
 337   hash_map<Value*, DSNodeHandle> &SM = G.getScalarMap();
 338   hash_map<Value*, DSNodeHandle>::iterator MII = SM.find(&MI);
 339
 340   // If we are modifying the original function, update the DSGraph...
 341   if (MII != SM.end()) {
 342     // V and Casted now point to whatever the original malloc did...
 343     SM.insert(std::make_pair(V, MII->second));
 344     SM.insert(std::make_pair(Casted, MII->second));
 345     SM.erase(MII);                     // The malloc is now destroyed
 346   } else {             // Otherwise, update the NewToOldValueMap
 347     std::map<Value*,const Value*>::iterator MII =
 348       FI.NewToOldValueMap.find(&MI);
 349     assert(MII != FI.NewToOldValueMap.end() && "MI not found in clone?");
 350     FI.NewToOldValueMap.insert(std::make_pair(V, MII->second));
 351     FI.NewToOldValueMap.insert(std::make_pair(Casted, MII->second));
 352     FI.NewToOldValueMap.erase(MII);
 353   }
 354 }
 355
 356 void FuncTransform::visitFreeInst(FreeInst &FI) {
 357   Value *Arg = FI.getOperand(0);
 358   Value *PH = getPoolHandle(Arg);  // Get the pool handle for this DSNode...
 359   if (PH == 0) return;
 360   // Insert a cast and a call to poolfree...
 361   Value *Casted = new CastInst(Arg, PointerType::get(Type::SByteTy),
 362                                Arg->getName()+".casted", &FI);
 363   new CallInst(PAInfo.PoolFree, make_vector(PH, Casted, 0), "", &FI);
 364
 365   // Delete the now obsolete free instruction...
 366   FI.getParent()->getInstList().erase(&FI);
 367 }
 368
 369 static void CalcNodeMapping(DSNode *Caller, DSNode *Callee,
 370                             std::map<DSNode*, DSNode*> &NodeMapping) {
 371   if (Callee == 0) return;
 372   assert(Caller && "Callee has node but caller doesn't??");
 373
 374   std::map<DSNode*, DSNode*>::iterator I = NodeMapping.find(Callee);
 375   if (I != NodeMapping.end()) {   // Node already in map...
 376     assert(I->second == Caller && "Node maps to different nodes on paths?");
 377   } else {
 378     NodeMapping.insert(I, std::make_pair(Callee, Caller));
 379
 380     // Recursively add pointed to nodes...
 381     for (unsigned i = 0, e = Callee->getNumLinks(); i != e; ++i)
 382       CalcNodeMapping(Caller->getLink(i << DS::PointerShift).getNode(),
 383                       Callee->getLink(i << DS::PointerShift).getNode(),
 384                       NodeMapping);
 385   }
 386 }
 387
 388 void FuncTransform::visitCallInst(CallInst &CI) {
 389   Function *CF = CI.getCalledFunction();
 390   assert(CF && "FIXME: Pool allocation doesn't handle indirect calls!");
 391
 392   FuncInfo *CFI = PAInfo.getFuncInfo(*CF);
 393   if (CFI == 0 || CFI->Clone == 0) return;  // Nothing to transform...
 394
 395   DEBUG(std::cerr << "  Handling call: " << CI);
 396
 397   DSGraph &CG = PAInfo.getBUDataStructures().getDSGraph(*CF);  // Callee graph
 398
 399   // We need to figure out which local pool descriptors correspond to the pool
 400   // descriptor arguments passed into the function call.  Calculate a mapping
 401   // from callee DSNodes to caller DSNodes.  We construct a partial isomophism
 402   // between the graphs to figure out which pool descriptors need to be passed
 403   // in.  The roots of this mapping is found from arguments and return values.
 404   //
 405   std::map<DSNode*, DSNode*> NodeMapping;
 406
 407   Function::aiterator AI = CF->abegin(), AE = CF->aend();
 408   unsigned OpNum = 1;
 409   for (; AI != AE; ++AI, ++OpNum)
 410     CalcNodeMapping(getDSNodeFor(CI.getOperand(OpNum)),
 411                     CG.getScalarMap()[AI].getNode(), NodeMapping);
 412   assert(OpNum == CI.getNumOperands() && "Varargs calls not handled yet!");
 413
 414   // Map the return value as well...
 415   CalcNodeMapping(getDSNodeFor(&CI), CG.getRetNode().getNode(), NodeMapping);
 416
 417
 418   // Okay, now that we have established our mapping, we can figure out which
 419   // pool descriptors to pass in...
 420   std::vector<Value*> Args;
 421
 422   // Add an argument for each pool which must be passed in...
 423   for (unsigned i = 0, e = CFI->ArgNodes.size(); i != e; ++i) {
 424     if (NodeMapping.count(CFI->ArgNodes[i])) {
 425       assert(NodeMapping.count(CFI->ArgNodes[i]) && "Node not in mapping!");
 426       DSNode *LocalNode = NodeMapping.find(CFI->ArgNodes[i])->second;
 427       assert(FI.PoolDescriptors.count(LocalNode) && "Node not pool allocated?");
 428       Args.push_back(FI.PoolDescriptors.find(LocalNode)->second);
 429     } else {
 430       Args.push_back(Constant::getNullValue(PoolDescPtr));
 431     }
 432   }
 433
 434   // Add the rest of the arguments...
 435   Args.insert(Args.end(), CI.op_begin()+1, CI.op_end());
 436
 437   std::string Name = CI.getName(); CI.setName("");
 438   Value *NewCall = new CallInst(CFI->Clone, Args, Name, &CI);
 439   CI.replaceAllUsesWith(NewCall);
 440
 441   DEBUG(std::cerr << "  Result Call: " << *NewCall);
 442   CI.getParent()->getInstList().erase(&CI);
 443 }