702 lines · plain
1==============================================2Kaleidoscope: Adding JIT and Optimizer Support3==============================================4 5.. contents::6 :local:7 8Chapter 4 Introduction9======================10 11Welcome to Chapter 4 of the "`Implementing a language with12LLVM <index.html>`_" tutorial. Chapters 1-3 described the implementation13of a simple language and added support for generating LLVM IR. This14chapter describes two new techniques: adding optimizer support to your15language, and adding JIT compiler support. These additions will16demonstrate how to get nice, efficient code for the Kaleidoscope17language.18 19Trivial Constant Folding20========================21 22Our demonstration for Chapter 3 is elegant and easy to extend.23Unfortunately, it does not produce wonderful code. The IRBuilder,24however, does give us obvious optimizations when compiling simple code:25 26::27 28 ready> def test(x) 1+2+x;29 Read function definition:30 define double @test(double %x) {31 entry:32 %addtmp = fadd double 3.000000e+00, %x33 ret double %addtmp34 }35 36This code is not a literal transcription of the AST built by parsing the37input. That would be:38 39::40 41 ready> def test(x) 1+2+x;42 Read function definition:43 define double @test(double %x) {44 entry:45 %addtmp = fadd double 2.000000e+00, 1.000000e+0046 %addtmp1 = fadd double %addtmp, %x47 ret double %addtmp148 }49 50Constant folding, as seen above, in particular, is a very common and51very important optimization: so much so that many language implementors52implement constant folding support in their AST representation.53 54With LLVM, you don't need this support in the AST. Since all calls to55build LLVM IR go through the LLVM IR builder, the builder itself checked56to see if there was a constant folding opportunity when you call it. If57so, it just does the constant fold and return the constant instead of58creating an instruction.59 60Well, that was easy :). In practice, we recommend always using61``IRBuilder`` when generating code like this. It has no "syntactic62overhead" for its use (you don't have to uglify your compiler with63constant checks everywhere) and it can dramatically reduce the amount of64LLVM IR that is generated in some cases (particular for languages with a65macro preprocessor or that use a lot of constants).66 67On the other hand, the ``IRBuilder`` is limited by the fact that it does68all of its analysis inline with the code as it is built. If you take a69slightly more complex example:70 71::72 73 ready> def test(x) (1+2+x)*(x+(1+2));74 ready> Read function definition:75 define double @test(double %x) {76 entry:77 %addtmp = fadd double 3.000000e+00, %x78 %addtmp1 = fadd double %x, 3.000000e+0079 %multmp = fmul double %addtmp, %addtmp180 ret double %multmp81 }82 83In this case, the LHS and RHS of the multiplication are the same value.84We'd really like to see this generate "``tmp = x+3; result = tmp*tmp;``"85instead of computing "``x+3``" twice.86 87Unfortunately, no amount of local analysis will be able to detect and88correct this. This requires two transformations: reassociation of89expressions (to make the adds lexically identical) and Common90Subexpression Elimination (CSE) to delete the redundant add instruction.91Fortunately, LLVM provides a broad range of optimizations that you can92use, in the form of "passes".93 94LLVM Optimization Passes95========================96 97LLVM provides many optimization passes, which do many different sorts of98things and have different tradeoffs. Unlike other systems, LLVM doesn't99hold to the mistaken notion that one set of optimizations is right for100all languages and for all situations. LLVM allows a compiler implementor101to make complete decisions about what optimizations to use, in which102order, and in what situation.103 104As a concrete example, LLVM supports both "whole module" passes, which105look across as large of body of code as they can (often a whole file,106but if run at link time, this can be a substantial portion of the whole107program). It also supports and includes "per-function" passes which just108operate on a single function at a time, without looking at other109functions. For more information on passes and how they are run, see the110`How to Write a Pass <../../WritingAnLLVMPass.html>`_ document and the111`List of LLVM Passes <../../Passes.html>`_.112 113For Kaleidoscope, we are currently generating functions on the fly, one114at a time, as the user types them in. We aren't shooting for the115ultimate optimization experience in this setting, but we also want to116catch the easy and quick stuff where possible. As such, we will choose117to run a few per-function optimizations as the user types the function118in. If we wanted to make a "static Kaleidoscope compiler", we would use119exactly the code we have now, except that we would defer running the120optimizer until the entire file has been parsed.121 122In addition to the distinction between function and module passes, passes can be123divided into transform and analysis passes. Transform passes mutate the IR, and124analysis passes compute information that other passes can use. In order to add125a transform pass, all analysis passes it depends upon must be registered in126advance.127 128In order to get per-function optimizations going, we need to set up a129`FunctionPassManager <../../WritingAnLLVMPass.html#what-passmanager-doesr>`_ to hold130and organize the LLVM optimizations that we want to run. Once we have131that, we can add a set of optimizations to run. We'll need a new132FunctionPassManager for each module that we want to optimize, so we'll133add to a function created in the previous chapter (``InitializeModule()``):134 135.. code-block:: c++136 137 void InitializeModuleAndManagers(void) {138 // Open a new context and module.139 TheContext = std::make_unique<LLVMContext>();140 TheModule = std::make_unique<Module>("KaleidoscopeJIT", *TheContext);141 TheModule->setDataLayout(TheJIT->getDataLayout());142 143 // Create a new builder for the module.144 Builder = std::make_unique<IRBuilder<>>(*TheContext);145 146 // Create new pass and analysis managers.147 TheFPM = std::make_unique<FunctionPassManager>();148 TheLAM = std::make_unique<LoopAnalysisManager>();149 TheFAM = std::make_unique<FunctionAnalysisManager>();150 TheCGAM = std::make_unique<CGSCCAnalysisManager>();151 TheMAM = std::make_unique<ModuleAnalysisManager>();152 ThePIC = std::make_unique<PassInstrumentationCallbacks>();153 TheSI = std::make_unique<StandardInstrumentations>(*TheContext,154 /*DebugLogging*/ true);155 TheSI->registerCallbacks(*ThePIC, TheMAM.get());156 ...157 158After initializing the global module ``TheModule`` and the FunctionPassManager,159we need to initialize other parts of the framework. The four AnalysisManagers160allow us to add analysis passes that run across the four levels of the IR161hierarchy. PassInstrumentationCallbacks and StandardInstrumentations are162required for the pass instrumentation framework, which allows developers to163customize what happens between passes.164 165Once these managers are set up, we use a series of "addPass" calls to add a166bunch of LLVM transform passes:167 168.. code-block:: c++169 170 // Add transform passes.171 // Do simple "peephole" optimizations and bit-twiddling optzns.172 TheFPM->addPass(InstCombinePass());173 // Reassociate expressions.174 TheFPM->addPass(ReassociatePass());175 // Eliminate Common SubExpressions.176 TheFPM->addPass(GVNPass());177 // Simplify the control flow graph (deleting unreachable blocks, etc).178 TheFPM->addPass(SimplifyCFGPass());179 180In this case, we choose to add four optimization passes.181The passes we choose here are a pretty standard set182of "cleanup" optimizations that are useful for a wide variety of code. I won't183delve into what they do but, believe me, they are a good starting place :).184 185Next, we register the analysis passes used by the transform passes.186 187.. code-block:: c++188 189 // Register analysis passes used in these transform passes.190 PassBuilder PB;191 PB.registerModuleAnalyses(*TheMAM);192 PB.registerFunctionAnalyses(*TheFAM);193 PB.crossRegisterProxies(*TheLAM, *TheFAM, *TheCGAM, *TheMAM);194 }195 196Once the PassManager is set up, we need to make use of it. We do this by197running it after our newly created function is constructed (in198``FunctionAST::codegen()``), but before it is returned to the client:199 200.. code-block:: c++201 202 if (Value *RetVal = Body->codegen()) {203 // Finish off the function.204 Builder.CreateRet(RetVal);205 206 // Validate the generated code, checking for consistency.207 verifyFunction(*TheFunction);208 209 // Optimize the function.210 TheFPM->run(*TheFunction, *TheFAM);211 212 return TheFunction;213 }214 215As you can see, this is pretty straightforward. The216``FunctionPassManager`` optimizes and updates the LLVM Function\* in217place, improving (hopefully) its body. With this in place, we can try218our test above again:219 220::221 222 ready> def test(x) (1+2+x)*(x+(1+2));223 ready> Read function definition:224 define double @test(double %x) {225 entry:226 %addtmp = fadd double %x, 3.000000e+00227 %multmp = fmul double %addtmp, %addtmp228 ret double %multmp229 }230 231As expected, we now get our nicely optimized code, saving a floating232point add instruction from every execution of this function.233 234LLVM provides a wide variety of optimizations that can be used in235certain circumstances. Some `documentation about the various236passes <../../Passes.html>`_ is available, but it isn't very complete.237Another good source of ideas can come from looking at the passes that238``Clang`` runs to get started. The "``opt``" tool allows you to239experiment with passes from the command line, so you can see if they do240anything.241 242Now that we have reasonable code coming out of our front-end, let's talk243about executing it!244 245Adding a JIT Compiler246=====================247 248Code that is available in LLVM IR can have a wide variety of tools249applied to it. For example, you can run optimizations on it (as we did250above), you can dump it out in textual or binary forms, you can compile251the code to an assembly file (.s) for some target, or you can JIT252compile it. The nice thing about the LLVM IR representation is that it253is the "common currency" between many different parts of the compiler.254 255In this section, we'll add JIT compiler support to our interpreter. The256basic idea that we want for Kaleidoscope is to have the user enter257function bodies as they do now, but immediately evaluate the top-level258expressions they type in. For example, if they type in "1 + 2;", we259should evaluate and print out 3. If they define a function, they should260be able to call it from the command line.261 262In order to do this, we first prepare the environment to create code for263the current native target and declare and initialize the JIT. This is264done by calling some ``InitializeNativeTarget\*`` functions and265adding a global variable ``TheJIT``, and initializing it in266``main``:267 268.. code-block:: c++269 270 static std::unique_ptr<KaleidoscopeJIT> TheJIT;271 ...272 int main() {273 InitializeNativeTarget();274 InitializeNativeTargetAsmPrinter();275 InitializeNativeTargetAsmParser();276 277 // Install standard binary operators.278 // 1 is lowest precedence.279 BinopPrecedence['<'] = 10;280 BinopPrecedence['+'] = 20;281 BinopPrecedence['-'] = 20;282 BinopPrecedence['*'] = 40; // highest.283 284 // Prime the first token.285 fprintf(stderr, "ready> ");286 getNextToken();287 288 TheJIT = std::make_unique<KaleidoscopeJIT>();289 290 // Run the main "interpreter loop" now.291 MainLoop();292 293 return 0;294 }295 296We also need to setup the data layout for the JIT:297 298.. code-block:: c++299 300 void InitializeModuleAndPassManager(void) {301 // Open a new context and module.302 TheContext = std::make_unique<LLVMContext>();303 TheModule = std::make_unique<Module>("my cool jit", TheContext);304 TheModule->setDataLayout(TheJIT->getDataLayout());305 306 // Create a new builder for the module.307 Builder = std::make_unique<IRBuilder<>>(*TheContext);308 309 // Create a new pass manager attached to it.310 TheFPM = std::make_unique<legacy::FunctionPassManager>(TheModule.get());311 ...312 313The KaleidoscopeJIT class is a simple JIT built specifically for these314tutorials, available inside the LLVM source code315at `llvm-src/examples/Kaleidoscope/include/KaleidoscopeJIT.h316<https://github.com/llvm/llvm-project/blob/main/llvm/examples/Kaleidoscope/include/KaleidoscopeJIT.h>`_.317In later chapters we will look at how it works and extend it with318new features, but for now we will take it as given. Its API is very simple:319``addModule`` adds an LLVM IR module to the JIT, making its functions320available for execution (with its memory managed by a ``ResourceTracker``); and321``lookup`` allows us to look up pointers to the compiled code.322 323We can take this simple API and change our code that parses top-level expressions to324look like this:325 326.. code-block:: c++327 328 static ExitOnError ExitOnErr;329 ...330 static void HandleTopLevelExpression() {331 // Evaluate a top-level expression into an anonymous function.332 if (auto FnAST = ParseTopLevelExpr()) {333 if (FnAST->codegen()) {334 // Create a ResourceTracker to track JIT'd memory allocated to our335 // anonymous expression -- that way we can free it after executing.336 auto RT = TheJIT->getMainJITDylib().createResourceTracker();337 338 auto TSM = ThreadSafeModule(std::move(TheModule), std::move(TheContext));339 ExitOnErr(TheJIT->addModule(std::move(TSM), RT));340 InitializeModuleAndPassManager();341 342 // Search the JIT for the __anon_expr symbol.343 auto ExprSymbol = ExitOnErr(TheJIT->lookup("__anon_expr"));344 assert(ExprSymbol && "Function not found");345 346 // Get the symbol's address and cast it to the right type (takes no347 // arguments, returns a double) so we can call it as a native function.348 double (*FP)() = ExprSymbol.getAddress().toPtr<double (*)()>();349 fprintf(stderr, "Evaluated to %f\n", FP());350 351 // Delete the anonymous expression module from the JIT.352 ExitOnErr(RT->remove());353 }354 355If parsing and codegen succeed, the next step is to add the module containing356the top-level expression to the JIT. We do this by calling addModule, which357triggers code generation for all the functions in the module, and accepts a358``ResourceTracker`` which can be used to remove the module from the JIT later. Once the module359has been added to the JIT it can no longer be modified, so we also open a new360module to hold subsequent code by calling ``InitializeModuleAndPassManager()``.361 362Once we've added the module to the JIT we need to get a pointer to the final363generated code. We do this by calling the JIT's ``lookup`` method, and passing364the name of the top-level expression function: ``__anon_expr``. Since we just365added this function, we assert that ``lookup`` returned a result.366 367Next, we get the in-memory address of the ``__anon_expr`` function by calling368``getAddress()`` on the symbol. Recall that we compile top-level expressions369into a self-contained LLVM function that takes no arguments and returns the370computed double. Because the LLVM JIT compiler matches the native platform ABI,371this means that you can just cast the result pointer to a function pointer of372that type and call it directly. This means, there is no difference between JIT373compiled code and native machine code that is statically linked into your374application.375 376Finally, since we don't support re-evaluation of top-level expressions, we377remove the module from the JIT when we're done to free the associated memory.378Recall, however, that the module we created a few lines earlier (via379``InitializeModuleAndPassManager``) is still open and waiting for new code to be380added.381 382With just these two changes, let's see how Kaleidoscope works now!383 384::385 386 ready> 4+5;387 Read top-level expression:388 define double @0() {389 entry:390 ret double 9.000000e+00391 }392 393 Evaluated to 9.000000394 395Well this looks like it is basically working. The dump of the function396shows the "no argument function that always returns double" that we397synthesize for each top-level expression that is typed in. This398demonstrates very basic functionality, but can we do more?399 400::401 402 ready> def testfunc(x y) x + y*2;403 Read function definition:404 define double @testfunc(double %x, double %y) {405 entry:406 %multmp = fmul double %y, 2.000000e+00407 %addtmp = fadd double %multmp, %x408 ret double %addtmp409 }410 411 ready> testfunc(4, 10);412 Read top-level expression:413 define double @1() {414 entry:415 %calltmp = call double @testfunc(double 4.000000e+00, double 1.000000e+01)416 ret double %calltmp417 }418 419 Evaluated to 24.000000420 421 ready> testfunc(5, 10);422 ready> LLVM ERROR: Program used external function 'testfunc' which could not be resolved!423 424 425Function definitions and calls also work, but something went very wrong on that426last line. The call looks valid, so what happened? As you may have guessed from427the API a Module is a unit of allocation for the JIT, and testfunc was part428of the same module that contained anonymous expression. When we removed that429module from the JIT to free the memory for the anonymous expression, we deleted430the definition of ``testfunc`` along with it. Then, when we tried to call431testfunc a second time, the JIT could no longer find it.432 433The easiest way to fix this is to put the anonymous expression in a separate434module from the rest of the function definitions. The JIT will happily resolve435function calls across module boundaries, as long as each of the functions called436has a prototype, and is added to the JIT before it is called. By putting the437anonymous expression in a different module we can delete it without affecting438the rest of the functions.439 440In fact, we're going to go a step further and put every function in its own441module. Doing so allows us to exploit a useful property of the KaleidoscopeJIT442that will make our environment more REPL-like: Functions can be added to the443JIT more than once (unlike a module where every function must have a unique444definition). When you look up a symbol in KaleidoscopeJIT it will always return445the most recent definition:446 447::448 449 ready> def foo(x) x + 1;450 Read function definition:451 define double @foo(double %x) {452 entry:453 %addtmp = fadd double %x, 1.000000e+00454 ret double %addtmp455 }456 457 ready> foo(2);458 Evaluated to 3.000000459 460 ready> def foo(x) x + 2;461 define double @foo(double %x) {462 entry:463 %addtmp = fadd double %x, 2.000000e+00464 ret double %addtmp465 }466 467 ready> foo(2);468 Evaluated to 4.000000469 470 471To allow each function to live in its own module we'll need a way to472re-generate previous function declarations into each new module we open:473 474.. code-block:: c++475 476 static std::unique_ptr<KaleidoscopeJIT> TheJIT;477 478 ...479 480 Function *getFunction(std::string Name) {481 // First, see if the function has already been added to the current module.482 if (auto *F = TheModule->getFunction(Name))483 return F;484 485 // If not, check whether we can codegen the declaration from some existing486 // prototype.487 auto FI = FunctionProtos.find(Name);488 if (FI != FunctionProtos.end())489 return FI->second->codegen();490 491 // If no existing prototype exists, return null.492 return nullptr;493 }494 495 ...496 497 Value *CallExprAST::codegen() {498 // Look up the name in the global module table.499 Function *CalleeF = getFunction(Callee);500 501 ...502 503 Function *FunctionAST::codegen() {504 // Transfer ownership of the prototype to the FunctionProtos map, but keep a505 // reference to it for use below.506 auto &P = *Proto;507 FunctionProtos[Proto->getName()] = std::move(Proto);508 Function *TheFunction = getFunction(P.getName());509 if (!TheFunction)510 return nullptr;511 512 513To enable this, we'll start by adding a new global, ``FunctionProtos``, that514holds the most recent prototype for each function. We'll also add a convenience515method, ``getFunction()``, to replace calls to ``TheModule->getFunction()``.516Our convenience method searches ``TheModule`` for an existing function517declaration, falling back to generating a new declaration from FunctionProtos if518it doesn't find one. In ``CallExprAST::codegen()`` we just need to replace the519call to ``TheModule->getFunction()``. In ``FunctionAST::codegen()`` we need to520update the FunctionProtos map first, then call ``getFunction()``. With this521done, we can always obtain a function declaration in the current module for any522previously declared function.523 524We also need to update HandleDefinition and HandleExtern:525 526.. code-block:: c++527 528 static void HandleDefinition() {529 if (auto FnAST = ParseDefinition()) {530 if (auto *FnIR = FnAST->codegen()) {531 fprintf(stderr, "Read function definition:");532 FnIR->print(errs());533 fprintf(stderr, "\n");534 ExitOnErr(TheJIT->addModule(535 ThreadSafeModule(std::move(TheModule), std::move(TheContext))));536 InitializeModuleAndPassManager();537 }538 } else {539 // Skip token for error recovery.540 getNextToken();541 }542 }543 544 static void HandleExtern() {545 if (auto ProtoAST = ParseExtern()) {546 if (auto *FnIR = ProtoAST->codegen()) {547 fprintf(stderr, "Read extern: ");548 FnIR->print(errs());549 fprintf(stderr, "\n");550 FunctionProtos[ProtoAST->getName()] = std::move(ProtoAST);551 }552 } else {553 // Skip token for error recovery.554 getNextToken();555 }556 }557 558In HandleDefinition, we add two lines to transfer the newly defined function to559the JIT and open a new module. In HandleExtern, we just need to add one line to560add the prototype to FunctionProtos.561 562.. warning::563 Duplication of symbols in separate modules is not allowed since LLVM-9. That means you can not redefine function in your Kaleidoscope as its shown below. Just skip this part.564 565 The reason is that the newer OrcV2 JIT APIs are trying to stay very close to the static and dynamic linker rules, including rejecting duplicate symbols. Requiring symbol names to be unique allows us to support concurrent compilation for symbols using the (unique) symbol names as keys for tracking.566 567With these changes made, let's try our REPL again (I removed the dump of the568anonymous functions this time, you should get the idea by now :) :569 570::571 572 ready> def foo(x) x + 1;573 ready> foo(2);574 Evaluated to 3.000000575 576 ready> def foo(x) x + 2;577 ready> foo(2);578 Evaluated to 4.000000579 580It works!581 582Even with this simple code, we get some surprisingly powerful capabilities -583check this out:584 585::586 587 ready> extern sin(x);588 Read extern:589 declare double @sin(double)590 591 ready> extern cos(x);592 Read extern:593 declare double @cos(double)594 595 ready> sin(1.0);596 Read top-level expression:597 define double @2() {598 entry:599 ret double 0x3FEAED548F090CEE600 }601 602 Evaluated to 0.841471603 604 ready> def foo(x) sin(x)*sin(x) + cos(x)*cos(x);605 Read function definition:606 define double @foo(double %x) {607 entry:608 %calltmp = call double @sin(double %x)609 %multmp = fmul double %calltmp, %calltmp610 %calltmp2 = call double @cos(double %x)611 %multmp4 = fmul double %calltmp2, %calltmp2612 %addtmp = fadd double %multmp, %multmp4613 ret double %addtmp614 }615 616 ready> foo(4.0);617 Read top-level expression:618 define double @3() {619 entry:620 %calltmp = call double @foo(double 4.000000e+00)621 ret double %calltmp622 }623 624 Evaluated to 1.000000625 626Whoa, how does the JIT know about sin and cos? The answer is surprisingly627simple: The KaleidoscopeJIT has a straightforward symbol resolution rule that628it uses to find symbols that aren't available in any given module: First629it searches all the modules that have already been added to the JIT, from the630most recent to the oldest, to find the newest definition. If no definition is631found inside the JIT, it falls back to calling "``dlsym("sin")``" on the632Kaleidoscope process itself. Since "``sin``" is defined within the JIT's633address space, it simply patches up calls in the module to call the libm634version of ``sin`` directly. But in some cases this even goes further:635as sin and cos are names of standard math functions, the constant folder636will directly evaluate the function calls to the correct result when called637with constants like in the "``sin(1.0)``" above.638 639In the future we'll see how tweaking this symbol resolution rule can be used to640enable all sorts of useful features, from security (restricting the set of641symbols available to JIT'd code), to dynamic code generation based on symbol642names, and even lazy compilation.643 644One immediate benefit of the symbol resolution rule is that we can now extend645the language by writing arbitrary C++ code to implement operations. For example,646if we add:647 648.. code-block:: c++649 650 #ifdef _WIN32651 #define DLLEXPORT __declspec(dllexport)652 #else653 #define DLLEXPORT654 #endif655 656 /// putchard - putchar that takes a double and returns 0.657 extern "C" DLLEXPORT double putchard(double X) {658 fputc((char)X, stderr);659 return 0;660 }661 662Note, that for Windows we need to actually export the functions because663the dynamic symbol loader will use ``GetProcAddress`` to find the symbols.664 665Now we can produce simple output to the console by using things like:666"``extern putchard(x); putchard(120);``", which prints a lowercase 'x'667on the console (120 is the ASCII code for 'x'). Similar code could be668used to implement file I/O, console input, and many other capabilities669in Kaleidoscope.670 671This completes the JIT and optimizer chapter of the Kaleidoscope672tutorial. At this point, we can compile a non-Turing-complete673programming language, optimize and JIT compile it in a user-driven way.674Next up we'll look into `extending the language with control flow675constructs <LangImpl05.html>`_, tackling some interesting LLVM IR issues676along the way.677 678Full Code Listing679=================680 681Here is the complete code listing for our running example, enhanced with682the LLVM JIT and optimizer. To build this example, use:683 684.. code-block:: bash685 686 # Compile687 clang++ -g toy.cpp `llvm-config --cxxflags --ldflags --system-libs --libs core orcjit native` -O3 -o toy688 # Run689 ./toy690 691If you are compiling this on Linux, make sure to add the "-rdynamic"692option as well. This makes sure that the external functions are resolved693properly at runtime.694 695Here is the code:696 697.. literalinclude:: ../../../examples/Kaleidoscope/Chapter4/toy.cpp698 :language: c++699 700`Next: Extending the language: control flow <LangImpl05.html>`_701 702