path: root/lib/Transforms/IPO/PoolAllocate.cpp

//===-- PoolAllocate.cpp - Pool Allocation Pass ---------------------------===//
//
// This transform changes programs so that disjoint data structures are
// allocated out of different pools of memory, increasing locality.
//
//===----------------------------------------------------------------------===//

#include "llvm/Transforms/PoolAllocate.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Analysis/DataStructure.h"
#include "llvm/Analysis/DSGraph.h"
#include "llvm/Module.h"
#include "llvm/DerivedTypes.h"
#include "llvm/Constants.h"
#include "llvm/Instructions.h"
#include "llvm/Target/TargetData.h"
#include "llvm/Support/InstVisitor.h"
#include "Support/Statistic.h"
#include "Support/VectorExtras.h"

using namespace PA;

namespace {
  const Type *VoidPtrTy = PointerType::get(Type::SByteTy);
  // The type to allocate for a pool descriptor: { sbyte*, uint }
  const Type *PoolDescType =
    StructType::get(make_vector<const Type*>(VoidPtrTy, Type::UIntTy, 0));
  const PointerType *PoolDescPtr = PointerType::get(PoolDescType);

  RegisterOpt<PoolAllocate>
  X("poolalloc", "Pool allocate disjoint data structures");
}

void PoolAllocate::getAnalysisUsage(AnalysisUsage &AU) const {
  AU.addRequired<BUDataStructures>();
  AU.addRequired<TargetData>();
}

bool PoolAllocate::run(Module &M) {
  if (M.begin() == M.end()) return false;
  CurModule = &M;
  
  AddPoolPrototypes();
  BU = &getAnalysis<BUDataStructures>();

  std::map<Function*, Function*> FuncMap;

  // Loop over only the function initially in the program, don't traverse newly
  // added ones.  If the function uses memory, make it's clone.
  Module::iterator LastOrigFunction = --M.end();
  for (Module::iterator I = M.begin(); ; ++I) {
    if (!I->isExternal())
      if (Function *R = MakeFunctionClone(*I))
        FuncMap[I] = R;
    if (I == LastOrigFunction) break;
  }

  ++LastOrigFunction;

  // Now that all call targets are available, rewrite the function bodies of the
  // clones.
  for (Module::iterator I = M.begin(); I != LastOrigFunction; ++I)
    if (!I->isExternal()) {
      std::map<Function*, Function*>::iterator FI = FuncMap.find(I);
      ProcessFunctionBody(*I, FI != FuncMap.end() ? *FI->second : *I);
    }

  FunctionInfo.clear();
  return true;
}


// AddPoolPrototypes - Add prototypes for the pool functions to the specified
// module and update the Pool* instance variables to point to them.
//
void PoolAllocate::AddPoolPrototypes() {
  CurModule->addTypeName("PoolDescriptor", PoolDescType);

  // Get poolinit function...
  FunctionType *PoolInitTy =
    FunctionType::get(Type::VoidTy,
                      make_vector<const Type*>(PoolDescPtr, Type::UIntTy, 0),
                      false);
  PoolInit = CurModule->getOrInsertFunction("poolinit", PoolInitTy);

  // Get pooldestroy function...
  std::vector<const Type*> PDArgs(1, PoolDescPtr);
  FunctionType *PoolDestroyTy =
    FunctionType::get(Type::VoidTy, PDArgs, false);
  PoolDestroy = CurModule->getOrInsertFunction("pooldestroy", PoolDestroyTy);

  // Get the poolalloc function...
  FunctionType *PoolAllocTy = FunctionType::get(VoidPtrTy, PDArgs, false);
  PoolAlloc = CurModule->getOrInsertFunction("poolalloc", PoolAllocTy);

  // Get the poolfree function...
  PDArgs.push_back(VoidPtrTy);       // Pointer to free
  FunctionType *PoolFreeTy = FunctionType::get(Type::VoidTy, PDArgs, false);
  PoolFree = CurModule->getOrInsertFunction("poolfree", PoolFreeTy);

#if 0
  Args[0] = Type::UIntTy;            // Number of slots to allocate
  FunctionType *PoolAllocArrayTy = FunctionType::get(VoidPtrTy, Args, true);
  PoolAllocArray = CurModule->getOrInsertFunction("poolallocarray",
                                                  PoolAllocArrayTy);
#endif
}


