1 #include "llvm/ADT/Triple.h"
2 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
3 #include "llvm/ExecutionEngine/Orc/CloneSubModule.h"
4 #include "llvm/IR/CallSite.h"
5 #include "llvm/IR/IRBuilder.h"
6 #include <set>
7 
8 using namespace llvm;
9 
10 namespace llvm {
11 
12 JITIndirections makeCallsSingleIndirect(
13     Module &M, const std::function<bool(const Function &)> &ShouldIndirect,
14     const char *JITImplSuffix, const char *JITAddrSuffix) {
15   std::vector<Function *> Worklist;
16   std::vector<std::string> FuncNames;
17 
18   for (auto &F : M)
19     if (ShouldIndirect(F) && (F.user_begin() != F.user_end())) {
20       Worklist.push_back(&F);
21       FuncNames.push_back(F.getName());
22     }
23 
24   for (auto *F : Worklist) {
25     GlobalVariable *FImplAddr = new GlobalVariable(
26         M, F->getType(), false, GlobalValue::ExternalLinkage,
27         Constant::getNullValue(F->getType()), F->getName() + JITAddrSuffix,
28         nullptr, GlobalValue::NotThreadLocal, 0, true);
29     FImplAddr->setVisibility(GlobalValue::HiddenVisibility);
30 
31     for (auto *U : F->users()) {
32       assert(isa<Instruction>(U) && "Cannot indirect non-instruction use");
33       IRBuilder<> Builder(cast<Instruction>(U));
34       U->replaceUsesOfWith(F, Builder.CreateLoad(FImplAddr));
35     }
36   }
37 
38   return JITIndirections(
39       FuncNames, [=](StringRef S) -> std::string { return std::string(S); },
40       [=](StringRef S)
41           -> std::string { return std::string(S) + JITAddrSuffix; });
42 }
43 
44 JITIndirections makeCallsDoubleIndirect(
45     Module &M, const std::function<bool(const Function &)> &ShouldIndirect,
46     const char *JITImplSuffix, const char *JITAddrSuffix) {
47 
48   std::vector<Function *> Worklist;
49   std::vector<std::string> FuncNames;
50 
51   for (auto &F : M)
52     if (!F.isDeclaration() && !F.hasAvailableExternallyLinkage() &&
53         ShouldIndirect(F))
54       Worklist.push_back(&F);
55 
56   for (auto *F : Worklist) {
57     std::string OrigName = F->getName();
58     F->setName(OrigName + JITImplSuffix);
59     FuncNames.push_back(OrigName);
60 
61     GlobalVariable *FImplAddr = new GlobalVariable(
62         M, F->getType(), false, GlobalValue::ExternalLinkage,
63         Constant::getNullValue(F->getType()), OrigName + JITAddrSuffix, nullptr,
64         GlobalValue::NotThreadLocal, 0, true);
65     FImplAddr->setVisibility(GlobalValue::HiddenVisibility);
66 
67     Function *FRedirect =
68         Function::Create(F->getFunctionType(), F->getLinkage(), OrigName, &M);
69 
70     F->replaceAllUsesWith(FRedirect);
71 
72     BasicBlock *EntryBlock =
73         BasicBlock::Create(M.getContext(), "entry", FRedirect);
74 
75     IRBuilder<> Builder(EntryBlock);
76     LoadInst *FImplLoadedAddr = Builder.CreateLoad(FImplAddr);
77 
78     std::vector<Value *> CallArgs;
79     for (Value &Arg : FRedirect->args())
80       CallArgs.push_back(&Arg);
81     CallInst *Call = Builder.CreateCall(FImplLoadedAddr, CallArgs);
82     Call->setTailCall();
83     Builder.CreateRet(Call);
84   }
85 
86   return JITIndirections(
87       FuncNames, [=](StringRef S)
88                      -> std::string { return std::string(S) + JITImplSuffix; },
89       [=](StringRef S)
90           -> std::string { return std::string(S) + JITAddrSuffix; });
91 }
92 
93 std::vector<std::unique_ptr<Module>>
94 explode(const Module &OrigMod,
95         const std::function<bool(const Function &)> &ShouldExtract) {
96 
97   std::vector<std::unique_ptr<Module>> NewModules;
98 
99   // Split all the globals, non-indirected functions, etc. into a single module.
100   auto ExtractGlobalVars = [&](GlobalVariable &New, const GlobalVariable &Orig,
101                                ValueToValueMapTy &VMap) {
102     copyGVInitializer(New, Orig, VMap);
103     if (New.getLinkage() == GlobalValue::PrivateLinkage) {
104       New.setLinkage(GlobalValue::ExternalLinkage);
105       New.setVisibility(GlobalValue::HiddenVisibility);
106     }
107   };
108 
109   auto ExtractNonImplFunctions =
110       [&](Function &New, const Function &Orig, ValueToValueMapTy &VMap) {
111         if (!ShouldExtract(New))
112           copyFunctionBody(New, Orig, VMap);
113       };
114 
115   NewModules.push_back(CloneSubModule(OrigMod, ExtractGlobalVars,
116                                       ExtractNonImplFunctions, true));
117 
118   // Preserve initializers for Common linkage vars, and make private linkage
119   // globals external: they are now provided by the globals module extracted
120   // above.
121   auto DropGlobalVars = [&](GlobalVariable &New, const GlobalVariable &Orig,
122                             ValueToValueMapTy &VMap) {
123     if (New.getLinkage() == GlobalValue::CommonLinkage)
124       copyGVInitializer(New, Orig, VMap);
125     else if (New.getLinkage() == GlobalValue::PrivateLinkage)
126       New.setLinkage(GlobalValue::ExternalLinkage);
127   };
128 
129   // Split each 'impl' function out in to its own module.
130   for (const auto &Func : OrigMod) {
131     if (Func.isDeclaration() || !ShouldExtract(Func))
132       continue;
133 
134     auto ExtractNamedFunction =
135         [&](Function &New, const Function &Orig, ValueToValueMapTy &VMap) {
136           if (New.getName() == Func.getName())
137             copyFunctionBody(New, Orig, VMap);
138         };
139 
140     NewModules.push_back(
141         CloneSubModule(OrigMod, DropGlobalVars, ExtractNamedFunction, false));
142   }
143 
144   return NewModules;
145 }
146 
147 std::vector<std::unique_ptr<Module>>
148 explode(const Module &OrigMod, const JITIndirections &Indirections) {
149   std::set<std::string> ImplNames;
150 
151   for (const auto &FuncName : Indirections.IndirectedNames)
152     ImplNames.insert(Indirections.GetImplName(FuncName));
153 
154   return explode(
155       OrigMod, [&](const Function &F) { return ImplNames.count(F.getName()); });
156 }
157 }
158