01. Particles: sort, cell lists, and data-dependent neighbour loops

An SPH density loss whose efficient form sorts particles into cells and loops over a data-dependent number of neighbours.

The problem

\(N\) particles in a periodic square box, each with an SPH density \[\rho_i = \sum_j m\, W(\lVert x_i - x_j\rVert, h),\] where \(W\) is the 2D cubic-spline kernel with support \(2h\), and the loss \[\mathcal{L}(x) = \sum_i (\rho_i - \rho_0)^2 .\] We want \(\nabla_x \mathcal{L}\), the gradient with respect to all particle positions. The box grows with \(N\) so every particle has about 31 neighbours.

This loss came up in a Slack discussion as “something JAX can’t do because it needs sort!”. That is not quite right: jax.lax.sort is differentiable. What matters is how the efficient algorithm works:

1. compute each particle’s cell (integer, zero derivative),
2. sort particles by cell so each cell is contiguous in memory,
3. for each particle, loop over the 3×3 neighbouring cells, each a range start[c]:start[c+1]-1 whose length depends on the data,
4. accumulate kernel-weighted contributions.

The gradient flows through the gather in step 2 (as a permutation) and through step 4. Steps 1–3 only decide which pairs interact.

Code

The Enzyme version is the algorithm above as plain loops, differentiated as written: on the CPU the whole pipeline, including the counting sort, is inside autodiff. JAX and Reactant cannot express step 3 (a loop whose bounds come from the data), so they regularise it. JAX uses jax-md’s padded neighbour list; Reactant uses a padded cell list. Both must pick a capacity from the data and check for overflow.

Enzyme.jl (CPU)

# The whole pipeline is differentiated: cell ids, counting sort, gather into sorted
# order, and neighbour loops whose ranges depend on the data.
function sph_loss(x::AbstractMatrix{T}, p::Params{T}) where {T}
    N, nc = size(x, 2), p.nc
    s, rc = p.L / nc, 2 * p.h

    # 1. cell id of every particle, and particles per cell
    cell  = Vector{Int}(undef, N)
    start = zeros(Int, nc * nc + 1)
    @inbounds for i in 1:N
        cx = mod(floor(Int, x[1, i] / s), nc)
        cy = mod(floor(Int, x[2, i] / s), nc)
        cell[i] = 1 + cx + nc * cy
        start[cell[i] + 1] += 1
    end

    # 2. counting sort by cell: start[c] is the first sorted index of cell c
    start[1] = 1
    for c in 1:nc * nc
        start[c + 1] += start[c]
    end
    cursor = start[1:end-1]
    perm = Vector{Int}(undef, N)
    @inbounds for i in 1:N
        perm[cursor[cell[i]]] = i
        cursor[cell[i]] += 1
    end

    # 3. gather positions into sorted order: each cell is contiguous in memory
    xs = similar(x)
    @inbounds for k in 1:N
        xs[1, k] = x[1, perm[k]]
        xs[2, k] = x[2, perm[k]]
    end

    # 4. density and loss over the 3×3 neighbouring cells; each cell is the range
    #    start[c]:start[c+1]-1, whose length depends on the data
    loss = zero(T)
    @inbounds for k in 1:N
        c = cell[perm[k]] - 1
        cx, cy = c % nc, c ÷ nc
        ρ = zero(T)
        for dy in -1:1, dx in -1:1
            cn = 1 + mod(cx + dx, nc) + nc * mod(cy + dy, nc)
            for l in start[cn]:start[cn + 1] - 1
                rx = minimg(xs[1, k] - xs[1, l], p.L)
                ry = minimg(xs[2, k] - xs[2, l], p.L)
                r2 = rx * rx + ry * ry
                if r2 < rc * rc
                    ρ += p.m * W(sqrt(r2), p.h)
                end
            end
        end
        loss += (ρ - p.ρ0)^2
    end
    return loss
end

sph_loss (generic function with 1 method)

Source: 01. Particles — Enzyme.jl

function sph_gradient(x::Matrix{T}, p::Params{T}) where {T}
    dx = zero(x)
    Enzyme.autodiff(Reverse, sph_loss, Active, Duplicated(x, dx), Const(p))
    return dx
