Add ThinLtoJIT example

[sgraenitz]

Summary

Prototype of a JIT compiler that utilizes ThinLTO summaries to compile modules ahead of time. This is an implementation of the concept I presented in my "ThinLTO Summaries in JIT Compilation" talk at the 2018 Developers' Meeting: http://llvm.org/devmtg/2018-10/talk-abstracts.html#lt8

Upfront the JIT first populates the *combined ThinLTO module index*, which provides fast access to the global call-graph and module paths by function. Next, it loads the main function's module and compiles it. All functions in the module will be emitted with two prolog instructions that *fire a discovery flag* once execution reaches them. In parallel, the *discovery thread* is busy-watching the existing flags. Once it detects one has fired, it uses the module index to find all functions that are reachable from it within a given number of calls and submits their defining modules to the compilation pipeline.

While execution continues, more flags are fired and further modules added. Ideally the JIT can be tuned in a way, so that the code on the execution path can always be compiled ahead of time. In cases where it doesn't work, the JIT has a *definition generator* in place that loads modules if missing functions are reached.

The example that was discussed in the presentation has a good size for debugging. Here's how to build it:

# Build LLVM
$ cd build-release
$ cmake -GNinja -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_PROJECTS=clang -DLLVM_TARGETS_TO_BUILD=host /path/to/llvm-project/llvm
$ ninja clang llvm-dis lli SpeculativeJIT ThinLtoJIT

# Get the example code
$ git clone https://github.com/weliveindetail/tinyexpr

# Unrepresentative minimal runtime comparison
$ /path/to/llvm-project/llvm/examples/ThinLtoJIT/bench
Usage: bench <path to llvm binaries> <path to c-sources> <main source file> [<override sysroot>]

$ /path/to/llvm-project/llvm/examples/ThinLtoJIT/bench $(pwd)/bin $(pwd)/tinyexpr example.c

Note that on macOS you may want to pass a sysroot override like /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX10.15.sdk.

I would like to gather some feedback. Do you think it's fine to have this as another JIT example? I am happy to maintain it and to continue working on it.

[sgraenitz]

Polishing

[sgraenitz]

Sample output:

> ./ThinLtoJIT -debug -debug-only="thinltojit" -compile-threads=2 example.bc tinyexpr.bc compile.bc free.bc eval.bc
1 new flags raised
Nudged 1 new functions into discovery
Instrumented 1 new functions in module example.bc
Generator: Added 1 new modules
Nudged 1 new functions into discovery
DiscoveryThread: 3 new modules (1 known modules)
1 new flags raised
DiscoveryThread: 0 new modules (3 known modules)
1 new flags raised
DiscoveryThread: 0 new modules (4 known modules)
Instrumented 3 new functions in module tinyexpr.bc
1 new flags raised
DiscoveryThread: 0 new modules (3 known modules)
Instrumented 23 new functions in module compile.bc
Generator: Added 0 new modules
11 new flags raised
DiscoveryThread: 0 new modules (3 known modules)
1 new flags raised
DiscoveryThread: 0 new modules (0 known modules)
Instrumented 1 new functions in module eval.bc
1 new flags raised
DiscoveryThread: 0 new modules (1 known modules)
Instrumented 2 new functions in module free.bc
Generator: Added 0 new modules
2 new flags raised
DiscoveryThread: 0 new modules (1 known modules)
2 new flags raised
DiscoveryThread: 0 new modules (0 known modules)
The expression:
	sqrt(5^2+7^2+11^2+(8-2)^2)
evaluates to:
	15.198684
Discovery flags stats
Alloc:  4096
Used:     32
Fired:    21

FlagOwnersMap has 11 remaining entries.

[lhames]

Hi Stefan. Thanks for posting this! I’m out on vacation this week, but looking forward to reviewing when I get back to the office next week.

[sgraenitz]

Pull in new modules also for global variables

[sgraenitz]

Workaround: parsed module requested but waiting in JITDylib::define()

[sgraenitz]

Make DiscoveryFlagsPerBatch and MemFence settings configurable

[sgraenitz]

Polishing, comments, cosmetics

[sgraenitz]

Make it configurable whether the symbol generator can nugde things into discovery

[sgraenitz]

Fix: Protect the set of parsed module IDs from concurrent access

[sgraenitz]

Use runAsMain() utility function from ExecutionUtils

[sgraenitz]

DataLayout can be owned by MangleWrapper

[sgraenitz]

In D72486#1814233, @lhames wrote:

Hi Stefan. Thanks for posting this! I’m out on vacation this week, but looking forward to reviewing when I get back to the office next week.

Hi Lang, that would be great! Looking forward to your comments.
It does pass the 403.gcc test case now, but can't beat lli orc-lazy per-module yet in terms of real-time performance.
There is one particular thing that I have to work on and would like to hear your recommendations. Please see the inline comment.

llvm/examples/ThinLtoJIT/ThinLtoJIT.cpp

101

I am struggling to find a solution for this: The discovery thread is about to add the required module, but it's waiting for the session lock in JITDylib::define(). As the lock is recursive, I can load it again and add it from here, but that's expensive and I cannot prevent the other one from getting into defineImpl() anymore! This causes a lot of duplicate symbol errors. They are not fatal, but add too much extra overhead.

I am currently working on parallel module parsing which makes this much more likely to happen. The only way I see so far is to change the design and avoid concurrent AddModule() calls. Would be possible, but I'd lose some speculation performance.

Happy about any suggestions. Thanks!

[sgraenitz]

Make LookaheadLevels and LookupOnAdd configurable

[sgraenitz]

llvm/examples/ThinLtoJIT/ThinLtoJIT.cpp

101

Ok, I should probably not add modules in parallel, but instead issue a lookup for one of their symbols from the DiscoveryThread. The lookup will end up in the definition generator and I can do the AddModule there. What do you think?

I have a first patch for it here, but it has some open todos:
https://reviews.llvm.org/differential/diff/237989/#change-u8RDZezBvAb6

[sgraenitz]

Avoid concurrent addModule() calls, improve ranking algorithm in discovery, add multi-threaded bitcode parsing

[sgraenitz]

Pardon, apparently the last update was missing the actual code changes. Trying again.

[sgraenitz]

Respect environment variables CFLAGS and EXEC_ARGS in the bench script

[sgraenitz]

Submit actual module keys and sort getAllModulePaths() by ID

[sgraenitz]

Fix GUID when instrumenting functions with local linkage

[sgraenitz]

Don't dispatch materialization for trivial modules to the thread pool

[sgraenitz]

Sort PathRank candidates by minimal distance

[sgraenitz]

If you want to build this patch, please apply D72301 before. I will try and make it work without it soon.

llvm/examples/ThinLtoJIT/ThinLtoJIT.cpp

259

@lhames: It looks like the session lock is a major bottleneck and not dispatching "trivial" modules to threads seems to help reducing the pressure. Do you think that's reasonable?

Are MaterializationUnit types meant to be distinguishable by type? Would it benefit from isa<>/cast<> support?

llvm/examples/ThinLtoJIT/ThinLtoModuleIndex.cpp

55

@tejohnson In which cases does the summary list for a global value have more than 1 entry? Is it reasonable to always refer to the first one (while this is a work in progress)? I am currently limited to sources in C and never see it happen.

[lhames]

I need to check this out and play around with it a bit to wrap my head around it, especially the performance tuning issues. This is very exciting though!

Heads up: I probably broke this with ce2207abaf9, but you just need to switch to using std::make_unique to construct the IRCompiler (see e.g. fixes in 98e55477558).

llvm/examples/ThinLtoJIT/ThinLtoJIT.cpp

259

Yes! I ran in to the same problem when developing our demo for the 2018 developer’s meeting.

I posted https://reviews.llvm.org/D39111 as a solution to distinguishing MUs by type, but haven’t had time to land it yet. Maybe it’s time to start making some noise on that review again.

[lhames]

llvm/examples/ThinLtoJIT/ThinLtoJIT.cpp

259

Side note: The bottle neck here may not be the session lock. I thought that was the culprit for our 2018 dev-meeting demo too, but it was really the thread pool: trivial materializes couldn’t get a thread to run on at all, rather than running but blocking on the session lock. Could that be what’s happening to you too?

