Revert "Fixes CPU slowdown with KernelAbstractions >= v0.8 (#3026)"

This reverts commit b37c00f.
CliMA · Mar 24, 2023 · 18d4eaa · 18d4eaa
1 parent b37c00f
commit 18d4eaa
Show file tree

Hide file tree

Showing 12 changed files with 106 additions and 113 deletions.
diff --git a/src/BoundaryConditions/apply_flux_bcs.jl b/src/BoundaryConditions/apply_flux_bcs.jl
@@ -33,28 +33,28 @@ Apply flux boundary conditions to a field `c` by adding the associated flux dive
 the source term `Gc` at the left and right.
 """
 apply_x_bcs!(Gc, grid::AbstractGrid, c, west_bc, east_bc, arch::AbstractArchitecture, dep, args...) =
-    launch!(arch, grid, :yz, _apply_x_bcs!, Gc, instantiated_location(Gc), grid, west_bc, east_bc, Tuple(args), dependencies=dep)
+    launch!(arch, grid, :yz, _apply_x_bcs!, Gc, instantiated_location(Gc), grid, west_bc, east_bc, args..., dependencies=dep)
 
 """
 Apply flux boundary conditions to a field `c` by adding the associated flux divergence to
 the source term `Gc` at the left and right.
 """
 apply_y_bcs!(Gc, grid::AbstractGrid, c, south_bc, north_bc, arch::AbstractArchitecture, dep, args...) =
-    launch!(arch, grid, :xz, _apply_y_bcs!, Gc, instantiated_location(Gc), grid, south_bc, north_bc, Tuple(args), dependencies=dep)
+    launch!(arch, grid, :xz, _apply_y_bcs!, Gc, instantiated_location(Gc), grid, south_bc, north_bc, args..., dependencies=dep)
 
 """
 Apply flux boundary conditions to a field `c` by adding the associated flux divergence to
 the source term `Gc` at the top and bottom.
 """
 apply_z_bcs!(Gc, grid::AbstractGrid, c, bottom_bc, top_bc, arch::AbstractArchitecture, dep, args...) =
-    launch!(arch, grid, :xy, _apply_z_bcs!, Gc, instantiated_location(Gc), grid, bottom_bc, top_bc, Tuple(args), dependencies=dep)
+    launch!(arch, grid, :xy, _apply_z_bcs!, Gc, instantiated_location(Gc), grid, bottom_bc, top_bc, args..., dependencies=dep)
 
 """
     _apply_x_bcs!(Gc, grid, west_bc, east_bc, args...)
 
 Apply a west and/or east boundary condition to variable `c`.
 """
-@kernel function _apply_x_bcs!(Gc, loc, grid, west_bc, east_bc, args) 
+@kernel function _apply_x_bcs!(Gc, loc, grid, west_bc, east_bc, args...)
     j, k = @index(Global, NTuple)
     apply_x_west_bc!(Gc, loc, west_bc, j, k, grid, args...)
     apply_x_east_bc!(Gc, loc, east_bc, j, k, grid, args...)
@@ -65,7 +65,7 @@ end
 
 Apply a south and/or north boundary condition to variable `c`.
 """
-@kernel function _apply_y_bcs!(Gc, loc, grid, south_bc, north_bc, args)
+@kernel function _apply_y_bcs!(Gc, loc, grid, south_bc, north_bc, args...)
     i, k = @index(Global, NTuple)
     apply_y_south_bc!(Gc, loc, south_bc, i, k, grid, args...)
     apply_y_north_bc!(Gc, loc, north_bc, i, k, grid, args...)