end

g_cpu = sph_gradient(x0, p0)
@printf "loss = %.10f   ‖∇loss‖ = %.6f\n" sph_loss(x0, p0) sqrt(sum(abs2, g_cpu))

loss = 12476.3617428793   ‖∇loss‖ = 7526.066606

Enzyme.jl (GPU)

# Same algorithm as KernelAbstractions kernels. The cell list is integer-valued (zero
# derivative), so it is built outside the differentiated function with CUDA.jl's sort.
# The gather and the neighbour loops are differentiated.

# Box size, smoothing length etc. are not differentiated; declaring that lets Params be
# passed to GPU kernels, which reject active (differentiable) scalar arguments.
Enzyme.EnzymeRules.inactive_type(::Type{<:Params}) = true

@kernel function cell_ids!(cell, counts, @Const(x), s, nc)
    i = @index(Global)
    @inbounds begin
        cx = mod(unsafe_trunc(Int32, floor(x[1, i] / s)), Int32(nc))
        cy = mod(unsafe_trunc(Int32, floor(x[2, i] / s)), Int32(nc))
        cell[i] = Int32(1) + cx + Int32(nc) * cy
        KernelAbstractions.@atomic counts[cell[i] + 1] += Int32(1)
    end
end

function build_cells(x::CuMatrix{T}, p::Params{T}) where {T}
    N, nc = size(x, 2), p.nc
    cell = CuVector{Int32}(undef, N)
    counts = CUDA.zeros(Int32, nc * nc + 1)
    cell_ids!(get_backend(x))(cell, counts, x, p.L / nc, nc; ndrange=N)
    counts[1:1] .= 1
    perm = sortperm(cell)
    return (; perm=Int32.(perm), cell=cell[perm], start=cumsum(counts))
end

@kernel function gather!(xs, @Const(x), @Const(perm))
    k = @index(Global)
    @inbounds xs[1, k] = x[1, perm[k]]
    @inbounds xs[2, k] = x[2, perm[k]]
end

@kernel function density_loss!(e, @Const(xs), @Const(cellk), @Const(start), p)
    k = @index(Global)
    T, nc = eltype(xs), Int32(p.nc)
    @inbounds begin
        c = cellk[k] - Int32(1)
        cx, cy = c % nc, c ÷ nc
        ρ = zero(T)
        for dy in Int32(-1):Int32(1), dx in Int32(-1):Int32(1)
            cn = Int32(1) + mod(cx + dx, nc) + nc * mod(cy + dy, nc)
            for l in start[cn]:start[cn + 1] - Int32(1)      # data-dependent range
                rx = minimg(xs[1, k] - xs[1, l], p.L)
                ry = minimg(xs[2, k] - xs[2, l], p.L)
                r2 = rx * rx + ry * ry
                if r2 < 4 * p.h * p.h
                    ρ += p.m * W(sqrt(r2), p.h)
                end
            end
        end
        e[k] = (ρ - p.ρ0)^2
    end
end

"Per-particle loss terms e (in sorted order); the loss is sum(e)."
function sph_terms!(e, xs, x, cells, p)
    backend, N = get_backend(x), size(x, 2)
    gather!(backend)(xs, x, cells.perm; ndrange=N)
    density_loss!(backend)(e, xs, cells.cell, cells.start, p; ndrange=N)
    return nothing
end

function sph_loss(x::CuMatrix{T}, p::Params{T}) where {T}
    e, xs = similar(x, size(x, 2)), similar(x)
    sph_terms!(e, xs, x, build_cells(x, p), p)
    return sum(e)
end

sph_loss (generic function with 2 methods)

"Gradient of sum(e): reverse mode seeded with de = 1."
function sph_gradient(x::CuMatrix{T}, p::Params{T}) where {T}
    N = size(x, 2)
    e, de = similar(x, N), CUDA.ones(T, N)
    xs, dxs = similar(x), CUDA.zeros(T, size(x))
    dx = CUDA.zeros(T, size(x))
    Enzyme.autodiff(Reverse, sph_terms!, Const,
                    Duplicated(e, de), Duplicated(xs, dxs), Duplicated(x, dx),
                    Const(build_cells(x, p)), Const(p))
    return dx
