Jit Header Image

Preface

This article references Call Go From Jit by Iskander Sharipov. It is recommended you read this first, even if this solution does not work in modern version of Go.

The method outlined in this article only works with Go versions using both ABI0 and ABI Internal (1.17 - ???). It has been tested on version between 1.17 and 1.25. If Go decides to write a breaking change: removing ABI0, adding ABI1 etc, it will not longer work.

Intro

I have started the process of adding a JIT compiler (Just in Time) to Guac, my Gameboy, Gameboy Advance, and (soon!) Nintendo DS Emulator, written in golang. This has been the most complicated and most rewarding project I have worked on. A JIT compiler allows for my emulated cpus to reach 2-5x the previous speed by batching instructions into blocks that run in pure assembly instead of decoding and handling in the native golang. This is achieved by dynamically generating machine code - like a train placing its own tracks while it is running.

The Problem

Golang (exe) -> Jit (handrolled asm) -> Golang function (exe)

A JIT itself can be difficult to implement, but the biggest complication in Guac has been the fact that golang must call the jit code, which then must call golang functions. for example a jit instruction sometimes must read the emulated virtual memory, using a “Read(addr uint32)” golang function. This in of itself can be fairly simple. However for more complicated golang functions, stack checks occur during stack growing, garbage collection etc. These checks require the runtime to walk up the stack to ensure it can adjust and update the stack properly. If the functions are not setup properly to match the ABI (application binary interface), then a crash will occur, usually stating unexcepted return pc, or unknown signal.

runtime: unexpected return pc for main.goFunc called from 0x7f9465f7c007

Most Jit and realtime compilers in Go refuse to let this occur by not letting your call go functions from jit code. Others such as nelhage/gojit and forks simply do not handle or setup properly for the abi (they may have worked in version 1.4 available at the time but I could still not get it to work). You can even find emulator developer Giovanni Bajo (rasky) asking for methods of handling this problem back in 2017. Simply put jit code is not supported by golang. BUT…

An Outdated Solution

Abi Image

I did run into one solution, written by Iskander Sharipov, a Go performance engineer who, through Intel and Huawai, added AVX512 support to Go, and other optimizations. Their method works flawlessly on Go 1.16 and earlier. This solution unfortunately no longer works - however I do recommend reading it to understand how we got here.

In Go 1.17 the ABI became split. ABI0 is the older stack based abi, where all function arguments and results were placed on the stack, while the new ABI Internal uses a register based abi, where function arguments and results are placed in registers instead of interacting with memory, argument/result count permitting. This lead to a ~5% speed increase. To not break Go Assembly code already in existence, the Go team decided to add autogenerated functions that convert between the abis.

This autogenerated function takes the assembly symbol name from the original function and gives the original a name change, for example callJIT becomes callJIT.abi0. In reference to Sharipov’s assembly call, this:

#include "funcdata.h"

TEXT ·callJIT(SB), 0, $8-8
    NO_LOCAL_POINTERS
    MOVQ code+0(FP), AX
    JMP AX
gocall:
    CALL CX
    JMP (SP)

No longer works since the callJIT symbol now points to a ABI Internal function that sets up the stack and calls the original callJIT (now called callJIT.abi0). This means when the offset to the “gocall” label is setup “j := funcAddr(callJIT) + 36” instead of pointing 36 bytes after the original callJIT function which we want, it instead points 36 bytes after the autogenerated callJIT function used to convert ABI Internal to ABI0.

The Fix

My solution is to build an additional asm function which gets the proper address of the callJIT(ABI0). This is then used as an alternative to “j := funcAddr(callJIT) + 36” to get the proper function base.

Asm:

TEXT ·callJITImplAddr(SB), 0, $0-8
    NO_LOCAL_POINTERS
    MOVQ $·callJIT(SB), AX  // address of ABI0 impl, not trampoline
    MOVQ AX, ret+0(FP)
    RET

j := callJITImplAddr() + 36

/* stub */
func callJITImplAddr() uintptr

Since both callJIT and callJITImplAddr are ABI0, there is not interaction with the ABI Internal autogenerated functions, allowing us to get the proper pointer value. We then just use this function to offset to the gocall label mentioned by Sharipov.

The code is available in the following repository: https://github.com/aabalke/jit-proof-of-concept/

Now, calling Go functions from Jit code is possible in 1.17+, including garbage collection, and recursion, which both cause stack checks. I have included a few additional adjustments, the offset is searched for initially, since OS and Go version will cause the offset to move slightly based on the instructions used. Additionally, I added functions to allow executable memory mapping on windows, replaced MOVQ with MOVABS for pointers and updated the funcAddr to use the modern reflect package instead.

With these changes I am hoping to build my own jit compiler / assembler based on nelhage/gojit but supporting go function calls. I will also be updating my emulator to include a jit with this functionality in the near future.

Resources

Call Go From Jit
nelhag/gojit
golang/go Issue 20123
Go internal ABI specification
Go 1.17 Release Notes

Calling Go Functions from Jit Code