Code generation

Alkahest can compile symbolic expressions to fast native or GPU code. Compiled code bypasses Python entirely during evaluation.

The compilation pipeline

Expressions lower through multiple IR levels:

ExprPool (hash-consed DAG)
    ↓  e-graph extraction + canonicalization
Canonical expression form
    ↓  alkahest MLIR dialect
High-level MLIR (math-aware ops: horner, poly_eval, interval_eval)
    ↓  lowering passes
Standard MLIR (arith, math, linalg, gpu)
    ↓
LLVM IR / PTX / StableHLO (depending on target)
    ↓
Native machine code / GPU kernel / XLA

The custom alkahest MLIR dialect is where math-aware optimizations happen: Horner’s method for polynomials, fused multiply-add emission, numerically stable rearrangements via StabilityCost.

compile_expr

compile_expr produces a callable from a symbolic expression and a list of input variables:

from alkahest import ExprPool, compile_expr, sin, cos

pool = ExprPool()
x = pool.symbol("x")
y = pool.symbol("y")

f = compile_expr(x**2 + sin(y), [x, y])
print(f([3.0, 0.0]))   # 9.0

The callable takes a list of floats (one per variable) and returns a float. For batch evaluation see numpy_eval below.

CPU compilation uses a three-tier dispatch chosen by CompileConfig (expression size and expected_evals, not feature-flag order):

1. Interpreter — small DAGs and few planned evaluations (zero compile latency)
2. Cranelift (--features cranelift) — pure Rust, ~10× faster compile than LLVM; no system LLVM required
3. LLVM (--features jit) — inkwell / LLVM 15 MCJIT; best throughput for large batch sweeps

Default PyPI wheels use the interpreter only. +jit / +full release wheels enable LLVM. Add cranelift to a from-source build for the fast-compile tier without LLVM.

Native JIT backends also emit a bulk entry point (alkahest_eval_bulk) for column-major batch evaluation; CompiledFn.call_batch / call_bulk use the same layout as numpy_eval.

CompileCache

Repeated compilation of the same expression is expensive. CompileCache memoizes by (ExprId, input variables) — hash-consing makes ExprId a stable content key:

from alkahest import ExprPool, compile_expr, CompileCache

pool = ExprPool()
x = pool.symbol("x")
cache = CompileCache()
f = cache.compile(x**2, [x], pool)   # JIT compiles on first call
g = cache.compile(x**2, [x], pool)   # cache hit — O(1)
print(cache.stats())                 # hits, compiles, hit_rate

eval_expr

For one-off evaluation without compiling:

from alkahest import eval_expr

result = eval_expr(x**2 + sin(y), {x: 3.0, y: 0.0})
print(result)   # 9.0

eval_expr is slower than a compiled function for repeated evaluation but has no compilation overhead.

numpy_eval

numpy_eval vectorises a compiled function over NumPy arrays via the batch path:

import numpy as np
from alkahest import numpy_eval

f = compile_expr(sin(x) * cos(x), [x])
xs = np.linspace(0, 2 * 3.14159, 1_000_000)
ys = numpy_eval(f, xs)   # vectorised, zero-copy

Also accepts PyTorch CPU tensors and JAX arrays via DLPack.

Parallel batch evaluation

With --features parallel, numpy_eval_par distributes evaluation across CPU cores (Rayon), releasing the GIL during computation:

from alkahest import numpy_eval_par

ys = numpy_eval_par(f, xs)   # same API as numpy_eval; multi-core

numpy_eval_par is only parallel when the extension was built with parallel. Every published wheel is — but parallel is not a Cargo default, so a source build that omits it gets a silent no-op alias: numpy_eval_par falls back to numpy_eval with the same correct result, the same single-threaded throughput, and no warning. Timing the two against each other on such a build measures one code path twice.

import alkahest as ak
ak.capabilities()["features"]["parallel"]   # True on PyPI wheels

If that prints False, rebuild with --features parallel — see Getting started.

A CompiledFn belongs to the thread that compiled it

numpy_eval_par fans out internally, so you do not need threads of your own to use it. If you add them anyway, compile inside each thread: the object returned by compile_expr owns JIT code pages and a backend handle, and is pinned to its creating thread. Touching it from another one is refused:

f = compile_expr(expr, [x])          # thread A
threading.Thread(target=lambda: numpy_eval(f, xs)).start()
# pyo3_runtime.PanicException: alkahest::PyCompiledFn is unsendable,
#                              but sent to another thread

Two things to know about that refusal. It has nothing to do with parallel — plain numpy_eval is refused identically. And PanicException derives from BaseException, not Exception, so a worker wrapped in a bare except Exception: will not catch it and the thread will die silently.

ExprPool and Expr have no such restriction: share the pool, build expressions anywhere, and call compile_expr per thread.

Horner-form emission

horner rewrites a polynomial expression into Horner’s form, which is numerically better conditioned and faster to evaluate:

from alkahest import horner

# x^3 + 2x^2 + 3x + 4 → x*(x*(x + 2) + 3) + 4
h = horner(x**3 + pool.integer(2)*x**2 + pool.integer(3)*x + pool.integer(4), x)

emit_c emits a C function string for embedding in other projects:

from alkahest import emit_c

c_code = emit_c(expr, [x, y], fn_name="f")
# → "double f(double x, double y) { return ...; }"

MLIR dialect

The alkahest-mlir crate exposes the custom MLIR dialect. The dialect ops are:

Three lowering targets are available:

• ArithMath — lowers to arith + math MLIR dialects; uses math.fma for Horner chains
• StableHlo — lowers to StableHLO ops for XLA/JAX integration
• Llvm — lowers to llvm dialect for LLVM IR / PTX emission

from alkahest import to_stablehlo

# Emit textual MLIR in the StableHLO dialect
mlir_text = to_stablehlo(expr, [x, y], fn_name="my_fn")
print(mlir_text)  # valid input to mlir-opt / XLA

GPU codegen (NVPTX)

With --features cuda and an LLVM 15 installation with NVPTX support — not in any published wheel, so pip install alkahest never has this. Full detail, including build prerequisites, the supported node set, error codes and the state of testing, is in GPU support (CUDA).

from alkahest import compile_cuda

f_gpu = compile_cuda(expr, [x, y])
result = f_gpu.call_batch(inputs)   # runs on CUDA device 0

Guard on the capability bit before reaching for it, since the name does not exist at all without the feature:

if alkahest.capabilities()["features"]["cuda"]:
    ...

The GPU compiler:

1. Lowers the expression through inkwell to NVPTX LLVM IR for sm_86 (Ampere)
2. Links libdevice.10.bc for transcendental functions (__nv_sin, etc.)
3. Emits PTX via LLVM’s target machine
4. Loads the PTX via the CUDA driver (cudarc)

The benchmark nvptx/nvptx_polynomial_1M shows 16.2× speedup over the CPU JIT on a 1M-point polynomial evaluation on an RTX 3090.

Upcoming (v1.1): AMD ROCm / amdgcn target (hardware-blocked pending RDNA3 availability).

Caching

Use CompileCache for explicit per-session memoization of compiled functions (see above). The persistent ExprPool (V1-14) can serialize expression DAGs across sessions; combine with CompileCache to avoid recompilation after reload.

Tier dispatch always falls back to the interpreter when native JIT features are unavailable or compilation fails.