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 its 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   if (DSNode *RetNode = G.getRetNode().getNode())
 136     if (RetNode->NodeType & DSNode::HeapNode)
 137       RetNode->markReachableNodes(MarkedNodes);
 138
 139   if (MarkedNodes.empty())   // We don't need to clone the function if there
 140     return 0;                // are no incoming arguments to be added.
 141
 142   // Figure out what the arguments are to be for the new version of the function
 143   const FunctionType *OldFuncTy = F.getFunctionType();
 144   std::vector<const Type*> ArgTys;
 145   ArgTys.reserve(OldFuncTy->getParamTypes().size() + MarkedNodes.size());
 146
 147   FI.ArgNodes.reserve(MarkedNodes.size());
 148   for (hash_set<DSNode*>::iterator I = MarkedNodes.begin(),
 149          E = MarkedNodes.end(); I != E; ++I)
 150     if ((*I)->NodeType & DSNode::Incomplete) {
 151       ArgTys.push_back(PoolDescPtr);      // Add the appropriate # of pool descs
 152       FI.ArgNodes.push_back(*I);
 153     }
 154   if (FI.ArgNodes.empty()) return 0;      // No nodes to be pool allocated!
 155
 156   ArgTys.insert(ArgTys.end(), OldFuncTy->getParamTypes().begin(),
 157                 OldFuncTy->getParamTypes().end());
 158
 159
 160   // Create the new function prototype
 161   FunctionType *FuncTy = FunctionType::get(OldFuncTy->getReturnType(), ArgTys,
 162                                            OldFuncTy->isVarArg());
 163   // Create the new function...
 164   Function *New = new Function(FuncTy, true, F.getName(), F.getParent());
 165
 166   // Set the rest of the new arguments names to be PDa<n> and add entries to the
 167   // pool descriptors map
 168   std::map<DSNode*, Value*> &PoolDescriptors = FI.PoolDescriptors;
 169   Function::aiterator NI = New->abegin();
 170   for (unsigned i = 0, e = FI.ArgNodes.size(); i != e; ++i, ++NI) {
 171     NI->setName("PDa");  // Add pd entry
 172     PoolDescriptors.insert(std::make_pair(FI.ArgNodes[i], NI));
 173   }
 174
 175   // Map the existing arguments of the old function to the corresponding
 176   // arguments of the new function.
 177   std::map<const Value*, Value*> ValueMap;
 178   for (Function::aiterator I = F.abegin(), E = F.aend(); I != E; ++I, ++NI) {
 179     ValueMap[I] = NI;
 180     NI->setName(I->getName());
 181   }
 182
 183   // Populate the value map with all of the globals in the program.
 184   // FIXME: This should be unneccesary!
 185   Module &M = *F.getParent();
 186   for (Module::iterator I = M.begin(), E=M.end(); I!=E; ++I)    ValueMap[I] = I;
 187   for (Module::giterator I = M.gbegin(), E=M.gend(); I!=E; ++I) ValueMap[I] = I;
 188
 189   // Perform the cloning.
 190   std::vector<ReturnInst*> Returns;
 191   CloneFunctionInto(New, &F, ValueMap, Returns);
 192
 193   // Invert the ValueMap into the NewToOldValueMap
 194   std::map<Value*, const Value*> &NewToOldValueMap = FI.NewToOldValueMap;
 195   for (std::map<const Value*, Value*>::iterator I = ValueMap.begin(),
 196          E = ValueMap.end(); I != E; ++I)
 197     NewToOldValueMap.insert(std::make_pair(I->second, I->first));
 198
 199   return FI.Clone = New;
 200 }
 201
 202
 203 // processFunction - Pool allocate any data structures which are contained in
 204 // the specified function...
 205 //
 206 void PoolAllocate::ProcessFunctionBody(Function &F, Function &NewF) {
 207   DSGraph &G = BU->getDSGraph(F);
 208   std::vector<DSNode*> &Nodes = G.getNodes();
 209   if (Nodes.empty()) return;     // Quick exit if nothing to do...
 210
 211   FuncInfo &FI = FunctionInfo[&F];   // Get FuncInfo for F
 212   hash_set<DSNode*> &MarkedNodes = FI.MarkedNodes;
 213
 214   DEBUG(std::cerr << "[" << F.getName() << "] Pool Allocate: ");
 215
 216   // Loop over all of the nodes which are non-escaping, adding pool-allocatable
 217   // ones to the NodesToPA vector.
 218   std::vector<DSNode*> NodesToPA;
 219   for (unsigned i = 0, e = Nodes.size(); i != e; ++i)
 220     if (Nodes[i]->NodeType & DSNode::HeapNode &&   // Pick nodes with heap elems
 221         !(Nodes[i]->NodeType & DSNode::Array) &&   // Doesn't handle arrays yet.
 222         !MarkedNodes.count(Nodes[i]))              // Can't be marked
 223       NodesToPA.push_back(Nodes[i]);
 224
 225   DEBUG(std::cerr << NodesToPA.size() << " nodes to pool allocate\n");
 226   if (!NodesToPA.empty()) {
 227     // Create pool construction/destruction code
 228     std::map<DSNode*, Value*> &PoolDescriptors = FI.PoolDescriptors;
 229     CreatePools(NewF, NodesToPA, PoolDescriptors);
 230   }
 231
 232   // Transform the body of the function now...
 233   TransformFunctionBody(NewF, G, FI);
 234 }
 235
 236
 237 // CreatePools - This creates the pool initialization and destruction code for
 238 // the DSNodes specified by the NodesToPA list.  This adds an entry to the
 239 // PoolDescriptors map for each DSNode.
 240 //
 241 void PoolAllocate::CreatePools(Function &F,
 242                                const std::vector<DSNode*> &NodesToPA,
 243                                std::map<DSNode*, Value*> &PoolDescriptors) {
 244   // Find all of the return nodes in the CFG...
 245   std::vector<BasicBlock*> ReturnNodes;
 246   for (Function::iterator I = F.begin(), E = F.end(); I != E; ++I)
 247     if (isa<ReturnInst>(I->getTerminator()))
 248       ReturnNodes.push_back(I);
 249
 250   TargetData &TD = getAnalysis<TargetData>();
 251
 252   // Loop over all of the pools, inserting code into the entry block of the
 253   // function for the initialization and code in the exit blocks for
 254   // destruction.
 255   //
 256   Instruction *InsertPoint = F.front().begin();
 257   for (unsigned i = 0, e = NodesToPA.size(); i != e; ++i) {
 258     DSNode *Node = NodesToPA[i];
 259
 260     // Create a new alloca instruction for the pool...
 261     Value *AI = new AllocaInst(PoolDescType, 0, "PD", InsertPoint);
 262
 263     Value *ElSize =
 264       ConstantUInt::get(Type::UIntTy, TD.getTypeSize(Node->getType()));
 265
 266     // Insert the call to initialize the pool...
 267     new CallInst(PoolInit, make_vector(AI, ElSize, 0), "", InsertPoint);
 268
 269     // Update the PoolDescriptors map
 270     PoolDescriptors.insert(std::make_pair(Node, AI));
 271
 272     // Insert a call to pool destroy before each return inst in the function
 273     for (unsigned r = 0, e = ReturnNodes.size(); r != e; ++r)
 274       new CallInst(PoolDestroy, make_vector(AI, 0), "",
 275                    ReturnNodes[r]->getTerminator());
 276   }
 277 }
 278
 279
 280 namespace {
 281   /// FuncTransform - This class implements transformation required of pool
 282   /// allocated functions.
 283   struct FuncTransform : public InstVisitor<FuncTransform> {
 284     PoolAllocate &PAInfo;
 285     DSGraph &G;
 286     FuncInfo &FI;
 287
 288     FuncTransform(PoolAllocate &P, DSGraph &g, FuncInfo &fi)
 289       : PAInfo(P), G(g), FI(fi) {}
 290
 291     void visitMallocInst(MallocInst &MI);
 292     void visitFreeInst(FreeInst &FI);
 293     void visitCallInst(CallInst &CI);
 294
 295   private:
 296     DSNode *getDSNodeFor(Value *V) {
 297       if (!FI.NewToOldValueMap.empty()) {
 298         // If the NewToOldValueMap is in effect, use it.
 299         std::map<Value*,const Value*>::iterator I = FI.NewToOldValueMap.find(V);
 300         if (I != FI.NewToOldValueMap.end())
 301           V = (Value*)I->second;
 302       }
 303
 304       return G.getScalarMap()[V].getNode();
 305     }
 306     Value *getPoolHandle(Value *V) {
 307       DSNode *Node = getDSNodeFor(V);
 308       // Get the pool handle for this DSNode...
 309       std::map<DSNode*, Value*>::iterator I = FI.PoolDescriptors.find(Node);
 310       return I != FI.PoolDescriptors.end() ? I->second : 0;
 311     }
 312   };
 313 }
 314
 315 void PoolAllocate::TransformFunctionBody(Function &F, DSGraph &G, FuncInfo &FI){
 316   FuncTransform(*this, G, FI).visit(F);
 317 }
 318
 319
 320 void FuncTransform::visitMallocInst(MallocInst &MI) {
 321   // Get the pool handle for the node that this contributes to...
 322   Value *PH = getPoolHandle(&MI);
 323   if (PH == 0) return;
 324
 325   // Insert a call to poolalloc
 326   Value *V = new CallInst(PAInfo.PoolAlloc, make_vector(PH, 0),
 327                           MI.getName(), &MI);
 328   MI.setName("");  // Nuke MIs name
 329
 330   // Cast to the appropriate type...
 331   Value *Casted = new CastInst(V, MI.getType(), V->getName(), &MI);
 332
 333   // Update def-use info
 334   MI.replaceAllUsesWith(Casted);
 335
 336   // Remove old malloc instruction
 337   MI.getParent()->getInstList().erase(&MI);
 338
 339   hash_map<Value*, DSNodeHandle> &SM = G.getScalarMap();
 340   hash_map<Value*, DSNodeHandle>::iterator MII = SM.find(&MI);
 341
 342   // If we are modifying the original function, update the DSGraph...
 343   if (MII != SM.end()) {
 344     // V and Casted now point to whatever the original malloc did...
 345     SM.insert(std::make_pair(V, MII->second));
 346     SM.insert(std::make_pair(Casted, MII->second));
 347     SM.erase(MII);                     // The malloc is now destroyed
 348   } else {             // Otherwise, update the NewToOldValueMap
 349     std::map<Value*,const Value*>::iterator MII =
 350       FI.NewToOldValueMap.find(&MI);
 351     assert(MII != FI.NewToOldValueMap.end() && "MI not found in clone?");
 352     FI.NewToOldValueMap.insert(std::make_pair(V, MII->second));
 353     FI.NewToOldValueMap.insert(std::make_pair(Casted, MII->second));
 354     FI.NewToOldValueMap.erase(MII);
 355   }
 356 }
 357
 358 void FuncTransform::visitFreeInst(FreeInst &FI) {
 359   Value *Arg = FI.getOperand(0);
 360   Value *PH = getPoolHandle(Arg);  // Get the pool handle for this DSNode...
 361   if (PH == 0) return;
 362   // Insert a cast and a call to poolfree...
 363   Value *Casted = new CastInst(Arg, PointerType::get(Type::SByteTy),
 364                                Arg->getName()+".casted", &FI);
 365   new CallInst(PAInfo.PoolFree, make_vector(PH, Casted, 0), "", &FI);
 366
 367   // Delete the now obsolete free instruction...
 368   FI.getParent()->getInstList().erase(&FI);
 369 }
 370
 371 static void CalcNodeMapping(DSNode *Caller, DSNode *Callee,
 372                             std::map<DSNode*, DSNode*> &NodeMapping) {
 373   if (Callee == 0) return;
 374   assert(Caller && "Callee has node but caller doesn't??");
 375
 376   std::map<DSNode*, DSNode*>::iterator I = NodeMapping.find(Callee);
 377   if (I != NodeMapping.end()) {   // Node already in map...
 378     assert(I->second == Caller && "Node maps to different nodes on paths?");
 379   } else {
 380     NodeMapping.insert(I, std::make_pair(Callee, Caller));
 381
 382     // Recursively add pointed to nodes...
 383     for (unsigned i = 0, e = Callee->getNumLinks(); i != e; ++i)
 384       CalcNodeMapping(Caller->getLink(i << DS::PointerShift).getNode(),
 385                       Callee->getLink(i << DS::PointerShift).getNode(),
 386                       NodeMapping);
 387   }
 388 }
 389
 390 void FuncTransform::visitCallInst(CallInst &CI) {
 391   Function *CF = CI.getCalledFunction();
 392   assert(CF && "FIXME: Pool allocation doesn't handle indirect calls!");
 393
 394   FuncInfo *CFI = PAInfo.getFuncInfo(*CF);
 395   if (CFI == 0 || CFI->Clone == 0) return;  // Nothing to transform...
 396
 397   DEBUG(std::cerr << "  Handling call: " << CI);
 398
 399   DSGraph &CG = PAInfo.getBUDataStructures().getDSGraph(*CF);  // Callee graph
 400
 401   // We need to figure out which local pool descriptors correspond to the pool
 402   // descriptor arguments passed into the function call.  Calculate a mapping
 403   // from callee DSNodes to caller DSNodes.  We construct a partial isomophism
 404   // between the graphs to figure out which pool descriptors need to be passed
 405   // in.  The roots of this mapping is found from arguments and return values.
 406   //
 407   std::map<DSNode*, DSNode*> NodeMapping;
 408
 409   Function::aiterator AI = CF->abegin(), AE = CF->aend();
 410   unsigned OpNum = 1;
 411   for (; AI != AE; ++AI, ++OpNum)
 412     CalcNodeMapping(getDSNodeFor(CI.getOperand(OpNum)),
 413                     CG.getScalarMap()[AI].getNode(), NodeMapping);
 414   assert(OpNum == CI.getNumOperands() && "Varargs calls not handled yet!");
 415
 416   // Map the return value as well...
 417   CalcNodeMapping(getDSNodeFor(&CI), CG.getRetNode().getNode(), NodeMapping);
 418
 419
 420   // Okay, now that we have established our mapping, we can figure out which
 421   // pool descriptors to pass in...
 422   std::vector<Value*> Args;
 423
 424   // Add an argument for each pool which must be passed in...
 425   for (unsigned i = 0, e = CFI->ArgNodes.size(); i != e; ++i) {
 426     if (NodeMapping.count(CFI->ArgNodes[i])) {
 427       assert(NodeMapping.count(CFI->ArgNodes[i]) && "Node not in mapping!");
 428       DSNode *LocalNode = NodeMapping.find(CFI->ArgNodes[i])->second;
 429       assert(FI.PoolDescriptors.count(LocalNode) && "Node not pool allocated?");
 430       Args.push_back(FI.PoolDescriptors.find(LocalNode)->second);
 431     } else {
 432       Args.push_back(Constant::getNullValue(PoolDescPtr));
 433     }
 434   }
 435
 436   // Add the rest of the arguments...
 437   Args.insert(Args.end(), CI.op_begin()+1, CI.op_end());
 438
 439   std::string Name = CI.getName(); CI.setName("");
 440   Value *NewCall = new CallInst(CFI->Clone, Args, Name, &CI);
 441   CI.replaceAllUsesWith(NewCall);
 442
 443   DEBUG(std::cerr << "  Result Call: " << *NewCall);
 444   CI.getParent()->getInstList().erase(&CI);
 445 }