Either way: I think I should add some logging to runSessionLocked to track how long threads spend holding the session lock, and how long they spend blocked waiting on it.

[tejohnson]

llvm/examples/ThinLtoJIT/ThinLtoModuleIndex.cpp

55

In C you can presumably still get duplicates on symbols declared with attribute((weak)), a GNU extension supported by gcc and clang.

Additionally, you could also get duplicates when there are same-named local symbols in same-named source files in different subdirectories and the user has not compiled with enough distinguishing source path. E.g. assume you have local symbol foo defined in path1/bar.c and path2/bar.c, and the user has done the following:

$ cd path1
$ clang -c -flto=thin bar.c
$ cd ../path2
$ clang -c -flto=thin bar.c

When you link those two bar.o together it will have two matching GUID since the globalized name for locals used to compute the GUID has the relative path (encoded in the module) prepended, and in this case that is both "." essentially. It happens, and we support this typically by looking for one in the same module as the caller:
http://llvm-cs.pcc.me.uk/lib/Transforms/IPO/FunctionImport.cpp#216

or in the current module when doing some handling in the ThinLTO backend, e.g.:
http://llvm-cs.pcc.me.uk/lib/Transforms/Utils/FunctionImportUtils.cpp#62

[lhames]

LGTM! I really like this. I'm all for adding it as another example if you're happy to maintain it.

The code could probably use some extra comments (including pointers to your talk), but I think that would be better dealt with in-tree.

[sgraenitz]

Adjust for emulated-TLS changes in ce2207abaf9a

[sgraenitz]

Drop requirement to apply D72301 before this patch

[sgraenitz]

In D72486#1839833, @lhames wrote:

LGTM! I really like this. I'm all for adding it as another example if you're happy to maintain it.

The code could probably use some extra comments (including pointers to your talk), but I think that would be better dealt with in-tree.

Thanks! I have a few small patches (polishing and fixing some todos) that I'd like to add next week, and a larger one (use call-through stubs for lazy module loading) that I will probably add in a separate commit. And sure, happy to add more comments step by step as the overall design consolidates.

I need to check this out and play around with it a bit to wrap my head around it, especially the performance tuning issues. This is very exciting though!

With the two last updates the code should build just fine from TOT. The tinyexpr example is a good starting point. It's small and I used it in my presentation. I split it up into multiple files in order to illustrate the principle with per-module lazy compilation (the original version has only two translation units): https://github.com/weliveindetail/tinyexpr

llvm/examples/ThinLtoJIT/ThinLtoJIT.cpp

259

I posted https://reviews.llvm.org/D39111 as a solution to distinguishing MUs by type

I see your point, that it would be convenient for clients to distinguish out-of-tree MU types using the same mechanism. For the moment, however, I could imagine the regular LLVM style isa for built-in MU types to be just fine. In case more and more clients utilize MUs as extension points to ORC, this would serve as a good argument for your open-class-hierarchy proposal (in addition to streamlining the llvm::ErrorInfo mechanism). I will put the former on my todo list and come back to you with my findings in a separate review.

trivial materializes couldn’t get a thread to run on at all, rather than running but blocking on the session lock. Could that be what’s happening to you too?

Possible. I will check that in more detail when going ahead with performance tuning. Thanks for the heads-up!

llvm/examples/ThinLtoJIT/ThinLtoModuleIndex.cpp

55

Ok, thanks for the explanation! I should check that all modules have unique paths upfront when building the index. Then I could turn the debug warning here into an assertion, saying that attribute((weak)) is not yet supported. I think that's fine for a first version and I can improve that once the overall design consolidates. Thanks for the pointers on how this is handled in other cases!

[sgraenitz]

bench script: link and disassemble global ThinLTO index file and respect LDFLAGS in static compilation

[sgraenitz]

Accept ThinLTO global index as input file

[sgraenitz]

Prefer std::atomic fetch_add over compare-exchange in reserveDiscoveryFlags()

[sgraenitz]

Don't risk "post-mortem" access violations. Keep DiscoveryThread attached and join it on shutdown.

[sgraenitz]

Drop DiscoveryThread module requests via the definition generator. Instead trigger asynchronous module adding from there.

[sgraenitz]

Let scheduleModuleParsing() process a range of paths with a single lock

[sgraenitz]

The blocking single-module parsing function parseModuleFromFile() should use the regular primitives

[sgraenitz]

Assume single entries in summary lists after checking module paths to be unique

[sgraenitz]

Split off ThinLtoInstrumentationLayer::dump() from its dtor

[sgraenitz]

Move ExplicitMemoryBarrier enum to ThinLtoJIT

[sgraenitz]

Add -print-stats command line flag

[sgraenitz]

Refine ThinLtoModuleIndex::selectNextPaths()

[sgraenitz]

Polish initialization, cleanup, formatting

llvm/examples/CMakeLists.txt

 add_subdirectory(BrainF)
 add_subdirectory(Fibonacci)
 add_subdirectory(HowToUseJIT)
 add_subdirectory(HowToUseLLJIT)
 add_subdirectory(IRTransforms)
 add_subdirectory(LLJITExamples)
 add_subdirectory(Kaleidoscope)
 add_subdirectory(ModuleMaker)
 add_subdirectory(SpeculativeJIT)
 add_subdirectory(Bye)
+add_subdirectory(ThinLtoJIT)
 
 if(LLVM_ENABLE_EH AND (NOT WIN32) AND (NOT "${LLVM_NATIVE_ARCH}" STREQUAL "ARM"))
     add_subdirectory(ExceptionDemo)
 endif()
 
 if( HAVE_PTHREAD_H )
   add_subdirectory(ParallelJIT)
 endif( HAVE_PTHREAD_H )

llvm/examples/ThinLtoJIT/CMakeLists.txt

+set(LLVM_LINK_COMPONENTS
+  Core
+  IRReader
+  OrcJIT
+  ExecutionEngine
+  Support
+  nativecodegen
+  Analysis
+  Passes
+  )
+
+# Disable warning for passing variable length argv to JITed main function.
+set_property(SOURCE main.cpp APPEND_STRING PROPERTY COMPILE_FLAGS "-Wno-vla-extension")
+
+add_llvm_example(ThinLtoJIT
+  main.cpp
+  ThinLtoJIT.cpp
+  ThinLtoModuleIndex.cpp
+  ThinLtoInstrumentationLayer.cpp
+  ThinLtoDiscoveryThread.cpp
+  )

llvm/examples/ThinLtoJIT/ThinLtoDiscoveryThread.h