end

g_gpu = Array(sph_gradient(CuArray(x0), p0))
@printf "max |∇_gpu − ∇_cpu| / max |∇_cpu| = %.1e\n" maximum(abs, g_gpu - g_cpu) / maximum(abs, g_cpu)

Reactant.jl

# Reactant traces Julia array code into one XLA program, so it has JAX's constraints:
# static shapes, no data-dependent loop bounds, no branches on traced values. The cell
# list is regularised like a padded JAX cell list: every neighbour cell is read
# through `cap` slots and slots past the cell's end are masked.

# `W` from common.jl branches (`if q < 1`), which cannot be traced. Select instead,
# evaluating both polynomials for every pair, as JAX's `jnp.where` does.
@inline function W_select(r, h::T) where {T}
    q = r / h
    σ = T(10) / (T(7) * T(π) * h * h)
    t = max(2 - q, zero(T))
    return σ * ifelse(q < 1, 1 - T(1.5) * q * q + T(0.75) * q * q * q, T(0.25) * t * t * t)
end

"""
First sorted index of each cell, using only sorts and gathers: merge the sorted cell
ids (keys 2c) with one probe per cell (keys 2c−1); the rank of probe c, minus the
c−1 probes before it, is the start of cell c.
"""
function cell_starts(cs, ncells)
    N = length(cs)
    probes = 2 .* Int32.(1:ncells+1) .- Int32(1)
    order = sortperm(vcat(2 .* cs, Reactant.promote_to(typeof(cs), probes)))
    rank = sortperm(order)                       # inverse permutation
    return rank[N+1:end] .- Int32.(0:ncells)
end

function sph_loss_padded(x::AbstractMatrix, p::Params{T}, ::Val{cap}) where {T,cap}
    N, nc = size(x, 2), p.nc
    s = p.L / nc

    cell = Int32(1) .+ mod.(floor.(Int32, x[1, :] ./ s), Int32(nc)) .+
           Int32(nc) .* mod.(floor.(Int32, x[2, :] ./ s), Int32(nc))
    perm = sortperm(cell)
    cs = cell[perm]
    start = cell_starts(cs, nc * nc)
    xs, ys = x[1, perm], x[2, perm]
    cxs, cys = (cs .- Int32(1)) .% Int32(nc), (cs .- Int32(1)) .÷ Int32(nc)
    slots = reshape(Int32.(0:cap-1), cap, 1)

    ρ = zero(xs)
    for dy in -1:1, dx in -1:1                   # unrolled at trace time
        cn = Int32(1) .+ mod.(cxs .+ Int32(dx), Int32(nc)) .+
             Int32(nc) .* mod.(cys .+ Int32(dy), Int32(nc))
        idx = reshape(start[cn], 1, N) .+ slots                  # (cap, N) candidates
        valid = idx .< reshape(start[cn .+ Int32(1)], 1, N)
        l = clamp.(idx, Int32(1), Int32(N))
        rx = minimg.(reshape(xs, 1, N) .- xs[l], p.L)
        ry = minimg.(reshape(ys, 1, N) .- ys[l], p.L)
        r2 = rx .^ 2 .+ ry .^ 2
        inside = valid .& (r2 .< 4 * p.h * p.h)
        r = sqrt.(ifelse.(inside, r2, one(T)))  # keep sqrt' finite on masked slots
        ρ = ρ .+ vec(sum(ifelse.(inside, p.m .* W_select.(r, p.h), zero(T)); dims=1))
    end
    return sum((ρ .- p.ρ0) .^ 2)
end

"Capacity needed: the largest cell occupancy. Too small silently drops neighbours."
function max_occupancy(x, p)
    s, nc = p.L / p.nc, p.nc
    counts = zeros(Int, nc * nc)
    for i in axes(x, 2)
        counts[1 + mod(floor(Int, x[1, i] / s), nc) + nc * mod(floor(Int, x[2, i] / s), nc)] += 1
    end
    return maximum(counts)
end

Main.Notebook.max_occupancy

