[ORC] Add an addObjectFile method to LLJIT.

The addObjectFile method adds the given object file to the JIT session, making
its code available for execution.

Support for the -extra-object flag is added to lli when operating in
-jit-kind=orc-lazy mode to support testing of this feature.

llvm-svn: 340870

llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h

@@ -51,6 +51,14 @@ public:
     return addIRModule(Main, std::move(M));
   }
 
+  /// Adds an object file to the given JITDylib.
+  Error addObjectFile(JITDylib &JD, std::unique_ptr<MemoryBuffer> Obj);
+
+  /// Adds an object file to the given JITDylib.
+  Error addObjectFile(std::unique_ptr<MemoryBuffer> Obj) {
+    return addObjectFile(Main, std::move(Obj));
+  }
+
   /// Look up a symbol in JITDylib JD by the symbol's linker-mangled name (to
   /// look up symbols based on their IR name use the lookup function instead).
   Expected<JITEvaluatedSymbol> lookupLinkerMangled(JITDylib &JD,

llvm/lib/ExecutionEngine/Orc/LLJIT.cpp

@@ -38,6 +38,13 @@ Error LLJIT::addIRModule(JITDylib &JD, std::unique_ptr<Module> M) {
   return CompileLayer.add(JD, K, std::move(M));
 }
 
+Error LLJIT::addObjectFile(JITDylib &JD, std::unique_ptr<MemoryBuffer> Obj) {
+  assert(Obj && "Can not add null object");
+
+  auto K = ES->allocateVModule();
+  return ObjLinkingLayer.add(JD, K, std::move(Obj));
+}
+
 Expected<JITEvaluatedSymbol> LLJIT::lookupLinkerMangled(JITDylib &JD,
                                                         StringRef Name) {
   return llvm::orc::lookup({&JD}, ES->getSymbolStringPool().intern(Name));

llvm/test/ExecutionEngine/OrcLazy/Inputs/basic-object-source.ll

@@ -0,0 +1,5 @@
+define i32 @foo() {
+entry:
+  ret i32 0
+}
+

llvm/test/ExecutionEngine/OrcLazy/basic-object-file-loading.ll

@@ -0,0 +1,13 @@
+; RUN: llc -filetype=obj -o %t %p/Inputs/basic-object-source.ll
+; RUN: lli -jit-kind=orc-lazy -extra-object %t %s
+;
+; Check that we can load an object file and call a function in it.
+
+declare i32 @foo()
+
+define i32 @main(i32 %argc, i8** %argv) {
+entry:
+  %0 = call i32 @foo()
+  ret i32 %0
+}
+

llvm/tools/lli/lli.cpp

@@ -337,8 +337,7 @@ static void reportError(SMDiagnostic Err, const char *ProgName) {
   exit(1);
 }
 
-int runOrcLazyJIT(LLVMContext &Ctx, std::vector<std::unique_ptr<Module>> Ms,
-                  const std::vector<std::string> &Args);
+int runOrcLazyJIT(const char *ProgName);
 
 //===----------------------------------------------------------------------===//
 // main Driver function
@@ -362,6 +361,9 @@ int main(int argc, char **argv, char * const *envp) {
   if (DisableCoreFiles)
     sys::Process::PreventCoreFiles();
 
+  if (UseJITKind == JITKind::OrcLazy)
+    return runOrcLazyJIT(argv[0]);
+
   LLVMContext Context;
 
   // Load the bitcode...
@@ -371,21 +373,6 @@ int main(int argc, char **argv, char * const *envp) {
   if (!Mod)
     reportError(Err, argv[0]);
 
-  if (UseJITKind == JITKind::OrcLazy) {
-    std::vector<std::unique_ptr<Module>> Ms;
-    Ms.push_back(std::move(Owner));
-    for (auto &ExtraMod : ExtraModules) {
-      Ms.push_back(parseIRFile(ExtraMod, Err, Context));
-      if (!Ms.back())
-        reportError(Err, argv[0]);
-    }
-    std::vector<std::string> Args;
-    Args.push_back(InputFile);
-    for (auto &Arg : InputArgv)
-      Args.push_back(Arg);
-    return runOrcLazyJIT(Context, std::move(Ms), Args);
-  }
-
   if (EnableCacheManager) {
     std::string CacheName("file:");
     CacheName.append(InputFile);
@@ -736,13 +723,10 @@ static orc::IRTransformLayer2::TransformFunction createDebugDumper() {
   llvm_unreachable("Unknown DumpKind");
 }
 
-int runOrcLazyJIT(LLVMContext &Ctx, std::vector<std::unique_ptr<Module>> Ms,
-                  const std::vector<std::string> &Args) {
-  // Bail out early if no modules loaded.
-  if (Ms.empty())
-    return 0;
+int runOrcLazyJIT(const char *ProgName) {
+  // Start setting up the JIT environment.
 
-  // Add lli's symbols into the JIT's search space.
+  // First add lli's symbols into the JIT's search space.
   std::string ErrMsg;
   sys::DynamicLibrary LibLLI =
       sys::DynamicLibrary::getPermanentLibrary(nullptr, &ErrMsg);
@@ -751,7 +735,14 @@ int runOrcLazyJIT(LLVMContext &Ctx, std::vector<std::unique_ptr<Module>> Ms,
     return 1;
   }
 
-  const auto &TT = Ms.front()->getTargetTriple();
+  // Parse the main module.
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  auto MainModule = parseIRFile(InputFile, Err, Ctx);
+  if (!MainModule)
+    reportError(Err, ProgName);
+
+  const auto &TT = MainModule->getTargetTriple();
   orc::JITTargetMachineBuilder TMD =
       TT.empty() ? ExitOnErr(orc::JITTargetMachineBuilder::detectHost())
                  : orc::JITTargetMachineBuilder(Triple(TT));
@@ -789,13 +780,35 @@ int runOrcLazyJIT(LLVMContext &Ctx, std::vector<std::unique_ptr<Module>> Ms,
   orc::LocalCXXRuntimeOverrides2 CXXRuntimeOverrides;
   ExitOnErr(CXXRuntimeOverrides.enable(J->getMainJITDylib(), Mangle));
 
-  for (auto &M : Ms) {
+  // Add the main module.
+  ExitOnErr(J->addLazyIRModule(std::move(MainModule)));
+
+  // Add any extra modules.
+  for (auto &ModulePath : ExtraModules) {
+    auto M = parseIRFile(ModulePath, Err, Ctx);
+    if (!M)
+      reportError(Err, ProgName);
+
     orc::makeAllSymbolsExternallyAccessible(*M);
     ExitOnErr(J->addLazyIRModule(std::move(M)));
   }
 
+  // Add the objects.
+  for (auto &ObjPath : ExtraObjects) {
+    auto Obj = ExitOnErr(errorOrToExpected(MemoryBuffer::getFile(ObjPath)));
+    ExitOnErr(J->addObjectFile(std::move(Obj)));
+  }
+
+  // Generate a argument string.
+  std::vector<std::string> Args;
+  Args.push_back(InputFile);
+  for (auto &Arg : InputArgv)
+    Args.push_back(Arg);
+
+  // Run any static constructors.
   ExitOnErr(J->runConstructors());
 
+  // Run main.
   auto MainSym = ExitOnErr(J->lookup("main"));
   typedef int (*MainFnPtr)(int, const char *[]);
   std::vector<const char *> ArgV;