1 //===-- PoolAllocate.cpp - Pool Allocation Pass ---------------------------===//
3 // This transform changes programs so that disjoint data structures are
4 // allocated out of different pools of memory, increasing locality.
6 //===----------------------------------------------------------------------===//
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"
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);
30 RegisterOpt<PoolAllocate>
31 X("poolalloc", "Pool allocate disjoint data structures");
34 void PoolAllocate::getAnalysisUsage(AnalysisUsage &AU) const {
35 AU.addRequired<BUDataStructures>();
36 AU.addRequired<TargetData>();
39 bool PoolAllocate::run(Module &M) {
40 if (M.begin() == M.end()) return false;
44 BU = &getAnalysis<BUDataStructures>();
46 std::map<Function*, Function*> FuncMap;
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) {
53 if (Function *R = MakeFunctionClone(*I))
55 if (I == LastOrigFunction) break;
60 // Now that all call targets are available, rewrite the function bodies of the
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);
73 // AddPoolPrototypes - Add prototypes for the pool functions to the specified
74 // module and update the Pool* instance variables to point to them.
76 void PoolAllocate::AddPoolPrototypes() {
77 CurModule->addTypeName("PoolDescriptor", PoolDescType);
79 // Get poolinit function...
80 FunctionType *PoolInitTy =
81 FunctionType::get(Type::VoidTy,
82 make_vector<const Type*>(PoolDescPtr, Type::UIntTy, 0),
84 PoolInit = CurModule->getOrInsertFunction("poolinit", PoolInitTy);
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);
92 // Get the poolalloc function...
93 FunctionType *PoolAllocTy = FunctionType::get(VoidPtrTy, PDArgs, false);
94 PoolAlloc = CurModule->getOrInsertFunction("poolalloc", PoolAllocTy);
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);
102 Args[0] = Type::UIntTy; // Number of slots to allocate
103 FunctionType *PoolAllocArrayTy = FunctionType::get(VoidPtrTy, Args, true);
104 PoolAllocArray = CurModule->getOrInsertFunction("poolallocarray",
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.
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
119 FuncInfo &FI = FunctionInfo[&F]; // Create a new entry for F
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;
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);
134 // Marked the returned node as alive...
135 if (DSNode *RetNode = G.getRetNode().getNode())
136 if (RetNode->NodeType & DSNode::HeapNode)
137 RetNode->markReachableNodes(MarkedNodes);
139 if (MarkedNodes.empty()) // We don't need to clone the function if there
140 return 0; // are no incoming arguments to be added.
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());
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);
154 if (FI.ArgNodes.empty()) return 0; // No nodes to be pool allocated!
156 ArgTys.insert(ArgTys.end(), OldFuncTy->getParamTypes().begin(),
157 OldFuncTy->getParamTypes().end());
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());
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));
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) {
180 NI->setName(I->getName());
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;
189 // Perform the cloning.
190 std::vector<ReturnInst*> Returns;
191 CloneFunctionInto(New, &F, ValueMap, Returns);
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));
199 return FI.Clone = New;
203 // processFunction - Pool allocate any data structures which are contained in
204 // the specified function...
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...
211 FuncInfo &FI = FunctionInfo[&F]; // Get FuncInfo for F
212 hash_set<DSNode*> &MarkedNodes = FI.MarkedNodes;
214 DEBUG(std::cerr << "[" << F.getName() << "] Pool Allocate: ");
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]);
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);
232 // Transform the body of the function now...
233 TransformFunctionBody(NewF, G, FI);
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.
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);
250 TargetData &TD = getAnalysis<TargetData>();
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
256 Instruction *InsertPoint = F.front().begin();
257 for (unsigned i = 0, e = NodesToPA.size(); i != e; ++i) {
258 DSNode *Node = NodesToPA[i];
260 // Create a new alloca instruction for the pool...
261 Value *AI = new AllocaInst(PoolDescType, 0, "PD", InsertPoint);
264 ConstantUInt::get(Type::UIntTy, TD.getTypeSize(Node->getType()));
266 // Insert the call to initialize the pool...
267 new CallInst(PoolInit, make_vector(AI, ElSize, 0), "", InsertPoint);
269 // Update the PoolDescriptors map
270 PoolDescriptors.insert(std::make_pair(Node, AI));
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());
281 /// FuncTransform - This class implements transformation required of pool
282 /// allocated functions.
283 struct FuncTransform : public InstVisitor<FuncTransform> {
284 PoolAllocate &PAInfo;
288 FuncTransform(PoolAllocate &P, DSGraph &g, FuncInfo &fi)
289 : PAInfo(P), G(g), FI(fi) {}
291 void visitMallocInst(MallocInst &MI);
292 void visitFreeInst(FreeInst &FI);
293 void visitCallInst(CallInst &CI);
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;
304 return G.getScalarMap()[V].getNode();
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;
315 void PoolAllocate::TransformFunctionBody(Function &F, DSGraph &G, FuncInfo &FI){
316 FuncTransform(*this, G, FI).visit(F);
320 void FuncTransform::visitMallocInst(MallocInst &MI) {
321 // Get the pool handle for the node that this contributes to...
322 Value *PH = getPoolHandle(&MI);
325 // Insert a call to poolalloc
326 Value *V = new CallInst(PAInfo.PoolAlloc, make_vector(PH, 0),
328 MI.setName(""); // Nuke MIs name
330 // Cast to the appropriate type...
331 Value *Casted = new CastInst(V, MI.getType(), V->getName(), &MI);
333 // Update def-use info
334 MI.replaceAllUsesWith(Casted);
336 // Remove old malloc instruction
337 MI.getParent()->getInstList().erase(&MI);
339 hash_map<Value*, DSNodeHandle> &SM = G.getScalarMap();
340 hash_map<Value*, DSNodeHandle>::iterator MII = SM.find(&MI);
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);
358 void FuncTransform::visitFreeInst(FreeInst &FI) {
359 Value *Arg = FI.getOperand(0);
360 Value *PH = getPoolHandle(Arg); // Get the pool handle for this DSNode...
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);
367 // Delete the now obsolete free instruction...
368 FI.getParent()->getInstList().erase(&FI);
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??");
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?");
380 NodeMapping.insert(I, std::make_pair(Callee, Caller));
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(),
390 void FuncTransform::visitCallInst(CallInst &CI) {
391 Function *CF = CI.getCalledFunction();
392 assert(CF && "FIXME: Pool allocation doesn't handle indirect calls!");
394 FuncInfo *CFI = PAInfo.getFuncInfo(*CF);
395 if (CFI == 0 || CFI->Clone == 0) return; // Nothing to transform...
397 DEBUG(std::cerr << " Handling call: " << CI);
399 DSGraph &CG = PAInfo.getBUDataStructures().getDSGraph(*CF); // Callee graph
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.
407 std::map<DSNode*, DSNode*> NodeMapping;
409 Function::aiterator AI = CF->abegin(), AE = CF->aend();
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!");
416 // Map the return value as well...
417 CalcNodeMapping(getDSNodeFor(&CI), CG.getRetNode().getNode(), NodeMapping);
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;
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);
432 Args.push_back(Constant::getNullValue(PoolDescPtr));
436 // Add the rest of the arguments...
437 Args.insert(Args.end(), CI.op_begin()+1, CI.op_end());
439 std::string Name = CI.getName(); CI.setName("");
440 Value *NewCall = new CallInst(CFI->Clone, Args, Name, &CI);
441 CI.replaceAllUsesWith(NewCall);
443 DEBUG(std::cerr << " Result Call: " << *NewCall);
444 CI.getParent()->getInstList().erase(&CI);