Stop the light attenuation kernels at the seafloor

[ali-ramadhan]

This PR resolves #414 by making integrate_light_attenuation! and update_MultiBandPhotosyntheticallyActiveRadiation! stop at the first immersed cell.

The multi-band kernel also no longer writes the top cell of fully immersed (land) columns. And _compute_euphotic_depth! stops at the seafloor too.

I suppose CliMA/Oceananigans.jl#5991 could also fix #414 but it feels correct to fix this too.

[jagoosw]

Similar to Simone's comment on CliMA/Oceananigans.jl#5991, doesn't this introduce branching?

Is the Oceananigans bug currently causing NaNs in the PAR or are you just saying something like that could cause this?

[ali-ramadhan]

Thanks for pointing that out. Yes, this PR currently introduces branch divergences on the GPU which did result in a measurable slowdown in that PR. I can refactor to avoid branch divergence.

Is the Oceananigans bug currently causing NaNs in the PAR or are you just saying something like that could cause this?

It is currently causing NaNs so I saw it as two issues worth fixing. Either fix will be enough, but perhaps this PR is more about answering the question: Should light attenuation be computed below the seafloor?

[jagoosw]

Is it okay if I merge the other PR now and we can discuss more?

[jagoosw]

Is there a good way to do this that doesn't introduce branching? I guess at the moment we are possibly also calculating in land columns which I think we could fix with the active cell map machinery?

[ali-ramadhan]

Is it okay if I merge the other PR now and we can discuss more?

Yes please go ahead! I just created the stacked PR to avoid merge conflicts, not necessarily to merge the two PRs together. Still figuring out if I like stacked PRs haha.

Is there a good way to do this that doesn't introduce branching? I guess at the moment we are possibly also calculating in land columns which I think we could fix with the active cell map machinery?

Lemme look at the code more closely. If we're worried about branch divergence in GPU kernels and CliMA/Oceananigans.jl#5991 avoids the NaNs for now, then we can save "not computing light attenuation below the seafloor" and related issues for a future more fundamental PR.

[jagoosw]

I was thinking about this more, if we know before the kernel launches where the bottom index is (e.g. by writing it to a field when the modifier is constructed) and loop Nz:-1:kf where kf is read from the field, then we don't get branching just some threads waiting while others loop further right?

[ali-ramadhan]

So I asked Claude to do something benchmarking (see below) and it seems like this PR does not introduce any slowdown with the break. Instead, this PR should actually speed things up a bit.

We can also do something like kᵇ = @inbounds k_bottom[i, j, 1] then iterate for k in grid.Nz:-1:kᵇ as you suggested and it's also faster (maybe a tiny bit faster)! If you think this is a better approach I can refactor the code here.

---

Does immersed_cell(i, j, k, grid) && break inside the light attenuation kernels cost anything on the GPU, and is there a branch-free formulation that would be better?

The divergence introduced by immersed_cell(i, j, k, grid) && break is the benign kind: idle lanes, not serialised code paths, so its worst case is the cost of the code it replaces, and realistic bathymetry runs 1.5–4.7× faster. The branch-free alternative is uniformly slower than break for every bathymetry tested. The only variant that beats break (hoisting the check into a precomputed per-column bottom index) has exactly the same lane-level divergence, so divergence is not what separates them.

What was measured

integrate_light_attenuation! for TwoBandPhotosyntheticallyActiveRadiation (the AbstractLight{N, Nothing} method) as it stands on the ali/light-skip-immersed-cells branch, on a 512 × 512 × 64 Float64 RectilinearGrid wrapped in an ImmersedBoundaryGrid with a GridFittedBottom, on one Tesla V100-PCIE-32GB. 262144 columns, one per thread, launched over :xy as in production. Each number is the best of 5 batches of 100 launches after 20 warm-up launches.

Four bathymetries, chosen to span the range of lane-level trip-count variation within a warp:

