#include "llvm/Pass.h" | |
#include "llvm/Function.h" | |
#include "llvm/Module.h" | |
#include "llvm/Instructions.h" | |
#include "llvm/Constants.h" | |
#include "llvm/GlobalVariable.h" | |
#include "llvm/Support/IRBuilder.h" | |
#include "llvm/ADT/StringMap.h" | |
#include "llvm/Analysis/LoopInfo.h" | |
#include "llvm/Transforms/Utils/BasicBlockUtils.h" | |
#include "llvm/DefaultPasses.h" | |
#include <string> | |
using namespace llvm; | |
using std::string; | |
using std::pair; | |
namespace | |
{ | |
class MemoiseExample : public ModulePass { | |
/// Module that we're currently optimising | |
Module *M; | |
/// Static cache. | |
llvm::StringMap<GlobalVariable*> statics; | |
// Lookup - call plus its argument | |
typedef std::pair<CallInst*,std::string> ExampleCall; | |
bool runOnFunction(Function &F); | |
public: | |
static char ID; | |
MemoiseExample() : ModulePass(ID) {} | |
virtual bool runOnModule(Module &Mod); | |
}; | |
RegisterPass<MemoiseExample> X("example-pass", | |
"Memoise example pass"); | |
} | |
char MemoiseExample::ID; | |
static void removeTerminator(BasicBlock *BB) { | |
TerminatorInst *BBTerm = BB->getTerminator(); | |
// Remove the BB as a predecessor from all of successors | |
for (unsigned i = 0, e = BBTerm->getNumSuccessors(); i != e; ++i) { | |
BBTerm->getSuccessor(i)->removePredecessor(BB); | |
} | |
BBTerm->replaceAllUsesWith(UndefValue::get(BBTerm->getType())); | |
// Remove the terminator instruction itself. | |
BBTerm->eraseFromParent(); | |
} | |
bool MemoiseExample::runOnFunction(Function &F) { | |
bool modified = false; | |
SmallVector<ExampleCall, 16> Lookups; | |
for (auto &i : F) { | |
for (auto &b : i) { | |
if (CallInst *c = dyn_cast<CallInst>(&b)) { | |
if (Function *func = c->getCalledFunction()) { | |
if (func->getName() == "example") { | |
ExampleCall lookup; | |
GlobalVariable *arg= dyn_cast<GlobalVariable>( | |
c->getOperand(0)->stripPointerCasts()); | |
if (0 == arg) { continue; } | |
ConstantDataArray *init = dyn_cast<ConstantDataArray>( | |
arg->getInitializer()); | |
if (0 == init || !init->isCString()) { continue; } | |
lookup.first = c; | |
lookup.second = init->getAsString(); | |
modified = true; | |
Lookups.push_back(lookup); | |
} | |
} | |
} | |
} | |
} | |
for (SmallVectorImpl<ExampleCall>::iterator i=Lookups.begin(), | |
e=Lookups.end() ; e!=i ; i++) { | |
llvm::Instruction *lookup = i->first; | |
std::string &arg= i->second; | |
Type *retTy = lookup->getType(); | |
GlobalVariable *cache = statics[arg]; | |
if (!cache) { | |
cache = new GlobalVariable(*M, retTy, false, | |
GlobalVariable::PrivateLinkage, Constant::getNullValue(retTy), | |
"._cache"); | |
statics[arg] = cache; | |
} | |
BasicBlock *beforeLookupBB = lookup->getParent(); | |
BasicBlock *lookupBB = SplitBlock(beforeLookupBB, lookup, this); | |
BasicBlock::iterator iter = lookup; | |
iter++; | |
BasicBlock *afterLookupBB = SplitBlock(iter->getParent(), iter, this); | |
// SplitBlock() adds an unconditional branch, which we don't want. | |
// Remove it. | |
removeTerminator(beforeLookupBB); | |
removeTerminator(lookupBB); | |
PHINode *phi = PHINode::Create(retTy, 2, arg, afterLookupBB->begin()); | |
// We replace all of the existing uses with the PHI node now, because | |
// we're going to add some more uses later that we don't want | |
// replaced. | |
lookup->replaceAllUsesWith(phi); | |
// In the original basic block, we test whether the cache is NULL, | |
// and skip the lookup if it isn't. | |
IRBuilder<> B(beforeLookupBB); | |
llvm::Value *cachedClass = | |
B.CreateBitCast(B.CreateLoad(cache), retTy); | |
llvm::Value *needsLookup = B.CreateIsNull(cachedClass); | |
B.CreateCondBr(needsLookup, lookupBB, afterLookupBB); | |
// In the lookup basic block, we just do the lookup, store it in the | |
// cache, and then jump to the continue block | |
B.SetInsertPoint(lookupBB); | |
B.CreateStore(lookup, cache); | |
B.CreateBr(afterLookupBB); | |
// Now we just need to set the PHI node to use the cache or the | |
// lookup result | |
phi->addIncoming(cachedClass, beforeLookupBB); | |
phi->addIncoming(lookup, lookupBB); | |
} | |
return modified; | |
} | |
bool MemoiseExample::runOnModule(Module &Mod) { | |
statics.empty(); | |
M = &Mod; | |
bool modified = false; | |
for (auto &F : Mod) { | |
if (F.isDeclaration()) { continue; } | |
modified |= runOnFunction(F); | |
} | |
return modified; | |
}; |