(0.90.7) Remove argument splatting in hydrostatic free surface tenden…

…cy kernel entry functions (#3477) * Dont splat * Simple benchmark * Simplify single column model benchmark * Update Project.toml --------- Co-authored-by: Navid C. Constantinou <[email protected]>
CliMA · Feb 16, 2024 · 0391b3a · 0391b3a · glwagner · Feb 16, 2024
1 parent c55878e
commit 0391b3a
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 19 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.90.6"
+version = "0.90.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl
@@ -122,7 +122,7 @@ function compute_hydrostatic_free_surface_tendency_contributions!(model, kernel_
                     c_tendency,
                     grid,
                     only_active_cells,
-                    args...;
+                    args;
                     only_active_cells)
         end
     end
@@ -134,11 +134,10 @@ end
 ##### Boundary condributions to hydrostatic free surface model
 #####
 
-function apply_flux_bcs!(Gcⁿ, c, arch, args...)
+function apply_flux_bcs!(Gcⁿ, c, arch, args)
     apply_x_bcs!(Gcⁿ, c, arch, args...)
     apply_y_bcs!(Gcⁿ, c, arch, args...)
     apply_z_bcs!(Gcⁿ, c, arch, args...)
-
     return nothing
 end
 
@@ -155,7 +154,7 @@ function compute_free_surface_tendency!(grid, model, kernel_parameters)
 
     launch!(arch, grid, kernel_parameters,
             compute_hydrostatic_free_surface_Gη!, model.timestepper.Gⁿ.η, 
-            grid, args...)
+            grid, args)
 
     return nothing
 end
@@ -189,12 +188,12 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para
     for parameters in kernel_parameters
         launch!(arch, grid, parameters,
                 compute_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid, 
-                only_active_cells, u_kernel_args...;
+                only_active_cells, u_kernel_args;
                 only_active_cells)
 
         launch!(arch, grid, parameters,
                 compute_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid, 
-                only_active_cells, v_kernel_args...;
+                only_active_cells, v_kernel_args;
                 only_active_cells)
     end
 
@@ -206,17 +205,19 @@ end
 """ Apply boundary conditions by adding flux divergences to the right-hand-side. """
 function compute_hydrostatic_boundary_tendency_contributions!(Gⁿ, arch, velocities, free_surface, tracers, args...)
 
+    args = Tuple(args)
+
     # Velocity fields
     for i in (:u, :v)
-        apply_flux_bcs!(Gⁿ[i], velocities[i], arch, args...)
+        apply_flux_bcs!(Gⁿ[i], velocities[i], arch, args)
     end
 
     # Free surface
-    apply_flux_bcs!(Gⁿ.η, displacement(free_surface), arch, args...)
+    apply_flux_bcs!(Gⁿ.η, displacement(free_surface), arch, args)
 
     # Tracer fields
     for i in propertynames(tracers)
-        apply_flux_bcs!(Gⁿ[i], tracers[i], arch, args...)
+        apply_flux_bcs!(Gⁿ[i], tracers[i], arch, args)
     end
 
     return nothing
@@ -227,24 +228,24 @@ end
 #####
 
 """ Calculate the right-hand-side of the u-velocity equation. """
-@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, interior_map, args...)
+@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, interior_map, args)
     i, j, k = @index(Global, NTuple)
     @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...)
 end
 
-@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid::ActiveCellsIBG, ::InteriorMap, args...)
+@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid::ActiveCellsIBG, ::InteriorMap, args)
     idx = @index(Global, Linear)
     i, j, k = active_linear_index_to_interior_tuple(idx, grid)
     @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...)
 end
 
 """ Calculate the right-hand-side of the v-velocity equation. """
-@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, interior_map, args...)
+@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, interior_map, args)
     i, j, k = @index(Global, NTuple)
     @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...)
 end
 
-@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid::ActiveCellsIBG, ::InteriorMap, args...)
+@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid::ActiveCellsIBG, ::InteriorMap, args)
     idx = @index(Global, Linear)
     i, j, k = active_linear_index_to_interior_tuple(idx, grid)
     @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...)
@@ -255,24 +256,24 @@ end
 #####
 
 """ Calculate the right-hand-side of the tracer advection-diffusion equation. """
-@kernel function compute_hydrostatic_free_surface_Gc!(Gc, grid, interior_map, args...)
+@kernel function compute_hydrostatic_free_surface_Gc!(Gc, grid, interior_map, args)
     i, j, k = @index(Global, NTuple)
     @inbounds Gc[i, j, k] = hydrostatic_free_surface_tracer_tendency(i, j, k, grid, args...)
 end
 
-@kernel function compute_hydrostatic_free_surface_Gc!(Gc, grid::ActiveCellsIBG, ::InteriorMap, args...)
+@kernel function compute_hydrostatic_free_surface_Gc!(Gc, grid::ActiveCellsIBG, ::InteriorMap, args)
     idx = @index(Global, Linear)
     i, j, k = active_linear_index_to_interior_tuple(idx, grid)
     @inbounds Gc[i, j, k] = hydrostatic_free_surface_tracer_tendency(i, j, k, grid, args...)
 end
 
 """ Calculate the right-hand-side of the subgrid scale energy equation. """
-@kernel function compute_hydrostatic_free_surface_Ge!(Ge, grid, interior_map, args...)
+@kernel function compute_hydrostatic_free_surface_Ge!(Ge, grid, interior_map, args)
     i, j, k = @index(Global, NTuple)
     @inbounds Ge[i, j, k] = hydrostatic_turbulent_kinetic_energy_tendency(i, j, k, grid, args...)
 end
 
-@kernel function compute_hydrostatic_free_surface_Ge!(Ge, grid::ActiveCellsIBG, ::InteriorMap, args...)
+@kernel function compute_hydrostatic_free_surface_Ge!(Ge, grid::ActiveCellsIBG, ::InteriorMap, args)
     idx = @index(Global, Linear)
     i, j, k = active_linear_index_to_interior_tuple(idx, grid)
     @inbounds Ge[i, j, k] = hydrostatic_turbulent_kinetic_energy_tendency(i, j, k, grid, args...)
@@ -283,7 +284,7 @@ end
 #####
 
 """ Calculate the right-hand-side of the free surface displacement (``η``) equation. """
-@kernel function compute_hydrostatic_free_surface_Gη!(Gη, grid, args...)
+@kernel function compute_hydrostatic_free_surface_Gη!(Gη, grid, args)
     i, j = @index(Global, NTuple)
     @inbounds Gη[i, j, grid.Nz+1] = free_surface_tendency(i, j, grid, args...)
 end
diff --git a/validation/vertical_mixing_closures/test_single_column_model.jl b/validation/vertical_mixing_closures/test_single_column_model.jl
@@ -0,0 +1,25 @@
+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
+
+