• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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", getGlobalContext());
135      TheModule->setDataLayout(TheJIT->getTargetMachine().createDataLayout());
136
137      // Create a new pass manager attached to it.
138      TheFPM = llvm::make_unique<FunctionPassManager>(TheModule.get());
139
140      // Provide basic AliasAnalysis support for GVN.
141      TheFPM.add(createBasicAliasAnalysisPass());
142      // Do simple "peephole" optimizations and bit-twiddling optzns.
143      TheFPM.add(createInstructionCombiningPass());
144      // Reassociate expressions.
145      TheFPM.add(createReassociatePass());
146      // Eliminate Common SubExpressions.
147      TheFPM.add(createGVNPass());
148      // Simplify the control flow graph (deleting unreachable blocks, etc).
149      TheFPM.add(createCFGSimplificationPass());
150
151      TheFPM.doInitialization();
152    }
153
154This code initializes the global module ``TheModule``, and the function pass
155manager ``TheFPM``, which is attached to ``TheModule``. Once the pass manager is
156set up, we use a series of "add" calls to add a bunch of LLVM passes.
157
158In this case, we choose to add five passes: one analysis pass (alias analysis),
159and four optimization passes. The passes we choose here are a pretty standard set
160of "cleanup" optimizations that are useful for a wide variety of code. I won't
161delve into what they do but, believe me, they are a good starting place :).
162
163Once the PassManager is set up, we need to make use of it. We do this by
164running it after our newly created function is constructed (in
165``FunctionAST::codegen()``), but before it is returned to the client:
166
167.. code-block:: c++
168
169      if (Value *RetVal = Body->codegen()) {
170        // Finish off the function.
171        Builder.CreateRet(RetVal);
172
173        // Validate the generated code, checking for consistency.
174        verifyFunction(*TheFunction);
175
176        // Optimize the function.
177        TheFPM->run(*TheFunction);
178
179        return TheFunction;
180      }
181
182As you can see, this is pretty straightforward. The
183``FunctionPassManager`` optimizes and updates the LLVM Function\* in
184place, improving (hopefully) its body. With this in place, we can try
185our test above again:
186
187::
188
189    ready> def test(x) (1+2+x)*(x+(1+2));
190    ready> Read function definition:
191    define double @test(double %x) {
192    entry:
193            %addtmp = fadd double %x, 3.000000e+00
194            %multmp = fmul double %addtmp, %addtmp
195            ret double %multmp
196    }
197
198As expected, we now get our nicely optimized code, saving a floating
199point add instruction from every execution of this function.
200
201LLVM provides a wide variety of optimizations that can be used in
202certain circumstances. Some `documentation about the various
203passes <../Passes.html>`_ is available, but it isn't very complete.
204Another good source of ideas can come from looking at the passes that
205``Clang`` runs to get started. The "``opt``" tool allows you to
206experiment with passes from the command line, so you can see if they do
207anything.
208
209Now that we have reasonable code coming out of our front-end, lets talk
210about executing it!
211
212Adding a JIT Compiler
213=====================
214
215Code that is available in LLVM IR can have a wide variety of tools
216applied to it. For example, you can run optimizations on it (as we did
217above), you can dump it out in textual or binary forms, you can compile
218the code to an assembly file (.s) for some target, or you can JIT
219compile it. The nice thing about the LLVM IR representation is that it
220is the "common currency" between many different parts of the compiler.
221
222In this section, we'll add JIT compiler support to our interpreter. The
223basic idea that we want for Kaleidoscope is to have the user enter
224function bodies as they do now, but immediately evaluate the top-level
225expressions they type in. For example, if they type in "1 + 2;", we
226should evaluate and print out 3. If they define a function, they should
227be able to call it from the command line.
228
229In order to do this, we first declare and initialize the JIT. This is
230done by adding a global variable ``TheJIT``, and initializing it in
231``main``:
232
233.. code-block:: c++
234
235    static std::unique_ptr<KaleidoscopeJIT> TheJIT;
236    ...
237    int main() {
238      ..
239      TheJIT = llvm::make_unique<KaleidoscopeJIT>();
240
241      // Run the main "interpreter loop" now.
242      MainLoop();
243
244      return 0;
245    }
246
247The KaleidoscopeJIT class is a simple JIT built specifically for these
248tutorials. In later chapters we will look at how it works and extend it with
249new features, but for now we will take it as given. Its API is very simple::
250``addModule`` adds an LLVM IR module to the JIT, making its functions
251available for execution; ``removeModule`` removes a module, freeing any
252memory associated with the code in that module; and ``findSymbol`` allows us
253to look up pointers to the compiled code.
254
255We can take this simple API and change our code that parses top-level expressions to
256look like this:
257
258.. code-block:: c++
259
260    static void HandleTopLevelExpression() {
261      // Evaluate a top-level expression into an anonymous function.
262      if (auto FnAST = ParseTopLevelExpr()) {
263        if (FnAST->codegen()) {
264
265          // JIT the module containing the anonymous expression, keeping a handle so
266          // we can free it later.
267          auto H = TheJIT->addModule(std::move(TheModule));
268          InitializeModuleAndPassManager();
269
270          // Search the JIT for the __anon_expr symbol.
271          auto ExprSymbol = TheJIT->findSymbol("__anon_expr");
272          assert(ExprSymbol && "Function not found");
273
274          // Get the symbol's address and cast it to the right type (takes no
275          // arguments, returns a double) so we can call it as a native function.
276          double (*FP)() = (double (*)())(intptr_t)ExprSymbol.getAddress();
277          fprintf(stderr, "Evaluated to %f\n", FP());
278
279          // Delete the anonymous expression module from the JIT.
280          TheJIT->removeModule(H);
281        }
282
283If parsing and codegen succeeed, the next step is to add the module containing
284the top-level expression to the JIT. We do this by calling addModule, which
285triggers code generation for all the functions in the module, and returns a
286handle that can be used to remove the module from the JIT later. Once the module
287has been added to the JIT it can no longer be modified, so we also open a new
288module to hold subsequent code by calling ``InitializeModuleAndPassManager()``.
289
290Once we've added the module to the JIT we need to get a pointer to the final
291generated code. We do this by calling the JIT's findSymbol method, and passing
292the name of the top-level expression function: ``__anon_expr``. Since we just
293added this function, we assert that findSymbol returned a result.
294
295Next, we get the in-memory address of the ``__anon_expr`` function by calling
296``getAddress()`` on the symbol. Recall that we compile top-level expressions
297into a self-contained LLVM function that takes no arguments and returns the
298computed double. Because the LLVM JIT compiler matches the native platform ABI,
299this means that you can just cast the result pointer to a function pointer of
300that type and call it directly. This means, there is no difference between JIT
301compiled code and native machine code that is statically linked into your
302application.
303
304Finally, since we don't support re-evaluation of top-level expressions, we
305remove the module from the JIT when we're done to free the associated memory.
306Recall, however, that the module we created a few lines earlier (via
307``InitializeModuleAndPassManager``) is still open and waiting for new code to be
308added.
309
310With just these two changes, lets see how Kaleidoscope works now!
311
312::
313
314    ready> 4+5;
315    Read top-level expression:
316    define double @0() {
317    entry:
318      ret double 9.000000e+00
319    }
320
321    Evaluated to 9.000000
322
323Well this looks like it is basically working. The dump of the function
324shows the "no argument function that always returns double" that we
325synthesize for each top-level expression that is typed in. This
326demonstrates very basic functionality, but can we do more?
327
328::
329
330    ready> def testfunc(x y) x + y*2;
331    Read function definition:
332    define double @testfunc(double %x, double %y) {
333    entry:
334      %multmp = fmul double %y, 2.000000e+00
335      %addtmp = fadd double %multmp, %x
336      ret double %addtmp
337    }
338
339    ready> testfunc(4, 10);
340    Read top-level expression:
341    define double @1() {
342    entry:
343      %calltmp = call double @testfunc(double 4.000000e+00, double 1.000000e+01)
344      ret double %calltmp
345    }
346
347    Evaluated to 24.000000
348
349    ready> testfunc(5, 10);
350    ready> LLVM ERROR: Program used external function 'testfunc' which could not be resolved!
351
352
353Function definitions and calls also work, but something went very wrong on that
354last line. The call looks valid, so what happened? As you may have guessed from
355the the API a Module is a unit of allocation for the JIT, and testfunc was part
356of the same module that contained anonymous expression. When we removed that
357module from the JIT to free the memory for the anonymous expression, we deleted
358the definition of ``testfunc`` along with it. Then, when we tried to call
359testfunc a second time, the JIT could no longer find it.
360
361The easiest way to fix this is to put the anonymous expression in a separate
362module from the rest of the function definitions. The JIT will happily resolve
363function calls across module boundaries, as long as each of the functions called
364has a prototype, and is added to the JIT before it is called. By putting the
365anonymous expression in a different module we can delete it without affecting
366the rest of the functions.
367
368In fact, we're going to go a step further and put every function in its own
369module. Doing so allows us to exploit a useful property of the KaleidoscopeJIT
370that will make our environment more REPL-like: Functions can be added to the
371JIT more than once (unlike a module where every function must have a unique
372definition). When you look up a symbol in KaleidoscopeJIT it will always return
373the most recent definition:
374
375::
376
377    ready> def foo(x) x + 1;
378    Read function definition:
379    define double @foo(double %x) {
380    entry:
381      %addtmp = fadd double %x, 1.000000e+00
382      ret double %addtmp
383    }
384
385    ready> foo(2);
386    Evaluated to 3.000000
387
388    ready> def foo(x) x + 2;
389    define double @foo(double %x) {
390    entry:
391      %addtmp = fadd double %x, 2.000000e+00
392      ret double %addtmp
393    }
394
395    ready> foo(2);
396    Evaluated to 4.000000
397
398
399To allow each function to live in its own module we'll need a way to
400re-generate previous function declarations into each new module we open:
401
402.. code-block:: c++
403
404    static std::unique_ptr<KaleidoscopeJIT> TheJIT;
405
406    ...
407
408    Function *getFunction(std::string Name) {
409      // First, see if the function has already been added to the current module.
410      if (auto *F = TheModule->getFunction(Name))
411        return F;
412
413      // If not, check whether we can codegen the declaration from some existing
414      // prototype.
415      auto FI = FunctionProtos.find(Name);
416      if (FI != FunctionProtos.end())
417        return FI->second->codegen();
418
419      // If no existing prototype exists, return null.
420      return nullptr;
421    }
422
423    ...
424
425    Value *CallExprAST::codegen() {
426      // Look up the name in the global module table.
427      Function *CalleeF = getFunction(Callee);
428
429    ...
430
431    Function *FunctionAST::codegen() {
432      // Transfer ownership of the prototype to the FunctionProtos map, but keep a
433      // reference to it for use below.
434      auto &P = *Proto;
435      FunctionProtos[Proto->getName()] = std::move(Proto);
436      Function *TheFunction = getFunction(P.getName());
437      if (!TheFunction)
438        return nullptr;
439
440
441To enable this, we'll start by adding a new global, ``FunctionProtos``, that
442holds the most recent prototype for each function. We'll also add a convenience
443method, ``getFunction()``, to replace calls to ``TheModule->getFunction()``.
444Our convenience method searches ``TheModule`` for an existing function
445declaration, falling back to generating a new declaration from FunctionProtos if
446it doesn't find one. In ``CallExprAST::codegen()`` we just need to replace the
447call to ``TheModule->getFunction()``. In ``FunctionAST::codegen()`` we need to
448update the FunctionProtos map first, then call ``getFunction()``. With this
449done, we can always obtain a function declaration in the current module for any
450previously declared function.
451
452We also need to update HandleDefinition and HandleExtern:
453
454.. code-block:: c++
455
456    static void HandleDefinition() {
457      if (auto FnAST = ParseDefinition()) {
458        if (auto *FnIR = FnAST->codegen()) {
459          fprintf(stderr, "Read function definition:");
460          FnIR->dump();
461          TheJIT->addModule(std::move(TheModule));
462          InitializeModuleAndPassManager();
463        }
464      } else {
465        // Skip token for error recovery.
466         getNextToken();
467      }
468    }
469
470    static void HandleExtern() {
471      if (auto ProtoAST = ParseExtern()) {
472        if (auto *FnIR = ProtoAST->codegen()) {
473          fprintf(stderr, "Read extern: ");
474          FnIR->dump();
475          FunctionProtos[ProtoAST->getName()] = std::move(ProtoAST);
476        }
477      } else {
478        // Skip token for error recovery.
479        getNextToken();
480      }
481    }
482
483In HandleDefinition, we add two lines to transfer the newly defined function to
484the JIT and open a new module. In HandleExtern, we just need to add one line to
485add the prototype to FunctionProtos.
486
487With these changes made, lets try our REPL again (I removed the dump of the
488anonymous functions this time, you should get the idea by now :) :
489
490::
491
492    ready> def foo(x) x + 1;
493    ready> foo(2);
494    Evaluated to 3.000000
495
496    ready> def foo(x) x + 2;
497    ready> foo(2);
498    Evaluated to 4.000000
499
500It works!
501
502Even with this simple code, we get some surprisingly powerful capabilities -
503check this out:
504
505::
506
507    ready> extern sin(x);
508    Read extern:
509    declare double @sin(double)
510
511    ready> extern cos(x);
512    Read extern:
513    declare double @cos(double)
514
515    ready> sin(1.0);
516    Read top-level expression:
517    define double @2() {
518    entry:
519      ret double 0x3FEAED548F090CEE
520    }
521
522    Evaluated to 0.841471
523
524    ready> def foo(x) sin(x)*sin(x) + cos(x)*cos(x);
525    Read function definition:
526    define double @foo(double %x) {
527    entry:
528      %calltmp = call double @sin(double %x)
529      %multmp = fmul double %calltmp, %calltmp
530      %calltmp2 = call double @cos(double %x)
531      %multmp4 = fmul double %calltmp2, %calltmp2
532      %addtmp = fadd double %multmp, %multmp4
533      ret double %addtmp
534    }
535
536    ready> foo(4.0);
537    Read top-level expression:
538    define double @3() {
539    entry:
540      %calltmp = call double @foo(double 4.000000e+00)
541      ret double %calltmp
542    }
543
544    Evaluated to 1.000000
545
546Whoa, how does the JIT know about sin and cos? The answer is surprisingly
547simple: The KaleidoscopeJIT has a straightforward symbol resolution rule that
548it uses to find symbols that aren't available in any given module: First
549it searches all the modules that have already been added to the JIT, from the
550most recent to the oldest, to find the newest definition. If no definition is
551found inside the JIT, it falls back to calling "``dlsym("sin")``" on the
552Kaleidoscope process itself. Since "``sin``" is defined within the JIT's
553address space, it simply patches up calls in the module to call the libm
554version of ``sin`` directly.
555
556In the future we'll see how tweaking this symbol resolution rule can be used to
557enable all sorts of useful features, from security (restricting the set of
558symbols available to JIT'd code), to dynamic code generation based on symbol
559names, and even lazy compilation.
560
561One immediate benefit of the symbol resolution rule is that we can now extend
562the language by writing arbitrary C++ code to implement operations. For example,
563if we add:
564
565.. code-block:: c++
566
567    /// putchard - putchar that takes a double and returns 0.
568    extern "C" double putchard(double X) {
569      fputc((char)X, stderr);
570      return 0;
571    }
572
573Now we can produce simple output to the console by using things like:
574"``extern putchard(x); putchard(120);``", which prints a lowercase 'x'
575on the console (120 is the ASCII code for 'x'). Similar code could be
576used to implement file I/O, console input, and many other capabilities
577in Kaleidoscope.
578
579This completes the JIT and optimizer chapter of the Kaleidoscope
580tutorial. At this point, we can compile a non-Turing-complete
581programming language, optimize and JIT compile it in a user-driven way.
582Next up we'll look into `extending the language with control flow
583constructs <LangImpl5.html>`_, tackling some interesting LLVM IR issues
584along the way.
585
586Full Code Listing
587=================
588
589Here is the complete code listing for our running example, enhanced with
590the LLVM JIT and optimizer. To build this example, use:
591
592.. code-block:: bash
593
594    # Compile
595    clang++ -g toy.cpp `llvm-config --cxxflags --ldflags --system-libs --libs core mcjit native` -O3 -o toy
596    # Run
597    ./toy
598
599If you are compiling this on Linux, make sure to add the "-rdynamic"
600option as well. This makes sure that the external functions are resolved
601properly at runtime.
602
603Here is the code:
604
605.. literalinclude:: ../../examples/Kaleidoscope/Chapter4/toy.cpp
606   :language: c++
607
608`Next: Extending the language: control flow <LangImpl5.html>`_
609
610