LLVM & Haskell (PART II)

In the first part we have seen an introduction to LLVM with Haskell, also how to transform the Abstract Syntax Tree into LLVM IR, some code generation, the Lexer and the Parser. In this part “2”, we will continue our discussion and go deeper into JIT and Control flow. I expect a third part and it will be the last one.

Optimization Passes:

LLVM provides many optimization passes, which do many different sorts of things and have different trade-offs. Unlike other systems, LLVM doesn’t hold to the mistaken notion that one set of optimizations is right for all languages and for all situations. LLVM allows a compiler implementor to make complete decisions about what optimizations to use, in which order, and in what situation.

As a concrete example, LLVM supports both “whole module” passes, which look across as large of body of code as they can (often a whole file, but if run at link time, this can be a substantial portion of the whole program). It also supports and includes “per-function” passes which just operate on a single function at a time, without looking at other functions. For more information on passes and how they are run, see the How to Write a Pass document and the List of LLVM Passes.

For Kaleidoscope, we are currently generating functions on the fly, one at a time, as the user types them in. We aren’t shooting for the ultimate optimization experience in this setting, but we also want to catch the easy and quick stuff where possible.

We won’t delve too much into the details of the passes since they are better described elsewhere. We will instead just invoke the default “curated passes” with an optimization level which will perform most of the common clean-ups and a few non-trivial optimizations.

passes :: PassSetSpec

passes = defaultCuratedPassSetSpec { optLevel = Just 3 }

To apply the passes we create a bracket for a PassManager and invoke runPassManager on our working module. Note that this modifies the module in-place.

runJIT :: AST.Module -> IO (Either String AST.Module) runJIT mod = do

. withContext $ \context ->

. runExceptT $ withModuleFromAST context mod $ \m ->

withPassManager passes $ \pm -> do runPassManager pm m

optmod <- moduleAST m

s <- moduleLLVMAssembly m putStrLn s

. return optmod

With this in place, we can try our test above again:

ready> def test(x) (1+2+x)*(x+(1+2)); ; ModuleID = ‘my cool jit’

Function Attrs: nounwind readnone define double @test(double %x) #0 { entry:
%0 = fadd double %x, 3.000000e+00
%1 = fmul double %0, %0
ret double %1
}

attributes #0 = { nounwind readnone }

As expected, we now get our nicely optimized code, saving a floating point add instruction from every execution of this function. We also see some extra metadata attached to our function, which we can ignore for now, but is indicating certain properties of the function that aid in later optimization.

LLVM provides a wide variety of optimizations that can be used in certain circumstances. Some documentation about the various passes is available, but it isn’t very complete. Another good source of ideas can come from looking at the passes that Clang runs to get started. The “opt” tool allows us to experiment with passes from the command line, so we can see if they do anything.

One important optimization pass is an “analysis pass” which will validate that the internal IR is well-formed. Since it quite possible (even easy!) to construct nonsensical or unsafe IR it is very good practice to validate our IR before attempting to optimize or execute it. To do so, we simply invoke the verify function with our active module.

runJIT :: AST.Module -> IO (Either String AST.Module) runJIT mod = do

…

withPassManager passes $ \pm -> do

runExceptT $ verify m

Now that we have reasonable code coming out of our front-end, let’s talk about executing it!

Adding a JIT Compiler:

Code that is available in LLVM IR can have a wide variety of tools applied to it. For example, we can run optimizations on it (as we did above), we can dump it out in textual or binary forms, we can compile the code to an assembly file (.s) for some target, or we can JIT compile it. The nice thing about the LLVM IR representation is that it is the “common currency” between many different parts of the compiler.

In this section, we’ll add JIT compiler support to our interpreter. The basic idea that we want for Kaleidoscope is to have the user enter function bodies as they do now, but immediately evaluate the top-level expressions they type in. For example, if they type in “1 + 2;”, we should evaluate and print out 3. If they define a function, they should be able to call it from the command line.

In order to do this, we add another function to bracket the creation of the JIT Execution Engine. There are two provided engines: jit and mcjit. The distinction is not important for us but we will opt to use the newer mcjit.

import qualified LLVM.ExecutionEngine as EE

jit :: Context -> (EE.MCJIT -> IO a) -> IO a

jit c = EE.withMCJIT c optlevel model ptrelim fastins

where

optlevel = Just 2 – optimization level

model = Nothing – code model ( Default ) ptrelim = Nothing – frame pointer elimination fastins = Nothing – fast instruction selection

The result of the JIT compiling our function will be a C function pointer which we can call from within the JIT’s process space. We need some (unsafe!) plumbing to coerce our foreign C function into a callable object from Haskell. Some care must be taken when performing these operations since we’re telling Haskell to “trust us” that the pointer we hand it is actually typed as we describe it. If we don’t take care with the casts we can expect undefined behavior.

External Functions:

The JIT provides a number of other more advanced interfaces for things like freeing allocated machine code, rejit’ing functions to update them, etc. However, even with this simple code, we get some surprisingly powerful capabilities – check this out:

ready> extern sin(x);

; ModuleID = ‘my cool jit’

declare double @sin(double) ready> extern cos(x);