// MakeFunctionClone - If the specified function needs to be modified for pool
// allocation support, make a clone of it, adding additional arguments as
// neccesary, and return it.  If not, just return null.
//
Function *PoolAllocate::MakeFunctionClone(Function &F) {
  DSGraph &G = BU->getDSGraph(F);
  std::vector<DSNode*> &Nodes = G.getNodes();
  if (Nodes.empty()) return 0;  // No memory activity, nothing is required

  FuncInfo &FI = FunctionInfo[&F];   // Create a new entry for F
  FI.Clone = 0;

  // Find DataStructure nodes which are allocated in pools non-local to the
  // current function.  This set will contain all of the DSNodes which require
  // pools to be passed in from outside of the function.
  hash_set<DSNode*> &MarkedNodes = FI.MarkedNodes;

  // Mark globals and incomplete nodes as live... (this handles arguments)
  if (F.getName() != "main")
    for (unsigned i = 0, e = Nodes.size(); i != e; ++i)
      if (Nodes[i]->NodeType & (DSNode::GlobalNode | DSNode::Incomplete) &&
          Nodes[i]->NodeType & (DSNode::HeapNode))
        Nodes[i]->markReachableNodes(MarkedNodes);

  // Marked the returned node as alive...
  G.getRetNode().getNode()->markReachableNodes(MarkedNodes);

  if (MarkedNodes.empty())   // We don't need to clone the function if there
    return 0;                // are no incoming arguments to be added.

  // Figure out what the arguments are to be for the new version of the function
  const FunctionType *OldFuncTy = F.getFunctionType();
  std::vector<const Type*> ArgTys;
  ArgTys.reserve(OldFuncTy->getParamTypes().size() + MarkedNodes.size());

  FI.ArgNodes.reserve(MarkedNodes.size());
  for (hash_set<DSNode*>::iterator I = MarkedNodes.begin(),
         E = MarkedNodes.end(); I != E; ++I)
    if ((*I)->NodeType & DSNode::Incomplete) {
      ArgTys.push_back(PoolDescPtr);      // Add the appropriate # of pool descs
      FI.ArgNodes.push_back(*I);
    }
  if (FI.ArgNodes.empty()) return 0;      // No nodes to be pool allocated!

  ArgTys.insert(ArgTys.end(), OldFuncTy->getParamTypes().begin(),
                OldFuncTy->getParamTypes().end());


  // Create the new function prototype
  FunctionType *FuncTy = FunctionType::get(OldFuncTy->getReturnType(), ArgTys,
                                           OldFuncTy->isVarArg());
  // Create the new function...
  Function *New = new Function(FuncTy, true, F.getName(), F.getParent());

  // Set the rest of the new arguments names to be PDa<n> and add entries to the
  // pool descriptors map
  std::map<DSNode*, Value*> &PoolDescriptors = FI.PoolDescriptors;
  Function::aiterator NI = New->abegin();
  for (unsigned i = 0, e = FI.ArgNodes.size(); i != e; ++i, ++NI) {
    NI->setName("PDa");  // Add pd entry
    PoolDescriptors.insert(std::make_pair(FI.ArgNodes[i], NI));
  }

  // Map the existing arguments of the old function to the corresponding
  // arguments of the new function.
  std::map<const Value*, Value*> ValueMap;
  for (Function::aiterator I = F.abegin(), E = F.aend(); I != E; ++I, ++NI) {
    ValueMap[I] = NI;
    NI->setName(I->getName());
  }

  // Populate the value map with all of the globals in the program.
  // FIXME: This should be unneccesary!
  Module &M = *F.getParent();
  for (Module::iterator I = M.begin(), E=M.end(); I!=E; ++I)    ValueMap[I] = I;
  for (Module::giterator I = M.gbegin(), E=M.gend(); I!=E; ++I) ValueMap[I] = I;

  // Perform the cloning.
  std::vector<ReturnInst*> Returns;