thautwarm / Diojit

DIO-JIT: General-purpose Python JIT

Important:

1. DIO-JIT now works for Python >= 3.8. We heavily rely on the LOAD_METHOD bytecode instruction.
2. DIO-JIT is not production-ready. a large number of specialisation rules are required to make DIO-JIT batteries-included.
3. This document is mainly provided for prospective developers. Users are not required to write any specialisation rules, which means that users need to learn nothing but @jit.jit and jit.spec_call.

Benchmark

The benchmark item "DNA READ" does not show a significant performance gain, this is because "DNA READ" heavily uses bytearray and bytes, whose specialised C-APIs are not exposed. In this case, although the JIT can infer the types, we have to fall back to CPython's default behaviour, or even worse: after all, the interpreter can access internal things, while we cannot.

P.S: DIO-JIT can do very powerful partial evaluation, which is disabled in default but you can leverage it in your domain specific tasks. Here is an example of achieving 500x speed up against pure Python: fibs.py

Install Instructions

$ pip install jill && jill install 1.6 --upstream Official

$ bash -ci "$(curl -fsSL https://raw.githubusercontent.com/abelsiqueira/jill/master/jill.sh)"

julia>
# press ]
pkg> add https://github.com/thautwarm/DIO.jl
# press backspace
julia> using DIO # precompile

$ pip install git+https://github.com/thautwarm/diojit

$ pip install -U diojit
$ julia -e "using Pkg; Pkg.update(string(:DIO));using DIO"

Usage from Python side is quite similar to that from Numba.

import diojit
from math import sqrt
# eagerjit: assuming all global references are fixed
@diojit.eagerjit
def fib(a):
    if a <= 2:
        return 1
    return fib(a + -1) + fib(a + -2)

jit_fib = diojit.spec_call(fib, diojit.oftype(int), diojit.oftype(int))
jit_fib(15) # 600% faster than pure python

It might look strange to you that we use a + -1 and a + -2 here.

Clever observation! And that's the point!

DIO-JIT relies on specilisation rules. We have written one for additions, more specifically, operator.__add__: specilisation for operator.__add__.

However, due to the bandwidth limitation, rules for operator.__sub__ is not implemented yet.

(P.S: why operator.__add__.)

Although specilisation is common in the scope of optimisation, unlike many other JIT attempts, DIO-JIT doesn't need to hard encode rules at compiler level. The DIO-JIT compiler implements the skeleton of abstract interpretation, but concrete rules for specialisation and other inferences can be added within Python itself in an extensible way!

See an example below.

Contribution Example: Add a specialisation rule for list.append

1. Python Side:

import diojit as jit
import timeit
jit.create_shape(list, oop=True)
@jit.register(list, attr="append")
def list_append_analysis(self: jit.Judge, *args: jit.AbsVal):
    if len(args) != 2:
        # rollback to CPython's default code
        return NotImplemented
    lst, elt = args

    return jit.CallSpec(
        instance=None,  # return value is not static
        e_call=jit.S(jit.intrinsic("PyList_Append"))(lst, elt),
        possibly_return_types=tuple({jit.S(type(None))}),
    )

jit.intrinsic("PyList_Append") mentioned in above code means the intrinsic provided by the Julia codegen backend. Usually it's calling a CPython C API, but sometimes may not.

No matter if it is an existing CPython C API, we can implement intrinsics in Julia.

• import PyList_Append symbol

• generate PyList_Append calling convention:

  @autoapi PyList_Append(PyPtr, PyPtr)::Cint != Cint(-1) cast(_cint2none) nocastexc
  

  As a consequence, we automatically generate an instrinsic function for DIO-JIT. This intrinsic function is capable of handling CPython exception and reference counting.

You can either do step 2) at Python side. It might looks more intuitive.

import diojit as jit
from diojit.runtime.julia_rt import jl_eval
jl_implemented_intrinsic = """
function PyList_Append(lst::Ptr, elt::PyPtr)
    if ccall(PyAPI.PyList_Append, Cint, (PyPtr, PyPtr), lst, elt) == -1
        return Py_NULL
    end
    nothing # automatically maps to a Python None
end
DIO.DIO_ExceptCode(::typeof(PyList_Append)) != Py_NULL
"""
jl_eval(jl_implemented_intrinsic)

You immediately get a >100% time speed up:

@jit.jit
def append3(xs, x):
    xs.append(x)
    xs.append(x)
    xs.append(x)

jit_append3 = jit.spec_call(append3, jit.oftype(list), jit.Top) # 'Top' means 'Any'
xs = [1]
jit_append3(xs, 3)

print("test jit_append3, [1] append 3 for 3 times:", xs)
# test jit func, [1] append 3 for 3 times: [1, 3, 3, 3]

xs = []
%timeit append3(xs, 1)
# 293 ns ± 26.2 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

xs = []
%timeit jit_append3(xs, 1)
# 142 ns ± 14.9 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

Why Julia?

We don't want to maintain a C compiler, and calling gcc or others will introduce cross-process IO, which is slow. We prefer compiling JITed code with LLVM, and Julia is quite a killer tool for this use case.

Current Limitations

1. Support for *varargs and **kwargs are not ready: we do can immediately support them with very tiny JIT performance gain, but considering backward compatibility we decide not to do this.

2. Exception handling is not yet supported inside JIT functions.

   Why?

   We haven't implemented the translation from exception handling bytecode to untyped DIO IR (jit.absint.abs.In_Stmt).

   Will support?

   Yes.

   In fact, now a callsite in any JIT function can raise an exception. It will not be handled by JIT functions, instead, it is lifted up to the root call, which is a pure Python call.

   Exception handling will be supported when we have efforts on translating CPython bytecode about exception handling into untyped DIO IR (jit.absint.abs.In_Stmt).

   P.S: This will be finished simultaneously with the support for for loop.

3. Support for for loop is missing.

   Why?

   Firstly, in CPython, for loop relies on exception handling, which is not supported yet.

   Secondly, we're considering a fast path for for loop, maybe proposing a __citer__ protocol for faster iteration for JIT functions, which requires communications with Python developers.

   Will support?

   Yes.

   This will be finished simultaneously with support for exception handling (faster for loop might come later).

4. Closure support is missing.

   Why?

   In imperative languages, closures use cell structures to achieve mutable free/cell variables.

   However, a writable cell makes it hard to optimise in a dynamic language.

   We recommend using types.MethodType to create immutable closures，which can be highly optimised in DIO-JIT(near future).

   import types
   def f(freevars, z):
           x, y = freevars
           return x + y + z
   
   def hof(x, y):
       return types.MethodType(f, (x, y))
   

   Will support?

   Still yes. However, don't expect much about the performance gain for Python's vanilla closures.

5. Specifying fixed global references(@diojit.jit(fixed_references=['isinstance', 'str', ...]) too annoying?

   Sorry, you have to. We are thinking about the possibility about automatic JIT covering all existing CPython code, but the biggest impediment is the volatile global variables.

   You might use @eagerjit, and in this case you'd be cautious in making global variables unchangeable.

   Possibility?

   Recently we found CPython's newly(:)) added feature Dict.ma_version_tag might be used to automatically notifying JITed functions to re-compile when the global references change.

   More research is required.

Contributions

1. Add more prescribed specialisation rules at jit.absint.prescr.
2. TODO

Benchmarks

Check benchmarks directory.