Implementing our own JIT compiler with LLJIT – JIT Compilation

by Nancy Rohan Exploring the lli tool, Extending the pass pipeline, ITCertification Exams, Simulating function calls

The lli tool is nothing more than a thin wrapper around LLVM APIs. In the first section, we learned that the ORC engine uses a layered approach. The ExecutionSession class represents a running JIT program. Besides other items, this class holds information such as used JITDylib instances. A JITDylib instance is a symbol table that maps symbol names to addresses. For example, these can be symbols defined in an LLVM IR file or the symbols of a loaded shared library.

For executing LLVM IR, we do not need to create a JIT stack on our own, as the LLJIT class provides this functionality. You can also make use of this class when migrating from the older MCJIT implementation, as this class essentially provides the same functionality.

To illustrate the functions of the LLJIT utility, we will be creating an interactive calculator application while incorporating JIT functionality. The main source code of our JIT calculator will be extended from the calc example from Chapter 2, The Structure of a Compiler.

The primary idea behind our interactive JIT calculator will be as follows:

  1. Allow the user to input a function definition, such as def f(x) = x*2.
  2. The function inputted by the user is then compiled by the LLJIT utility into a function – in this case, f.
  3. Allow the user to call the function they have defined with a numerical value: f(3).
  4. Evaluate the function with the provided argument, and print the result to the console: 6.

Before we discuss incorporating JIT functionality into the calculator source code, there are a few main differences to point out with respect to the original calculator example:

  • Firstly, we previously only input and parsed functions beginning with the with keyword, rather than the def keyword described previously. For this chapter, we instead only accept function definitions beginning with def, and this is represented as a particular node in our abstract syntax tree (AST) class, known as DefDecl. The DefDecl class is aware of the arguments and their names it is defined with, and the function name is also stored within this class.
  • Secondly, we also need our AST to be aware of function calls, to represent the functions that the LLJIT utility has consumed or JIT’ted. Whenever a user inputs the name of a function, followed by arguments enclosed in parentheses, the AST recognizes these as FuncCallFromDef nodes. This class essentially is aware of the same information as the DefDecl class.

Due to the addition of these two AST classes, it is obvious to expect that the semantic analysis, parser, and code generation classes will be adapted accordingly to handle the changes in our AST. One additional thing to note is the addition of a new data structure, called JITtedFunctions, which all these classes are aware of. This data structure is a map with the defined function names as keys, and the number of arguments a function is defined with is stored as values within the map. We will see later how this data structure will be utilized in our JIT calculator.

For more details on the changes we have made to the calc example, the full source containing the changes from calc and this section’s JIT implementation can be found within the lljit source directory.