@@ -76,7 +76,7 @@ end
 
 Apply a top and/or bottom boundary condition to variable `c`.
 """
-@kernel function _apply_z_bcs!(Gc, loc, grid, bottom_bc, top_bc, args) 
+@kernel function _apply_z_bcs!(Gc, loc, grid, bottom_bc, top_bc, args...)
     i, j = @index(Global, NTuple)
     apply_z_bottom_bc!(Gc, loc, bottom_bc, i, j, grid, args...)
        apply_z_top_bc!(Gc, loc, top_bc,    i, j, grid, args...)

diff --git a/src/BoundaryConditions/fill_halo_regions.jl b/src/BoundaryConditions/fill_halo_regions.jl
@@ -170,23 +170,23 @@ fill_first(bc1, bc2)               = true
 ##### General fill_halo! kernels
 #####
 
-@kernel function _fill_west_and_east_halo!(c, west_bc, east_bc, offset, loc, grid, args) 
+@kernel function _fill_west_and_east_halo!(c, west_bc, east_bc, offset, loc, grid, args...)
     j, k = @index(Global, NTuple)
     j′ = j + offset[1]
     k′ = k + offset[2]
     _fill_west_halo!(j′, k′, grid, c, west_bc, loc, args...)
     _fill_east_halo!(j′, k′, grid, c, east_bc, loc, args...)
 end
 
-@kernel function _fill_south_and_north_halo!(c, south_bc, north_bc, offset, loc, grid, args)
+@kernel function _fill_south_and_north_halo!(c, south_bc, north_bc, offset, loc, grid, args...)
     i, k = @index(Global, NTuple)
     i′ = i + offset[1]
     k′ = k + offset[2]
     _fill_south_halo!(i′, k′, grid, c, south_bc, loc, args...)
     _fill_north_halo!(i′, k′, grid, c, north_bc, loc, args...)
 end
 
-@kernel function _fill_bottom_and_top_halo!(c, bottom_bc, top_bc, offset, loc, grid, args)
+@kernel function _fill_bottom_and_top_halo!(c, bottom_bc, top_bc, offset, loc, grid, args...)
     i, j = @index(Global, NTuple)
     i′ = i + offset[1]
     j′ = j + offset[2]
@@ -200,7 +200,7 @@ end
 
 import Oceananigans.Utils: @constprop
 
-@kernel function _fill_west_and_east_halo!(c::NTuple, west_bc, east_bc, offset, loc, grid, args)
+@kernel function _fill_west_and_east_halo!(c::NTuple, west_bc, east_bc, offset, loc, grid, args...)
     j, k = @index(Global, NTuple)
     ntuple(Val(length(west_bc))) do n
         Base.@_inline_meta
@@ -212,7 +212,7 @@ import Oceananigans.Utils: @constprop
     end
 end
 
-@kernel function _fill_south_and_north_halo!(c::NTuple, south_bc, north_bc, offset, loc, grid, args) 
+@kernel function _fill_south_and_north_halo!(c::NTuple, south_bc, north_bc, offset, loc, grid, args...)
     i, k = @index(Global, NTuple)
     ntuple(Val(length(south_bc))) do n
         Base.@_inline_meta
@@ -224,7 +224,7 @@ end
     end
 end
 
-@kernel function _fill_bottom_and_top_halo!(c::NTuple, bottom_bc, top_bc, offset, loc, grid, args) 
+@kernel function _fill_bottom_and_top_halo!(c::NTuple, bottom_bc, top_bc, offset, loc, grid, args...)
     i, j = @index(Global, NTuple)
     ntuple(Val(length(bottom_bc))) do n
         Base.@_inline_meta
@@ -237,13 +237,13 @@ end
 end
 
 fill_west_and_east_halo!(c, west_bc, east_bc, size, offset, loc, arch, dep, grid, args...; kwargs...) =
-    launch!(arch, grid, size, _fill_west_and_east_halo!, c, west_bc, east_bc, offset, loc, grid, Tuple(args); dependencies=dep, kwargs...)
+    launch!(arch, grid, size, _fill_west_and_east_halo!, c, west_bc, east_bc, offset, loc, grid, args...; dependencies=dep, kwargs...)
 
 fill_south_and_north_halo!(c, south_bc, north_bc, size, offset, loc, arch, dep, grid, args...; kwargs...) =
-    launch!(arch, grid, size, _fill_south_and_north_halo!, c, south_bc, north_bc, offset, loc, grid, Tuple(args); dependencies=dep, kwargs...)
+    launch!(arch, grid, size, _fill_south_and_north_halo!, c, south_bc, north_bc, offset, loc, grid, args...; dependencies=dep, kwargs...)
 
 fill_bottom_and_top_halo!(c, bottom_bc, top_bc, size, offset, loc, arch, dep, grid, args...; kwargs...) =
-    launch!(arch, grid, size, _fill_bottom_and_top_halo!, c, bottom_bc, top_bc, offset, loc, grid, Tuple(args); dependencies=dep, kwargs...)
+    launch!(arch, grid, size, _fill_bottom_and_top_halo!, c, bottom_bc, top_bc, offset, loc, grid, args...; dependencies=dep, kwargs...)
 
 #####
 ##### Calculate kernel size and offset for Windowed and Sliced Fields

diff --git a/src/BoundaryConditions/fill_halo_regions_open.jl b/src/BoundaryConditions/fill_halo_regions_open.jl
@@ -9,33 +9,33 @@
 # because the boundary-normal index can vary (and array boundary conditions need to be
 # 3D in general).
 
-@kernel function set_west_or_east_u!(u, offset, i_boundary, bc, grid, args) 
+@kernel function set_west_or_east_u!(u, offset, i_boundary, bc, grid, args...)
     j, k = @index(Global, NTuple)
     j′ = j + offset[1]
     k′ = k + offset[2]
 @inbounds u[i_boundary, j′, k′] = getbc(bc, j′, k′, grid, args...)
 end
 
-@kernel function set_south_or_north_v!(v, offset, j_boundary, bc, grid, args)
+@kernel function set_south_or_north_v!(v, offset, j_boundary, bc, grid, args...)
     i, k = @index(Global, NTuple)
     i′ = i + offset[1]
     k′ = k + offset[2]
 @inbounds v[i′, j_boundary, k′] = getbc(bc, i′, k′, grid, args...)
 end
 
-@kernel function set_bottom_or_top_w!(w, offset, k_boundary, bc, grid, args) 
+@kernel function set_bottom_or_top_w!(w, offset, k_boundary, bc, grid, args...)
     i, j = @index(Global, NTuple)
     i′ = i + offset[1]
     j′ = j + offset[2]
 @inbounds w[i′, j′, k_boundary] = getbc(bc, i′, j′, grid, args...)
 end
 
-@inline   fill_west_halo!(u, bc::OBC, kernel_size, offset, loc, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, kernel_size, set_west_or_east_u!,   u, offset,           1, bc, grid, Tuple(args); dependencies=dep, kwargs...)
-@inline   fill_east_halo!(u, bc::OBC, kernel_size, offset, loc, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, kernel_size, set_west_or_east_u!,   u, offset, grid.Nx + 1, bc, grid, Tuple(args); dependencies=dep, kwargs...)
-@inline  fill_south_halo!(v, bc::OBC, kernel_size, offset, loc, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, kernel_size, set_south_or_north_v!, v, offset,           1, bc, grid, Tuple(args); dependencies=dep, kwargs...)
-@inline  fill_north_halo!(v, bc::OBC, kernel_size, offset, loc, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, kernel_size, set_south_or_north_v!, v, offset, grid.Ny + 1, bc, grid, Tuple(args); dependencies=dep, kwargs...)
-@inline fill_bottom_halo!(w, bc::OBC, kernel_size, offset, loc, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, kernel_size, set_bottom_or_top_w!,  w, offset,           1, bc, grid, Tuple(args); dependencies=dep, kwargs...)
-@inline    fill_top_halo!(w, bc::OBC, kernel_size, offset, loc, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, kernel_size, set_bottom_or_top_w!,  w, offset, grid.Nz + 1, bc, grid, Tuple(args); dependencies=dep, kwargs...)
+@inline   fill_west_halo!(u, bc::OBC, kernel_size, offset, loc, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, kernel_size, set_west_or_east_u!,   u, offset,           1, bc, grid, args...; dependencies=dep, kwargs...)
+@inline   fill_east_halo!(u, bc::OBC, kernel_size, offset, loc, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, kernel_size, set_west_or_east_u!,   u, offset, grid.Nx + 1, bc, grid, args...; dependencies=dep, kwargs...)
+@inline  fill_south_halo!(v, bc::OBC, kernel_size, offset, loc, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, kernel_size, set_south_or_north_v!, v, offset,           1, bc, grid, args...; dependencies=dep, kwargs...)
+@inline  fill_north_halo!(v, bc::OBC, kernel_size, offset, loc, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, kernel_size, set_south_or_north_v!, v, offset, grid.Ny + 1, bc, grid, args...; dependencies=dep, kwargs...)
+@inline fill_bottom_halo!(w, bc::OBC, kernel_size, offset, loc, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, kernel_size, set_bottom_or_top_w!,  w, offset,           1, bc, grid, args...; dependencies=dep, kwargs...)
+@inline    fill_top_halo!(w, bc::OBC, kernel_size, offset, loc, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, kernel_size, set_bottom_or_top_w!,  w, offset, grid.Nz + 1, bc, grid, args...; dependencies=dep, kwargs...)
 
 @inline   _fill_west_halo!(j, k, grid, c, bc::OBC, loc, args...) = @inbounds c[1, j, k]           = getbc(bc, j, k, grid, args...)
 @inline   _fill_east_halo!(j, k, grid, c, bc::OBC, loc, args...) = @inbounds c[grid.Nx + 1, j, k] = getbc(bc, j, k, grid, args...)

diff --git a/src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_tendencies.jl
@@ -104,28 +104,26 @@ function calculate_hydrostatic_free_surface_interior_tendency_contributions!(mod
                                                                                          model.closure,
                                                                                          model.diffusivity_fields)
 
-        args = tuple(Val(tracer_index),
-                     Val(tracer_name),
-                     c_advection,
-                     closure,
-                     c_immersed_bc,
-                     model.buoyancy,
-                     model.biogeochemistry,
-                     model.velocities,
-                     model.free_surface,
-                     model.tracers,
-                     top_tracer_bcs,
-                     diffusivity_fields,
-                     model.auxiliary_fields,
-                     c_forcing,
-                     model.clock)
-
         Gc_event = launch!(arch, grid, :xyz,
                            calculate_hydrostatic_free_surface_Gc!,
                            c_tendency,
                            c_kernel_function,
                            grid,
-                           args;
+                           Val(tracer_index),
+                           Val(tracer_name),
+                           c_advection,
+                           closure,
+                           c_immersed_bc,
+                           model.buoyancy,
+                           model.biogeochemistry,
+                           model.velocities,
+                           model.free_surface,
+                           model.tracers,
+                           top_tracer_bcs,
+                           diffusivity_fields,
+                           model.auxiliary_fields,
+                           c_forcing,
+                           model.clock;
                            dependencies = barrier, 
                            only_active_cells)
 
@@ -155,17 +153,15 @@ function calculate_free_surface_tendency!(grid, model, dependencies)
 
     arch = architecture(grid)
 
-    args = tuple(model.velocities,
-                 model.free_surface,
-                 model.tracers,
-                 model.auxiliary_fields,
-                 model.forcing,
-                 model.clock)
-
     Gη_event = launch!(arch, grid, :xy,
                        calculate_hydrostatic_free_surface_Gη!, model.timestepper.Gⁿ.η,
                        grid,
-                       args;
+                       model.velocities,
+                       model.free_surface,
+                       model.tracers,
+                       model.auxiliary_fields,
+                       model.forcing,
+                       model.clock;
                        dependencies = dependencies)
 
     return Gη_event
@@ -181,7 +177,8 @@ function calculate_hydrostatic_momentum_tendencies!(model, velocities; dependenc
     u_immersed_bc = immersed_boundary_condition(velocities.u)
     v_immersed_bc = immersed_boundary_condition(velocities.v)
 
-    start_momentum_kernel_args = (model.advection.momentum,
+    start_momentum_kernel_args = (grid,
+                                  model.advection.momentum,
                                   model.coriolis,
                                   model.closure)
 
@@ -201,11 +198,11 @@ function calculate_hydrostatic_momentum_tendencies!(model, velocities; dependenc
     only_active_cells = use_only_active_cells(grid)
 
     Gu_event = launch!(arch, grid, :xyz,
-                       calculate_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid, u_kernel_args;
+                       calculate_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, u_kernel_args...;
                        dependencies = dependencies, only_active_cells)
 
     Gv_event = launch!(arch, grid, :xyz,
-                       calculate_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid, v_kernel_args;
+                       calculate_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, v_kernel_args...;
                        dependencies = dependencies, only_active_cells)
 
     Gη_event = calculate_free_surface_tendency!(grid, model, dependencies)
@@ -246,24 +243,24 @@ end
 #####
 
 """ Calculate the right-hand-side of the u-velocity equation. """
-@kernel function calculate_hydrostatic_free_surface_Gu!(Gu, grid, args)
+@kernel function calculate_hydrostatic_free_surface_Gu!(Gu, grid, 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 calculate_hydrostatic_free_surface_Gu!(Gu, grid::ActiveCellsIBG, args)
+@kernel function calculate_hydrostatic_free_surface_Gu!(Gu, grid::ActiveCellsIBG, args...)
     idx = @index(Global, Linear)
     i, j, k = active_linear_index_to_ntuple(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 calculate_hydrostatic_free_surface_Gv!(Gv, grid, args)
+@kernel function calculate_hydrostatic_free_surface_Gv!(Gv, grid, 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 calculate_hydrostatic_free_surface_Gv!(Gv, grid::ActiveCellsIBG, args)
+@kernel function calculate_hydrostatic_free_surface_Gv!(Gv, grid::ActiveCellsIBG, args...)
     idx = @index(Global, Linear)
     i, j, k = active_linear_index_to_ntuple(idx, grid)
     @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...)
@@ -274,12 +271,12 @@ end
 #####
 
 """ Calculate the right-hand-side of the tracer advection-diffusion equation. """
-@kernel function calculate_hydrostatic_free_surface_Gc!(Gc, tendency_kernel_function, grid, args)
+@kernel function calculate_hydrostatic_free_surface_Gc!(Gc, tendency_kernel_function, grid, args...)
     i, j, k = @index(Global, NTuple)
     @inbounds Gc[i, j, k] = tendency_kernel_function(i, j, k, grid, args...)
 end
 
-@kernel function calculate_hydrostatic_free_surface_Gc!(Gc, tendency_kernel_function, grid::ActiveCellsIBG, args)
+@kernel function calculate_hydrostatic_free_surface_Gc!(Gc, tendency_kernel_function, grid::ActiveCellsIBG, args...)
     idx = @index(Global, Linear)
     i, j, k = active_linear_index_to_ntuple(idx, grid)
     @inbounds Gc[i, j, k] = tendency_kernel_function(i, j, k, grid, args...)
@@ -290,7 +287,7 @@ end
 #####
 
 """ Calculate the right-hand-side of the free surface displacement (``η``) equation. """
-@kernel function calculate_hydrostatic_free_surface_Gη!(Gη, grid, args)
+@kernel function calculate_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