From b381a532ac920a4b7b9bf937c99489d1f092d0b4 Mon Sep 17 00:00:00 2001 From: Yahya Jabary Date: Wed, 21 Oct 2026 18:55:29 +0200 Subject: [PATCH] blog: add ExoJIT post --- docs/blog/posts/blog/2026-10-21-exojit.md | 97 +++++++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100644 docs/blog/posts/blog/2026-10-21-exojit.md diff --git a/docs/blog/posts/blog/2026-10-21-exojit.md b/docs/blog/posts/blog/2026-10-21-exojit.md new file mode 100644 index 000000000000..2a3086f81c5d --- /dev/null +++ b/docs/blog/posts/blog/2026-10-21-exojit.md @@ -0,0 +1,97 @@ +--- +date: 2026-10-21 +title: ExoJIT +tags: + - technical +--- + +User-schedulable languages like the [Exo](https://github.com/exo-lang/exo) +Python-DSL let users write numerical kernels and declaratively optimize their +scheduling (e.g. loop unrolling). + +So far, benchmarking any given schedule in the Python-DSL and invoking it in a +CPython runtime required several tedious steps. + +This is because Exo's compiler emits C and header files, requiring separate +compilation, linking process and FFI calls or CPython extensions. + +We introduce [exojit](https://github.com/sueszli/exojit), a just-in-time +compiler for Exo procedures that alleviates all these pain-points, by letting +users annotate their Exo functions with `@jit(optimize=...)` and invoke them +like regular Python functions. + +The Exo compiler applies the schedule and verifies that invariances are held. +Then we translate Exo's LoopIR through xDSL to LLVM, use the MCJIT engine to +generate the binary and bind it with CFFI. + +Users can now write, schedule, execute, benchmark and verify a kernel all in the +same Python script, making experimentation with Exo a lot more pleasant to +integrate into their existing workflow and opening up the door to faster +iteration cycles for humans, LLMs and auto-schedulers. + +The following demonstrates loop fission, reordering, splitting and unrolling on +a matmul GEMM, that applies the `@jit` decorator: + +```python +# /// script +# requires-python = "==3.14.*" +# dependencies = ["exojit @ git+https://github.com/sueszli/exojit"] +# /// + +import math +import random + +from exo import * +from exo.API import Procedure +from exo.stdlib.scheduling import divide_loop, fission, reorder_loops, simplify, unroll_loop + +from exojit import jit + + +# The schedule transforms matmul into: +# +# for i in range(32): +# for j in range(32): +# out[i, j] = 0.0 +# for i in range(32): +# for k in range(32): +# for jo in range(8): +# j = 4 * jo +# out[i, j] += A[i, k] * B[k, j] +# out[i, j + 1] += A[i, k] * B[k, j + 1] +# out[i, j + 2] += A[i, k] * B[k, j + 2] +# out[i, j + 3] += A[i, k] * B[k, j + 3] +def optimize_matmul(p: Procedure) -> Procedure: + p = fission(p, p.find("for k in _: _").before(), n_lifts=2) + p = reorder_loops(p, "j k") + p = divide_loop(p, "j #1", 4, ["jo", "ji"], perfect=True) + p = unroll_loop(p, "ji") + return simplify(p) + + +@jit(optimize=optimize_matmul) +def matmul( + out: f32[32, 32] @ DRAM, + A: f32[32, 32] @ DRAM, + B: f32[32, 32] @ DRAM, +): + for i in seq(0, 32): + for j in seq(0, 32): + out[i, j] = 0.0 + for k in seq(0, 32): + out[i, j] += A[i, k] * B[k, j] + + +def test_matmul(): + A = [[random.uniform(-1.0, 1.0) for _ in range(32)] for _ in range(32)] + B = [[random.uniform(-1.0, 1.0) for _ in range(32)] for _ in range(32)] + out = [[0.0] * 32 for _ in range(32)] + expected = [[sum(A[i][k] * B[k][j] for k in range(32)) for j in range(32)] for i in range(32)] + + matmul(out, A, B) + assert all(math.isclose(out[i][j], expected[i][j], rel_tol=1e-5, abs_tol=1e-5) for i in range(32) for j in range(32)) + + +if __name__ == "__main__": + test_matmul() +```