bathymetry	variation within a warp
no seafloor (every column full depth)	none — every lane runs Nz iterations
uniform shelf at 20 % of the grid depth	none — every lane runs 0.2 Nz iterations
shelf–slope, 10 % → 100 % of grid depth across x	small — neighbouring columns differ slightly
independent random depth per column	maximal — almost every warp holds a near-full-depth column

The four variants

1. walk all — what main did before this PR: descend the whole column unconditionally.

@inbounds for k in grid.Nz:-1:1
    ...
    PAR[i, j, k] = - PAR⁰ * total_cell_average(K, t, la)
    K = attenuate(K, t, la)
end

2. break — what this PR does: stop at the first immersed cell.

@inbounds for k in grid.Nz:-1:1
    immersed_cell(i, j, k, grid) && break
    ...
end

3. predicated — branch-free: uniform control flow, mask the store.

@inbounds for k in grid.Nz:-1:1
    ...
    PAR[i, j, k] = ifelse(immersed_cell(i, j, k, grid), zero(PAR⁰), - PAR⁰ * total_cell_average(K, t, la))
    K = attenuate(K, t, la)
end

4. bounded — hoist the check out of the loop: a counted loop down to a per-column bottom index, precomputed once with Sediments.calculate_bottom_indices (which is already "first wet k from the top" for any static immersed boundary).

kᵇ = @inbounds k_bottom[i, j, 1]
@inbounds for k in grid.Nz:-1:kᵇ
    ...
end

Results

Kernel time per launch, and the mean fraction of each column that is wet:

bathymetry	walk all	break	predicated	bounded	wet fraction
no seafloor (all wet)	1.331 ms	1.374 ms	1.370 ms	1.330 ms	1.00
uniform shelf, 20 % of grid	1.330 ms	0.286 ms	1.376 ms	0.276 ms	0.20
shelf–slope, 10 % → 100 % in x	1.332 ms	0.910 ms	1.371 ms	0.873 ms	0.55
random depth per column	1.332 ms	1.316 ms	1.370 ms	1.264 ms	0.50

Interpretation

The break costs at most ~3 %, and usually saves multiples of that. The check itself (immersed_cell → one rnode and one compare against bottom_height[i, j, 1]) is the only overhead, visible as 1.374 vs 1.331 ms in the all-wet row where it can never fire. Everything else is a saving: 1.5× on the slope, 4.7× on the shelf.

The divergence is lanes going idle, not code being run twice. There is no else branch to execute — a lane that hits the seafloor is simply masked off for the rest of the loop. A warp therefore takes as long as its deepest column, which is precisely what walk all paid for every warp. That is why the adversarial random-depth row is a wash (1.316 vs 1.332 ms) rather than a regression: the worst case for divergence is the old cost, and there is no bathymetry where break can be meaningfully slower. The divergence to worry about is the kind where a warp serially executes both sides of a branch; this is not that.

Branch-free predication is the worst option. It buys uniform control flow at the price of doing the full-depth work unconditionally, so it lands within noise of walk all (+3 % for the check) in every row — including the ones where break is 1.5–4.7× faster. Uniform control flow is not the objective; doing less work is.

Hoisting the check into a bottom index is the only real win, and it is not about divergence. bounded is consistently fastest because it drops the per-iteration check, recovering the ~3 %. Its lane-level trip-count variation is identical to break — same idle lanes, same warp-max cost — which is the clearest evidence that the branch, not the divergence, is what break pays for. It would mean carrying a bottom_indices field on each light model (the sediment models already do exactly this). Worth a follow-up if the per-iteration branch is a concern; not worth blocking this PR on, since it is a ~3 % refinement of a change that is already a 1.5–4.7× improvement.

Genuinely divergence-free would be a different algorithm. Replacing the serial descent with a parallel scan along z (accumulate optical depth τ, then PAR = PAR⁰ exp(-τ_top) (1 - t) / x per cell with immersed cells masked) removes the per-column loop altogether, at the cost of a scratch field and an extra pass. The kernel is already parallel over 262144 columns and finishes in ~1 ms, so there is nothing to win.