1============================================== 2Kaleidoscope: Adding JIT and Optimizer Support 3============================================== 4 5.. contents:: 6 :local: 7 8Chapter 4 Introduction 9====================== 10 11Welcome to Chapter 4 of the "`Implementing a language with 12LLVM <index.html>`_" tutorial. Chapters 1-3 described the implementation 13of a simple language and added support for generating LLVM IR. This 14chapter describes two new techniques: adding optimizer support to your 15language, and adding JIT compiler support. These additions will 16demonstrate how to get nice, efficient code for the Kaleidoscope 17language. 18 19Trivial Constant Folding 20======================== 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, %x 33 ret double %addtmp 34 } 35 36This code is not a literal transcription of the AST built by parsing the 37input. 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+00 46 %addtmp1 = fadd double %addtmp, %x 47 ret double %addtmp1 48 } 49 50Constant folding, as seen above, in particular, is a very common and 51very important optimization: so much so that many language implementors 52implement constant folding support in their AST representation. 53 54With LLVM, you don't need this support in the AST. Since all calls to 55build LLVM IR go through the LLVM IR builder, the builder itself checked 56to see if there was a constant folding opportunity when you call it. If 57so, it just does the constant fold and return the constant instead of 58creating an instruction. 59 60Well, that was easy :). In practice, we recommend always using 61``IRBuilder`` when generating code like this. It has no "syntactic 62overhead" for its use (you don't have to uglify your compiler with 63constant checks everywhere) and it can dramatically reduce the amount of 64LLVM IR that is generated in some cases (particular for languages with a 65macro preprocessor or that use a lot of constants). 66 67On the other hand, the ``IRBuilder`` is limited by the fact that it does 68all of its analysis inline with the code as it is built. If you take a 69slightly 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, %x 78 %addtmp1 = fadd double %x, 3.000000e+00 79 %multmp = fmul double %addtmp, %addtmp1 80 ret double %multmp 81 } 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 and 88correct this. This requires two transformations: reassociation of 89expressions (to make the add's lexically identical) and Common 90Subexpression Elimination (CSE) to delete the redundant add instruction. 91Fortunately, LLVM provides a broad range of optimizations that you can 92use, in the form of "passes". 93 94LLVM Optimization Passes 95======================== 96 97LLVM provides many optimization passes, which do many different sorts of 98things and have different tradeoffs. Unlike other systems, LLVM doesn't 99hold to the mistaken notion that one set of optimizations is right for 100all languages and for all situations. LLVM allows a compiler implementor 101to make complete decisions about what optimizations to use, in which 102order, and in what situation. 103 104As a concrete example, LLVM supports both "whole module" passes, which 105look 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 whole 107program). It also supports and includes "per-function" passes which just 108operate on a single function at a time, without looking at other 109functions. For more information on passes and how they are run, see the 110`How to Write a Pass <../WritingAnLLVMPass.html>`_ document and the 111`List of LLVM Passes <../Passes.html>`_. 112 113For Kaleidoscope, we are currently generating functions on the fly, one 114at a time, as the user types them in. We aren't shooting for the 115ultimate optimization experience in this setting, but we also want to 116catch the easy and quick stuff where possible. As such, we will choose 117to run a few per-function optimizations as the user types the function 118in. If we wanted to make a "static Kaleidoscope compiler", we would use 119exactly the code we have now, except that we would defer running the 120optimizer until the entire file has been parsed. 121 122In order to get per-function optimizations going, we need to set up a 123`FunctionPassManager <../WritingAnLLVMPass.html#what-passmanager-doesr>`_ to hold 124and organize the LLVM optimizations that we want to run. Once we have 125that, we can add a set of optimizations to run. We'll need a new 126FunctionPassManager for each module that we want to optimize, so we'll 127write a function to create and initialize both the module and pass manager 128for us: 129 130.. code-block:: c++ 131 132 void InitializeModuleAndPassManager(void) { 133 // Open a new module. 134 TheModule = llvm::make_unique<Module>("my cool jit", TheContext); 135 136 // Create a new pass manager attached to it. 137 TheFPM = llvm::make_unique<FunctionPassManager>(TheModule.get()); 138 139 // Do simple "peephole" optimizations and bit-twiddling optzns. 140 TheFPM->add(createInstructionCombiningPass()); 141 // Reassociate expressions. 142 TheFPM->add(createReassociatePass()); 143 // Eliminate Common SubExpressions. 144 TheFPM->add(createGVNPass()); 145 // Simplify the control flow graph (deleting unreachable blocks, etc). 146 TheFPM->add(createCFGSimplificationPass()); 147 148 TheFPM->doInitialization(); 149 } 150 151This code initializes the global module ``TheModule``, and the function pass 152manager ``TheFPM``, which is attached to ``TheModule``. Once the pass manager is 153set up, we use a series of "add" calls to add a bunch of LLVM passes. 154 155In this case, we choose to add four optimization passes. 156The passes we choose here are a pretty standard set 157of "cleanup" optimizations that are useful for a wide variety of code. I won't 158delve into what they do but, believe me, they are a good starting place :). 159 160Once the PassManager is set up, we need to make use of it. We do this by 161running it after our newly created function is constructed (in 162``FunctionAST::codegen()``), but before it is returned to the client: 163 164.. code-block:: c++ 165 166 if (Value *RetVal = Body->codegen()) { 167 // Finish off the function. 168 Builder.CreateRet(RetVal); 169 170 // Validate the generated code, checking for consistency. 171 verifyFunction(*TheFunction); 172 173 // Optimize the function. 174 TheFPM->run(*TheFunction); 175 176 return TheFunction; 177 } 178 179As you can see, this is pretty straightforward. The 180``FunctionPassManager`` optimizes and updates the LLVM Function\* in 181place, improving (hopefully) its body. With this in place, we can try 182our test above again: 183 184:: 185 186 ready> def test(x) (1+2+x)*(x+(1+2)); 187 ready> Read function definition: 188 define double @test(double %x) { 189 entry: 190 %addtmp = fadd double %x, 3.000000e+00 191 %multmp = fmul double %addtmp, %addtmp 192 ret double %multmp 193 } 194 195As expected, we now get our nicely optimized code, saving a floating 196point add instruction from every execution of this function. 197 198LLVM provides a wide variety of optimizations that can be used in 199certain circumstances. Some `documentation about the various 200passes <../Passes.html>`_ is available, but it isn't very complete. 201Another good source of ideas can come from looking at the passes that 202``Clang`` runs to get started. The "``opt``" tool allows you to 203experiment with passes from the command line, so you can see if they do 204anything. 205 206Now that we have reasonable code coming out of our front-end, let's talk 207about executing it! 208 209Adding a JIT Compiler 210===================== 211 212Code that is available in LLVM IR can have a wide variety of tools 213applied to it. For example, you can run optimizations on it (as we did 214above), you can dump it out in textual or binary forms, you can compile 215the code to an assembly file (.s) for some target, or you can JIT 216compile it. The nice thing about the LLVM IR representation is that it 217is the "common currency" between many different parts of the compiler. 218 219In this section, we'll add JIT compiler support to our interpreter. The 220basic idea that we want for Kaleidoscope is to have the user enter 221function bodies as they do now, but immediately evaluate the top-level 222expressions they type in. For example, if they type in "1 + 2;", we 223should evaluate and print out 3. If they define a function, they should 224be able to call it from the command line. 225 226In order to do this, we first prepare the environment to create code for 227the current native target and declare and initialize the JIT. This is 228done by calling some ``InitializeNativeTarget\*`` functions and 229adding a global variable ``TheJIT``, and initializing it in 230``main``: 231 232.. code-block:: c++ 233 234 static std::unique_ptr<KaleidoscopeJIT> TheJIT; 235 ... 236 int main() { 237 InitializeNativeTarget(); 238 InitializeNativeTargetAsmPrinter(); 239 InitializeNativeTargetAsmParser(); 240 241 // Install standard binary operators. 242 // 1 is lowest precedence. 243 BinopPrecedence['<'] = 10; 244 BinopPrecedence['+'] = 20; 245 BinopPrecedence['-'] = 20; 246 BinopPrecedence['*'] = 40; // highest. 247 248 // Prime the first token. 249 fprintf(stderr, "ready> "); 250 getNextToken(); 251 252 TheJIT = llvm::make_unique<KaleidoscopeJIT>(); 253 254 // Run the main "interpreter loop" now. 255 MainLoop(); 256 257 return 0; 258 } 259 260We also need to setup the data layout for the JIT: 261 262.. code-block:: c++ 263 264 void InitializeModuleAndPassManager(void) { 265 // Open a new module. 266 TheModule = llvm::make_unique<Module>("my cool jit", TheContext); 267 TheModule->setDataLayout(TheJIT->getTargetMachine().createDataLayout()); 268 269 // Create a new pass manager attached to it. 270 TheFPM = llvm::make_unique<FunctionPassManager>(TheModule.get()); 271 ... 272 273The KaleidoscopeJIT class is a simple JIT built specifically for these 274tutorials, available inside the LLVM source code 275at llvm-src/examples/Kaleidoscope/include/KaleidoscopeJIT.h. 276In later chapters we will look at how it works and extend it with 277new features, but for now we will take it as given. Its API is very simple: 278``addModule`` adds an LLVM IR module to the JIT, making its functions 279available for execution; ``removeModule`` removes a module, freeing any 280memory associated with the code in that module; and ``findSymbol`` allows us 281to look up pointers to the compiled code. 282 283We can take this simple API and change our code that parses top-level expressions to 284look like this: 285 286.. code-block:: c++ 287 288 static void HandleTopLevelExpression() { 289 // Evaluate a top-level expression into an anonymous function. 290 if (auto FnAST = ParseTopLevelExpr()) { 291 if (FnAST->codegen()) { 292 293 // JIT the module containing the anonymous expression, keeping a handle so 294 // we can free it later. 295 auto H = TheJIT->addModule(std::move(TheModule)); 296 InitializeModuleAndPassManager(); 297 298 // Search the JIT for the __anon_expr symbol. 299 auto ExprSymbol = TheJIT->findSymbol("__anon_expr"); 300 assert(ExprSymbol && "Function not found"); 301 302 // Get the symbol's address and cast it to the right type (takes no 303 // arguments, returns a double) so we can call it as a native function. 304 double (*FP)() = (double (*)())(intptr_t)ExprSymbol.getAddress(); 305 fprintf(stderr, "Evaluated to %f\n", FP()); 306 307 // Delete the anonymous expression module from the JIT. 308 TheJIT->removeModule(H); 309 } 310 311If parsing and codegen succeeed, the next step is to add the module containing 312the top-level expression to the JIT. We do this by calling addModule, which 313triggers code generation for all the functions in the module, and returns a 314handle that can be used to remove the module from the JIT later. Once the module 315has been added to the JIT it can no longer be modified, so we also open a new 316module to hold subsequent code by calling ``InitializeModuleAndPassManager()``. 317 318Once we've added the module to the JIT we need to get a pointer to the final 319generated code. We do this by calling the JIT's findSymbol method, and passing 320the name of the top-level expression function: ``__anon_expr``. Since we just 321added this function, we assert that findSymbol returned a result. 322 323Next, we get the in-memory address of the ``__anon_expr`` function by calling 324``getAddress()`` on the symbol. Recall that we compile top-level expressions 325into a self-contained LLVM function that takes no arguments and returns the 326computed double. Because the LLVM JIT compiler matches the native platform ABI, 327this means that you can just cast the result pointer to a function pointer of 328that type and call it directly. This means, there is no difference between JIT 329compiled code and native machine code that is statically linked into your 330application. 331 332Finally, since we don't support re-evaluation of top-level expressions, we 333remove the module from the JIT when we're done to free the associated memory. 334Recall, however, that the module we created a few lines earlier (via 335``InitializeModuleAndPassManager``) is still open and waiting for new code to be 336added. 337 338With just these two changes, let's see how Kaleidoscope works now! 339 340:: 341 342 ready> 4+5; 343 Read top-level expression: 344 define double @0() { 345 entry: 346 ret double 9.000000e+00 347 } 348 349 Evaluated to 9.000000 350 351Well this looks like it is basically working. The dump of the function 352shows the "no argument function that always returns double" that we 353synthesize for each top-level expression that is typed in. This 354demonstrates very basic functionality, but can we do more? 355 356:: 357 358 ready> def testfunc(x y) x + y*2; 359 Read function definition: 360 define double @testfunc(double %x, double %y) { 361 entry: 362 %multmp = fmul double %y, 2.000000e+00 363 %addtmp = fadd double %multmp, %x 364 ret double %addtmp 365 } 366 367 ready> testfunc(4, 10); 368 Read top-level expression: 369 define double @1() { 370 entry: 371 %calltmp = call double @testfunc(double 4.000000e+00, double 1.000000e+01) 372 ret double %calltmp 373 } 374 375 Evaluated to 24.000000 376 377 ready> testfunc(5, 10); 378 ready> LLVM ERROR: Program used external function 'testfunc' which could not be resolved! 379 380 381Function definitions and calls also work, but something went very wrong on that 382last line. The call looks valid, so what happened? As you may have guessed from 383the API a Module is a unit of allocation for the JIT, and testfunc was part 384of the same module that contained anonymous expression. When we removed that 385module from the JIT to free the memory for the anonymous expression, we deleted 386the definition of ``testfunc`` along with it. Then, when we tried to call 387testfunc a second time, the JIT could no longer find it. 388 389The easiest way to fix this is to put the anonymous expression in a separate 390module from the rest of the function definitions. The JIT will happily resolve 391function calls across module boundaries, as long as each of the functions called 392has a prototype, and is added to the JIT before it is called. By putting the 393anonymous expression in a different module we can delete it without affecting 394the rest of the functions. 395 396In fact, we're going to go a step further and put every function in its own 397module. Doing so allows us to exploit a useful property of the KaleidoscopeJIT 398that will make our environment more REPL-like: Functions can be added to the 399JIT more than once (unlike a module where every function must have a unique 400definition). When you look up a symbol in KaleidoscopeJIT it will always return 401the most recent definition: 402 403:: 404 405 ready> def foo(x) x + 1; 406 Read function definition: 407 define double @foo(double %x) { 408 entry: 409 %addtmp = fadd double %x, 1.000000e+00 410 ret double %addtmp 411 } 412 413 ready> foo(2); 414 Evaluated to 3.000000 415 416 ready> def foo(x) x + 2; 417 define double @foo(double %x) { 418 entry: 419 %addtmp = fadd double %x, 2.000000e+00 420 ret double %addtmp 421 } 422 423 ready> foo(2); 424 Evaluated to 4.000000 425 426 427To allow each function to live in its own module we'll need a way to 428re-generate previous function declarations into each new module we open: 429 430.. code-block:: c++ 431 432 static std::unique_ptr<KaleidoscopeJIT> TheJIT; 433 434 ... 435 436 Function *getFunction(std::string Name) { 437 // First, see if the function has already been added to the current module. 438 if (auto *F = TheModule->getFunction(Name)) 439 return F; 440 441 // If not, check whether we can codegen the declaration from some existing 442 // prototype. 443 auto FI = FunctionProtos.find(Name); 444 if (FI != FunctionProtos.end()) 445 return FI->second->codegen(); 446 447 // If no existing prototype exists, return null. 448 return nullptr; 449 } 450 451 ... 452 453 Value *CallExprAST::codegen() { 454 // Look up the name in the global module table. 455 Function *CalleeF = getFunction(Callee); 456 457 ... 458 459 Function *FunctionAST::codegen() { 460 // Transfer ownership of the prototype to the FunctionProtos map, but keep a 461 // reference to it for use below. 462 auto &P = *Proto; 463 FunctionProtos[Proto->getName()] = std::move(Proto); 464 Function *TheFunction = getFunction(P.getName()); 465 if (!TheFunction) 466 return nullptr; 467 468 469To enable this, we'll start by adding a new global, ``FunctionProtos``, that 470holds the most recent prototype for each function. We'll also add a convenience 471method, ``getFunction()``, to replace calls to ``TheModule->getFunction()``. 472Our convenience method searches ``TheModule`` for an existing function 473declaration, falling back to generating a new declaration from FunctionProtos if 474it doesn't find one. In ``CallExprAST::codegen()`` we just need to replace the 475call to ``TheModule->getFunction()``. In ``FunctionAST::codegen()`` we need to 476update the FunctionProtos map first, then call ``getFunction()``. With this 477done, we can always obtain a function declaration in the current module for any 478previously declared function. 479 480We also need to update HandleDefinition and HandleExtern: 481 482.. code-block:: c++ 483 484 static void HandleDefinition() { 485 if (auto FnAST = ParseDefinition()) { 486 if (auto *FnIR = FnAST->codegen()) { 487 fprintf(stderr, "Read function definition:"); 488 FnIR->print(errs()); 489 fprintf(stderr, "\n"); 490 TheJIT->addModule(std::move(TheModule)); 491 InitializeModuleAndPassManager(); 492 } 493 } else { 494 // Skip token for error recovery. 495 getNextToken(); 496 } 497 } 498 499 static void HandleExtern() { 500 if (auto ProtoAST = ParseExtern()) { 501 if (auto *FnIR = ProtoAST->codegen()) { 502 fprintf(stderr, "Read extern: "); 503 FnIR->print(errs()); 504 fprintf(stderr, "\n"); 505 FunctionProtos[ProtoAST->getName()] = std::move(ProtoAST); 506 } 507 } else { 508 // Skip token for error recovery. 509 getNextToken(); 510 } 511 } 512 513In HandleDefinition, we add two lines to transfer the newly defined function to 514the JIT and open a new module. In HandleExtern, we just need to add one line to 515add the prototype to FunctionProtos. 516 517With these changes made, let's try our REPL again (I removed the dump of the 518anonymous functions this time, you should get the idea by now :) : 519 520:: 521 522 ready> def foo(x) x + 1; 523 ready> foo(2); 524 Evaluated to 3.000000 525 526 ready> def foo(x) x + 2; 527 ready> foo(2); 528 Evaluated to 4.000000 529 530It works! 531 532Even with this simple code, we get some surprisingly powerful capabilities - 533check this out: 534 535:: 536 537 ready> extern sin(x); 538 Read extern: 539 declare double @sin(double) 540 541 ready> extern cos(x); 542 Read extern: 543 declare double @cos(double) 544 545 ready> sin(1.0); 546 Read top-level expression: 547 define double @2() { 548 entry: 549 ret double 0x3FEAED548F090CEE 550 } 551 552 Evaluated to 0.841471 553 554 ready> def foo(x) sin(x)*sin(x) + cos(x)*cos(x); 555 Read function definition: 556 define double @foo(double %x) { 557 entry: 558 %calltmp = call double @sin(double %x) 559 %multmp = fmul double %calltmp, %calltmp 560 %calltmp2 = call double @cos(double %x) 561 %multmp4 = fmul double %calltmp2, %calltmp2 562 %addtmp = fadd double %multmp, %multmp4 563 ret double %addtmp 564 } 565 566 ready> foo(4.0); 567 Read top-level expression: 568 define double @3() { 569 entry: 570 %calltmp = call double @foo(double 4.000000e+00) 571 ret double %calltmp 572 } 573 574 Evaluated to 1.000000 575 576Whoa, how does the JIT know about sin and cos? The answer is surprisingly 577simple: The KaleidoscopeJIT has a straightforward symbol resolution rule that 578it uses to find symbols that aren't available in any given module: First 579it searches all the modules that have already been added to the JIT, from the 580most recent to the oldest, to find the newest definition. If no definition is 581found inside the JIT, it falls back to calling "``dlsym("sin")``" on the 582Kaleidoscope process itself. Since "``sin``" is defined within the JIT's 583address space, it simply patches up calls in the module to call the libm 584version of ``sin`` directly. But in some cases this even goes further: 585as sin and cos are names of standard math functions, the constant folder 586will directly evaluate the function calls to the correct result when called 587with constants like in the "``sin(1.0)``" above. 588 589In the future we'll see how tweaking this symbol resolution rule can be used to 590enable all sorts of useful features, from security (restricting the set of 591symbols available to JIT'd code), to dynamic code generation based on symbol 592names, and even lazy compilation. 593 594One immediate benefit of the symbol resolution rule is that we can now extend 595the language by writing arbitrary C++ code to implement operations. For example, 596if we add: 597 598.. code-block:: c++ 599 600 #ifdef _WIN32 601 #define DLLEXPORT __declspec(dllexport) 602 #else 603 #define DLLEXPORT 604 #endif 605 606 /// putchard - putchar that takes a double and returns 0. 607 extern "C" DLLEXPORT double putchard(double X) { 608 fputc((char)X, stderr); 609 return 0; 610 } 611 612Note, that for Windows we need to actually export the functions because 613the dynamic symbol loader will use GetProcAddress to find the symbols. 614 615Now we can produce simple output to the console by using things like: 616"``extern putchard(x); putchard(120);``", which prints a lowercase 'x' 617on the console (120 is the ASCII code for 'x'). Similar code could be 618used to implement file I/O, console input, and many other capabilities 619in Kaleidoscope. 620 621This completes the JIT and optimizer chapter of the Kaleidoscope 622tutorial. At this point, we can compile a non-Turing-complete 623programming language, optimize and JIT compile it in a user-driven way. 624Next up we'll look into `extending the language with control flow 625constructs <LangImpl05.html>`_, tackling some interesting LLVM IR issues 626along the way. 627 628Full Code Listing 629================= 630 631Here is the complete code listing for our running example, enhanced with 632the LLVM JIT and optimizer. To build this example, use: 633 634.. code-block:: bash 635 636 # Compile 637 clang++ -g toy.cpp `llvm-config --cxxflags --ldflags --system-libs --libs core mcjit native` -O3 -o toy 638 # Run 639 ./toy 640 641If you are compiling this on Linux, make sure to add the "-rdynamic" 642option as well. This makes sure that the external functions are resolved 643properly at runtime. 644 645Here is the code: 646 647.. literalinclude:: ../../examples/Kaleidoscope/Chapter4/toy.cpp 648 :language: c++ 649 650`Next: Extending the language: control flow <LangImpl05.html>`_ 651 652