02. Gauss–Seidel: a sequential in-place sweep

Each update reads values updated a moment earlier in the same sweep. Natural as a loop, awkward as an array program.

The problem

Solve \(-\Delta u = \theta\) on the unit square (\(n \times n\) interior points, \(u = 0\) on the boundary) with \(K\) Gauss–Seidel sweeps from \(u = 0\), and differentiate \[\mathcal{L}(\theta) = \lVert u_K(\theta) - u_{\text{obs}} \rVert^2\] with respect to the source \(\theta\). Each update is \[u_{ij} \leftarrow \tfrac14\left(h^2\theta_{ij} + u_{i-1,j} + u_{i+1,j} + u_{i,j-1} + u_{i,j+1}\right),\] in place: \(u_{i-1,j}\) and \(u_{i,j-1}\) are values already updated earlier in the same sweep. That is what makes Gauss–Seidel converge about twice as fast as Jacobi, and it is what makes it a sequential loop.

The same pattern appears in SOR, incomplete factorisations, the Thomas algorithm, upwind sweeps and line relaxation.

Two variants:

Lexicographic, the textbook sweep: a triple loop.
Red-black, which colours the grid like a chessboard so that each colour can be updated in parallel. It is a different iteration, with a different result, and it is the formulation a GPU or JAX user writes. Gradients are compared within each variant.

Code

Lexicographic, CPU:

# Lexicographic Gauss–Seidel, in place: u[i, j] uses the *new* u[i-1, j] and u[i, j-1].
function gs_loss(θ::AbstractMatrix{T}, uobs, K) where {T}
    n = size(θ, 1)
    h2 = one(T) / (n + 1)^2
    u = zeros(T, n + 2, n + 2)
    for _ in 1:K, j in 2:n+1, i in 2:n+1
        u[i, j] = (h2 * θ[i-1, j-1] + u[i-1, j] + u[i+1, j] + u[i, j-1] + u[i, j+1]) / 4
    end
    loss = zero(T)
    for j in 1:n, i in 1:n
        loss += (u[i+1, j+1] - uobs[i, j])^2
    end
    return loss
end

gs_loss (generic function with 1 method)

Source: 02. Gauss–Seidel — Enzyme.jl

function gs_gradient(θ, uobs, K)
    dθ = zero(θ)
    Enzyme.autodiff(Reverse, gs_loss, Active, Duplicated(θ, dθ), Const(uobs), Const(K))
    return dθ
end

g_lex = gs_gradient(θ0, uobs0, K_CHECK)
@printf "loss = %.12f   ‖∇loss‖ = %.6e\n" gs_loss(θ0, uobs0, K_CHECK) sqrt(sum(abs2, g_lex))

loss = 0.105860373742   ‖∇loss‖ = 6.409859e-04

Red-black, KernelAbstractions (CPU and GPU):

# Red-black Gauss–Seidel as a KernelAbstractions kernel (runs on CPU and GPU): points
# of one colour read only the other colour, so all threads of a colour are independent.
@kernel function rb_half_sweep!(u, @Const(θ), colour)
    I = @index(Global, Cartesian)
    i, j = I[1] + 1, I[2] + 1
    h2 = one(eltype(u)) / (size(θ, 1) + 1)^2   # computed here: GPU kernels take no active scalars
    @inbounds if (i + j) % 2 == colour
        u[i, j] = (h2 * θ[i-1, j-1] + u[i-1, j] + u[i+1, j] + u[i, j-1] + u[i, j+1]) / 4
    end
end

@kernel function sqdiff!(e, @Const(u), @Const(uobs))
    I = @index(Global, Cartesian)
    @inbounds e[I] = (u[I[1]+1, I[2]+1] - uobs[I])^2
end

"Loss terms e (the loss is sum(e)); u is the in-place solution buffer."
function rb_terms!(e, u, θ, uobs, K)
    backend, n = get_backend(u), size(θ, 1)
    fill!(u, 0)
    for _ in 1:K, colour in (0, 1)
        rb_half_sweep!(backend)(u, θ, colour; ndrange=(n, n))
    end
    sqdiff!(backend)(e, u, uobs; ndrange=(n, n))
    return nothing
end

function rb_loss(θ, uobs, K)
    n = size(θ, 1)
    e, u = similar(θ), similar(θ, n + 2, n + 2)
    rb_terms!(e, u, θ, uobs, K)
    return sum(e)
end

rb_loss (generic function with 1 method)

function rb_gradient(θ, uobs, K)
    n = size(θ, 1)
    e, de = similar(θ), fill!(similar(θ), 1)
    u, du = similar(θ, n + 2, n + 2), fill!(similar(θ, n + 2, n + 2), 0)
    dθ = fill!(similar(θ), 0)
    Enzyme.autodiff(Reverse, rb_terms!, Const, Duplicated(e, de), Duplicated(u, du),
                    Duplicated(θ, dθ), Const(uobs), Const(K))
    return dθ
end

