scramjit

Python at native speed.

Live demo · Docs

_{The same function on one clock, native on arm64. scramjit is capped at one fractal frame per GIF frame; uncapped, it does 21x the iterations. Script.}

Usage

import numpy as np
from scramjit import scram

@scram
def ema(x, y, a: float):
    y[0] = x[0]
    for i in range(1, len(x)):
        y[i] = a * x[i] + (1 - a) * y[i - 1]

x = np.random.rand(1_000_000)
y = np.zeros_like(x)
ema(x, y, 0.1)     # first call: compiles for these argument types, about 0.1 s
ema(x, y, 0.1)     # 3 ms; undecorated, 250 ms

scramjit is a JIT written in pure Python. On the first call for a combination of argument types, it generates C specialized to those types and compiles it in-process with cimo, a C compiler that is also pure Python. It needs no LLVM, C compiler or binary wheels. Arrays are passed without copying, and compiled code releases the GIL.

Install

pip install scramjit numpy

Requires CPython 3.9 or newer. numpy is only needed for your arrays.

Live demo

The demo runs Pyodide, scramjit and cimo in the browser, with no compile server, and targets WebAssembly. The code is editable and recompiles as you type.

Examples

examples/ has runnable scripts: n-body gravity, edit distance, union-find, Karplus-Strong audio, a Mandelbrot renderer, Monte Carlo option pricing, threads, and what falling back looks like.

Benchmarks

Arm64 Linux, from bench/kernels.py and bench/m1.py. The hand-written C is compiled by cimo too.

Viewshed pays for exact math semantics; with parity=False it takes 7.3 ms. On Gray-Scott and flood fill, scramjit's C compiled by gcc -O3 matches hand-written C compiled the same way, so the remaining gap there is cimo's backend (bench/gcc.py). Four threads ran four calls 3.5x faster than running them one after another.

How it works

function ──▶ syntax tree ──▶ type inference ──▶ C ──▶ cimo ──▶ machine code
                             from the first call's             arm64, x86-64,
                             argument types                    WebAssembly

The generated C implements Python's semantics, including floor division, negative indices and bounds checks that raise IndexError. range, zip, enumerate, map, filter and comprehensions are lowered to plain loops. Later calls with the same types go straight to the compiled code.

What compiles

** compiles. Where Python's result would change type, a negative k in n ** k (Python gives a float) or a negative float x in x ** 0.5 (a complex), the call runs as Python if the function writes nothing, and raises ValueError if it does.

A function that can't be compiled runs as Python, with a FallbackWarning giving the reason:

import warnings
import numpy as np
from scramjit import scram, CompileError

@scram
def count(a):
    seen = {}
    for i in range(len(a)):
        seen[a[i]] = i
    return len(seen)

with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter('always')
    print(count(np.arange(3.0)))     # 3
print(caught[0].message)             # count(float64[:]) runs as Python: the Dict expression

@scram(fallback=False)               # raise CompileError instead
def strict(a):
    return {0: a[0]}[0]

try:
    strict(np.arange(3.0))
except CompileError as e:
    print(e)

Preflight

scramjit preflight shows what compiles and which lines stop it, without importing your code:

$ scramjit preflight src/
src/sim.py
  ✓ ema        compiled
  ✓ grid       compiled
  ✗ summarize  runs as Python
      22 │     names = {}
         │             ^^ dicts aren't compiled
         │ hint: if the keys are small integers, use an array indexed by key; otherwise do this
         │       outside the compiled function

2 compiled, 1 runs as Python

It also lists undecorated functions with a loop that would compile if decorated, assuming their arguments are numbers or arrays. scramjit preflight --run app.py (or --run -m pytest) reports what happened with the real types: calls, compile time, and time in native code and in Python. Both exit non-zero if a function isn't compiled, so they work as a CI check, and --json gives the report to tools. Add --counts for how often each line ran and each branch was taken, and --live to watch it happen. scramjit preflight FILE:FUNCTION or f.preflight() annotates one function line by line, including which bounds checks were proven, hoisted out of the loop or kept. scramjit c FILE:FUNCTION prints the generated C, and SCRAMJIT_LOG=1 logs each compile and fallback. scramjit --help has the rest.

Compiled code is cached on disk: compiling took 0.53 s in the first process and 0.06 s in later ones on a Gray-Scott kernel. @scram(cache=False) writes nothing.

CPython parity

Compiled functions are tested against the same functions running undecorated, including on randomly generated expressions. This covers floor division and modulo on negatives, ZeroDivisionError, IndexError, math domain errors and numpy's scalar rules.

There are three exceptions: integers are 64-bit and wrap, and the two ** cases above raise ValueError in a function that writes. Compiled calls also skip numpy's warnings and errstate, are invisible to debuggers and coverage, and cost about 1.4 µs each from Python, so a loop should be compiled too, not only a small helper. See Gotchas.

Compiled code runs with the full privileges of the Python process, like an extension module. There is no sandbox.

Status

Pre-release: the API may still change.

• Releases. 0.0.3 is on PyPI. New since 0.0.2: calls between @scram functions, default parameter values, enums, nested tuples, calling C from cimo.compile, arrays made inside a function, module-level constants, arithmetic on float32, uint8, uint64 and the other narrow dtypes, << and >>, exact 64-bit hashes, and int sums that vectorize.
• Tests. About 2,500 differential tests compare each compiled function with CPython and numpy, bit for bit. Developed and benchmarked on arm64 Linux and in Pyodide. Run by hand on arm64 Windows and x86-64 Windows (emulated), where a NaN's sign can differ. cimo also targets x86-64 Linux and macOS, where scramjit's tests haven't been run yet, and nothing runs in CI yet.
• A function that doesn't compile whole has its loops compiled on their own, where they fit.
• Data structures: dicts passed in with keys that don't change, tuples to any depth, enums, local lists, and module-level numbers, tuples, arrays and dicts compile. Lists passed in, nested lists and classes are next.
• Next: see TODO.md.

Development

pip install -e .[dev]       # with cimo installed, or checked out next to this repository
python -m pytest tests
python bench/kernels.py

MIT licensed.