lib/ExecutionEngine/Orc/IndirectionUtils.cpp

   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 }