Source: 01. Particles — Reactant.jl

sph_gradient_padded(x, p, cap) =
    Enzyme.gradient(Reverse, Const(sph_loss_padded), x, Const(p), Const(cap))[1]

cap0 = Val(max_occupancy(x0, p0))
xr0 = Reactant.to_rarray(x0)
loss_c = @compile sph_loss_padded(xr0, p0, cap0)
grad_c = @compile sph_gradient_padded(xr0, p0, cap0)
g = Array(grad_c(xr0, p0, cap0))
@printf "cap = %d   loss = %.10f   ‖∇loss‖ = %.6f\n" max_occupancy(x0, p0) Float64(loss_c(xr0, p0, cap0)) sqrt(sum(abs2, g))

I0000 00:00:1791991487.502186  561125 se_gpu_pjrt_compiler.cc:422] Found GPU target config and a PjRtClient with matching configuration. Performing a JIT compilation.
I0000 00:00:1791991488.027130  561125 cuda_dnn.cc:478] Loaded cuDNN version 92100
cap = 23   loss = 12476.3617428792   ‖∇loss‖ = 7526.066606

JAX

# Idiomatic JAX: jax-md's neighbour list. It is a fixed-capacity padded pair list;
# `allocate` (outside jit) sizes it from the data, `update` (inside jit) refills it,
# and `did_buffer_overflow` must be checked after every call.
NUMBER_DENSITY, H, MASS = 10.0, 0.5, 1.0
RC, RHO0 = 2 * H, NUMBER_DENSITY

def box_size(n):
    return math.sqrt(n / NUMBER_DENSITY)

def W(r, h=H):
    """2D cubic spline. No `if`: both branches are evaluated and selected."""
    q = r / h
    sigma = 10.0 / (7.0 * math.pi * h * h)
    t = jnp.maximum(2.0 - q, 0.0)
    return sigma * jnp.where(q < 1.0, 1.0 - 1.5 * q**2 + 0.75 * q**3, 0.25 * t**3)

def make_loss(n, displacement):
    w0 = MASS * float(W(0.0))  # self-contribution; jax-md's pair list excludes i == j

    def loss(x, nbrs):
        nbrs = nbrs.update(x)  # rebuild the cell list from the current positions
        i, j = nbrs.idx
        mask = i < n  # padded slots hold index n
        dR = jax.vmap(displacement)(x[i], x[j])
        dR = jnp.where(mask[:, None], dR, 1.0)  # padded dR = 0 would make ∇‖dR‖ NaN
        r = jnp.sqrt(jnp.sum(dR**2, axis=1))
        w = jnp.where(mask, MASS * W(r), 0.0)
        rho = w0 + jax.ops.segment_sum(w, i, num_segments=n)
        return jnp.sum((rho - RHO0) ** 2), nbrs

    return loss

class Particles:
    """Jitted loss and gradient with jax-md's overflow-and-reallocate protocol."""

    def __init__(self, x, capacity_multiplier=1.25):
        self.n = x.shape[0]
        L = box_size(self.n)
        displacement, _ = space.periodic(L)
        self.neighbor_fn = partition.neighbor_list(
            displacement, L, RC, dr_threshold=0.0,
            capacity_multiplier=capacity_multiplier, format=partition.Sparse)
        self.nbrs = self.neighbor_fn.allocate(x)
        loss = make_loss(self.n, displacement)
        self._loss = jax.jit(loss)
        self._grad = jax.jit(jax.grad(loss, has_aux=True))

    def _checked(self, f, x):
        out, nbrs = f(x, self.nbrs)
        if nbrs.did_buffer_overflow:  # host sync, as in jax-md's examples
            self.nbrs = self.neighbor_fn.allocate(x)
            out, nbrs = f(x, self.nbrs)
        self.nbrs = nbrs
        return out

    def loss(self, x):
        return self._checked(self._loss, x)

    def grad(self, x):
        return self._checked(self._grad, x)

Source: 01. Particles — JAX

P0 = Particles(x0)
g = P0.grad(x0)
print(f"loss = {float(P0.loss(x0)):.10f}   ‖∇loss‖ = {float(jnp.linalg.norm(g)):.6f}")

