POC JIT With Go (plugins)

I recently implemented a very rudimentary JIT compiler in go. This compiler generates go source code, compiles it by invoking the go compiler and loading the resulting shared object into the current process and invoking the compiled function.

The scope is currently limited to arithmetic expressions. I implemented a tree-walk interpreter as well as a byte code compiler and virtual machine to compare this JIT to. The front-end (lexer, parser) is shared, the back-end is specific to the approach of evaluation.

As mentioned before the scope I choose for this compiler is limited to arithmetic expressions, thus I want to accept the following:

11+1.2-5

Tokenizing

We convert our character stream to tokens:

1NUMBER 1
2PLUS
3NUMBER 1.2
4MINUS
5NUMBER 5

Parsing

We parse the tokens and produce an abstract syntax tree via recursive descent:

 1Binary {
 2    Type: MINUS,
 3    Left: Binary {
 4        Type: PLUS,
 5        Left: Number {
 6            Value: 1
 7        }
 8        Right: Number {
 9            Value: 1.2
10        }
11    }
12    Right: Number {
13       Value: 5
14    }
15}

Code-generation

I want to generate go code from the abstract syntax tree, I know for arithmetics its pretty idiotic, but i want to evaluate the pipeline in a comparable way to byte code compilation+evaluation and tree-walk interpreting. Thus we generate the following go source code:

1package main
2func Main()float64{return 1E+00+1.2E+00-5E+00}

Compilation

The go compiler sits in an internal go package, see its repo here and I could not include it in my JIT, thus I had to use os/exec for invoking the compiler:

 1package main
 2
 3import (
 4    "os"
 5    "os/exec"
 6    "fmt"
 7)
 8
 9func JIT() (func() float64, error) {
10    generatedCode := `package main
11func Main() float64 {
12    return 1E+00+1.2E+00-5E+00
13}
14   `
15    err := os.WriteFile("jit_output.go", []byte(generatedCode), 0777)
16    if err != nil {
17        return nil, err
18    }
19    cmd := exec.Command("go", "build", "-buildmode=plugin", "-o", "jit_output.so", "jit_output.go")
20    err := cmd.Run()
21    if err != nil {
22        return nil, err
23    }
24
25
26
27    return nil, nil
28}
29
30func main() {
31    function, err := JIT()
32    if err != nil {
33        panic(err)
34    }
35    fmt.Println(function())
36}

This first step calls the compiler and wants it to build a go plugin.

1[...]
2
3	-buildmode mode
4		build mode to use. See 'go help buildmode' for more.
5
6[...]

 1The 'go build' and 'go install' commands take a -buildmode argument which
 2indicates which kind of object file is to be built. Currently supported values
 3are:
 4
 5[...]
 6
 7	-buildmode=plugin
 8		Build the listed main packages, plus all packages that they
 9		import, into a Go plugin. Packages not named main are ignored.
10[...]

Now a go plugin called jit_output.so sits in our project root.

Go plugins

Go features a package called plugin for loading and resolving go symbols of go plugins. There are some drawbacks such as missing portability due to no windows support and easily exploitable bugs in plugin loaders - since this is a POC JIT requiring the go compiler toolchain in the path I’m not even going to walk down the road of portability.

Since we know the location of the plugin and the symbols in it our interactions with the plugin packages should be minimal - I need to open the plugin file, locate the Main function, cast it to func() float64 and return the result.

 1package main
 2
 3import (
 4    "os"
 5    "os/exec"
 6    "fmt"
 7    "plugin"
 8    "errors"
 9)
10
11func JIT() (func() float64, error) {
12    generatedCode := `package main
13func Main() float64 {
14    return 1E+00+1.2E+00-5E+00
15}
16   `
17    err := os.WriteFile("jit_output.go", []byte(generatedCode), 0777)
18    if err != nil {
19        return nil, err
20    }
21    cmd := exec.Command("go", "build", "-buildmode=plugin", "-o", "jit_output.so", "jit_output.go")
22    err := cmd.Run()
23    if err != nil {
24        return nil, err
25    }
26
27    plug, err := plugin.Open("jit_output.so")
28    if err != nil {
29        return nil, err
30    }
31
32	symbol, err := plug.Lookup("Main")
33	if err != nil {
34		fmt.Println(sharedObjectPath)
35		return nil, err
36	}
37
38	Main, ok := symbol.(func() float64)
39	if !ok {
40		return nil, errors.New("Error while accessing jit compiled symbols")
41	}
42
43	return Main, nil
44}
45
46func main() {
47    function, err := JIT()
48    if err != nil {
49        panic(err)
50    }
51    fmt.Println(function())
52}

I know this JIT compiler is very rudimentary and uses the go tool chain, however I plan on adding a bytecode compiler and a bytecode vm to the sophia programming language. This vm will include meta tracing capabilities and a JIT compiler for turning hot paths into go source code and compiling it on the fly.

What to do with this JIT

My plan is to compare the three different approaches to evaluation techniques in a blog article in the next few weeks, expect some in depth benchmarks.