+#ifndef LLVM_EXAMPLES_THINLTOJIT_DISCOVERYTHREAD_H
+#define LLVM_EXAMPLES_THINLTOJIT_DISCOVERYTHREAD_H
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/IR/ModuleSummaryIndex.h"
+
+#include "ThinLtoJIT.h"
+
+#include <atomic>
+#include <set>
+#include <vector>
+
+namespace llvm {
+namespace orc {
+
+class ThinLtoModuleIndex;
+class ThinLtoInstrumentationLayer;
+
+class ThinLtoDiscoveryThread {
+public:
+  ThinLtoDiscoveryThread(std::atomic<bool> &RunningFlag,
+                         ThinLtoInstrumentationLayer &L,
+                         ThinLtoModuleIndex &GlobalIndex,
+                         unsigned LookaheadLevels,
+                         ThinLtoJIT::AddModuleFunction AddModule)
+      : KeepRunning(RunningFlag), Layer(L), GlobalIndex(GlobalIndex),
+        AddModule(std::move(AddModule)), LookaheadLevels(LookaheadLevels) {}
+
+  void operator()();
+
+private:
+  std::atomic<bool> &KeepRunning;
+  ThinLtoInstrumentationLayer &Layer;
+  ThinLtoModuleIndex &GlobalIndex;
+  ThinLtoJIT::AddModuleFunction AddModule;
+  unsigned LookaheadLevels;
+
+  std::set<StringRef> discoverCalleeModulePaths(FunctionSummary *S,
+                                                unsigned Levels);
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif

llvm/examples/ThinLtoJIT/ThinLtoDiscoveryThread.cpp

+#include "ThinLtoDiscoveryThread.h"
+
+#include "llvm/IR/GlobalValue.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/Error.h"
+
+#include "ThinLtoInstrumentationLayer.h"
+#include "ThinLtoModuleIndex.h"
+
+#define DEBUG_TYPE "thinltojit"
+
+namespace llvm {
+namespace orc {
+
+void ThinLtoDiscoveryThread::operator()() {
+  while (KeepRunning.load()) {
+    std::vector<unsigned> Indexes = Layer.takeFlagsThatFired();
+
+    // TODO: Dispatch actual module loading into a thread pool.
+    if (!Indexes.empty()) {
+      LLVM_DEBUG(dbgs() << Indexes.size() << " new flags raised\n");
+      auto ReachedFunctions = Layer.takeFlagOwners(std::move(Indexes));
+
+      // While traversing the call graph, collect the modules we cross.
+      std::set<StringRef> Paths;
+      for (GlobalValue::GUID F : ReachedFunctions) {
+        if (GlobalValueSummary *S = GlobalIndex.getSummary(F)) {
+          if (isa<FunctionSummary>(S)) {
+            std::set<StringRef> NewPaths = discoverCalleeModulePaths(
+                cast<FunctionSummary>(S), LookaheadLevels - 1);
+            Paths.insert(NewPaths.begin(), NewPaths.end());
+          } else {
+            LLVM_DEBUG(dbgs()
+                       << "Reached symbol is not a function: " << F << "\n");
+          }
+        }
+      }
+
+#ifndef NDEBUG
+      unsigned Added = 0;
+#endif
+      for (StringRef M : Paths) {
+        Expected<Optional<ThreadSafeModule>> TSM =
+            GlobalIndex.parseNewModuleFromFile(M);
+        if (!TSM) {
+          // Failed to parse the module.
+          Layer.getExecutionSession().reportError(TSM.takeError());
+          continue;
+        }
+
+        if (!*TSM)
+          // This module was added already.
+          continue;
+
+        if (Error LoadErr = AddModule(std::move(**TSM))) {
+          // Failed to load the module.
+          Layer.getExecutionSession().reportError(std::move(LoadErr));
+        }
+
+#ifndef NDEBUG
+        ++Added;
+#endif
+      }
+
+      LLVM_DEBUG(dbgs() << "DiscoveryThread: " << Added << " new modules "
+                        << "(" << Paths.size() - Added << " known modules)\n");
+    }
+  }
+}
+
+// We don't filter visited functions here. Discovery will often be retriggered
+// from the middle of already visited functions and aims to reach a little
+// further each time.
+std::set<StringRef>
+ThinLtoDiscoveryThread::discoverCalleeModulePaths(FunctionSummary *S,
+                                                  unsigned Levels) {
+  // Summaries for function callees must be FunctionSummaries.
+  auto getCalleeSummary = [](const ValueInfo &VI) -> FunctionSummary * {
+    const auto &SummaryList = VI.getSummaryList();
+    if (SummaryList.empty())
+      return nullptr;
+    return cast<FunctionSummary>(SummaryList.front().get()->getBaseObject());
+  };
+
+  // TODO: The paths we discover should be weighed, e.g. by number of edges that
+  // lead there and the minimal distance to go?
+  std::set<StringRef> Paths;
+  bool VisitChildren = (Levels > 0);
+
+  for (const auto &Edge : S->calls()) {
+    if (FunctionSummary *CalleeSummary = getCalleeSummary(Edge.first)) {
+      Paths.insert(CalleeSummary->modulePath());
+      if (VisitChildren) {
+        std::set<StringRef> ChildPaths =
+            discoverCalleeModulePaths(CalleeSummary, Levels - 1);
+        Paths.insert(ChildPaths.begin(), ChildPaths.end());
+      }
+    }
+  }
+  return Paths;
+}
+
+} // namespace orc
+} // namespace llvm

llvm/examples/ThinLtoJIT/ThinLtoInstrumentationLayer.h

+#ifndef LLVM_EXAMPLES_THINLTOJIT_DISCOVERYLAYER_H
+#define LLVM_EXAMPLES_THINLTOJIT_DISCOVERYLAYER_H
+
+#include "llvm/ExecutionEngine/JITSymbol.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/Layer.h"
+#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
+#include "llvm/IR/GlobalValue.h"
+#include "llvm/IR/IRBuilder.h"
+
+#include <atomic>
+#include <cstdint>
+#include <map>
+#include <mutex>
+#include <vector>
+
+namespace llvm {
+namespace orc {
+
+class ThinLtoInstrumentationLayer : public IRLayer {
+public:
+  enum ExplicitMemoryBarrier {
+    Never = 0,
+    StaticCode = 1,
+    JITedCode = 2,
+    Always = 3
+  };
+
+  ThinLtoInstrumentationLayer(ExecutionSession &ES, IRCompileLayer &BaseLayer,
+                              ExplicitMemoryBarrier InsertMemBarrier,
+                              unsigned FlagsPerBatch)
+      : IRLayer(ES), BaseLayer(BaseLayer), InsertMemBarrier(InsertMemBarrier) {
+    // TODO: So far we only allocate one batch.
+    allocateDiscoveryFlags(FlagsPerBatch);
+  }
+
+  ~ThinLtoInstrumentationLayer() override;
+
+  void emit(MaterializationResponsibility R, ThreadSafeModule TSM) override;
+
+  unsigned reserveDiscoveryFlags(unsigned Count);
+  void registerDiscoveryFlagOwners(std::vector<GlobalValue::GUID> Guids,
+                                   unsigned FirstIdx);
+
+  void nudgeIntoDiscovery(std::vector<GlobalValue::GUID> Functions);
+
+  std::vector<unsigned> takeFlagsThatFired();
+  std::vector<GlobalValue::GUID> takeFlagOwners(std::vector<unsigned> Indexes);
+
+private:
+  IRCompileLayer &BaseLayer;
+  ExplicitMemoryBarrier InsertMemBarrier;
+
+  enum Flag : uint8_t { Clear = 0, Fired = 1 };
+
+  // Lock-free read access.
+  uint8_t *FlagsStorage;
+  Flag *FlagsIncoming; // lock-free write by design
+  Flag *FlagsHandled;
+  unsigned NumFlagsAllocated;
+  std::atomic<unsigned> NumFlagsUsed; // spin-lock
+
+  // Acquire/release sync between writers and reader
+  std::atomic<uint64_t> FlagsSync;
+
+  // STL container requires locking for both, read and write access.
+  mutable std::mutex DiscoveryFlagsInfoLock;
+  std::map<unsigned, GlobalValue::GUID> FlagOwnersMap;
+
+  void allocateDiscoveryFlags(unsigned MinFlags);
+  void compileFunctionReachedFlagSetter(IRBuilder<> &B, Flag *F);
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif

llvm/examples/ThinLtoJIT/ThinLtoInstrumentationLayer.cpp

+#include "ThinLtoInstrumentationLayer.h"
+
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Type.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/Process.h"
+
+#include <cstdlib>
+
+#define DEBUG_TYPE "thinltojit"
+
+namespace llvm {
+namespace orc {
+
+// TODO: Fixed set of flags may not always be enough. Make this expandable.
+void ThinLtoInstrumentationLayer::allocateDiscoveryFlags(unsigned MinFlags) {
+  // Round up to full memory pages.
+  unsigned PageSize = sys::Process::getPageSizeEstimate();
+  unsigned NumPagesEach = (MinFlags + (PageSize - 1)) / PageSize;
+  unsigned NumPagesTotal = 2 * NumPagesEach;
+  assert(isPowerOf2_64(PageSize) && "Adjust aligned memory alloc below");
+
+  // Allocate one more page to make up for size loss due to alignment.
+  void *Storage = std::calloc(NumPagesTotal + 1, PageSize);
+  uint64_t StorageAddr = reinterpret_cast<uint64_t>(Storage);
+  uint64_t PageSizeDecr = PageSize - 1;
+  uint64_t AlignedAddr = ((StorageAddr + PageSizeDecr) & ~PageSizeDecr);
+  uint64_t Diff = AlignedAddr - StorageAddr;
+
+  // For each flag we allocate one byte in each location: Incoming and Handled.
+  // TODO: 'Handled' could be a bitset, but size must be dynamic
+  NumFlagsUsed.store(0);
+  NumFlagsAllocated = NumPagesEach * PageSize;
+  FlagsStorage = static_cast<uint8_t *>(Storage);
+  FlagsIncoming = reinterpret_cast<Flag *>(FlagsStorage + Diff);
+  FlagsHandled = FlagsIncoming + NumFlagsAllocated;
+
+  static_assert(sizeof(FlagsIncoming[0]) == sizeof(uint8_t), "Flags are bytes");
+  assert(reinterpret_cast<uint64_t>(FlagsIncoming) % PageSize == 0);
+  assert(reinterpret_cast<uint64_t>(FlagsHandled) % PageSize == 0);
+  assert(NumFlagsAllocated >= MinFlags);
+}
+
+unsigned ThinLtoInstrumentationLayer::reserveDiscoveryFlags(unsigned Count) {
+  assert(Count > 0);
+  unsigned Before, After;
+  do {
+    Before = NumFlagsUsed.load();
+    After = Before + Count;
+  } while (!NumFlagsUsed.compare_exchange_weak(Before, After));
+
+#ifndef NDEBUG
+  for (unsigned i = Before; i < After; i++) {
+    assert(FlagsIncoming[i] == Clear);
+  }
+#endif
+
+  return Before; // First reserved index
+}
+
+void ThinLtoInstrumentationLayer::registerDiscoveryFlagOwners(
+    std::vector<GlobalValue::GUID> Guids, unsigned FirstIdx) {
+  unsigned Count = Guids.size();
+
+  std::lock_guard<std::mutex> Lock(DiscoveryFlagsInfoLock);
+  for (unsigned i = 0; i < Count; i++) {
+    assert(!FlagOwnersMap.count(FirstIdx + i) &&
+           "Flag should not have an owner at this point");
+    FlagOwnersMap[FirstIdx + i] = Guids[i];
+  }
+}
+
+std::vector<unsigned> ThinLtoInstrumentationLayer::takeFlagsThatFired() {
+  // This is only effective with the respective Release.
+  FlagsSync.load(std::memory_order_acquire);
+
+  std::vector<unsigned> Indexes;
+  unsigned NumIndexesUsed = NumFlagsUsed.load();
+  for (unsigned i = 0; i < NumIndexesUsed; i++) {
+    if (FlagsIncoming[i] == Fired && FlagsHandled[i] == Clear) {
+      FlagsHandled[i] = Fired;
+      Indexes.push_back(i);
+    }
+  }
+
+  return Indexes;
+}
+
+std::vector<GlobalValue::GUID>
+ThinLtoInstrumentationLayer::takeFlagOwners(std::vector<unsigned> Indexes) {
+  std::vector<GlobalValue::GUID> ReachedFunctions;
+  std::lock_guard<std::mutex> Lock(DiscoveryFlagsInfoLock);
+
+  for (unsigned i : Indexes) {
+    auto KV = FlagOwnersMap.find(i);
+    assert(KV != FlagOwnersMap.end());
+    ReachedFunctions.push_back(KV->second);
+    FlagOwnersMap.erase(KV);
+  }
+
+  return ReachedFunctions;
+}
+
+void ThinLtoInstrumentationLayer::nudgeIntoDiscovery(
+    std::vector<GlobalValue::GUID> Functions) {
+  unsigned Count = Functions.size();
+
+  // Registering synthetic flags in advance. We expect them to get processed
+  // before the respective functions get emitted. If not, the emit() function
+  unsigned FirstFlagIdx = reserveDiscoveryFlags(Functions.size());
+  registerDiscoveryFlagOwners(std::move(Functions), FirstFlagIdx);
+
+  // Initialize the flags as fired and force a cache sync, so discovery will
+  // pick them up as soon as possible.
+  for (unsigned i = FirstFlagIdx; i < FirstFlagIdx + Count; i++) {
+    FlagsIncoming[i] = Fired;
+  }
+  if (InsertMemBarrier & ExplicitMemoryBarrier::StaticCode) {
+    FlagsSync.store(0, std::memory_order_release);
+  }
+
+  LLVM_DEBUG(dbgs() << "Nudged " << Count << " new functions into discovery\n");
+}
+
+void ThinLtoInstrumentationLayer::emit(MaterializationResponsibility R,
+                                       ThreadSafeModule TSM) {
+  TSM.withModuleDo([this](Module &M) {
+    std::vector<Function *> FunctionsToInstrument;
+
+    // We may have discovered ahead of some functions already, but we still
+    // instrument them all. Their notifications steer the future direction of
+    // discovery.
+    for (Function &F : M.getFunctionList())
+      if (!F.isDeclaration())
+        FunctionsToInstrument.push_back(&F);
+
+    if (!FunctionsToInstrument.empty()) {
+      IRBuilder<> B(M.getContext());
+      std::vector<GlobalValue::GUID> NewDiscoveryRoots;
+
+      // Flags that fire must have owners registered. We will do it below and
+      // that's fine, because they can only be reached once the code is emitted.
+      unsigned FirstFlagIdx =
+          reserveDiscoveryFlags(FunctionsToInstrument.size());
+
+      unsigned NextFlagIdx = FirstFlagIdx;
+      for (Function *F : FunctionsToInstrument) {
+
+        BasicBlock *E = &F->getEntryBlock();
+        B.SetInsertPoint(BasicBlock::Create(
+            M.getContext(), "NotifyFunctionReachedProlog", F, E));
+        compileFunctionReachedFlagSetter(B, FlagsIncoming + NextFlagIdx);
+        B.CreateBr(E);
+
+        NewDiscoveryRoots.push_back(GlobalValue::getGUID(F->getName()));
+        ++NextFlagIdx;
+      }
+
+      LLVM_DEBUG(dbgs() << "Instrumented " << NewDiscoveryRoots.size()
+                        << " new functions in module " << M.getName() << "\n");
+
+      // Submit owner info, so the DiscoveryThread can evaluate the flags.
+      registerDiscoveryFlagOwners(std::move(NewDiscoveryRoots), FirstFlagIdx);
+    }
+  });
+
+  BaseLayer.emit(std::move(R), std::move(TSM));
+}
+
+void ThinLtoInstrumentationLayer::compileFunctionReachedFlagSetter(
+    IRBuilder<> &B, Flag *F) {
+  assert(*F == Clear);
+  Type *Int64Ty = Type::getInt64Ty(B.getContext());
+
+  // Write one immediate 8bit value to a fixed location in memory.
+  auto FlagAddr = pointerToJITTargetAddress(F);
+  Type *FlagTy = Type::getInt8Ty(B.getContext());
+  B.CreateStore(ConstantInt::get(FlagTy, Fired),
+                B.CreateIntToPtr(ConstantInt::get(Int64Ty, FlagAddr),
+                                 FlagTy->getPointerTo()));
+
+  if (InsertMemBarrier & ExplicitMemoryBarrier::JITedCode) {
+    // Overwrite the sync value with Release ordering. The discovery thread
+    // reads it with Acquire ordering. The actual value doesn't matter.
+    static constexpr bool IsVolatile = true;
+    static constexpr Instruction *NoInsertBefore = nullptr;
+    auto SyncFlagAddr = pointerToJITTargetAddress(&FlagsSync);
+
+    B.Insert(
+        new StoreInst(ConstantInt::get(Int64Ty, 0),
+                      B.CreateIntToPtr(ConstantInt::get(Int64Ty, SyncFlagAddr),
+                                       Int64Ty->getPointerTo()),
+                      IsVolatile, MaybeAlign(64), AtomicOrdering::Release,
+                      SyncScope::System, NoInsertBefore));
+  }
+}
+
+ThinLtoInstrumentationLayer::~ThinLtoInstrumentationLayer() {
+  LLVM_DEBUG({
+    dbgs() << "Discovery flags stats\n";
+
+    unsigned NumFlagsFired = 0;
+    for (unsigned i = 0; i < NumFlagsAllocated; i++) {
+      if (FlagsIncoming[i] == Fired)
+        ++NumFlagsFired;
+    }
+    dbgs() << "Alloc:  " << format("%6.d", NumFlagsAllocated) << "\n";
+    dbgs() << "Issued: " << format("%6.d", NumFlagsUsed.load()) << "\n";
+    dbgs() << "Fired:  " << format("%6.d", NumFlagsFired) << "\n";
+
+    unsigned RemainingFlagOwners = 0;
+    for (const auto &_ : FlagOwnersMap) {
+      ++RemainingFlagOwners;
+      (void)_;
+    }
+    dbgs() << "\nFlagOwnersMap has " << RemainingFlagOwners
+           << " remaining entries.\n";
+  });
+
+  std::free(FlagsStorage);
+}
+
+} // namespace orc
+} // namespace llvm

llvm/examples/ThinLtoJIT/ThinLtoJIT.h

+#ifndef LLVM_EXAMPLES_THINLTOJIT_THINLTOJIT_H
+#define LLVM_EXAMPLES_THINLTOJIT_THINLTOJIT_H
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/Triple.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/SymbolStringPool.h"
+#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/ThreadPool.h"
+
+#include "ThinLtoInstrumentationLayer.h"
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace llvm {
+namespace orc {
+
+class ThinLtoModuleIndex;
+class ThinLtoDiscoveryThread;
+
+class RTDyldObjectLinkingLayer;
+class IRCompileLayer;
+class CompileOnDemandLayer;
+
+class JITDylib;
+class MangleAndInterner;
+class LazyCallThroughManager;
+
+class ThinLtoJIT {
+public:
+  using AddModuleFunction = std::function<Error(ThreadSafeModule)>;
+
+  ThinLtoJIT(std::vector<StringRef> ModuleFiles, StringRef MainFunctionName,
+             unsigned NumCompileThreads, unsigned DiscoveryFlagsPerBatch,
+             ThinLtoInstrumentationLayer::ExplicitMemoryBarrier MemFence,
+             bool AllowNudgeIntoDiscovery,
+             Error &Err);
+  ~ThinLtoJIT();
+
+  ThinLtoJIT(const ThinLtoJIT &) = delete;
+  ThinLtoJIT &operator=(const ThinLtoJIT &) = delete;
+  ThinLtoJIT(ThinLtoJIT &&) = delete;
+  ThinLtoJIT &operator=(ThinLtoJIT &&) = delete;
+
+  Expected<int> main(int argc, const char *argv[]) {
+    auto MainSym = ES.lookup({MainJD}, MainFunctionMangled);
+    if (!MainSym)
+      return MainSym.takeError();
+
+    auto Main = jitTargetAddressToFunction<int (*)(int, const char *[])>(
+        MainSym->getAddress());
+    return Main(argc, argv);
+  }
+
+private:
+  ExecutionSession ES;
+  std::unique_ptr<DataLayout> DL;
+
+  // Wrapper class to allow late construction of the mangler while preserving
+  // the conventional Mangle(SymbolName) syntax.
+  // TODO: Alternatively the original MangleAndInterner could be turned into a
+  // regular class (default-constructible and assignable).
+  struct MangleWrapper {
+    SymbolStringPtr operator()(StringRef S) { return Impl->operator()(S); }
+    std::unique_ptr<MangleAndInterner> Impl{nullptr};
+  };
+
+  MangleWrapper Mangle;
+
+  JITDylib *MainJD;
+  SymbolStringPtr MainFunctionMangled;
+  std::unique_ptr<ThreadPool> CompileThreads;
+  std::unique_ptr<ThinLtoModuleIndex> GlobalIndex;
+
+  AddModuleFunction AddModule;
+  AddModuleFunction AddModuleAndLookup;
+  std::unique_ptr<RTDyldObjectLinkingLayer> ObjLinkingLayer;
+  std::unique_ptr<IRCompileLayer> CompileLayer;
+  std::unique_ptr<ThinLtoInstrumentationLayer> InstrumentationLayer;
+  std::unique_ptr<CompileOnDemandLayer> OnDemandLayer;
+
+  std::atomic<bool> JitRunning;
+  std::unique_ptr<ThinLtoDiscoveryThread> DiscoveryThreadWorker;
+  std::unique_ptr<LazyCallThroughManager> CallThroughManager;
+
+  Error
+  setupLayers(Triple TT, unsigned DiscoveryFlagsPerBatch,
+              ThinLtoInstrumentationLayer::ExplicitMemoryBarrier MemFence);
+  Error setupJITDylib(JITDylib *JD, bool AllowNudge);
+  Error setupDiscovery(unsigned NumCompileThreads);
+  Expected<ThreadSafeModule> setupMainModule(StringRef MainFunction);
+
+  static void exitOnLazyCallThroughFailure() {
+    errs() << "Compilation failed. Aborting.\n";
+    exit(1);
+  }
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif

llvm/examples/ThinLtoJIT/ThinLtoJIT.cpp

+#include "ThinLtoJIT.h"
+
+#include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
+#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
+#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
+#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
+#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/Support/Debug.h"
+
+#include "ThinLtoDiscoveryThread.h"
+#include "ThinLtoInstrumentationLayer.h"
+#include "ThinLtoModuleIndex.h"
+
+#include <set>
+#include <string>
+#include <thread>
+
+#ifndef NDEBUG
+#include <chrono>
+#endif
+
+#define DEBUG_TYPE "thinltojit"
+
+namespace llvm {
+namespace orc {
+
+class ThinLtoDefinitionGenerator : public JITDylib::DefinitionGenerator {
+public:
+  ThinLtoDefinitionGenerator(ThinLtoModuleIndex &GlobalIndex,
+                             ThinLtoInstrumentationLayer &InstrumentationLayer,
+                             ThinLtoJIT::AddModuleFunction AddModule,
+                             char ManglePrefix, bool AllowNudge)
+      : GlobalIndex(GlobalIndex), InstrumentationLayer(InstrumentationLayer),
+        AddModule(std::move(AddModule)), ManglePrefix(ManglePrefix),
+        AllowNudgeIntoDiscovery(AllowNudge) {}
+
+  Error tryToGenerate(LookupKind K, JITDylib &JD,
+                      JITDylibLookupFlags JDLookupFlags,
+                      const SymbolLookupSet &Symbols) override;
+
+private:
+  ThinLtoModuleIndex &GlobalIndex;
+  ThinLtoInstrumentationLayer &InstrumentationLayer;
+  ThinLtoJIT::AddModuleFunction AddModule;
+  char ManglePrefix;
+  bool AllowNudgeIntoDiscovery;
+
+  // ThinLTO summaries encode unprefixed names.
+  StringRef stripGlobalManglePrefix(StringRef Symbol) const {
+    bool Strip = (ManglePrefix != '\0' && Symbol[0] == ManglePrefix);
+    return Strip ? StringRef(Symbol.data() + 1, Symbol.size() - 1) : Symbol;
+  }
+};
+
+Error ThinLtoDefinitionGenerator::tryToGenerate(
+    LookupKind K, JITDylib &JD, JITDylibLookupFlags JDLookupFlags,
+    const SymbolLookupSet &Symbols) {
+#ifndef NDEBUG
+  unsigned Added = 0;
+#endif
+
+  std::vector<GlobalValue::GUID> NewDiscoveryRoots;
+  StringMap<std::vector<StringRef>> SymbolNamesByModulePath;
+
+  for (const auto &KV : Symbols) {
+    StringRef Name = stripGlobalManglePrefix(*KV.first);
+    auto Guid = GlobalValue::getGUID(Name);
+
+    if (GlobalValueSummary *S = GlobalIndex.getSummary(Guid)) {
+      std::vector<StringRef> &Names = SymbolNamesByModulePath[S->modulePath()];
+      Names.push_back(Name);
+      if (AllowNudgeIntoDiscovery && isa<FunctionSummary>(S)) {
+        NewDiscoveryRoots.push_back(GlobalValue::getGUID(Name));
+      }
+    }
+  }
+
+  for (const auto &KV : SymbolNamesByModulePath) {
+    Expected<Optional<ThreadSafeModule>> TSM =
+        GlobalIndex.parseNewModuleFromFile(KV.first());
+    if (!TSM) {
+      // Parsing the module from disk failed, after we successfully obtained
+      // ValueInfos for its symbols from ThinLTO summaries.
+      return TSM.takeError();
+    }
+
+    // We did parse the module already, but the add request is waiting "outside"
+    // for the lock that will be freed once this request is done. There appears
+    // to be no way to temporarily suspend the request and get the symbol ready.
+    // Thus, we must parse the module here again and submit it before returning.
+    // It's quite expensive as we are actively blocking execution at this point.
+    if (!*TSM) {
+      InstrumentationLayer.getExecutionSession().reportError(
+          createStringError(inconvertibleErrorCode(),
+                            "Module %s required for symbol %s was added while "
+                            "request for it was in progress. Reparsing!",
+                            KV.first().data(), KV.second.front().data()));
+
+      static constexpr bool ForceLoad = true;
+      TSM = GlobalIndex.parseNewModuleFromFile(KV.first(), ForceLoad);
+      if (!TSM) {
+        // Parsing the module from disk failed. This may happen any time.
+        return TSM.takeError();
+      }
+      assert(*TSM && "We forced the load operation");
+    }
+
+    if (Error LoadErr = AddModule(std::move(**TSM)))
+      // Found a module but failed to add it.
+      return LoadErr;
+
+#ifndef NDEBUG
+    ++Added;
+#endif
+  }
+
+  LLVM_DEBUG(dbgs() << "Generator: Added " << Added << " new modules\n");
+
+  // We can anticipate to run into the requested functions as soon as execution
+  // continues. Thus, we may trigger discovery flags for them already now to
+  // initiate discovery behind them. We will probably compile a few unnecessary
+  // things in this case.
+  if (!NewDiscoveryRoots.empty()) {
+    // The registration involves locking a mutex, so better do it in a
+    // separate thread.
+    std::thread(
+        [this](std::vector<GlobalValue::GUID> Rs) {
+          InstrumentationLayer.nudgeIntoDiscovery(std::move(Rs));
+        },
+        std::move(NewDiscoveryRoots))
+        .detach();
+  }
+
+  return Error::success();
+}
+
+ThinLtoJIT::ThinLtoJIT(
+    std::vector<StringRef> ModuleFiles, StringRef MainFunctionName,
+    unsigned NumCompileThreads, unsigned DiscoveryFlagsPerBatch,
+    ThinLtoInstrumentationLayer::ExplicitMemoryBarrier MemFence,
+    bool AllowNudgeIntoDiscovery, Error &Err) {
+  ErrorAsOutParameter ErrAsOutParam(&Err);
+
+  GlobalIndex = std::make_unique<ThinLtoModuleIndex>(ES);
+  for (StringRef F : ModuleFiles) {
+    if (auto Err = GlobalIndex->add(F))
+      ES.reportError(std::move(Err));
+  }
+
+  auto TSM = setupMainModule(MainFunctionName);
+  if (!TSM) {
+    Err = TSM.takeError();
+    return;
+  }
+
+  ThreadSafeModule MainModule = std::move(*TSM);
+  DL = std::make_unique<DataLayout>(MainModule.getModuleUnlocked());
+
+  // Now that we know the target data layout we can setup the mangler.
+  Mangle.Impl = std::make_unique<MangleAndInterner>(ES, *DL);
+  MainFunctionMangled = Mangle(MainFunctionName);
+
+  Err = setupLayers(Triple(MainModule.getModuleUnlocked()->getTargetTriple()),
+                    DiscoveryFlagsPerBatch, MemFence);
+  if (Err)
+    return;
+
+  Err = setupDiscovery(NumCompileThreads);
+  if (Err)
+    return;
+
+  MainJD = &ES.createJITDylib("main");
+  Err = setupJITDylib(MainJD, AllowNudgeIntoDiscovery);
+  if (Err)
+    return;
+
+  Err = AddModule(std::move(MainModule));
+  if (Err)
+    return;
+
+  if (AllowNudgeIntoDiscovery) {
+    auto MainFunctionGuid = GlobalValue::getGUID(MainFunctionName);
+    InstrumentationLayer->nudgeIntoDiscovery({MainFunctionGuid});
+  }
+
+#ifndef NDEBUG
+  // Uncomment to give the discovery thread some time do things.
+  // std::this_thread::sleep_for(std::chrono::milliseconds(1000));
+#endif
+}
+
+Expected<ThreadSafeModule> ThinLtoJIT::setupMainModule(StringRef MainFunction) {
+  Optional<StringRef> M = GlobalIndex->getModulePathForSymbol(MainFunction);
+  if (!M) {
+    std::string Buffer;
+    raw_string_ostream OS(Buffer);
+    OS << "No ValueInfo for symbol '" << MainFunction;
+    OS << "' in provided modules: ";
+    for (StringRef P : GlobalIndex->getAllModulePaths())
+      OS << P << " ";
+    OS << "\n";
+    return createStringError(inconvertibleErrorCode(), OS.str());
+  }
+  Expected<Optional<ThreadSafeModule>> TSM =
+      GlobalIndex->parseNewModuleFromFile(*M);
+  if (!TSM)
+    return TSM.takeError();
+  assert(*TSM && "This is the first module, it cannot exist yet");
+  return std::move(**TSM);
+}
+
+Error ThinLtoJIT::setupLayers(
+    Triple TT, unsigned DiscoveryFlagsPerBatch,
+    ThinLtoInstrumentationLayer::ExplicitMemoryBarrier MemFence) {
+  ObjLinkingLayer = std::make_unique<RTDyldObjectLinkingLayer>(
+      ES, []() { return std::make_unique<SectionMemoryManager>(); });
+  CompileLayer = std::make_unique<IRCompileLayer>(
+      ES, *ObjLinkingLayer, ConcurrentIRCompiler(JITTargetMachineBuilder(TT)));
+
+  InstrumentationLayer = std::make_unique<ThinLtoInstrumentationLayer>(
+      ES, *CompileLayer, MemFence, DiscoveryFlagsPerBatch);
+
+  auto ISMB = createLocalIndirectStubsManagerBuilder(TT);
+  auto LCTM = createLocalLazyCallThroughManager(
+      TT, ES, pointerToJITTargetAddress(exitOnLazyCallThroughFailure));
+  if (!LCTM)
+    return LCTM.takeError();
+
+  CallThroughManager = std::move(*LCTM);
+  OnDemandLayer = std::make_unique<CompileOnDemandLayer>(
+      ES, *InstrumentationLayer, *CallThroughManager, std::move(ISMB));
+  // Don't break up modules. Insert stubs on module boundaries.
+  OnDemandLayer->setPartitionFunction(CompileOnDemandLayer::compileWholeModule);
+
+  AddModule = [this](ThreadSafeModule TSM) -> Error {
+    assert(MainJD && "Setup MainJD JITDylib before calling");
+    return OnDemandLayer->add(*MainJD, std::move(TSM));
+  };
+
+  return Error::success();
+}
+
+Error ThinLtoJIT::setupDiscovery(unsigned NumCompileThreads) {
+  // Delegate compilation to the thread pool.
+  CompileThreads = std::make_unique<ThreadPool>(NumCompileThreads);
+  ES.setDispatchMaterialization(
+      [this](JITDylib &JD, std::shared_ptr<MaterializationUnit> MU) {
+        CompileThreads->async([MU, &JD]() { MU->doMaterialize(JD); });
+      });
+
+  unsigned LookaheadLevels = 3;
+
+#ifndef NDEBUG
+  // Uncomment to avoid discovering all at once when debugging small examples.
+  // LookaheadLevels = 1;
+#endif
+
+  // Spawn discovery thread and let it add newly discovered modules to the JIT.
+  // We lookup one symbol right away to force immediately materialization.
+  JitRunning.store(true);
+  DiscoveryThreadWorker = std::make_unique<ThinLtoDiscoveryThread>(
+      JitRunning, *InstrumentationLayer, *GlobalIndex, LookaheadLevels,
+      [this](ThreadSafeModule TSM) -> Error {
+        std::string FunctionNameForLookup =
+            TSM.getModuleUnlocked()->getFunctionList().front().getName().str();
+        Error SubmitErr = AddModule(std::move(TSM));
+        if (SubmitErr)
+          return SubmitErr;
+        // TODO: This is quite workaroundish. CompileOnDemandLayer could have a
+        // flag to force materialization without an extra lookup like this.
+        auto LookupRes = ES.lookup({MainJD}, Mangle(FunctionNameForLookup));
+        if (!LookupRes)
+          return LookupRes.takeError();
+        assert(LookupRes->getAddress() && "Function should be emitted now");
+        return Error::success();
+      });
+
+  std::thread(std::ref(*DiscoveryThreadWorker)).detach();
+  return Error::success();
+}
+
+Error ThinLtoJIT::setupJITDylib(JITDylib *JD, bool AllowNudge) {
+  // Register symbols for C++ static destructors.
+  LocalCXXRuntimeOverrides CXXRuntimeoverrides;
+  Error Err = CXXRuntimeoverrides.enable(*JD, *Mangle.Impl);
+  if (Err)
+    return Err;
+
+  // Lookup symbol names in the global ThinLTO module index first
+  char Prefix = DL->getGlobalPrefix();
+  JD->addGenerator(std::make_unique<ThinLtoDefinitionGenerator>(
+      *GlobalIndex, *InstrumentationLayer, AddModule, Prefix, AllowNudge));
+  // Then try lookup in the host process.
+  auto HostLookup = DynamicLibrarySearchGenerator::GetForCurrentProcess(Prefix);
+  if (!HostLookup)
+    return HostLookup.takeError();
+  JD->addGenerator(std::move(*HostLookup));
+
+  return Error::success();
+}
+
+ThinLtoJIT::~ThinLtoJIT() {
+  // Signal the DiscoveryThread to shut down.
+  JitRunning.store(false);
+  // Wait for potential compile actions to finish.
+  CompileThreads->wait();
+}
+
+} // namespace orc
+} // namespace llvm

llvm/examples/ThinLtoJIT/ThinLtoModuleIndex.h

+#ifndef LLVM_EXAMPLES_THINLTOJIT_THINLTOJITMODULEINDEX_H
+#define LLVM_EXAMPLES_THINLTOJIT_THINLTOJITMODULEINDEX_H
+
+#include "llvm/ADT/Optional.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
+#include "llvm/IR/GlobalValue.h"
+#include "llvm/IR/ModuleSummaryIndex.h"
+#include "llvm/Support/Error.h"
+
+#include <cstdint>
+#include <mutex>
+#include <set>
+#include <vector>
+
+namespace llvm {
+namespace orc {
+
+class SymbolStringPtr;
+
+class ThinLtoModuleIndex {
+  static constexpr bool HaveGVs = false;
+
+public:
+  ThinLtoModuleIndex(ExecutionSession &ES)
+      : ES(ES), CombinedSummaryIndex(HaveGVs), NextModuleId(0) {}
+
+  Error add(StringRef ModulePath);
+  GlobalValueSummary *getSummary(GlobalValue::GUID Function) const;
+  std::vector<StringRef> getAllModulePaths() const;
+  Optional<StringRef> getModulePathForSymbol(StringRef Name) const;
+  Expected<Optional<ThreadSafeModule>>
+  parseNewModuleFromFile(StringRef Path, bool ForceLoad = false);
+
+private:
+  ExecutionSession &ES;
+  ModuleSummaryIndex CombinedSummaryIndex;
+  uint64_t NextModuleId;
+
+  std::mutex ParsedModulesLock;
+  std::set<SymbolStringPtr> ParsedModules;
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif

llvm/examples/ThinLtoJIT/ThinLtoModuleIndex.cpp

+#include "ThinLtoModuleIndex.h"
+
+#include "llvm/Bitcode/BitcodeReader.h"
+#include "llvm/ExecutionEngine/Orc/SymbolStringPool.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <memory>
+#include <string>
+
+#define DEBUG_TYPE "thinltojit"
+
+namespace llvm {
+namespace orc {
+
+Error ThinLtoModuleIndex::add(StringRef ModulePath) {
+  auto Buffer = errorOrToExpected(MemoryBuffer::getFile(ModulePath));
+  if (!Buffer)
+    return Buffer.takeError();
+
+  Error ParseErr = readModuleSummaryIndex((*Buffer)->getMemBufferRef(),
+                                          CombinedSummaryIndex, NextModuleId);
+  if (ParseErr)
+    return ParseErr;
+
+  ++NextModuleId;
+  return Error::success();
+}
+
+std::vector<StringRef> ThinLtoModuleIndex::getAllModulePaths() const {
+  std::vector<StringRef> Paths;
+  for (const auto &KV : CombinedSummaryIndex.modulePaths()) {
+    StringRef Path = KV.first();
+    Paths.push_back(Path);
+  }
+  return Paths;
+}
+
+GlobalValueSummary *
+ThinLtoModuleIndex::getSummary(GlobalValue::GUID Function) const {
+  ValueInfo VI = CombinedSummaryIndex.getValueInfo(Function);
+  if (!VI || VI.getSummaryList().empty())
+    return nullptr;
+#ifndef NDEBUG
+  if (VI.getSummaryList().size() > 1) {
+    LLVM_DEBUG(dbgs() << "SummaryList with multiple entries!\n");
+  }
+#endif
+  return VI.getSummaryList().front().get()->getBaseObject();
+}
+
+Optional<StringRef>
+ThinLtoModuleIndex::getModulePathForSymbol(StringRef Name) const {
+  if (GlobalValueSummary *S = getSummary(GlobalValue::getGUID(Name)))
+    return S->modulePath();
+  return None; // We don't know the symbol.
+}
+
+Expected<Optional<ThreadSafeModule>>
+ThinLtoModuleIndex::parseNewModuleFromFile(StringRef Path, bool ForceLoad) {
+  if (!ForceLoad) {
+    std::lock_guard<std::mutex> Lock(ParsedModulesLock);
+
+    SymbolStringPtr PathId = ES.intern(Path);
+    auto It = ParsedModules.find(PathId);
+    if (It != ParsedModules.end() && !ForceLoad)
+      // This is not a new module.
+      return None;
+
+    ParsedModules.insert(PathId);
+  }
+
+  // TODO: make a SMDiagnosticError class for this
+  SMDiagnostic Err;
+  auto Ctx = std::make_unique<LLVMContext>();
+  auto M = parseIRFile(Path, Err, *Ctx);
+  if (!M) {
+    std::string ErrDescription;
+    {
+      raw_string_ostream S(ErrDescription);
+      Err.print("ThinLtoJIT", S);
+    }
+    return createStringError(inconvertibleErrorCode(),
+                             "Failed load module from file '%s' (%s)",
+                             Path.data(), ErrDescription.c_str());
+  }
+
+  return ThreadSafeModule(std::move(M), std::move(Ctx));
+}
+
+} // namespace orc
+} // namespace llvm

llvm/examples/ThinLtoJIT/bench

+#!/bin/bash
+#set -x
+
+if [ $# -gt 2 ]; then
+  TOOLS_DIR="$1"
+  SOURCE_DIR="$2"
+  MAIN_SOURCE_FILE="$3"
+else
+  echo "Usage: bench <path to llvm binaries> <path to c-sources> <main source file> [<override sysroot>]"
+  exit 1
+fi
+
+if [ $# -gt 3 ]; then
+  SYS_ROOT="$4"
+else
+  SYS_ROOT="/"
+fi
+
+function check_tool ()
+{
+  if [ -e "${TOOLS_DIR}/$1" ]; then
+    echo "Found: $1"
+  else
+    echo "!!! Cannot find required tool, please provide it in the LLVM binaries folder: $1"
+  fi
+}
+
+check_tool lli
+check_tool SpeculativeJIT
+check_tool ThinLtoJIT
+
+SKIP_BITCODE_GEN=0
+if [[ -e bc-default || -e bc-thinlto || -e ll-default || -e ll-thinlto ]]; then
+  echo "Skipping bitcode generation: output directories existing"
+  echo "Please clean up manually: rm -R bc-default bc-thinlto ll-default ll-thinlto"
+  SKIP_BITCODE_GEN=1
+else
+  check_tool clang
+  check_tool llvm-dis
+  mkdir bc-default
+  mkdir bc-thinlto
+  mkdir ll-default
+  mkdir ll-thinlto
+fi
+
+ROOT_DIR=$(pwd)
+ALL_BITCODE_FILES=""
+
+MAIN_FILE_BASENAME=$(basename "${MAIN_SOURCE_FILE%.c*}")
+LLI_EXTRA_MODULES=""
+
+for f in ${SOURCE_DIR}/*.c* ; do
+  BASE_NAME=$(basename "${f%.c*}")
+
+  if [ ${SKIP_BITCODE_GEN} -eq 0 ]; then
+    echo "Compile: $f -> ${BASE_NAME}.bc"
+
+    "${TOOLS_DIR}/clang" -c -I ${SOURCE_DIR} -isysroot ${SYS_ROOT} -emit-llvm \
+                        -o "bc-default/${BASE_NAME}.bc" "$f"
+    "${TOOLS_DIR}/clang" -c -I ${SOURCE_DIR} -isysroot ${SYS_ROOT} -flto=thin \
+                        -o "bc-thinlto/${BASE_NAME}.bc" "$f"
+
+    echo "Disassemble ${BASE_NAME}.bc -> ${BASE_NAME}.ll"
+    ${TOOLS_DIR}/llvm-dis bc-default/${BASE_NAME}.bc -o ll-default/${BASE_NAME}.ll
+    ${TOOLS_DIR}/llvm-dis bc-thinlto/${BASE_NAME}.bc -o ll-thinlto/${BASE_NAME}.ll
+  fi
+
+  ALL_BITCODE_FILES="${ALL_BITCODE_FILES} ${BASE_NAME}.bc"
+  if [ "${BASE_NAME}" != "${MAIN_FILE_BASENAME}" ]; then
+    LLI_EXTRA_MODULES="${LLI_EXTRA_MODULES} -extra-module=${BASE_NAME}.bc"
+  fi
+done
+
+set -x
+cd ${ROOT_DIR}/bc-default
+time (${TOOLS_DIR}/clang -o ${MAIN_FILE_BASENAME} -O0 ${ALL_BITCODE_FILES} && ./${MAIN_FILE_BASENAME} 1>/dev/null)
+time ${TOOLS_DIR}/lli ${LLI_EXTRA_MODULES} -jit-kind=mcjit "${MAIN_FILE_BASENAME}.bc" 1>/dev/null
+time ${TOOLS_DIR}/lli ${LLI_EXTRA_MODULES} -jit-kind=orc-mcjit "${MAIN_FILE_BASENAME}.bc" 1>/dev/null
+time ${TOOLS_DIR}/lli ${LLI_EXTRA_MODULES} -jit-kind=orc-lazy "${MAIN_FILE_BASENAME}.bc" 1>/dev/null
+time ${TOOLS_DIR}/lli ${LLI_EXTRA_MODULES} -jit-kind=orc-lazy -compile-threads=2 "${MAIN_FILE_BASENAME}.bc" 1>/dev/null
+time ${TOOLS_DIR}/lli ${LLI_EXTRA_MODULES} -jit-kind=orc-lazy -per-module-lazy "${MAIN_FILE_BASENAME}.bc" 1>/dev/null
+time ${TOOLS_DIR}/lli ${LLI_EXTRA_MODULES} -jit-kind=orc-lazy -per-module-lazy -compile-threads=2 "${MAIN_FILE_BASENAME}.bc" 1>/dev/null
+time ${TOOLS_DIR}/lli ${LLI_EXTRA_MODULES} -jit-kind=orc-lazy -per-module-lazy -compile-threads=2 -O1 "${MAIN_FILE_BASENAME}.bc" 1>/dev/null
+time ${TOOLS_DIR}/lli ${LLI_EXTRA_MODULES} -jit-kind=orc-lazy -per-module-lazy -compile-threads=2 -O0 "${MAIN_FILE_BASENAME}.bc" 1>/dev/null
+time ${TOOLS_DIR}/SpeculativeJIT -num-threads=2 ${ALL_BITCODE_FILES} 1>/dev/null
+
+cd ${ROOT_DIR}/bc-thinlto
+#time (${TOOLS_DIR}/clang -flto=thin -o test ${ALL_BITCODE_FILES} && ./test 1>/dev/null)
+time ${TOOLS_DIR}/ThinLtoJIT -compile-threads=2 ${ALL_BITCODE_FILES} 1>/dev/null

llvm/examples/ThinLtoJIT/main.cpp

+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/TargetSelect.h"
+
+#include "ThinLtoInstrumentationLayer.h"
+#include "ThinLtoJIT.h"
+
+#include <string>
+#include <vector>
+
+using namespace llvm;
+
+static cl::list<std::string> InputFiles(cl::Positional, cl::OneOrMore,
+                                        cl::desc("<bitcode files>"));
+
+static cl::list<std::string> InputArgs("args", cl::Positional,
+                                       cl::desc("<program arguments>..."),
+                                       cl::ZeroOrMore, cl::PositionalEatsArgs);
+
+static cl::opt<unsigned> CompileThreads("compile-threads", cl::Optional,
+                                        cl::desc("Number of compile threads"),
+                                        cl::init(4));
+
+static cl::opt<unsigned> DiscoveryFlagsBatchSize(
+    "discovery-flag-batch-size", cl::Optional,
+    cl::desc("Number of discovery flags allocated in one go"), cl::init(4096));
+
+static cl::opt<orc::ThinLtoInstrumentationLayer::ExplicitMemoryBarrier>
+    MemFence(
+        "mem-fence",
+        cl::desc(
+            "Choose where to install memory fences for cache synchronization"),
+        cl::init(orc::ThinLtoInstrumentationLayer::Always),
+        cl::values(clEnumValN(orc::ThinLtoInstrumentationLayer::Never, "never",
+                              "No use of memory fences"),
+                   clEnumValN(orc::ThinLtoInstrumentationLayer::StaticCode,
+                              "static",
+                              "Use of memory fences in static code only"),
+                   clEnumValN(orc::ThinLtoInstrumentationLayer::JITedCode,
+                              "jited",
+                              "Install memory fences in JITed code only"),
+                   clEnumValN(orc::ThinLtoInstrumentationLayer::Always,
+                              "always", "Always use of memory fences")));
+
+static cl::opt<bool> AllowNudge("allow-nudge",
+                                cl::desc("Allow the symbol generator to nudge symbols into discovery although they haven't been reached"),
+                                cl::init(true));
+
+int main(int argc, char *argv[]) {
+  InitLLVM X(argc, argv);
+  InitializeNativeTarget();
+  InitializeNativeTargetAsmPrinter();
+  cl::ParseCommandLineOptions(argc, argv, "ThinLtoJIT");
+
+  std::vector<StringRef> ModuleFiles;
+  for (const std::string &F : InputFiles) {
+    ModuleFiles.push_back(F);
+  }
+
+  Error Err = Error::success();
+  unsigned Threads = std::max(1u, static_cast<unsigned>(CompileThreads));
+
+  orc::ThinLtoJIT Jit(std::move(ModuleFiles), "main", Threads,
+                      DiscoveryFlagsBatchSize, MemFence, AllowNudge, Err);
+  if (Err) {
+    logAllUnhandledErrors(std::move(Err), errs(), "ThinLtoJIT: ");
+    exit(1);
+  }
+
+  unsigned JitMainArgc = InputArgs.size();
+  const char *JitMainArgv[JitMainArgc];
+  for (unsigned i = 0; i < JitMainArgc; i++) {
+    JitMainArgv[i] = InputArgs[i].c_str();
+  }
+
+  ExitOnError ExitOnErr;
+  ExitOnErr.setBanner("ThinLtoJIT: ");
+
+  return ExitOnErr(Jit.main(JitMainArgc, JitMainArgv));
+}