Clone Function in LLVM

LLVM is a powerful compiler infrastructure, but sometimes compiler engineers need to design their own functions for their own purposes. This Blog will introduce how to clone an existing function in LLVM to the same module.

Background

In my current project, CoSense (hope will soon be published and open-sourced), I’m going to clone some functions first and optimize each of these cloned functions with specific algorithms. Then my compiler will eliminate the redundant functions.

There’s no existing API to clone functions quickly and conveniently, so I did a bit of research and wrote this blog.

Related Work

While browsing on Google, there are lots of good experiences.

KeksimvsMaximvs^[1] asked how to clone function in another module, and gave a good solution to it by CloneFunctionInto. Scott A. Carr^[2] shared his experience of how to search and fix this problem, which is really helpful, but didn’t provide a usable source code. And both of them didn’t mention that when a callee function is cloned, the callInst node in the caller function should change the called name. In LLVM’s mailing list, Zhang^[3] answered this question with CallInst::setCalledFunction. But the problems still didn’t have a good solution.

Solution

In my solution, I use a std::map<std::string, int> to store the cloned function’s ID with the same base function, as they are separate functions with their own names. And the IDs are used to generate the new name of the cloned functions.

Then I create the new function with the new name and set arguments to it. Next, I use CloneFunctionInto to generate the new function’s implementation, set visibility and linkage, and copy metadata from the base function.

The last step is to insert the new function into the module (also the base function’s parent node), and changed the call instruction.

Code

Here’s the partial but sufficient key code.

std::map<std::string, int> calleeCounter;

/*  
 * collect the inner bound info for each callee function.
 */
 auto llvmIrCallInstruction = dyn_cast<CallInst>(&llvmIrInstruction);
 Function * calledFunction = llvmIrCallInstruction->getCalledFunction();
 auto calleeCounterIt = calleeCounter.find(calledFunction->getName().str());  
 if (calleeCounterIt != calleeCounter.end()) {  
     calleeCounterIt->second++;  
 } else {  
     calleeCounter.emplace(calledFunction->getName().str(), 0);
 }
 std::string newFuncName = calledFunction->getName().str() + '_' + std::to_string(calleeCounterIt->second);
 /*  
  * rename the llvmIrCallInstruction to the new function name 
  */
 ValueToValueMapTy vMap;  
auto overloadFunc = Function::Create(calledFunction->getFunctionType(),  
                                     calledFunction->getLinkage(),  
                                     calledFunction->getAddressSpace(),  
                                     newFuncName);  
auto *newFuncArgIt = overloadFunc->arg_begin();  
for (auto &arg : calledFunction->args()) {  
    auto argName = arg.getName();  
    newFuncArgIt->setName(argName);  
    vMap[&arg] = &(*newFuncArgIt++);  
}  
SmallVector<ReturnInst*, 8> Returns;  
CloneFunctionInto(overloadFunc, calledFunction, vMap,  
                  CloneFunctionChangeType::LocalChangesOnly, Returns);  
// Set the linkage and visibility late as CloneFunctionInto has some  
// implicit requirements.  
overloadFunc->setVisibility(GlobalValue::DefaultVisibility);  
overloadFunc->setLinkage(GlobalValue::PrivateLinkage);  
  
// Copy metadata  
SmallVector<std::pair<unsigned, MDNode *>, 1> MDs;  
calledFunction->getAllMetadata(MDs);  
for (auto MDIt : MDs) {  
    if (!overloadFunc->hasMetadata()) {  
        overloadFunc->addMetadata(MDIt.first, *MDIt.second);  
    }  
}  
  
Module &funcModule = *calledFunction->getParent();  
funcModule.getFunctionList().insert(calledFunction->getIterator(), overloadFunc);  
overloadFunc->setDSOLocal(true);  
llvmIrCallInstruction->setCalledFunction(overloadFunc);

PS: I used LLVM-13, and the function CloneFunctionInto may change in different LLVM versions.

Conclusion

So this is how to clone functions in LLVM. I will introduce how to eliminate the same (redundant) function after cloning. It will also be fun.

Copyright Notice

Anyone is free to use it, and please indicate the reference when using or publishing. Thank you!

References

[1] https://stackoverflow.com/questions/61189720/llvm-clone-function-pass-to-different-module
[2] http://scottcarr.github.io/2016/01/24/my-week-in-llvm.html
[3] https://lists.llvm.org/pipermail/llvm-dev/2018-March/122114.html