docs/tutorial/BuildingAJIT3.rst

   1 =============================================
   2 Building a JIT: Per-function Lazy Compilation
   3 =============================================
   4
   5 .. contents::
   6    :local:
   7
   8 **This tutorial is under active development. It is incomplete and details may
   9 change frequently.** Nonetheless we invite you to try it out as it stands, and
  10 we welcome any feedback.
  11
  12 Chapter 3 Introduction
  13 ======================
  14
  15 Welcome to Chapter 3 of the "Building an ORC-based JIT in LLVM" tutorial. This
  16 chapter discusses lazy JITing and shows you how to enable it by adding an ORC
  17 CompileOnDemand layer the JIT from `Chapter 2 <BuildingAJIT2.html>`_.
  18
  19 Lazy Compilation
  20 ================
  21
  22 When we add a module to the KaleidoscopeJIT class from Chapter 2 it is
  23 immediately optimized, compiled and linked for us by the IRTransformLayer,
  24 IRCompileLayer and RTDyldObjectLinkingLayer respectively. This scheme, where all the
  25 work to make a Module executable is done up front, is simple to understand and
  26 its performance characteristics are easy to reason about. However, it will lead
  27 to very high startup times if the amount of code to be compiled is large, and
  28 may also do a lot of unnecessary compilation if only a few compiled functions
  29 are ever called at runtime. A truly "just-in-time" compiler should allow us to
  30 defer the compilation of any given function until the moment that function is
  31 first called, improving launch times and eliminating redundant work. In fact,
  32 the ORC APIs provide us with a layer to lazily compile LLVM IR:
  33 *CompileOnDemandLayer*.
  34
  35 The CompileOnDemandLayer class conforms to the layer interface described in
  36 Chapter 2, but its addModule method behaves quite differently from the layers
  37 we have seen so far: rather than doing any work up front, it just scans the
  38 Modules being added and arranges for each function in them to be compiled the
  39 first time it is called. To do this, the CompileOnDemandLayer creates two small
  40 utilities for each function that it scans: a *stub* and a *compile
  41 callback*. The stub is a pair of a function pointer (which will be pointed at
  42 the function's implementation once the function has been compiled) and an
  43 indirect jump through the pointer. By fixing the address of the indirect jump
  44 for the lifetime of the program we can give the function a permanent "effective
  45 address", one that can be safely used for indirection and function pointer
  46 comparison even if the function's implementation is never compiled, or if it is
  47 compiled more than once (due to, for example, recompiling the function at a
  48 higher optimization level) and changes address. The second utility, the compile
  49 callback, represents a re-entry point from the program into the compiler that
  50 will trigger compilation and then execution of a function. By initializing the
  51 function's stub to point at the function's compile callback, we enable lazy
  52 compilation: The first attempted call to the function will follow the function
  53 pointer and trigger the compile callback instead. The compile callback will
  54 compile the function, update the function pointer for the stub, then execute
  55 the function. On all subsequent calls to the function, the function pointer
  56 will point at the already-compiled function, so there is no further overhead
  57 from the compiler. We will look at this process in more detail in the next
  58 chapter of this tutorial, but for now we'll trust the CompileOnDemandLayer to
  59 set all the stubs and callbacks up for us. All we need to do is to add the
  60 CompileOnDemandLayer to the top of our stack and we'll get the benefits of
  61 lazy compilation. We just need a few changes to the source:
  62
  63 .. code-block:: c++
  64
  65   ...
  66   #include "llvm/ExecutionEngine/SectionMemoryManager.h"
  67   #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
  68   #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
  69   ...
  70
  71   ...
  72   class KaleidoscopeJIT {
  73   private:
  74     std::unique_ptr<TargetMachine> TM;
  75     const DataLayout DL;
  76     RTDyldObjectLinkingLayer ObjectLayer;
  77     IRCompileLayer<decltype(ObjectLayer), SimpleCompiler> CompileLayer;
  78
  79     using OptimizeFunction =
  80         std::function<std::shared_ptr<Module>(std::shared_ptr<Module>)>;
  81
  82     IRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer;
  83
  84     std::unique_ptr<JITCompileCallbackManager> CompileCallbackManager;
  85     CompileOnDemandLayer<decltype(OptimizeLayer)> CODLayer;
  86
  87   public:
  88     using ModuleHandle = decltype(CODLayer)::ModuleHandleT;
  89
  90 First we need to include the CompileOnDemandLayer.h header, then add two new
  91 members: a std::unique_ptr<JITCompileCallbackManager> and a CompileOnDemandLayer,
  92 to our class. The CompileCallbackManager member is used by the CompileOnDemandLayer
  93 to create the compile callback needed for each function.
  94
  95 .. code-block:: c++
  96
  97   KaleidoscopeJIT()
  98       : TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
  99         ObjectLayer([]() { return std::make_shared<SectionMemoryManager>(); }),
 100         CompileLayer(ObjectLayer, SimpleCompiler(*TM)),
 101         OptimizeLayer(CompileLayer,
 102                       [this](std::shared_ptr<Module> M) {
 103                         return optimizeModule(std::move(M));
 104                       }),
 105         CompileCallbackManager(
 106             orc::createLocalCompileCallbackManager(TM->getTargetTriple(), 0)),
 107         CODLayer(OptimizeLayer,
 108                  [this](Function &F) { return std::set<Function*>({&F}); },
 109                  *CompileCallbackManager,
 110                  orc::createLocalIndirectStubsManagerBuilder(
 111                    TM->getTargetTriple())) {
 112     llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
 113   }
 114
 115 Next we have to update our constructor to initialize the new members. To create
 116 an appropriate compile callback manager we use the
 117 createLocalCompileCallbackManager function, which takes a TargetMachine and a
 118 JITTargetAddress to call if it receives a request to compile an unknown
 119 function.  In our simple JIT this situation is unlikely to come up, so we'll
 120 cheat and just pass '0' here. In a production quality JIT you could give the
 121 address of a function that throws an exception in order to unwind the JIT'd
 122 code's stack.
 123
 124 Now we can construct our CompileOnDemandLayer. Following the pattern from
 125 previous layers we start by passing a reference to the next layer down in our
 126 stack -- the OptimizeLayer. Next we need to supply a 'partitioning function':
 127 when a not-yet-compiled function is called, the CompileOnDemandLayer will call
 128 this function to ask us what we would like to compile. At a minimum we need to
 129 compile the function being called (given by the argument to the partitioning
 130 function), but we could also request that the CompileOnDemandLayer compile other
 131 functions that are unconditionally called (or highly likely to be called) from
 132 the function being called. For KaleidoscopeJIT we'll keep it simple and just
 133 request compilation of the function that was called. Next we pass a reference to
 134 our CompileCallbackManager. Finally, we need to supply an "indirect stubs
 135 manager builder": a utility function that constructs IndirectStubManagers, which
 136 are in turn used to build the stubs for the functions in each module. The
 137 CompileOnDemandLayer will call the indirect stub manager builder once for each
 138 call to addModule, and use the resulting indirect stubs manager to create
 139 stubs for all functions in all modules in the set. If/when the module set is
 140 removed from the JIT the indirect stubs manager will be deleted, freeing any
 141 memory allocated to the stubs. We supply this function by using the
 142 createLocalIndirectStubsManagerBuilder utility.
 143
 144 .. code-block:: c++
 145
 146   // ...
 147           if (auto Sym = CODLayer.findSymbol(Name, false))
 148   // ...
 149   return cantFail(CODLayer.addModule(std::move(Ms),
 150                                      std::move(Resolver)));
 151   // ...
 152
 153   // ...
 154   return CODLayer.findSymbol(MangledNameStream.str(), true);
 155   // ...
 156
 157   // ...
 158   CODLayer.removeModule(H);
 159   // ...
 160
 161 Finally, we need to replace the references to OptimizeLayer in our addModule,
 162 findSymbol, and removeModule methods. With that, we're up and running.
 163
 164 **To be done:**
 165
 166 ** Chapter conclusion.**
 167
 168 Full Code Listing
 169 =================
 170
 171 Here is the complete code listing for our running example with a CompileOnDemand
 172 layer added to enable lazy function-at-a-time compilation. To build this example, use:
 173
 174 .. code-block:: bash
 175
 176     # Compile
 177     clang++ -g toy.cpp `llvm-config --cxxflags --ldflags --system-libs --libs core orcjit native` -O3 -o toy
 178     # Run
 179     ./toy
 180
 181 Here is the code:
 182
 183 .. literalinclude:: ../../examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h
 184    :language: c++
 185
 186 `Next: Extreme Laziness -- Using Compile Callbacks to JIT directly from ASTs <BuildingAJIT4.html>`_