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.

foreign import ccall "dynamic" haskFun :: FunPtr (IO Double) -> (IO Double)
run :: FunPtr a -> IO Double
run fn = haskFun (castFunPtr fn :: FunPtr (IO Double))

Integrating this with our function from above we can now manifest our IR as executable code inside the ExecutionEngine and pass the resulting native types to and from the Haskell runtime.

runJIT :: AST.Module -> IO (Either String ())
runJIT mod = do
    ...
    jit context $ \executionEngine ->
        ...
        EE.withModuleInEngine executionEngine m $ \ee -> do
          mainfn <- EE.getFunction ee (AST.Name "main")
          case mainfn of
            Just fn -> do
              res <- run fn
              putStrLn $ "Evaluated to: " ++ show res
            Nothing -> return ()

Having to statically declare our function pointer type is rather inflexible. If we wish to extend to this to be more flexible, a library like libffi is very useful for calling functions with argument types that can be determined at runtime.