g_rb = rb_gradient(θ0, uobs0, K_CHECK)
g_rb_gpu = Array(rb_gradient(CuArray(θ0), CuArray(uobs0), K_CHECK))
@printf "loss = %.12f   ‖∇loss‖ = %.6e   GPU vs CPU: %.1e\n" rb_loss(θ0, uobs0, K_CHECK) sqrt(sum(abs2, g_rb)) maximum(abs, g_rb_gpu - g_rb) / maximum(abs, g_rb)

loss = 0.105843754705   ‖∇loss‖ = 6.366981e-04   GPU vs CPU: 3.2e-16

Lexicographic:

# Lexicographic Gauss–Seidel. Arrays are immutable and loop bodies must be pure, so
# the sweep is a scan over columns, each a scan over rows carrying the new u[i-1, j].
# Same update order as the Julia loop (i fastest, then j).
@functools.partial(jax.jit, static_argnums=2)
def gs_loss(theta, uobs, K):
    n = theta.shape[0]
    h2 = 1.0 / (n + 1) ** 2

    def column(u, j):
        left, old, right = u[:, j], u[:, j + 1], u[:, j + 2]

        def point(below, i):  # below = new u[i-1, j]
            new = (h2 * theta[i, j] + below + old[i + 2] + left[i + 1] + right[i + 1]) / 4
            return new, new

        _, col = jax.lax.scan(point, old[0], jnp.arange(n))
        return u.at[1:n + 1, j + 1].set(col), None

    def sweep(u, _):
        return jax.lax.scan(column, u, jnp.arange(n))[0], None

    u, _ = jax.lax.scan(sweep, jnp.zeros((n + 2, n + 2), theta.dtype), None, length=K)
    return jnp.sum((u[1:-1, 1:-1] - uobs) ** 2)

Source: 02. Gauss–Seidel — JAX

gs_grad = jax.jit(jax.grad(gs_loss), static_argnums=2)
g_lex = gs_grad(theta0, uobs0, K_CHECK)
print(f"loss = {float(gs_loss(theta0, uobs0, K_CHECK)):.12f}   ‖∇loss‖ = {float(jnp.linalg.norm(g_lex)):.6e}")

loss = 0.105860373742   ‖∇loss‖ = 6.409859e-04

Red-black:

# Red-black Gauss–Seidel: the reformulation JAX users actually write. Each colour is
# one vectorised stencil, selected with a mask.
@functools.partial(jax.jit, static_argnums=2)
def rb_loss(theta, uobs, K):
    n = theta.shape[0]
    h2 = 1.0 / (n + 1) ** 2
    i, j = jnp.meshgrid(jnp.arange(2, n + 2), jnp.arange(2, n + 2), indexing="ij")
    red = (i + j) % 2 == 0  # 1-based indices, as in the Julia kernel

    def half(u, colour):
        new = (h2 * theta + u[:-2, 1:-1] + u[2:, 1:-1] + u[1:-1, :-2] + u[1:-1, 2:]) / 4
        return u.at[1:-1, 1:-1].set(jnp.where(colour, new, u[1:-1, 1:-1]))

    def sweep(u, _):
        return half(half(u, red), ~red), None

    u, _ = jax.lax.scan(sweep, jnp.zeros((n + 2, n + 2), theta.dtype), None, length=K)
    return jnp.sum((u[1:-1, 1:-1] - uobs) ** 2)

rb_grad = jax.jit(jax.grad(rb_loss), static_argnums=2)
g_rb = rb_grad(theta0, uobs0, K_CHECK)
print(f"loss = {float(rb_loss(theta0, uobs0, K_CHECK)):.12f}   ‖∇loss‖ = {float(jnp.linalg.norm(g_rb)):.6e}")

loss = 0.105843754705   ‖∇loss‖ = 6.366981e-04

Lexicographic:

# The same triple loop as the Enzyme version. Loops over traced data must be marked
# @trace, and indexing a traced array at a traced position must be explicitly allowed.
# Each scalar update becomes one step of an XLA while loop.
function gs_loss(θ, uobs, K)
    n = size(θ, 1)
    h2 = 1.0 / (n + 1)^2
    u = Reactant.Ops.fill(0.0, (n + 2, n + 2))
    @trace for _ in 1:K
        @trace for j in 2:n+1
            @trace for i in 2:n+1
                Reactant.@allowscalar u[i, j] =
                    (h2 * θ[i-1, j-1] + u[i-1, j] + u[i+1, j] + u[i, j-1] + u[i, j+1]) / 4
            end
        end
    end
    return sum((u[2:end-1, 2:end-1] .- uobs) .^ 2)
end

gs_loss (generic function with 1 method)

Source: 02. Gauss–Seidel — Reactant.jl

Red-black:

