(0.90.7) Remove argument splatting in hydrostatic free surface tendency kernel entry functions

[glwagner]

This PR removes argument splatting in intermediate functions that are called to compute the hydrostatic free surface tendencies.

Argument splatting was removed in a prior PR (that'd be great if someone can remember), but was reinstated in #3360.

This PR re-removes splatting. It yields a 2x performance gain for a column model:

main

julia> include("test_single_column_model.jl")
┌ Info: Running a simulation of HydrostaticFreeSurfaceModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
│ ├── grid: 1×1×64 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
│ ├── timestepper: QuasiAdamsBashforth2TimeStepper
│ ├── tracers: (b, e)
│ ├── closure: CATKEVerticalDiffusivity{VerticallyImplicitTimeDiscretization}
│ ├── buoyancy: BuoyancyTracer with ĝ = NegativeZDirection()
└ └── coriolis: FPlane{Float64}...
  0.398369 seconds (727.30 k allocations: 706.103 MiB, 30.55% gc time)

This PR

julia> include("test_single_column_model.jl")
┌ Info: Running a simulation of HydrostaticFreeSurfaceModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
│ ├── grid: 1×1×64 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
│ ├── timestepper: QuasiAdamsBashforth2TimeStepper
│ ├── tracers: (b, e)
│ ├── closure: CATKEVerticalDiffusivity{VerticallyImplicitTimeDiscretization}
│ ├── buoyancy: BuoyancyTracer with ĝ = NegativeZDirection()
└ └── coriolis: FPlane{Float64}...
  0.214935 seconds (258.50 k allocations: 195.374 MiB, 10.57% gc time)

It also reduces memory allocation.

Note that the nonhydrostatic model was not changed (it does not splat).

[glwagner]

@wsmoses @jlk9

[glwagner]

Here's a script for generating timings

using Oceananigans
using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity
using Oceananigans.TimeSteppers: time_step!

grid = RectilinearGrid(size=64, z=(-256, 0), topology=(Flat, Flat, Bounded))
coriolis = FPlane(f=1e-4)
closure = CATKEVerticalDiffusivity()

boundary_conditions = (b = FieldBoundaryConditions(top = FluxBoundaryCondition(1e-8)),
                       u = FieldBoundaryConditions(top = FluxBoundaryCondition(-2e-4)))

model = HydrostaticFreeSurfaceModel(; grid, closure, coriolis, boundary_conditions,
                                    tracers = (:b, :e), buoyancy = BuoyancyTracer())

bᵢ(z) = 1e-6 * z
set!(model, b=bᵢ, e=1e-6)

# Compile
time_step!(model, 600)

@time for n = 1:100
    time_step!(model, 600)
end

[glwagner]

@simone-silvestri note that the tridiagonal solver also did not re-introduce splatting, so no change needs to be made there.

[navidcy]

Argument splatting was removed in a prior PR (that'd be great if someone can remember), but was reinstated in #3360.

Did you mean #3026 and #2996?

[glwagner]

That's right @navidcy, nice detective work

[wsmoses]

as mentioned on slack here's @jlk9 's Enzyme integration test PR #3480

We should check that the tuple doesn't also break that (or try to find the correct mutually happy solution like maybe vararg type)

[glwagner]

as mentioned on slack here's @jlk9 's Enzyme integration test PR #3480

We should check that the tuple doesn't also break that (or try to find the correct mutually happy solution like maybe vararg type)

@jlk9 , the test on #3480 does not pass right now, is that correct?

[glwagner]

I think we should merge this now, since Oceananigans is barely useable at the moment, and then pick up getting the Enzyme tests passing in @jlk9's PR (which also needs to involve performance benchmarking to ensure we maintain performance)