loss = 12476.3617428792   ‖∇loss‖ = 7526.066606

Shared definitions (Params, the kernel W, the minimum image, the input) are in common.jl.

Same gradient from all three

Each implementation computes the Float64 gradient for the same 1000 particles, checks it against finite differences, and saves it. The build fails unless all three agree.

implementation	‖g‖₂	max rel. diff vs enzyme
enzyme	7526.066606	0.0e+00
reactant	7526.066606	6.6e-16
jax	7526.066606	1.1e-15

Timings

Float64, median of 3–10 runs after warm-up, GPU synchronised. Each gradient includes rebuilding the cell or neighbour list. NVIDIA RTX 5000 Ada Laptop GPU and Intel i9-13980HX, with Enzyme on the CPU single-threaded and JAX on the CPU using XLA’s default threading.

Table 1: N = 1,000,000 particles

implementation	forward	gradient	gradient / forward	pair slots evaluated	compile [s]
enzyme (cpu)	1.42 s	4.78 s	3×	91 M	0.4
enzyme (gpu)	30.6 ms	2.57 s	84×	91 M	0.1
jax (cpu)	5.46 s	5.98 s	1×	39 M	13.9
reactant (gpu)	39.6 ms	196.1 ms	5×	234 M	17.7

jax (gpu), N = 1,000,000: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.07GiB with allocator GPU_0_bfc on device 0.

“Pair slots evaluated” counts the candidate pairs each method computes a distance for. Enzyme visits exactly the particles in the 3×3 cells. JAX’s pair list is padded to 1.25× the pairs found at allocation. Reactant’s padded cell list reads cap slots per cell, with cap the largest cell occupancy.

What this shows

Differentiating the irregular algorithm works. Enzyme differentiates the mutating, sort-based, data-dependent loop code directly and gets the same gradient as JAX and Reactant to machine precision. No padding, no capacity, no overflow check, no where tricks to keep sqrt finite. On the CPU at \(N = 10^6\), single-threaded Enzyme takes 4.78 s for the gradient, against 5.98 s for multithreaded JAX. The forward pass alone takes 1.42 s against 5.46 s.

The padded version runs out of memory on the GPU. In Float64, jax-md’s neighbour list for \(10^6\) particles does not fit on the 16 GB GPU: building it needs a further 4 GiB intermediate on top of a pair list of 39 million slots. Enzyme’s cell list needs no capacity at all.

Reactant behaves like JAX. Reactant is Julia, but compiled through XLA, and it needs the same regularisation: padded cells, masks, a branch-free kernel. At \(N = 10^5\) its gradient takes 21.1 ms against JAX’s 10.3 ms. Its padded cell list is leaner than jax-md’s pair list (no stored pair indices), so it still fits at \(N = 10^6\) (196.1 ms). The constraint comes from the static-graph execution model, not from the language.

Where JAX wins / where Enzyme hurts

WarningOn the GPU, today, the padded versions are much faster than Enzyme’s

At \(N = 10^5\) the Enzyme GPU gradient takes 119.1 ms, against 10.3 ms for JAX and 21.1 ms for Reactant. At \(N = 10^6\) it takes 2.57 s, 84× its own forward pass.

The reason is how Enzyme compiles reverse-mode GPU kernels with a data-dependent inner loop: it stores values for every inner iteration in device malloc memory, in both the augmented forward kernel and the reverse kernel. With CUDA’s default 8 MB device heap this crashes above about 30 000 particles, so the notebook raises the heap to 4 GiB. A branch-free kernel and a fixed-trip-count kernel showed the same behaviour.

What we did not do:

• No hand-written adjoint on either side. A JAX expert could write the neighbour loop as a Pallas kernel with a hand-written backward pass. An Enzyme user could restructure the GPU gradient as a forward kernel plus a gather kernel that applies Enzyme only to the scalar pair function. Both are faster, and both give up “the adjoint for free”.
• No multithreaded Enzyme CPU version. The CPU numbers are single-threaded Julia against multithreaded XLA.
• jax-md’s capacity_multiplier left at 1.25, its default.