Inference problem on the CPU

[simone-silvestri]

Working with @milankl I noticed that a stripped-down version of ClimaOcean (with only the HydrostaticFreeSurfaceModel) was terribly slow compared to a SpeedyWeather simulation with the same gridsize.

After some investigating I found out that the problem was the forcing added as an argument to the momentum kernels after the splatting of the other args....
I am not sure what caused this loss of performance, but the benchmark of the code:

using ClimaOcean
using Oceananigans
using Oceananigans.Units
using BenchmarkTools

arch    = Oceananigans.CPU()
r_faces = ClimaOcean.exponential_z_faces(; Nz = 8, depth = 5000)

grid = LatitudeLongitudeGrid(arch, 
                             size=(360, 152, 8), 
                             latitude=(-75.5, 75.5), 
                             longitude=(0, 360),
                             halo = (6, 6, 2),
                             z = r_faces)

ocean = ClimaOcean.ocean_simulation(grid; 
                                    closure=nothing,
                                    momentum_advection=VectorInvariant(),
                                    tracer_advection=Centered(),
                                    free_surface=SplitExplicitFreeSurface(grid; substeps=20))

time_step!(ocean)
time_step!(ocean)

@benchmark time_step!(ocean)

on main:

and on this branch with the correction:

This should not impact GPU execution.

[glwagner]

Oh. I think this is basically #2996 coming back. This keeps happening so I wonder if there is a way to test it.

[simone-silvestri]

We might need to add some performance tests. Possibly very simple ones, for example tests that check that there is no extra allocation that comes from types not being inferred correctly.