# Red-black as array operations: one masked stencil per colour. The colour mask is built
# from traced iota arrays rather than captured as a Julia constant: with a constant mask
# compilation took minutes at n = 64.
function rb_loss(θ, uobs, K)
    n = size(θ, 1)
    h2 = 1.0 / (n + 1)^2
    i = Reactant.Ops.iota(Int64, [n, n]; iota_dimension=1)
    j = Reactant.Ops.iota(Int64, [n, n]; iota_dimension=2)
    red = iseven.(i .+ j)
    u = Reactant.Ops.fill(0.0, (n + 2, n + 2))
    @trace for _ in 1:K
        for colour in (red, .!red)
            new = (h2 .* θ .+ u[1:end-2, 2:end-1] .+ u[3:end, 2:end-1] .+
                   u[2:end-1, 1:end-2] .+ u[2:end-1, 3:end]) ./ 4
            u[2:end-1, 2:end-1] = ifelse.(colour, new, u[2:end-1, 2:end-1])
        end
    end
    return sum((u[2:end-1, 2:end-1] .- uobs) .^ 2)
end

rb_loss (generic function with 1 method)

gradient(f, θ, uobs, K) = Enzyme.gradient(Reverse, Const(f), θ, Const(uobs), Const(K))[1]

g_lex = Array((@compile gradient(gs_loss, θ0, uobs0, K_CHECK))(gs_loss, θ0, uobs0, K_CHECK))
g_rb = Array((@compile gradient(rb_loss, θ0, uobs0, K_CHECK))(rb_loss, θ0, uobs0, K_CHECK))
@printf "‖∇loss‖: lexicographic %.6e, red-black %.6e\n" sqrt(sum(abs2, g_lex)) sqrt(sum(abs2, g_rb))

I0000 00:00:1792007279.908700   31484 se_gpu_pjrt_compiler.cc:422] Found GPU target config and a PjRtClient with matching configuration. Performing a JIT compilation.
I0000 00:00:1792007280.026785   31484 cuda_dnn.cc:478] Loaded cuDNN version 92100
‖∇loss‖: lexicographic 6.409859e-04, red-black 6.366981e-04

Same gradient

\(n = 32\), \(K = 20\), Float64. Each gradient is also checked against finite differences.

Table 1: lexicographic

implementation	‖g‖₂	max rel. diff vs enzyme-lex
enzyme-lex	6.409859e-04	0.0e+00
reactant-lex	6.409859e-04	4.3e-16
jax-lex	6.409859e-04	2.1e-16

Table 2: red-black

implementation	‖g‖₂	max rel. diff vs enzyme-rb
enzyme-rb	6.366981e-04	0.0e+00
reactant-rb	6.366981e-04	5.4e-16
jax-rb	6.366981e-04	4.3e-16

Timings

\(K = 100\) sweeps, Float64, median of 3–10 runs after warm-up. A sweep stops once a single gradient takes more than 5 s.

What this shows

The sequential sweep on the CPU. At \(n = 512\) the lexicographic gradient takes 325.7 ms with Enzyme and 19.7 s with JAX, although their forward passes are close (177.6 ms and 517.8 ms). Enzyme’s gradient costs 1.8× its forward pass; JAX’s costs 38×. JAX expresses the sweep as nested lax.scans, and the transposed scan of a scan does much more work per step than the scalar loop it came from.

The sequential sweep on the GPU. At \(n = 64\) the lexicographic gradient takes 7.3 s with JAX and 8.3 s with Reactant, against 4.1 ms for Enzyme on one CPU core. Every scalar update becomes a step of an XLA loop, which on the GPU means a device round trip per grid point. Reactant traces the same Julia triple loop as Enzyme, so the language is not the issue: the loop is slow once it goes through XLA.

Memory does not grow with K, for either tool. We expected reverse mode through \(K\) in-place sweeps to store \(K\) copies of \(u\). It doesn’t, because each update is linear: its derivative is the constant \(\tfrac14\), so the backward pass needs no stored values. Both Enzyme’s activity analysis and JAX’s linearisation see this.

Table 3: Gradient memory [MiB], lexicographic, n = 256 (Enzyme: bytes allocated by the call; JAX: XLA temp buffers)

impl	enzyme	jax
K
10	1.52	2.04
100	1.52	2.04
1000	1.52	2.04

A nonlinear sweep, such as nonlinear Gauss–Seidel with a local Newton step, would need stored values in both tools and make memory a real issue. That is not tested here.

Where JAX wins / what we did not do

Note

Red-black on the GPU is where the XLA tools shine. At \(n = 1024\) the red-black gradient takes 16.8 ms with Reactant, 20.0 ms with JAX and 21.9 ms with Enzyme through KernelAbstractions; on small grids Enzyme is several times slower. Enzyme launches two kernels per sweep, forward and reverse, while XLA compiles the whole loop into one program. That launch overhead dominates on small grids.

A JAX expert can rewrite the lexicographic sweep without a scalar loop. Along each column the update is a first-order linear recurrence, which jax.lax.associative_scan evaluates in parallel. That is a different algorithm, with different rounding and \(O(\log n)\) depth, and you have to know the trick.
No Enzyme multithreading and no fused multi-sweep GPU kernel.