; ModuleID = ‘my cool jit’ declare double @sin(double) declare double @cos(double) ready> sin(1.0);

; ModuleID = ‘my cool jit’ declare double @sin(double) declare double @cos(double)

define double @main() { entry:
%0 = call double @sin(double 1.000000e+00)
ret double %0 }

Evaluated to: 0.8414709848078965

Whoa, how does the JIT know about sin and cos? The answer is surprisingly simple: in this example, the JIT started execution of a function and got to a function call. It realized that the function was not yet JIT compiled and invoked the standard set of routines to resolve the function. In this case, there is no body defined for the function, so the JIT ended up calling dlsym(“sin”) on the Kaleidoscope process itself. Since “sin” is defined within the JIT’s address space, it simply patches up calls in the module to call the libm version of sin directly.

The LLVM JIT provides a number of interfaces for controlling how unknown functions get resolved. It allows us to establish explicit mappings between IR objects and addresses (useful for LLVM global variables that we want to map to static tables, for example), allows us to dynamically decide on the fly based on the function name, and even allows us JIT compile functions lazily the first time they’re called.

Control flow:

‘if’ Expressions

Extending Kaleidoscope to support if/then/else is quite straightforward. It basically requires adding lexer support for this “new” concept to the lexer, parser, AST, and LLVM code emitter. This example is nice, because it shows how easy it is to “grow” a language over time, incrementally extending it as new ideas are discovered.

Before we get going on “how” we add this extension, let’s talk about “what” we want. The basic idea is that we want to be able to write this sort of thing:

def fib(x)
    if x < 3 then
       1
    else
       fib(x-1) + fib(x-2);

In Kaleidoscope, every construct is an expression: there are no statements. As such, the if/then/else expression needs to return a value like any other. Since we’re using a mostly functional form, we’ll have it evaluate its conditional, then return the ‘then’ or ‘else’ value based on how the condition was resolved. This is very similar to the C “?:” expression.

The semantics of the if/then/else expression is that it evaluates the condition to a boolean equality value: 0.0 is considered to be false and everything else is considered to be true. If the condition is true, the first subexpression is evaluated and returned, if the condition is false, the second subexpression is evaluated and returned. Since Kaleidoscope allows side-effects, this behavior is important to nail down.

Now that we know what we “want”, let’s break this down into its constituent pieces.

data Expr

…

| If Expr Expr Expr deriving (Eq, Ord, Show)

We also extend our lexer definition with the new reserved names.

lexer :: Tok.TokenParser ()

lexer = Tok.makeTokenParser style

where

ops = [“+”,”*”,”-”,”/”,”;”,”,”,”<”]

names = [“def”,”extern”,”if”,”then”,”else”]

style = emptyDef {
Tok.commentLine = “#”
Tok.reservedOpNames = ops
Tok.reservedNames = names
}

Now that we have it parsing and building the AST, the final piece is adding LLVM code generation support. This is the most interesting part of the if/then/else example, because this is where it starts to introduce new concepts. All of the code above has been thoroughly described in previous chapters.

To motivate the code we want to produce, let’s take a look at a simple example. Consider:

extern foo();
extern bar();
def baz(x) if x then foo() else bar();
declare double @foo()
declare double @bar()
define double @baz(double %x) {
entry:
%ifcond = fcmp one double %x, 0.000000e+00
br i1 %ifcond, label %then, label %else
then: ; preds = %entry
%calltmp = call double @foo()
br label %ifcont
else: ; preds = %entry %calltmp1 = call double @bar() br label %ifcont
ifcont: ; preds = %else, %then
%iftmp = phi double [ %calltmp, %then ], [ %calltmp1, %else ]
ret double %iftmp
}

To visualize the control flow graph, we can use a nifty feature of the LLVM opt tool. If we put this LLVM IR into “t.ll” and run

$ llvm-as < t.ll | opt -analyze -view-cfg A window will pop up and we’ll see this graph:

LLVM has many nice features for visualizing various graphs, but note that these are available only if your LLVM was built with Graphviz support (accomplished by having Graphviz and Ghostview installed when building LLVM).

Getting back to the generated code, it is fairly simple: the entry block evaluates the conditional expression (“x” in our case here) and compares the result to 0.0 with the fcmp one instruction (one is “Ordered and Not Equal”). Based on the result of this expression, the code jumps to either the “then” or “else” blocks, which contain the expressions for the true/false cases.

Once the then/else blocks are finished executing, they both branch back to the if.exit block to execute the code that happens after the if/then/else. In this case the only thing left to do is to return to the caller of the function. The question then becomes: how does the code know which expression to return?

The answer to this question involves an important SSA operation: the Phi operation. If you’re not familiar with SSA, the Wikipedia article is a good introduction and there are various other introductions to it available on your favorite search engine.

In order to generate code for this, we implement the Codegen method for If node:

cgen (S.If cond tr fl) = do ifthen <- addBlock “if.then” ifelse <- addBlock “if.else” ifexit <- addBlock “if.exit”
– %entry
– – – – – – – – –
cond <- cgen cond
test <- fcmp FP.ONE false cond
cbr test ifthen ifelse – Branch based on the condition
– if.then
– – – – – – – – –

