277 lines · plain
1=====================================================================2Building a JIT: Adding Optimizations -- An introduction to ORC Layers3=====================================================================4 5.. contents::6 :local:7 8**This tutorial is under active development. It is incomplete and details may9change frequently.** Nonetheless we invite you to try it out as it stands, and10we welcome any feedback.11 12Chapter 2 Introduction13======================14 15**Warning: This tutorial is currently being updated to account for ORC API16changes. Only Chapters 1 and 2 are up-to-date.**17 18**Example code from Chapters 3 to 5 will compile and run, but has not been19updated**20 21Welcome to Chapter 2 of the "Building an ORC-based JIT in LLVM" tutorial. In22`Chapter 1 <BuildingAJIT1.html>`_ of this series we examined a basic JIT23class, KaleidoscopeJIT, that could take LLVM IR modules as input and produce24executable code in memory. KaleidoscopeJIT was able to do this with relatively25little code by composing two off-the-shelf *ORC layers*: IRCompileLayer and26ObjectLinkingLayer, to do much of the heavy lifting.27 28In this layer we'll learn more about the ORC layer concept by using a new layer,29IRTransformLayer, to add IR optimization support to KaleidoscopeJIT.30 31Optimizing Modules using the IRTransformLayer32=============================================33 34In `Chapter 4 <LangImpl04.html>`_ of the "Implementing a language with LLVM"35tutorial series the llvm *FunctionPassManager* is introduced as a means for36optimizing LLVM IR. Interested readers may read that chapter for details, but37in short: to optimize a Module we create an llvm::FunctionPassManager38instance, configure it with a set of optimizations, then run the PassManager on39a Module to mutate it into a (hopefully) more optimized but semantically40equivalent form. In the original tutorial series the FunctionPassManager was41created outside the KaleidoscopeJIT and modules were optimized before being42added to it. In this Chapter we will make optimization a phase of our JIT43instead. For now this will provide us a motivation to learn more about ORC44layers, but in the long term making optimization part of our JIT will yield an45important benefit: When we begin lazily compiling code (i.e. deferring46compilation of each function until the first time it's run) having47optimization managed by our JIT will allow us to optimize lazily too, rather48than having to do all our optimization up-front.49 50To add optimization support to our JIT we will take the KaleidoscopeJIT from51Chapter 1 and compose an ORC *IRTransformLayer* on top. We will look at how the52IRTransformLayer works in more detail below, but the interface is simple: the53constructor for this layer takes a reference to the execution session and the54layer below (as all layers do) plus an *IR optimization function* that it will55apply to each Module that is added via addModule:56 57.. code-block:: c++58 59 class KaleidoscopeJIT {60 private:61 ExecutionSession ES;62 RTDyldObjectLinkingLayer ObjectLayer;63 IRCompileLayer CompileLayer;64 IRTransformLayer TransformLayer;65 66 DataLayout DL;67 MangleAndInterner Mangle;68 ThreadSafeContext Ctx;69 70 public:71 72 KaleidoscopeJIT(JITTargetMachineBuilder JTMB, DataLayout DL)73 : ObjectLayer(ES,74 []() { return std::make_unique<SectionMemoryManager>(); }),75 CompileLayer(ES, ObjectLayer, ConcurrentIRCompiler(std::move(JTMB))),76 TransformLayer(ES, CompileLayer, optimizeModule),77 DL(std::move(DL)), Mangle(ES, this->DL),78 Ctx(std::make_unique<LLVMContext>()) {79 ES.getMainJITDylib().addGenerator(80 cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(DL.getGlobalPrefix())));81 }82 83Our extended KaleidoscopeJIT class starts out the same as it did in Chapter 1,84but after the CompileLayer we introduce a new member, TransformLayer, which sits85on top of our CompileLayer. We initialize our OptimizeLayer with a reference to86the ExecutionSession and output layer (standard practice for layers), along with87a *transform function*. For our transform function we supply our classes88optimizeModule static method.89 90.. code-block:: c++91 92 // ...93 return cantFail(OptimizeLayer.addModule(std::move(M),94 std::move(Resolver)));95 // ...96 97Next we need to update our addModule method to replace the call to98``CompileLayer::add`` with a call to ``OptimizeLayer::add`` instead.99 100.. code-block:: c++101 102 static Expected<ThreadSafeModule>103 optimizeModule(ThreadSafeModule M, const MaterializationResponsibility &R) {104 // Create a function pass manager.105 auto FPM = std::make_unique<legacy::FunctionPassManager>(M.get());106 107 // Add some optimizations.108 FPM->add(createInstructionCombiningPass());109 FPM->add(createReassociatePass());110 FPM->add(createGVNPass());111 FPM->add(createCFGSimplificationPass());112 FPM->doInitialization();113 114 // Run the optimizations over all functions in the module being added to115 // the JIT.116 for (auto &F : *M)117 FPM->run(F);118 119 return M;120 }121 122At the bottom of our JIT we add a private method to do the actual optimization:123*optimizeModule*. This function takes the module to be transformed as input (as124a ThreadSafeModule) along with a reference to a reference to a new class:125``MaterializationResponsibility``. The MaterializationResponsibility argument126can be used to query JIT state for the module being transformed, such as the set127of definitions in the module that JIT'd code is actively trying to call/access.128For now we will ignore this argument and use a standard optimization129pipeline. To do this we set up a FunctionPassManager, add some passes to it, run130it over every function in the module, and then return the mutated module. The131specific optimizations are the same ones used in `Chapter 4 <LangImpl04.html>`_132of the "Implementing a language with LLVM" tutorial series. Readers may visit133that chapter for a more in-depth discussion of these, and of IR optimization in134general.135 136And that's it in terms of changes to KaleidoscopeJIT: When a module is added via137addModule the OptimizeLayer will call our optimizeModule function before passing138the transformed module on to the CompileLayer below. Of course, we could have139called optimizeModule directly in our addModule function and not gone to the140bother of using the IRTransformLayer, but doing so gives us another opportunity141to see how layers compose. It also provides a neat entry point to the *layer*142concept itself, because IRTransformLayer is one of the simplest layers that143can be implemented.144 145.. code-block:: c++146 147 // From IRTransformLayer.h:148 class IRTransformLayer : public IRLayer {149 public:150 using TransformFunction = std::function<Expected<ThreadSafeModule>(151 ThreadSafeModule, const MaterializationResponsibility &R)>;152 153 IRTransformLayer(ExecutionSession &ES, IRLayer &BaseLayer,154 TransformFunction Transform = identityTransform);155 156 void setTransform(TransformFunction Transform) {157 this->Transform = std::move(Transform);158 }159 160 static ThreadSafeModule161 identityTransform(ThreadSafeModule TSM,162 const MaterializationResponsibility &R) {163 return TSM;164 }165 166 void emit(MaterializationResponsibility R, ThreadSafeModule TSM) override;167 168 private:169 IRLayer &BaseLayer;170 TransformFunction Transform;171 };172 173 // From IRTransformLayer.cpp:174 175 IRTransformLayer::IRTransformLayer(ExecutionSession &ES,176 IRLayer &BaseLayer,177 TransformFunction Transform)178 : IRLayer(ES), BaseLayer(BaseLayer), Transform(std::move(Transform)) {}179 180 void IRTransformLayer::emit(MaterializationResponsibility R,181 ThreadSafeModule TSM) {182 assert(TSM.getModule() && "Module must not be null");183 184 if (auto TransformedTSM = Transform(std::move(TSM), R))185 BaseLayer.emit(std::move(R), std::move(*TransformedTSM));186 else {187 R.failMaterialization();188 getExecutionSession().reportError(TransformedTSM.takeError());189 }190 }191 192This is the whole definition of IRTransformLayer, from193``llvm/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h`` and194``llvm/lib/ExecutionEngine/Orc/IRTransformLayer.cpp``. This class is concerned195with two very simple jobs: (1) Running every IR Module that is emitted via this196layer through the transform function object, and (2) implementing the ORC197``IRLayer`` interface (which itself conforms to the general ORC Layer concept,198more on that below). Most of the class is straightforward: a typedef for the199transform function, a constructor to initialize the members, a setter for the200transform function value, and a default no-op transform. The most important201method is ``emit`` as this is half of our IRLayer interface. The emit method202applies our transform to each module that it is called on and, if the transform203succeeds, passes the transformed module to the base layer. If the transform204fails, our emit function calls205``MaterializationResponsibility::failMaterialization`` (this JIT clients who206may be waiting on other threads know that the code they were waiting for has207failed to compile) and logs the error with the execution session before bailing208out.209 210The other half of the IRLayer interface we inherit unmodified from the IRLayer211class:212 213.. code-block:: c++214 215 Error IRLayer::add(JITDylib &JD, ThreadSafeModule TSM, VModuleKey K) {216 return JD.define(std::make_unique<BasicIRLayerMaterializationUnit>(217 *this, std::move(K), std::move(TSM)));218 }219 220This code, from ``llvm/lib/ExecutionEngine/Orc/Layer.cpp``, adds a221ThreadSafeModule to a given JITDylib by wrapping it up in a222``MaterializationUnit`` (in this case a ``BasicIRLayerMaterializationUnit``).223Most layers that derived from IRLayer can rely on this default implementation224of the ``add`` method.225 226These two operations, ``add`` and ``emit``, together constitute the layer227concept: A layer is a way to wrap a part of a compiler pipeline (in this case228the "opt" phase of an LLVM compiler) whose API is opaque to ORC with an229interface that ORC can call as needed. The add method takes an230module in some input program representation (in this case an LLVM IR module)231and stores it in the target ``JITDylib``, arranging for it to be passed back232to the layer's emit method when any symbol defined by that module is requested.233Each layer can complete its own work by calling the ``emit`` method of its base234layer. For example, in this tutorial our IRTransformLayer calls through to235our IRCompileLayer to compile the transformed IR, and our IRCompileLayer in236turn calls our ObjectLayer to link the object file produced by our compiler.237 238So far we have learned how to optimize and compile our LLVM IR, but we have239not focused on when compilation happens. Our current REPL optimizes and240compiles each function as soon as it is referenced by any other code,241regardless of whether it is ever called at runtime. In the next chapter we242will introduce a fully lazy compilation, in which functions are not compiled243until they are first called at run-time. At this point the trade-offs get much244more interesting: the lazier we are, the quicker we can start executing the245first function, but the more often we will have to pause to compile newly246encountered functions. If we only code-gen lazily, but optimize eagerly, we247will have a longer startup time (as everything is optimized at that time) but248relatively short pauses as each function just passes through code-gen. If we249both optimize and code-gen lazily we can start executing the first function250more quickly, but we will have longer pauses as each function has to be both251optimized and code-gen'd when it is first executed. Things become even more252interesting if we consider interprocedural optimizations like inlining, which253must be performed eagerly. These are complex trade-offs, and there is no254one-size-fits all solution to them, but by providing composable layers we leave255the decisions to the person implementing the JIT, and make it easy for them to256experiment with different configurations.257 258`Next: Adding Per-function Lazy Compilation <BuildingAJIT3.html>`_259 260Full Code Listing261=================262 263Here is the complete code listing for our running example with an264IRTransformLayer added to enable optimization. To build this example, use:265 266.. code-block:: bash267 268 # Compile269 clang++ -g toy.cpp `llvm-config --cxxflags --ldflags --system-libs --libs core orcjit native` -O3 -o toy270 # Run271 ./toy272 273Here is the code:274 275.. literalinclude:: ../../examples/Kaleidoscope/BuildingAJIT/Chapter2/KaleidoscopeJIT.h276 :language: c++277