[WIP][ORC][ThinLTO] Early ThinLTO-JIT prototype and basic tests for next-gen ORC APIs

[sgraenitz]

Summary

Mimimal implementation that I could use to demonstrate speedups. Will add inline comments where I think this is not satisfying yet. Apart from that, what do you think? Do you have proposals for improvements?

Originally I wanted a real ThinLTOLayer with a proper ThinLTOLayerMaterializationUnit, integration with llvm::orc::lookup(), etc. So far I failed to get the SymbolFlags handling right, when passing the MaterializationResponsibility down to the underlying IRLayer. Would be great to fix that and upstream it, but it's no hurry.

[sgraenitz]

Not sure what's the best way to get rid of the binary object files. IIUC the ThinLTO tests are lit-based, so they build their summary-enhanced objects with clang and test via llvm-lto2 tool. Maybe I can turn all the unit tests into lit-test and add a hook to llvm-lto2? Don't know if there is a way to emit object files with summaries using yaml2obj or so? Any ideas?

lib/ExecutionEngine/Orc/LLThinLTOJIT.cpp

154

It's not currently possible to obtain the SearchOrder to check for existing entries. That's on purpose?

unittests/ExecutionEngine/Orc/LLThinLTOJITTest.cpp

230

I think I could do much better here with an actual layer. It currently dumps:

[ RUN      ] LLThinLTOJITTest.Incomplete
JIT session error: Symbols not found: { "_bar" }
Failed as expected: Failed to materialize symbols: { "_foo" }
Failed again: Symbols not found: { "_foo" }
[       OK ] LLThinLTOJITTest.Incomplete (11 ms)

With control over the MaterializationUnit, I could invoke the code I currently call prepare for not-yet-discovered callees internally and successfully materialize already with the first lookup("foo") above.

[lhames]

You definitely want to do this as a real layer and materialization unit. It is beneficial even without lazy compilation, but the ability to compose it with laziness would be extra cool.

What aspect of the symbol flags gave you trouble? This is one of the most challenging parts of the new API, so I'm keen to hear where it is tripping people up so I can improve on it.

[lhames]

lib/ExecutionEngine/Orc/LLThinLTOJIT.cpp

154

There is no getSearchOrder() method as that would not be thread safe. You can use withSearchOrderDo:

auto &NewJD = // Get dylib to add to search order.
JD.withSearchOrderDo(
  [&](const JITDylibList &SearchOrder) {
    // Search to see if NewJD is already in the search order, bail out if it is.
    for (auto *E : SearchOrder)
      if (E == &NewJD)
        return;
    // NewJD was not in the search order. Add it.
    addToSearchOrder(NewJD);
  });

If this add-without-duplication operation turns out to be common we could add a DiscardDuplicates argument to addToSearchOrder and use that instead.

[dexonsmith]

In D51497#1219698, @sgraenitz wrote:

Not sure what's the best way to get rid of the binary object files. IIUC the ThinLTO tests are lit-based, so they build their summary-enhanced objects with clang and test via llvm-lto2 tool. Maybe I can turn all the unit tests into lit-test and add a hook to llvm-lto2? Don't know if there is a way to emit object files with summaries using yaml2obj or so? Any ideas?

You should check in LLVM textual IR files (.ll) instead of C++ source files or object files. Using clang -S -emit-llvm (or llvm-dis some-thinlto-bitcode-file.o) should get you started.

[sgraenitz]

In D51497#1219835, @dexonsmith wrote:

You should check in LLVM textual IR files (.ll) instead of C++ source files or object files. Using clang -S -emit-llvm (or llvm-dis some-thinlto-bitcode-file.o) should get you started.

Yes that would be great, but wouldn't it require IR representation for summaries? I didn't double-check myself, but there's a note in the language ref, "that temporarily the summary entries are skipped when parsing the assembly". Not sure if I understand that correctly, but I was under the impression that we can't have summaries in IR files yet.
https://llvm.org/docs/LangRef.html#thinlto-summary

[dexonsmith]

In D51497#1220621, @sgraenitz wrote:

In D51497#1219835, @dexonsmith wrote:

You should check in LLVM textual IR files (.ll) instead of C++ source files or object files. Using clang -S -emit-llvm (or llvm-dis some-thinlto-bitcode-file.o) should get you started.

