[Orc] Update the Orc indirection utils and refactor the CompileOnDemand layer.

This patch replaces most of the Orc indirection utils API with a new class:
JITCompileCallbackManager, which creates and manages JIT callbacks.
Exposing this functionality directly allows the user to create callbacks that
are associated with user supplied compilation actions. For example, you can
create a callback to lazyily IR-gen something from an AST. (A kaleidoscope
example demonstrating this will be committed shortly).

This patch also refactors the CompileOnDemand layer to use the
JITCompileCallbackManager API.

llvm-svn: 229461

1 files changed, 76 insertions, 124 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/IndirectionUtils.cpp b/llvm/lib/ExecutionEngine/Orc/IndirectionUtils.cpp
index 2fcfb82237c..57616a54840 100644
--- a/llvm/lib/ExecutionEngine/Orc/IndirectionUtils.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/IndirectionUtils.cpp
@@ -9,149 +9,101 @@ using namespace llvm;
 
 namespace llvm {
 
-JITIndirections makeCallsSingleIndirect(
-    Module &M, const std::function<bool(const Function &)> &ShouldIndirect,
-    const char *JITImplSuffix, const char *JITAddrSuffix) {
-  std::vector<Function *> Worklist;
-  std::vector<std::string> FuncNames;
-
-  for (auto &F : M)
-    if (ShouldIndirect(F) && (F.user_begin() != F.user_end())) {
-      Worklist.push_back(&F);
-      FuncNames.push_back(F.getName());
-    }
-
-  for (auto *F : Worklist) {
-    GlobalVariable *FImplAddr = new GlobalVariable(
-        M, F->getType(), false, GlobalValue::ExternalLinkage,
-        Constant::getNullValue(F->getType()), F->getName() + JITAddrSuffix,
-        nullptr, GlobalValue::NotThreadLocal, 0, true);
-    FImplAddr->setVisibility(GlobalValue::HiddenVisibility);
-
-    for (auto *U : F->users()) {
-      assert(isa<Instruction>(U) && "Cannot indirect non-instruction use");
-      IRBuilder<> Builder(cast<Instruction>(U));
-      U->replaceUsesOfWith(F, Builder.CreateLoad(FImplAddr));
-    }
-  }
-
-  return JITIndirections(
-      FuncNames, [=](StringRef S) -> std::string { return std::string(S); },
-      [=](StringRef S)
-          -> std::string { return std::string(S) + JITAddrSuffix; });
+GlobalVariable* createImplPointer(Function &F, const Twine &Name,
+                                  Constant *Initializer) {
+  assert(F.getParent() && "Function isn't in a module.");
+  if (!Initializer)
+    Initializer = Constant::getNullValue(F.getType());
+  Module &M = *F.getParent();
+  return new GlobalVariable(M, F.getType(), false, GlobalValue::ExternalLinkage,
+                            Initializer, Name, nullptr,
+                            GlobalValue::NotThreadLocal, 0, true);
 }
 
-JITIndirections makeCallsDoubleIndirect(
-    Module &M, const std::function<bool(const Function &)> &ShouldIndirect,
-    const char *JITImplSuffix, const char *JITAddrSuffix) {
-
-  std::vector<Function *> Worklist;
-  std::vector<std::string> FuncNames;
-
-  for (auto &F : M)
-    if (!F.isDeclaration() && !F.hasAvailableExternallyLinkage() &&
-        ShouldIndirect(F))
-      Worklist.push_back(&F);
-
-  for (auto *F : Worklist) {
-    std::string OrigName = F->getName();
-    F->setName(OrigName + JITImplSuffix);
-    FuncNames.push_back(OrigName);
-
-    GlobalVariable *FImplAddr = new GlobalVariable(
-        M, F->getType(), false, GlobalValue::ExternalLinkage,
-        Constant::getNullValue(F->getType()), OrigName + JITAddrSuffix, nullptr,
-        GlobalValue::NotThreadLocal, 0, true);
-    FImplAddr->setVisibility(GlobalValue::HiddenVisibility);
+void makeStub(Function &F, GlobalVariable &ImplPointer) {
+  assert(F.isDeclaration() && "Can't turn a definition into a stub.");
+  assert(F.getParent() && "Function isn't in a module.");
+  Module &M = *F.getParent();
+  BasicBlock *EntryBlock = BasicBlock::Create(M.getContext(), "entry", &F);
+  IRBuilder<> Builder(EntryBlock);
+  LoadInst *ImplAddr = Builder.CreateLoad(&ImplPointer);
+  std::vector<Value*> CallArgs;
+  for (auto &A : F.args())
+    CallArgs.push_back(&A);
+  CallInst *Call = Builder.CreateCall(ImplAddr, CallArgs);
+  Call->setTailCall();
+  Builder.CreateRet(Call);
+}
 
-    Function *FRedirect =
-        Function::Create(F->getFunctionType(), F->getLinkage(), OrigName, &M);
+void partition(Module &M, const ModulePartitionMap &PMap) {
 
-    F->replaceAllUsesWith(FRedirect);
+  for (auto &KVPair : PMap) {
 
-    BasicBlock *EntryBlock =
-        BasicBlock::Create(M.getContext(), "entry", FRedirect);
+    auto ExtractGlobalVars =
+      [&](GlobalVariable &New, const GlobalVariable &Orig,
+          ValueToValueMapTy &VMap) {
+        if (KVPair.second.count(&Orig)) {
+          copyGVInitializer(New, Orig, VMap);
+        }
+        if (New.getLinkage() == GlobalValue::PrivateLinkage) {
+          New.setLinkage(GlobalValue::ExternalLinkage);
+          New.setVisibility(GlobalValue::HiddenVisibility);
+        }
+      };
 
-    IRBuilder<> Builder(EntryBlock);
-    LoadInst *FImplLoadedAddr = Builder.CreateLoad(FImplAddr);
+    auto ExtractFunctions =
+      [&](Function &New, const Function &Orig, ValueToValueMapTy &VMap) {
+        if (KVPair.second.count(&Orig))
+          copyFunctionBody(New, Orig, VMap);
+        if (New.getLinkage() == GlobalValue::InternalLinkage) {
+          New.setLinkage(GlobalValue::ExternalLinkage);
+          New.setVisibility(GlobalValue::HiddenVisibility);
+        }
+      };
 
-    std::vector<Value *> CallArgs;
-    for (Value &Arg : FRedirect->args())
-      CallArgs.push_back(&Arg);
-    CallInst *Call = Builder.CreateCall(FImplLoadedAddr, CallArgs);
-    Call->setTailCall();
-    Builder.CreateRet(Call);
+    CloneSubModule(*KVPair.first, M, ExtractGlobalVars, ExtractFunctions,
+                   false);
   }
-
-  return JITIndirections(
-      FuncNames, [=](StringRef S)
-                     -> std::string { return std::string(S) + JITImplSuffix; },
-      [=](StringRef S)
-          -> std::string { return std::string(S) + JITAddrSuffix; });
 }
 
-std::vector<std::unique_ptr<Module>>
-explode(const Module &OrigMod,
-        const std::function<bool(const Function &)> &ShouldExtract) {
-
-  std::vector<std::unique_ptr<Module>> NewModules;
+FullyPartitionedModule fullyPartition(Module &M) {
+  FullyPartitionedModule MP;
 
-  // Split all the globals, non-indirected functions, etc. into a single module.
-  auto ExtractGlobalVars = [&](GlobalVariable &New, const GlobalVariable &Orig,
-                               ValueToValueMapTy &VMap) {
-    copyGVInitializer(New, Orig, VMap);
-    if (New.getLinkage() == GlobalValue::PrivateLinkage) {
-      New.setLinkage(GlobalValue::ExternalLinkage);
-      New.setVisibility(GlobalValue::HiddenVisibility);
-    }
-  };
+  ModulePartitionMap PMap;
 
-  auto ExtractNonImplFunctions =
-      [&](Function &New, const Function &Orig, ValueToValueMapTy &VMap) {
-        if (!ShouldExtract(New))
-          copyFunctionBody(New, Orig, VMap);
-      };
+  for (auto &F : M) {
 
-  NewModules.push_back(CloneSubModule(OrigMod, ExtractGlobalVars,
-                                      ExtractNonImplFunctions, true));
-
-  // Preserve initializers for Common linkage vars, and make private linkage
-  // globals external: they are now provided by the globals module extracted
-  // above.
-  auto DropGlobalVars = [&](GlobalVariable &New, const GlobalVariable &Orig,
-                            ValueToValueMapTy &VMap) {
-    if (New.getLinkage() == GlobalValue::CommonLinkage)
-      copyGVInitializer(New, Orig, VMap);
-    else if (New.getLinkage() == GlobalValue::PrivateLinkage)
-      New.setLinkage(GlobalValue::ExternalLinkage);
-  };
-
-  // Split each 'impl' function out in to its own module.
-  for (const auto &Func : OrigMod) {
-    if (Func.isDeclaration() || !ShouldExtract(Func))
+    if (F.isDeclaration())
       continue;
 
-    auto ExtractNamedFunction =
-        [&](Function &New, const Function &Orig, ValueToValueMapTy &VMap) {
-          if (New.getName() == Func.getName())
-            copyFunctionBody(New, Orig, VMap);
-        };
-
-    NewModules.push_back(
-        CloneSubModule(OrigMod, DropGlobalVars, ExtractNamedFunction, false));
+    std::string NewModuleName = (M.getName() + "." + F.getName()).str();
+    MP.Functions.push_back(
+      llvm::make_unique<Module>(NewModuleName, M.getContext()));
+    MP.Functions.back()->setDataLayout(M.getDataLayout());
+    PMap[MP.Functions.back().get()].insert(&F);
   }
 
-  return NewModules;
-}
+  MP.GlobalVars =
+    llvm::make_unique<Module>((M.getName() + ".globals_and_stubs").str(),
+                              M.getContext());
+  MP.GlobalVars->setDataLayout(M.getDataLayout());
 
-std::vector<std::unique_ptr<Module>>
-explode(const Module &OrigMod, const JITIndirections &Indirections) {
-  std::set<std::string> ImplNames;
+  MP.Commons =
+    llvm::make_unique<Module>((M.getName() + ".commons").str(), M.getContext());
+  MP.Commons->setDataLayout(M.getDataLayout());
 
-  for (const auto &FuncName : Indirections.IndirectedNames)
-    ImplNames.insert(Indirections.GetImplName(FuncName));
+  // Make sure there's at least an empty set for the stubs map or we'll fail
+  // to clone anything for it (including the decls).
+  PMap[MP.GlobalVars.get()] = ModulePartitionMap::mapped_type();
+  for (auto &GV : M.globals())
+    if (GV.getLinkage() == GlobalValue::CommonLinkage)
+      PMap[MP.Commons.get()].insert(&GV);
+    else
+      PMap[MP.GlobalVars.get()].insert(&GV);
 
-  return explode(
-      OrigMod, [&](const Function &F) { return ImplNames.count(F.getName()); });
+  partition(M, PMap);
+
+  return MP;
 }
+
 }