02. Gauss–Seidel — Reactant.jl

In [1]:

using WEANJ, Reactant, Enzyme, Printf
include("common.jl")
Reactant.set_default_backend("gpu")

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

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1792007231.147872   31484 service.cc:155] XLA service 0x2eacfbf0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1792007231.147924   31484 service.cc:171]   StreamExecutor [0]: NVIDIA RTX 5000 Ada Generation Laptop GPU, Compute Capability 8.9 (Driver: 13.0.0[580.178.4]; Runtime: 13.1.0; Toolkit: 13.1.0; DNN: 9.21.0)
I0000 00:00:1792007231.148509   31484 se_gpu_pjrt_client.cc:1335] Using BFC allocator.
I0000 00:00:1792007231.148545   31484 gpu_helpers.cc:162] XLA backend allocating 11.68GiB (12538871808 bytes) on device 0 for BFCAllocator.

(Reactant.ConcretePJRTArray{Float64, 2, 1}([-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]), Reactant.ConcretePJRTArray{Float64, 2, 1}([-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]:

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

In [1]:

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

In [1]:

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

In [1]:

lex_c = @compile gs_loss(θ0, uobs0, K_CHECK)
rb_c = @compile rb_loss(θ0, uobs0, K_CHECK)
θh = Array(θ0)
e1 = fd_check(t -> Float64(lex_c(Reactant.to_rarray(t), uobs0, K_CHECK)), θh, g_lex; n=10, h=1e-2)
e2 = fd_check(t -> Float64(rb_c(Reactant.to_rarray(t), uobs0, K_CHECK)), θh, g_rb; n=10, h=1e-2)
save_gradient(EXAMPLE, "reactant-lex", g_lex)
save_gradient(EXAMPLE, "reactant-rb", g_rb)
@printf "finite differences: lexicographic %.1e, red-black %.1e\n" e1 e2

finite differences: lexicographic 8.3e-10, red-black 3.3e-10

In [1]:

ns = [32, 64, 128, 256, 512, 1024]
K = 100
budget_s = 5.0
records = []
for (variant, f) in (("lexicographic", gs_loss), ("red-black", rb_loss))
    for n in ns
        θ, uobs = Reactant.to_rarray.(inputs(Float64, n))
        compile_s = @elapsed begin
            lf = @compile f(θ, uobs, K)
            gf = @compile gradient(f, θ, uobs, K)
        end
        samples = n >= 512 ? 3 : 10
        r = (; dtype="Float64", variant, N=n^2, n, K,
             forward_s = time_median(() -> lf(θ, uobs, K); sync=Reactant.synchronize, samples),
             grad_s    = time_median(() -> gf(f, θ, uobs, K); sync=Reactant.synchronize, samples),
             compile_s, peak_mem_bytes = nothing)
        push!(records, r)
        r.grad_s > budget_s && break
    end
end
write_results(EXAMPLE, "reactant", "gpu", records)
for r in records
    @printf "gpu %-13s n=%-5d forward %9.3f ms  gradient %9.3f ms  compile %5.1f s\n" r.variant r.n 1e3r.forward_s 1e3r.grad_s r.compile_s
end

gpu lexicographic n=32    forward  1113.409 ms  gradient  2121.248 ms  compile   1.1 s
gpu lexicographic n=64    forward  4421.362 ms  gradient  8339.931 ms  compile   1.2 s
gpu red-black     n=32    forward     1.519 ms  gradient     2.330 ms  compile   0.8 s
gpu red-black     n=64    forward     1.484 ms  gradient     2.546 ms  compile   0.9 s
gpu red-black     n=128   forward     1.403 ms  gradient     2.576 ms  compile   1.0 s
gpu red-black     n=256   forward     1.815 ms  gradient     2.955 ms  compile   0.9 s
gpu red-black     n=512   forward     3.167 ms  gradient     6.099 ms  compile   0.9 s
gpu red-black     n=1024  forward     6.715 ms  gradient    16.799 ms  compile   1.0 s