03. Compaction and scatter-add: work proportional to the data

Select the few particles that matter at each step (x[mask], push!) and scatter them into a histogram at data-dependent indices.

The problem

\(N\) particles move with fixed velocities on the periodic interval \([0, 1)\). At each of \(T = 50\) steps, a narrow window sweeps across the domain. The particles inside it, about 1 %, are deposited with linear (cloud-in-cell) weights into a histogram \(H\) of 1024 bins. A taper \(w(d) = (1 - |d|/h_w)^2\) fades the contribution to zero at the window edges, which keeps the loss differentiable as particles enter and leave. We differentiate \[\mathcal{L}(x_0) = \lVert H(x_0) - H_{\text{target}} \rVert^2\] with respect to the initial positions.

Two irregular patterns are involved:

  • Compaction. The work worth doing is proportional to the number of selected particles, which changes with the data: x[mask], findall, push! into a buffer.
  • Scatter-add at data-dependent indices: H[j] += w.

The same patterns appear in event detection, active-set methods, particle-in-cell deposits, histogramming, and resampling.

Code

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)
Source: 03. Compaction and scatter-add — Enzyme.jl
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
# 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)
"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
def minimg(d):
    return d - jnp.round(d)

def taper(d):
    return jnp.maximum(1.0 - jnp.abs(d) / HW, 0.0) ** 2

def deposit(H, x, w):
    """Cloud-in-cell scatter-add into the histogram."""
    g = x * NB - 0.5
    j = jnp.floor(g).astype(jnp.int32)
    f = g - j
    return H.at[j % NB].add(w * (1 - f)).at[(j + 1) % NB].add(w * f)

# Masked: every particle is deposited at every step, with weight 0 outside the window.
@jax.jit
def loss_masked(x0, v, target):
    def step(H, t):
        x = x0 + v * (t * DT)
        d = minimg(x - t / T)
        w = jnp.where(jnp.abs(d) < HW, taper(d), 0.0)
        return deposit(H, x - jnp.floor(x), w), None
    H, _ = jax.lax.scan(step, jnp.zeros(NB, x0.dtype), jnp.arange(1, T + 1))
    return jnp.sum((H - target) ** 2)

# Compacted: the closest JAX gets to x[mask]. `nonzero` needs a static `size`, so the
# selection is padded to a capacity chosen from the data, checked outside jit.
@functools.partial(jax.jit, static_argnums=3)
def loss_compact(x0, v, target, cap):
    n = x0.shape[0]
    def step(H, t):
        inside = jnp.abs(minimg(x0 + v * (t * DT) - t / T)) < HW
        idx = jnp.nonzero(inside, size=cap, fill_value=n)[0]
        i = jnp.minimum(idx, n - 1)
        x = x0[i] + v[i] * (t * DT)
        w = jnp.where(idx < n, taper(minimg(x - t / T)), 0.0)
        return deposit(H, x - jnp.floor(x), w), None
    H, _ = jax.lax.scan(step, jnp.zeros(NB, x0.dtype), jnp.arange(1, T + 1))
    return jnp.sum((H - target) ** 2)

def needed_capacity(x0, v):
    """Largest number of selected particles at any step (host-side, outside jit)."""
    t = np.arange(1, T + 1)[:, None]
    d = np.asarray(x0)[None] + np.asarray(v)[None] * t * DT - t / T
    return int((np.abs(d - np.round(d)) < HW).sum(axis=1).max())
Source: 03. Compaction and scatter-add — JAX
grad_masked = jax.jit(jax.grad(loss_masked))
grad_compact = jax.jit(jax.grad(loss_compact), static_argnums=3)

cap0 = needed_capacity(x0, v0)
g = grad_compact(x0, v0, target0, cap0)
print(f"capacity = {cap0}   loss = {float(loss_compact(x0, v0, target0, cap0)):.10f}   "
      f"‖∇loss‖ = {float(jnp.linalg.norm(g)):.6f}")
capacity = 30   loss = 559.5605712700   ‖∇loss‖ = 52854.055962

Same gradient

implementation ‖g‖₂ max rel. diff vs enzyme
enzyme 52854.055962 0.0e+00
jax 52854.055962 4.9e-14

The gradients differ at the 1e-14 level, not 1e-16, because scatter-adds into the same bin are summed in a different order.

Timings

Float64, \(T = 50\) steps, median of 3–10 runs after warm-up.

What this shows

CPU, \(N = 10^7\). The gradient takes 3.18 s with Enzyme’s compacting loop, 4.82 s with JAX’s nonzero(size=cap), and 11.99 s with JAX’s masked version, which deposits all \(N\) particles at every step with zero weight. Compaction pays off in JAX too, but there it needs a capacity, a host-side check, and a recompilation whenever the capacity changes.

Memory is where masking really hurts. Reverse mode through the 50-step scan keeps each step’s residuals. For the masked version that is all \(N\) particles per step; for the compacted version and for Enzyme, only the selected ones.

Table 1: Gradient memory [MiB]. Enzyme: host bytes allocated by the call. JAX: XLA temp buffers.
  N = 1e+04 N = 1e+05 N = 1e+06 N = 1e+07
implementation        
enzyme (cpu) 0.1 0.8 7.9 78.0
jax / compacted (cpu) 0.6 5.8 57.5 567.3
jax / compacted (gpu) 0.3 2.8 27.9 271.2
jax / masked (cpu) 24.1 240.7 2,407.1 24,070.7
jax / masked (gpu) 23.9 240.3 2,403.3 24,032.6

On the GPU (16 GB) the masked JAX gradient runs out of memory at \(N = 10^7\):

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

Where JAX wins / what we did not do

Note

On the GPU, JAX’s compacted version is the fastest. At \(N = 10^7\) the gradient takes 52.6 ms with JAX against 146.9 ms with Enzyme. On a GPU every implementation tests all \(N\) particles at every step. Enzyme’s kernel lets threads outside the window return immediately and scatters the rest with atomic adds, while XLA’s nonzero compaction followed by a dense deposit on ~1 % of the particles is faster here. Getting there in JAX required knowing the nonzero(size=…) idiom and computing the capacity outside jit.

  • Neither side sorts particles by position, which would make the selection itself \(O(\text{selected})\) instead of \(O(N)\) per step.
  • No Enzyme GPU compaction (a stream-compaction kernel before the deposit), which is what a CUDA programmer would write.
  • No Reactant column on this page yet.