Yes that would be great, but wouldn't it require IR representation for summaries? I didn't double-check myself, but there's a note in the language ref, "that temporarily the summary entries are skipped when parsing the assembly". Not sure if I understand that correctly, but I was under the impression that we can't have summaries in IR files yet.
https://llvm.org/docs/LangRef.html#thinlto-summary

Looking at existing testcases in test/LTO, it seems like the existing practice is to check in a .ll file without a summary, and generate the summary using opt. Something like:

; RUN: opt -module-summary -o %t1.bc %s

That's probably better than my original suggestion anyway.

[tejohnson]

In D51497#1220621, @sgraenitz wrote:

In D51497#1219835, @dexonsmith wrote:

You should check in LLVM textual IR files (.ll) instead of C++ source files or object files. Using clang -S -emit-llvm (or llvm-dis some-thinlto-bitcode-file.o) should get you started.

Yes that would be great, but wouldn't it require IR representation for summaries? I didn't double-check myself, but there's a note in the language ref, "that temporarily the summary entries are skipped when parsing the assembly". Not sure if I understand that correctly, but I was under the impression that we can't have summaries in IR files yet.
https://llvm.org/docs/LangRef.html#thinlto-summary

This is stale - Sorry, I totally forgot to remove when I implemented the parsing side, which went in a couple months ago (r335602). For an example with a summary, see test/Assembler/thinlto-summary.ll. You need to use llvm-as to assembly it (not opt). Eventually other tools such as llvm-lto2 should be able to read assembly with summary directly, but I haven't had a chance to do that yet.

[sgraenitz]

Thanks for your input. For now I worked more on internals than on testing, but at least I managed to check in .ll files instead of .cpp and use .bc file in my tests:

clang -c -flto=thin -o Foo.o Foo.cpp
llvm-dis -o Foo.ll Foo.o
llvm-as -o Foo.bc Foo.ll

I will close this review in favour of https://reviews.llvm.org/D51744, where I use this approach.

lib/ExecutionEngine/Orc/LLThinLTOJIT.cpp

154

Sure right, I saw that earlier but didn't remember. Thanks!

include/llvm/ExecutionEngine/Orc/LLThinLTOJIT.h

+//===- LLThinLTOJIT.h - ThinLTO Module Summaries for JITing -----*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_LLTHINLTOJIT_H
+#define LLVM_EXECUTIONENGINE_ORC_LLTHINLTOJIT_H
+
+#include "llvm/Bitcode/BitcodeReader.h"
+#include "llvm/ExecutionEngine/JITSymbol.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/LLJIT.h"
+#include "llvm/IR/Mangler.h"
+#include "llvm/IR/ModuleSummaryIndex.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/IPO/FunctionImport.h"
+
+#include <memory>
+#include <string>
+
+namespace llvm {
+namespace orc {
+
+/// An extended version of LLJIT that supports lazy function-at-a-time
+/// compilation of LLVM IR.
+class LLThinLTOJIT : public LLJIT {
+  struct ModuleInfo {
+    std::string Path;
+    JITDylib *Dylib;
+    bool StateEmitted;
+  };
+
+public:
+  LLThinLTOJIT(std::unique_ptr<ExecutionSession> ES,
+               std::unique_ptr<TargetMachine> TM, DataLayout DL,
+               CompileFtorFactory MakeCompileFtor, LLVMContext &Ctx);
+
+  /// Add prefix paths for searching ThinLTO IR Modules.
+  void addSearchPath(std::string Dir) {
+    ModuleSearchPaths.push_back(std::move(Dir));
+  }
+  
+  /// Add a module to be compiled to JITDylib JD.
+  Expected<VModuleKey> addThinLTOIRModule(JITDylib &JD, std::string Path);
+  
+  /// Add a module to be compiled to the main JITDylib.
+  Expected<VModuleKey> addThinLTOIRModule(std::string Path) {
+    return addThinLTOIRModule(Main, std::move(Path));
+  }
+
+  /// Emit modules for the given function name and discovered callees.
+  Expected<JITDylib *> prepare(std::string Name, unsigned DiscoveryLimit = -1);
+  
+private:
+  LLVMContext &Context;
+  ModuleSummaryIndex CombinedIndex;
+  std::vector<std::string> ModuleSearchPaths;
+  std::map<VModuleKey, ModuleInfo> RegisteredModules;
+  
+  Expected<bool> submit(ModuleInfo &Info);
+  Expected<std::string> findFullModulePath(std::string Path) const;
+  FunctionImporter::ImportMapTy discoverCallees(GlobalValue::GUID Guid,
+                                                StringRef ModulePath,
+                                                unsigned DiscoveryLevelLimit);
+};
+
+} // end namespace orc
+} // end namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_LLTHINLTOJIT_H

