`RRTMGP` CUDA kernels slow when called from `ClimaAtmos`

[sriharshakandala]

RRTMGP CUDA kernels take much longer when called from ClimaAtmos compared to the corresponding RRTMGP benchmark for an equivalent problem size.

Please see below for the corresponding ClimaAtmos and RRTMGP builds for the equivalent all-sky problem with aerosols. The number of aerosols is reduced to two to match with the RRTMGP benchmark.

ClimaAtmos PR: #3740

ClimaAtmos build: gpu_aquaplanet_dyamond - strong scaling - 1 GPU - No MPI @
https://buildkite.com/clima/climaatmos-target-gpu-simulations/builds/414#0195ce32-92da-478d-9f29-282e94952871

RRTMGP build: GPU all-sky with aerosols DYAMOND benchmark @ https://buildkite.com/clima/rrtmgp-clima-a100-pipeline/builds/46#_

In this example, the CUDA kernels for the longwave and shortwave problems take about 1.5x and 1.6x the time taken by the full longwave or shortwave solver in RRTMGP.

[sriharshakandala]

cc: @szy21 @akshaysridhar @charleskawczynski @dennisYatunin

[sriharshakandala]

I updated the grid for ClimaAtmos simulation to use helems = 28 (ncol = 75264), zelem = 72.
The RRTMGP benchmark uses ncol = 75748, nlev = 73.
The updated ClimaAtmos build is here: https://buildkite.com/clima/climaatmos-target-gpu-simulations/builds/415#0195cea9-ab0f-4e79-a93e-490f2c404831

After this update, the CUDA kernels for the longwave and shortwave problems take about 1.42x and 1.58x the time taken by the full longwave or shortwave solver in RRTMGP.

[sriharshakandala]

I see the following results when using BenchmarkTools from inside ClimaAtmos:

@show "benchmarking longwave solver"
    trial_lw = @benchmark CUDA.@sync RRTMGP.RTESolver.solve_lw!(
        $(model.lw_solver),
        $(model.as),
        $(model.lookups.lookup_lw),
        $(model.lookups.lookup_lw_cld),
        $(model.lookups.lookup_lw_aero),
    )
    @show "median lw solver time = ", Statistics.median(trial_lw)
    @show eltype(model.lw_solver.flux.flux_up)

    @show "Benchmarking shortwave solver"
    trial_sw = @benchmark CUDA.@sync RRTMGP.RTESolver.solve_sw!(
        $(model.sw_solver),
        $(model.as),
        $(model.lookups.lookup_sw),
        $(model.lookups.lookup_sw_cld),
        $(model.lookups.lookup_sw_aero),
    )
    @show "median sw solver time = ", Statistics.median(trial_sw)
    @show eltype(model.sw_solver.flux.flux_up)

model.radiation_mode = ClimaAtmos.RRTMGPInterface.AllSkyRadiationWithClearSkyDiagnostics{ClimaAtmos.InteractiveCloudInRadiation}(false, false, ClimaAtmos.InteractiveCloudInRadiation(), true, true, false)
"benchmarking longwave solver" = "benchmarking longwave solver"
("median lw solver time = ", Statistics.median(trial_lw)) = ("median lw solver time = ", TrialEstimate(1.347 s))
eltype(model.lw_solver.flux.flux_up) = Float32
"Benchmarking shortwave solver" = "Benchmarking shortwave solver"
("median sw solver time = ", Statistics.median(trial_sw)) = ("median sw solver time = ", TrialEstimate(1.060 s))
eltype(model.sw_solver.flux.flux_up) = Float32

These numbers are consistent with what is seen from the Nsight Systems profile presented above.

[sriharshakandala]

With add_isothermal_boundary_layer: false, we see a tangible improvement:

"benchmarking longwave solver" = "benchmarking longwave solver"
("median lw solver time = ", Statistics.median(trial_lw)) = ("median lw solver time = ", TrialEstimate(1.175 s))
eltype(model.lw_solver.flux.flux_up) = Float32
"Benchmarking shortwave solver" = "Benchmarking shortwave solver"
("median sw solver time = ", Statistics.median(trial_sw)) = ("median sw solver time = ", TrialEstimate(890.091 ms))
eltype(model.sw_solver.flux.flux_up) = Float32

My guess is that the performance degradation is due to views.
For reference, the corresponding RRTMGP times are 959.424 ms for longwave solver and 673.576 ms for shortwave solver respectively!

[charleskawczynski]

My guess is that the performance degradation is due to views.

Are any views actually passed to RRTMGP?

I wonder if the extra vertical layer creates cache alignment issues. Can you try add_isothermal_boundary_layer: false and add_isothermal_boundary_layer: true for Spaces.nlevels(ᶜspace) == 62/Spaces.nlevels(ᶠspace) == 63 (I think that this should result in nlevels = 64 passed to RRTMGP.)