blog: add ExoJIT post

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()
+```