02. Gauss–Seidel — Enzyme.jl

In [1]:

using WEANJ, Enzyme, CUDA, KernelAbstractions, Printf
include("common.jl")   # problem size, shared inputs

d = ensure_inputs()
θ0, uobs0 = d["theta"], d["uobs"]

([-0.36335748145177754 1.9207540740525255 … 0.20732993930891253 0.007365170012581141; 0.2517372155742292 0.8006488606751957 … 1.221741766000742 -0.4708761069990038; … ; 0.17127065460959456 -0.4400172451059574 … 0.27999985275170924 -0.7484878004000803; 1.1199182981069984 -1.7799429142766356 … 0.5500269297887572 -0.25196386505341467], [-0.00019879962436752713 -0.00558994605387216 … -0.00801162607331382 -0.01157992248121622; 0.002182456557295407 0.0005819751744237116 … 0.0002654612932840401 -0.0028831162758658426; … ; 0.011481650345351322 -0.023480482251067098 … 0.009341572856865593 -0.010644550873280046; -0.001636400736785318 -0.005467296315401388 … 0.003665309993473059 -0.004358409328340684])

In [1]:

# 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)

In [1]:

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

In [1]:

# 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)

In [1]:

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

In [1]:

# The loss is quadratic in θ, so a large finite-difference step is exact up to rounding.
e1 = fd_check(t -> gs_loss(t, uobs0, K_CHECK), θ0, g_lex; n=10, h=1e-2)
e2 = fd_check(t -> rb_loss(t, uobs0, K_CHECK), θ0, g_rb; n=10, h=1e-2)
@assert maximum(abs, g_rb_gpu - g_rb) <= 1e-12 * maximum(abs, g_rb)
save_gradient(EXAMPLE, "enzyme-lex", g_lex)
save_gradient(EXAMPLE, "enzyme-rb", g_rb)
@printf "finite differences: lexicographic %.1e, red-black %.1e\n" e1 e2

finite differences: lexicographic 3.2e-09, red-black 3.6e-10

In [1]:

ns = [32, 64, 128, 256, 512, 1024]
K = 100
budget_s = 5.0   # stop a sweep once one gradient takes longer than this
gpu_sync = _ -> CUDA.synchronize()

function sweep(variant, device, loss, grad; prep=identity, sync=identity, allocated)
    records = []
    for (k, n) in enumerate(ns)
        θ, uobs = prep.(inputs(Float64, n))
        compile_s = k == 1 ? @elapsed(sync(grad(θ, uobs, K))) : nothing
        samples = n >= 512 ? 3 : 10
        r = (; dtype="Float64", variant, N=n^2, n, K,
             forward_s = time_median(() -> loss(θ, uobs, K); sync, samples),
             grad_s    = time_median(() -> grad(θ, uobs, K); sync, samples),
             compile_s, peak_mem_bytes = allocated(() -> grad(θ, uobs, K)))
        push!(records, r)
        r.grad_s > budget_s && break
    end
    return records
end

# Memory allocated by one gradient call (host for CPU, device for GPU).
host_alloc(f) = @allocated f()
dev_alloc(f) = CUDA.@allocated f()

cpu = vcat(sweep("lexicographic", "cpu", gs_loss, gs_gradient; allocated=host_alloc),
           sweep("red-black", "cpu", rb_loss, rb_gradient; allocated=host_alloc))
gpu = sweep("red-black", "gpu", rb_loss, rb_gradient; prep=CuArray, sync=gpu_sync,
            allocated=dev_alloc)

# Gradient memory against the number of sweeps K, at n = 256.
θm, uobsm = inputs(Float64, 256)
for Km in (10, 100, 1000)
    push!(cpu, (; dtype="Float64", variant="lexicographic", N=256^2, n=256, K=Km,
                forward_s=nothing, grad_s=nothing, compile_s=nothing,
                peak_mem_bytes=@allocated(gs_gradient(θm, uobsm, Km))))
end

write_results(EXAMPLE, "enzyme", "cpu", cpu)
write_results(EXAMPLE, "enzyme", "gpu", gpu)
for (dev, rs) in (("cpu", cpu), ("gpu", gpu)), r in rs
    r.grad_s === nothing && continue
    @printf "%s %-13s n=%-5d forward %9.3f ms  gradient %9.3f ms  memory %8.2f MiB\n" dev r.variant r.n 1e3r.forward_s 1e3r.grad_s r.peak_mem_bytes / 2^20
end

cpu lexicographic n=32    forward     0.435 ms  gradient     0.876 ms  memory     0.03 MiB
cpu lexicographic n=64    forward     2.132 ms  gradient     4.140 ms  memory     0.10 MiB
cpu lexicographic n=128   forward     9.953 ms  gradient    18.652 ms  memory     0.38 MiB
cpu lexicographic n=256   forward    42.688 ms  gradient    78.594 ms  memory     1.52 MiB
cpu lexicographic n=512   forward   177.599 ms  gradient   325.750 ms  memory     6.03 MiB
cpu lexicographic n=1024  forward   730.356 ms  gradient  1318.426 ms  memory    24.06 MiB
cpu red-black     n=32    forward     0.345 ms  gradient     2.704 ms  memory     0.30 MiB
cpu red-black     n=64    forward     1.224 ms  gradient     8.139 ms  memory     0.49 MiB
cpu red-black     n=128   forward     4.395 ms  gradient    32.837 ms  memory     1.27 MiB
cpu red-black     n=256   forward    16.652 ms  gradient   167.094 ms  memory     4.35 MiB
cpu red-black     n=512   forward    65.567 ms  gradient   517.138 ms  memory    16.64 MiB
cpu red-black     n=1024  forward   263.518 ms  gradient  3476.960 ms  memory    65.76 MiB
gpu red-black     n=32    forward     1.199 ms  gradient    10.198 ms  memory     0.04 MiB
gpu red-black     n=64    forward     1.170 ms  gradient    10.425 ms  memory     0.16 MiB
gpu red-black     n=128   forward     1.197 ms  gradient    10.554 ms  memory     0.63 MiB
gpu red-black     n=256   forward     1.556 ms  gradient    10.579 ms  memory     2.52 MiB
gpu red-black     n=512   forward     2.944 ms  gradient    12.674 ms  memory    10.03 MiB
gpu red-black     n=1024  forward    10.316 ms  gradient    21.928 ms  memory    40.06 MiB