lib/ExecutionEngine/Orc/CMakeLists.txt

 add_llvm_library(LLVMOrcJIT
   CompileOnDemandLayer.cpp
   Core.cpp
   ExecutionUtils.cpp
   IndirectionUtils.cpp
   IRCompileLayer.cpp
   IRTransformLayer.cpp
   Legacy.cpp
   Layer.cpp
   LLJIT.cpp
+  LLThinLTOJIT.cpp
   NullResolver.cpp
   ObjectTransformLayer.cpp
   OrcABISupport.cpp
   OrcCBindings.cpp
   OrcError.cpp
   OrcMCJITReplacement.cpp
   RPCUtils.cpp
   RTDyldObjectLinkingLayer.cpp

lib/ExecutionEngine/Orc/LLThinLTOJIT.cpp

+//===--------- LLJIT.cpp - An ORC-based JIT for compiling LLVM IR ---------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/LLThinLTOJIT.h"
+
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/ExecutionEngine/Orc/OrcError.h"
+#include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/IR/Mangler.h"
+#include "llvm/Support/FileSystem.h"
+
+namespace llvm {
+namespace orc {
+
+/// Local shortcut to create a formatted string error.
+template <typename... Ts>
+static Error error(char const *Fmt, const Ts &... Vals) {
+  return createStringError(inconvertibleErrorCode(), Fmt, Vals...);
+}
+
+/// We're reading summaries from BC sources.
+constexpr static bool PerformingIRModuleAnalysis = false;
+
+LLThinLTOJIT::LLThinLTOJIT(std::unique_ptr<ExecutionSession> ES,
+                           std::unique_ptr<TargetMachine> TM,
+                           DataLayout DL, CompileFtorFactory MakeCompileFtor,
+                           LLVMContext &Ctx)
+  : LLJIT(std::move(ES), std::move(TM), std::move(DL),
+          std::move(MakeCompileFtor)),
+    Context(Ctx), CombinedIndex(PerformingIRModuleAnalysis) {}
+
+Expected<std::string> LLThinLTOJIT::findFullModulePath(std::string P) const {
+  if (sys::fs::exists(P))
+    return P;
+  
+  for (const std::string &SearchPath : ModuleSearchPaths) {
+    std::string FullPath = SearchPath + P;
+    if (sys::fs::exists(FullPath))
+      return FullPath;
+  }
+
+  auto MSPStr = join(ModuleSearchPaths.begin(), ModuleSearchPaths.end(), ", ");
+  return error("Can't find file %s.\nModule search paths are: ",
+               P.c_str(), MSPStr.c_str());
+}
+  
+Expected<VModuleKey> LLThinLTOJIT::addThinLTOIRModule(JITDylib &JD,
+                                                      std::string Path) {
+  // Find the full path of the module on disk.
+  if (auto FullPath = findFullModulePath(std::move(Path)))
+    Path = *FullPath;
+  else
+    return FullPath.takeError();
+
+  // Module registered already? Then keep the existing registration. We assume
+  // files never change on disk and we don't attempt to reload them.
+  if (VModuleKey ExistingKey = CombinedIndex.getModuleId(Path)) {
+    const ModuleInfo &Info = RegisteredModules[ExistingKey];
+    if (Info.Dylib != &JD)
+      return createStringError(
+          inconvertibleErrorCode(),
+          "Path %s can't be added to Dylib %s. "
+          "It's already associated with Dylib %s.",
+          Path.c_str(), JD.getName().c_str(),
+          Info.Dylib->getName().c_str());
+    
+    return ExistingKey;
+  }
+
+  // Memory-map bitcode file.
+  // TODO Cache it?
+  auto B = MemoryBuffer::getFile(Path);
+  if (!B)
+    return errorCodeToError(B.getError());
+
+  // Read summary section and merge it into the global CombinedIndex.
+  VModuleKey K = ES->allocateVModule();
+  MemoryBufferRef BufferRef = (*B)->getMemBufferRef();
+  if (Error Err = readModuleSummaryIndex(BufferRef, CombinedIndex, K))
+    return std::move(Err);
+
+  // VModuleKey is the Orc equivalent for ThinLTO's ModuleID
+  assert(CombinedIndex.getModuleId(Path) == K);
+
+  RegisteredModules[K] = { std::move(Path), &JD, false };
+  return K;
+};
+
+Expected<JITDylib *> LLThinLTOJIT::prepare(std::string FunctionName,
+                                           unsigned DiscoveryLimit) {
+  // ValueInfo tells us which module(s) to compile.
+  auto Guid = GlobalValue::getGUID(FunctionName);
+  ValueInfo VI = CombinedIndex.getValueInfo(Guid);
+  if (!VI || VI.getSummaryList().empty())
+    return error("Can't find ValueInfo for GUID %" PRIx64, Guid);
+
+  // Find the matching summary.
+  // TODO Usually there is only one, right?
+  const GlobalValueSummary *Summary = nullptr;
+  for (const auto &S : VI.getSummaryList()) {
+    // TODO Any more conditions here?
+    if (S->getSummaryKind() == GlobalValueSummary::FunctionKind) {
+      Summary = S.get();
+      break;
+    }
+  }
+
+  if (Summary == nullptr)
+    return error("Summary for name %s with GUID %" PRIx64 " is not a function",
+                 FunctionName.c_str(), Guid);
+
+  // If we found the ValueInfo, the module must be registered.
+  StringRef FullPath = Summary->modulePath();
+  VModuleKey K = CombinedIndex.getModuleId(FullPath);
+  ModuleInfo &DefiningModuleInfo = RegisteredModules[K];
+  
+  // Pass it on to the next layer.
+  Expected<bool> Result = submit(DefiningModuleInfo);
+  if (!Result)
+    return Result.takeError();
+
+  // Discover functions that are referenced from the given one.
+  FunctionImporter::ImportMapTy DiscoveredFunctions =
+      discoverCallees(Guid, FullPath, DiscoveryLimit);
+
+#if !defined(NDEBUG)
+  for (const auto &Entry : DiscoveredFunctions) {
+    dbgs() << "Discovered functions from " << Entry.first() << "\n";
+    for (const auto &GUID : Entry.second) {
+      dbgs() << "  " << GUID << "\n";
+    }
+  }
+#endif
+
+  // Add the modules of all discovered functions.
+  for (const auto &Entry : DiscoveredFunctions) {
+    VModuleKey K = CombinedIndex.getModuleId(Entry.first());
+    ModuleInfo &RelatedModuleInfo = RegisteredModules[K];
+
+    // Pass it on to the next layer.
+    Expected<bool> Result = submit(RelatedModuleInfo);
+    if (!Result)
+      return Result.takeError();
+
+    // If we materialized into a new Dylib, add it to our SearchPath.
+    bool AddedModule = *Result;
+    if (AddedModule && RelatedModuleInfo.Dylib != DefiningModuleInfo.Dylib) {
+      DefiningModuleInfo.Dylib->addToSearchOrder(*RelatedModuleInfo.Dylib);
+    }
+  }
+
+  return DefiningModuleInfo.Dylib;
+};
+
+FunctionImporter::ImportMapTy
+LLThinLTOJIT::discoverCallees(GlobalValue::GUID Guid, StringRef ModulePath,
+                              unsigned DiscoveryLevelLimit) {
+  if (DiscoveryLevelLimit == 0)
+    return {};
+
+  GlobalValueSummary *Summary =
+      CombinedIndex.findSummaryInModule(Guid, ModulePath);
+
+  // The list of functions we are interested in.
+  GVSummaryMapTy FunctionSummaryMap;
+  FunctionSummaryMap.try_emplace(Guid, Summary);
+
+  // Limit on instruction count for edges we are following (ThinLTO default).
+  // TODO Make it configurable.
+  unsigned DiscoveryInstrLimit = 100;
+  FunctionImporter::ImportThresholdsTy ImportThresholds;
+
+  // List of discovered functions that need to be processed.
+  SmallVector<EdgeInfo, 128> Worklist;
+
+  FunctionSummary *FuncSummary = dyn_cast<FunctionSummary>(Summary);
+  assert(FuncSummary && "FunctionKind was checked earlier");
+
+  // Discovery level 1:
+  FunctionImporter::ImportMapTy DiscoveredFunctions;
+  computeImportForFunction(*FuncSummary, CombinedIndex, DiscoveryInstrLimit,
+                           FunctionSummaryMap, Worklist, DiscoveredFunctions,
+                           nullptr, ImportThresholds);
+
+  // Process the newly imported functions and add callees to the worklist.
+  // TODO Add discovery level info to worklist items.
+  while (!Worklist.empty()) {
+    auto FuncInfo = Worklist.pop_back_val();
+    auto *Summary = std::get<0>(FuncInfo);
+    auto Threshold = std::get<1>(FuncInfo);
+    computeImportForFunction(*Summary, CombinedIndex, Threshold,
+                             FunctionSummaryMap, Worklist,
+                             DiscoveredFunctions, nullptr, ImportThresholds);
+  }
+  
+  return DiscoveredFunctions;
+};
+
+Expected<bool> LLThinLTOJIT::submit(ModuleInfo &Info) {
+  if (Info.StateEmitted)
+    return false; // Not added
+  
+  // If we cached the buffer, we'd reuse it here.
+  VModuleKey K = CombinedIndex.getModuleId(Info.Path);
+  auto B = MemoryBuffer::getFile(Info.Path);
+  if (!B)
+    return errorCodeToError(B.getError());
+
+  // Parse the full module.
+  MemoryBufferRef R = (*B)->getMemBufferRef();
+  auto M = parseBitcodeFile(R, Context);
+  if (!M)
+    return M.takeError();
+
+  // Pass the module on to the base layer.
+  if (Error Err = CompileLayer.add(*Info.Dylib, K, std::move(*M)))
+    return std::move(Err);
+
+  RegisteredModules[K].StateEmitted = true;
+  return true; // Added successfully
+}
+  
+} // End namespace orc.
+} // End namespace llvm.

unittests/ExecutionEngine/Orc/CMakeLists.txt

 
 set(LLVM_LINK_COMPONENTS
   Core
   ExecutionEngine
   IRReader
   Object
   OrcJIT
+  IPO
   RuntimeDyld
   Support
   native
   )
 
 add_llvm_unittest(OrcJITTests
   CompileOnDemandLayerTest.cpp
   CoreAPIsTest.cpp
   IndirectionUtilsTest.cpp
   GlobalMappingLayerTest.cpp
   LazyEmittingLayerTest.cpp
   LegacyAPIInteropTest.cpp
   LLJITTest.cpp
+  LLThinLTOJITTest.cpp
   ObjectTransformLayerTest.cpp
   OrcCAPITest.cpp
   OrcTestCommon.cpp
   QueueChannel.cpp
   RemoteObjectLayerTest.cpp
   RPCUtilsTest.cpp
   RTDyldObjectLinkingLayerTest.cpp
   RTDyldObjectLinkingLayer2Test.cpp
   SymbolStringPoolTest.cpp
   )
 
 target_link_libraries(OrcJITTests PRIVATE ${ORC_JIT_TEST_LIBS})

unittests/ExecutionEngine/Orc/Inputs/Bar.cpp

+extern "C" int bar() {
+  return 1;
+}

unittests/ExecutionEngine/Orc/Inputs/Foo.cpp

+extern "C" int bar();
+
+extern "C" int foo() {
+  return bar();
+}

unittests/ExecutionEngine/Orc/LLThinLTOJITTest.cpp

+//===- LLThinLTOJITTest.cpp - Unit tests for the ThinLTO layer ------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/LLThinLTOJIT.h"
+#include "OrcTestCommon.h"
+
+#include "llvm/ADT/ScopeExit.h"
+#include "llvm/ExecutionEngine/ExecutionEngine.h"
+#include "llvm/ExecutionEngine/JITSymbol.h"
+#include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
+#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
+#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
+#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
+#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/TargetRegistry.h"
+#include "llvm/Support/TargetSelect.h"
+
+#include "gtest/gtest.h"
+
+#include <memory>
+#include <thread>
+#include <vector>
+
+using namespace llvm;
+using namespace llvm::orc;
+
+namespace {
+
+// TODO That's not gonna work for others..
+// TODO Test input objects must be compiled with clang:
+// clang -c -flto=thin Foo.cpp Bar.cpp
+static const std::string TestInputsDir =
+  "/Users/sgranitz/Develop/orcjit/llvm/unittests/ExecutionEngine/Orc/Inputs/";
+
+class LLThinLTOJITTest : public testing::Test {
+  struct Init {
+    Init() {
+      InitializeNativeTarget();
+      InitializeNativeTargetAsmParser();
+      InitializeNativeTargetAsmPrinter();
+    }
+    
+    // TODO Make a proper LLThinLTOJIT::Create()
+    static std::unique_ptr<LLThinLTOJIT>
+    makeJit(std::unique_ptr<ExecutionSession> ES, LLVMContext &Ctx,
+            bool Multithreaded) {
+      std::unique_ptr<TargetMachine> TM;
+      LLThinLTOJIT::CompileFtorFactory MakeCompileFtor;
+
+      if (Multithreaded) {
+        auto JTMB = cantFail(JITTargetMachineBuilder::detectHost());
+        TM = cantFail(JTMB.createTargetMachine());
+        MakeCompileFtor = [JTMB](TargetMachine &) -> LLThinLTOJIT::CompileFtor {
+          return MultiThreadedSimpleCompiler(std::move(JTMB));
+        };
+      } else {
+        TM = std::unique_ptr<TargetMachine>(EngineBuilder().selectTarget());
+        MakeCompileFtor = [](TargetMachine &TM) -> LLThinLTOJIT::CompileFtor {
+          return SimpleCompiler(TM);
+        };
+      }
+
+      auto DL = TM->createDataLayout();
+      return llvm::make_unique<LLThinLTOJIT>(std::move(ES), std::move(TM),
+                                             std::move(DL),
+                                             std::move(MakeCompileFtor), Ctx);
+    }
+  };
+
+  Init StaticInit;
+  LLVMContext Ctx;
+  std::shared_ptr<SymbolStringPool> SharedStringPool;
+  DataLayout MandIDL;
+  ExecutionSession MandIES;
+  MangleAndInterner MandI;
+
+public:
+  static int Exec(JITEvaluatedSymbol FuncSym) {
+    auto FuncPtr = (int (*)())FuncSym.getAddress();
+    return FuncPtr();
+  }
+
+  static void JoinAll(std::vector<std::thread> &Ts) {
+    for (std::thread &T : Ts)
+      if (T.joinable())
+        T.join();
+  }
+
+  LLThinLTOJITTest()
+    : StaticInit(), SharedStringPool(std::make_shared<SymbolStringPool>()),
+      MandIDL(EngineBuilder().selectTarget()->createDataLayout()),
+      MandIES(SharedStringPool), MandI(MandIES, MandIDL) {}
+
+  std::unique_ptr<LLThinLTOJIT> createSinglethreadedJit() {
+    auto ES = llvm::make_unique<ExecutionSession>(SharedStringPool);
+    auto Jit = Init::makeJit(std::move(ES), Ctx, false);
+    Jit->addSearchPath(TestInputsDir);
+    return Jit;
+  }
+
+  std::unique_ptr<LLThinLTOJIT> createMultithreadedJit() {
+    auto ES = llvm::make_unique<ExecutionSession>(SharedStringPool);
+    auto Jit = Init::makeJit(std::move(ES), Ctx, true);
+    Jit->addSearchPath(TestInputsDir);
+    return Jit;
+  }
+  
+  SymbolStringPtr mangledPoolStr(StringRef Name) {
+    return MandI(Name);
+  }
+};
+
+//===----------------------------------------------------------------------===//
+  
+TEST_F(LLThinLTOJITTest, UnthreadedSingleDylib) {
+  auto Jit = createSinglethreadedJit();
+
+  cantFail(Jit->addThinLTOIRModule("Foo.o"));
+  cantFail(Jit->addThinLTOIRModule("Bar.o"));
+  cantFail(Jit->prepare("foo"));
+  
+  JITEvaluatedSymbol FooSym = cantFail(Jit->lookup("foo"));
+  EXPECT_EQ(1, Exec(FooSym));
+}
+
+//===----------------------------------------------------------------------===//
+  
+TEST_F(LLThinLTOJITTest, UnthreadedMultiDylib) {
+  auto Jit = createSinglethreadedJit();
+  ExecutionSession &ES = Jit->getExecutionSession();
+
+  JITDylib &DylibFoo = ES.createJITDylib("Foo");
+  cantFail(Jit->addThinLTOIRModule(DylibFoo, "Foo.o"));
+
+  JITDylib &DylibBar = ES.createJITDylib("Bar");
+  cantFail(Jit->addThinLTOIRModule(DylibBar, "Bar.o"));
+
+  JITDylib *OwnerDylib = cantFail(Jit->prepare("foo"));
+  EXPECT_EQ(OwnerDylib, &DylibFoo);
+
+  JITEvaluatedSymbol FooSym = cantFail(Jit->lookup(DylibFoo, "foo"));
+  EXPECT_EQ(1, Exec(FooSym));
+  
+  JITEvaluatedSymbol BarSym = cantFail(Jit->lookup(DylibBar, "bar"));
+  EXPECT_EQ(1, Exec(BarSym));
+}
+
+//===----------------------------------------------------------------------===//
+  
+TEST_F(LLThinLTOJITTest, ThreadedSingleDylib) {
+  auto Jit = createMultithreadedJit();
+  cantFail(Jit->addThinLTOIRModule("Foo.o"));
+  cantFail(Jit->addThinLTOIRModule("Bar.o"));
+
+  // Dispatch compile jobs to separate threads. Below we invoke lookups from
+  // worker threads. Eventually execution arrives here and spawns additional
+  // materialization threads. These must be joined back to the original runner
+  // threads, which is quite hard to track. Thus, we use thread-local scope-exit
+  // guards here. It's not a good idea when calling from the main thread though!
+  std::thread::id MainThreadId = std::this_thread::get_id();
+  ExecutionSession &ES = Jit->getExecutionSession();
+  ES.setDispatchMaterialization([&](JITDylib &JD,
+                                    std::unique_ptr<MaterializationUnit> MU) {
+    assert(std::this_thread::get_id() != MainThreadId &&
+           "Don't join via thread-local scope-exit "
+           "when calling from main thread!!!");
+
+    auto SMU = std::shared_ptr<MaterializationUnit>(std::move(MU));
+    auto CompilerThread =
+        std::make_shared<std::thread>([SMU, &JD]() { SMU->doMaterialize(JD); });
+
+    thread_local auto JoinCompilerThread =
+        llvm::make_scope_exit([CompilerThread]() { CompilerThread->join(); });
+  });
+
+  // Compile and run functions in different threads.
+  std::vector<std::thread> Workers;
+  
+  cantFail(Jit->prepare("foo"));
+  Workers.emplace_back([&Jit]() {
+    JITEvaluatedSymbol FooSym = cantFail(Jit->lookup("foo"));
+    EXPECT_EQ(1, Exec(FooSym));
+  });
+
+  cantFail(Jit->prepare("bar"));
+  Workers.emplace_back([&Jit]() {
+    JITEvaluatedSymbol BarSym = cantFail(Jit->lookup("bar"));
+    EXPECT_EQ(1, Exec(BarSym));
+  });
+
+  JoinAll(Workers);
+}
+
+//===----------------------------------------------------------------------===//
+  
+TEST_F(LLThinLTOJITTest, Incomplete) {
+  auto Jit = createMultithreadedJit();
+  cantFail(Jit->addThinLTOIRModule("Foo.o"));
+  cantFail(Jit->addThinLTOIRModule("Bar.o"));
+
+  // DiscoveryLimit zero means we don't follow any edge. This way "bar" will not
+  // be discovered. This will cause an error when looking up "foo" below.
+  // It would be great to get the symbol query back, but I didn't manage to wire
+  // my experimental ThinLTOLayer up that far yet..
+  unsigned DiscoveryLimit = 0;
+  cantFail(Jit->prepare("foo", DiscoveryLimit));
+
+  Expected<JITEvaluatedSymbol> NoFooSym = Jit->lookup("foo");
+  EXPECT_FALSE(NoFooSym);
+  logAllUnhandledErrors(NoFooSym.takeError(), dbgs(), "Failed as expected: ");
+
+  // Manually adding the module that defines the only call from "foo". I assume
+  // it would be possible to do this automatically with a real ThinLTOLayer.
+  cantFail(Jit->prepare("bar", DiscoveryLimit));
+
+  // Now it should work, but it doesn't because the MaterializationUnit that
+  // defined "foo" was abandoned and I don't know how to recover it.
+  Expected<JITEvaluatedSymbol> FooSym = Jit->lookup("foo");
+  EXPECT_FALSE(FooSym);
+  logAllUnhandledErrors(FooSym.takeError(), dbgs(), "Failed again: ");
+}
+
+} // namespace