03. Compaction and scatter-add — Enzyme.jl

In [1]:

using WEANJ, Enzyme, CUDA, KernelAbstractions, Printf
include("common.jl")   # Setup, wrap, minimg, taper, inputs

d = ensure_inputs()
x0, v0, target0 = d["x0"], d["v"], d["target"]
s0 = Setup()

Setup{Float64}(50, 0.01, 0.005, 1024)

In [1]:

function deposit_loss(x0::Vector{T}, v, target, s::Setup) where {T}
    N, nb = length(x0), s.nb
    H = zeros(T, nb)
    sel = Int[]                                   # grows and shrinks with the data
    for t in 1:s.T
        c = T(t) / s.T                            # window centre
        empty!(sel)
        for i in 1:N                              # compaction: indices of x[mask]
            d = minimg(x0[i] + v[i] * (t * s.dt) - c)
            abs(d) < s.halfwidth && push!(sel, i)
        end
        for i in sel                              # work only for the selected particles
            x = wrap(x0[i] + v[i] * (t * s.dt))
            w = taper(minimg(x - c), s.halfwidth)
            g = x * nb - T(0.5)                   # cloud-in-cell deposit
            j = floor(Int, g)
            f = g - j
            H[mod(j, nb) + 1] += w * (1 - f)      # scatter-add at a data-dependent index
            H[mod(j + 1, nb) + 1] += w * f
        end
    end
    return sum(abs2, H .- target)
end

deposit_loss (generic function with 1 method)

In [1]:

function deposit_gradient(x0::Vector, v, target, s)
    dx = zero(x0)
    Enzyme.autodiff(Reverse, deposit_loss, Active, Duplicated(x0, dx), Const(v), Const(target), Const(s))
    return dx
end

g_cpu = deposit_gradient(x0, v0, target0, s0)
@printf "loss = %.10f   ‖∇loss‖ = %.6f   nonzero gradient entries: %d of %d\n" deposit_loss(x0, v0, target0, s0) sqrt(sum(abs2, g_cpu)) count(!iszero, g_cpu) length(x0)

loss = 559.5605712700   ‖∇loss‖ = 52854.055962   nonzero gradient entries: 983 of 2000

In [1]:

# On the GPU the natural form is one thread per particle and step: threads outside the
# window return immediately, the rest scatter into H with atomic adds.
Enzyme.EnzymeRules.inactive_type(::Type{<:Setup}) = true   # constants, never differentiated

@kernel function deposit_step!(H, @Const(x0), @Const(v), s, t)
    i = @index(Global)
    T = eltype(H)
    @inbounds begin
        x = wrap(x0[i] + v[i] * (t * s.dt))
        d = minimg(x - T(t) / s.T)
        if abs(d) < s.halfwidth
            w = taper(d, s.halfwidth)
            g = x * s.nb - T(0.5)
            j = unsafe_trunc(Int32, floor(g))
            f = g - j
            KernelAbstractions.@atomic H[mod(j, Int32(s.nb)) + 1] += w * (1 - f)
            KernelAbstractions.@atomic H[mod(j + Int32(1), Int32(s.nb)) + 1] += w * f
        end
    end
end

function deposit!(H, x0, v, s)
    fill!(H, 0)
    for t in 1:s.T
        deposit_step!(get_backend(H))(H, x0, v, s, t; ndrange=length(x0))
    end
    return nothing
end

function deposit_loss(x0::CuVector, v, target, s)
    H = similar(x0, s.nb)
    deposit!(H, x0, v, s)
    return sum(abs2, H .- target)
end

deposit_loss (generic function with 2 methods)

In [1]:

"Reverse through deposit!, seeded with ∂loss/∂H = 2(H − target)."
function deposit_gradient(x0::CuVector{T}, v, target, s) where {T}
    H = similar(x0, s.nb)
    deposit!(H, x0, v, s)
    dH = 2 .* (H .- target)
    dx = CUDA.zeros(T, length(x0))
    Enzyme.autodiff(Reverse, deposit!, Const, Duplicated(H, dH), Duplicated(x0, dx), Const(v), Const(s))
    return dx
end

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

max |∇_gpu − ∇_cpu| / max |∇_cpu| = 4.0e-16

In [1]:

err = fd_check(y -> deposit_loss(y, v0, target0, s0), x0, g_cpu; n=20, h=1e-7)
@assert maximum(abs, g_gpu - g_cpu) <= 1e-10 * maximum(abs, g_cpu)   # atomics reorder sums
save_gradient(EXAMPLE, "enzyme", g_cpu)
@printf "finite differences: max rel. error %.1e on 20 components\n" err

finite differences: max rel. error 7.2e-06 on 20 components

In [1]:

Ns = [10^4, 10^5, 10^6, 10^7]
gpu_sync = _ -> CUDA.synchronize()
records = Dict("cpu" => [], "gpu" => [])
for (device, prep, sync) in (("cpu", identity, identity), ("gpu", CuArray, gpu_sync))
    for (k, N) in enumerate(Ns)
        x, v, tg = prep.(inputs(N))
        compile_s = k == 1 ? @elapsed(sync(deposit_gradient(x, v, tg, s0))) : nothing
        samples = N >= 10^6 ? 3 : 10
        push!(records[device], (; dtype="Float64", N,
            forward_s = time_median(() -> deposit_loss(x, v, tg, s0); sync, samples),
            grad_s    = time_median(() -> deposit_gradient(x, v, tg, s0); sync, samples),
            compile_s,
            # host bytes allocated by one gradient (CPU only): Enzyme's tape stores the
            # selected particles of every step, not all N
            peak_mem_bytes = device == "cpu" ? @allocated(deposit_gradient(x, v, tg, s0)) : nothing))
    end
    write_results(EXAMPLE, "enzyme", device, records[device])
end
for d in ("cpu", "gpu"), r in records[d]
    @printf "%s  N=%-9d forward %9.3f ms   gradient %9.3f ms   %s\n" d r.N 1e3r.forward_s 1e3r.grad_s (r.peak_mem_bytes === nothing ? "" : @sprintf("allocated %.1f MiB", r.peak_mem_bytes / 2^20))
end

cpu  N=10000     forward     0.884 ms   gradient     0.988 ms   allocated 0.1 MiB
cpu  N=100000    forward     8.875 ms   gradient    23.716 ms   allocated 0.8 MiB
cpu  N=1000000   forward    94.128 ms   gradient   275.931 ms   allocated 7.9 MiB
cpu  N=10000000  forward  1253.366 ms   gradient  3180.572 ms   allocated 78.0 MiB
gpu  N=10000     forward     0.404 ms   gradient     3.407 ms   
gpu  N=100000    forward     0.788 ms   gradient     3.832 ms   
gpu  N=1000000   forward     4.978 ms   gradient    16.430 ms   
gpu  N=10000000  forward    41.608 ms   gradient   146.902 ms