setBlock ifthen

Get Idriss Chaouch’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

trval <- cgen tr

br ifexit

ifthen <- getBlock

– if.else

– – – – – – – – –

setBlock ifelse

flval <- cgen fl

br ifexit

ifelse <- getBlock

– if.exit

– – – – – – – – –

– Generate code for the true branch

– Branch to the merge block

– Generate code for the false branch

– Branch to the merge block

setBlock ifexit

phi double [(trval, ifthen), (flval, ifelse)]

We start by creating three blocks:

ifthen <- addBlock “if.then”

ifelse <- addBlock “if.else”

ifexit <- addBlock “if.exit”

Next emit the expression for the condition, then compare that value to zero to get a truth value as a 1-bit (i.e. bool) value. We end this entry block by emitting the conditional branch that chooses between the two cases.

‘for’ Loop Expressions

Now that we know how to add basic control flow constructs to the language, we have the tools to add more powerful things. Let’s add something more aggressive, a ‘for’ expression:

extern putchard(char);
def printstar(n)
for i = 1, i < n, 1.0 in
putchard(42); # ascii 42 = ‘*’
# print 100 ‘*’ characters
printstar(100);

This expression defines a new variable (i in this case) which iterates from a starting value, while the condition (i < n in this case) is true, incrementing by an optional step value (1.0 in this case). While the loop is true, it executes its body expression. Because we don’t have anything better to return, we’ll just define the loop as always returning 0.0. In the future when we have mutable variables, it will get more useful.

To get started, we again extend our lexer with new reserved names “for” and “in”.

lexer :: Tok.TokenParser ()

lexer = Tok.makeTokenParser style

where

ops = [“+”,”*”,”-”,”/”,”;”,”,”,”<”]

names = [“def”,”extern”,”if”,”then”,”else”,”in”,”for”]

style = emptyDef {

Tok.commentLine = “#”

, Tok.reservedOpNames = ops

, Tok.reservedNames = names

}

As before, let’s talk about the changes that we need to Kaleidoscope to support this. The AST node is just as simple. It basically boils down to capturing the variable name and the constituent expressions in the node.

data Expr

…

| For Name Expr Expr Expr Expr deriving (Eq, Ord, Show)

The parser code captures a named value for the iterator variable and the four expressions objects for the parameters of the loop parameters.

for :: Parser Expr for = do

reserved “for”

var <- identifier

reservedOp “=”

start <- expr

reservedOp “,”

cond <- expr

reservedOp “,”

step <- expr

reserved “in”

body <- expr

return $ For var start cond step body

Now we get to the good part: the LLVM IR we want to generate for this thing. With the simple example above, we get this LLVM IR (note that this dump is generated with optimizations disabled for clarity):

declare double @putchard(double)
define double @printstar(double %n) {
entry:
br label %loop
loop:
%i = phi double [ 1.000000e+00, %entry ], [ %nextvar, %loop ]
%calltmp = call double @putchard(double 4.200000e+01)
%nextvar = fadd double %i, 1.000000e+00
%cmptmp = fcmp ult double %i, %n
%booltmp = uitofp i1 %cmptmp to double
%loopcond = fcmp one double %booltmp, 0.000000e+00
br i1 %loopcond, label %loop, label %afterloop
afterloop:
ret double 0.000000e+00
}

The code to generate this is only slightly more complicated than the above “if” statement.

cgen (S.For ivar start cond step body) = do forloop <- addBlock “for.loop”

forexit <- addBlock “for.exit”

– %entry

– – – – – – – – –

i <- alloca double

istart <- cgen start

stepval <- cgen step

– Generate loop variable initial value

– Generate loop variable step

store i istart

. assign ivar i

. br forloop

– for.loop

. – – – – – – – – –

. setBlock forloop

. cgen body

. ival <- load i

. inext <- fadd ival stepval

. store i inext

cond <- cgen cond

test <- fcmp FP.ONE false cond – Test if the loop condition is True ( 1.0 ) cbr test forloop forexit – Generate the loop condition

The first step is to set up the LLVM basic block for the start of the loop body. In the case above, the whole loop body is one block, but remember that the generating code for the body of the loop could consist of multiple blocks (e.g. if it contains an if/then/else or a for/in expression).

. forloop <- addBlock “for.loop”

. forexit <- addBlock “for.exit”

Next we allocate the iteration variable and generate the code for the constant initial value and step.

i <- alloca double

istart <- cgen start – Generate loop variable initial value stepval <- cgen step – Generate loop variable step

Now the code starts to get more interesting. Our ‘for’ loop introduces a new variable to the symbol table. This means that our symbol table can now contain either function arguments or loop variables. Once the loop variable is set into the symbol table, the code recursively codegen’s the body. This allows the body to use the loop variable: any references to it will naturally find it in the symbol table.

store i istart

assign ivar i

br forloop

– Store the loop variable initial value

– Assign loop variable to the variable name

– Branch to the loop body block

Now that the “preheader” for the loop is set up, we switch to emitting code for the loop body.

Have a great day, keep safe.

TO BE CONTINUED..