Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 18d4eaa

Browse files
navidcynavidcy
authoredMar 24, 2023
Revert "Fixes CPU slowdown with KernelAbstractions >= v0.8 (#3026)"
This reverts commit b37c00f.
1 parent b37c00f commit 18d4eaa

File tree

12 files changed

+106
-113
lines changed

12 files changed

+106
-113
lines changed
 

‎src/BoundaryConditions/apply_flux_bcs.jl

+6-6
Original file line numberDiff line numberDiff line change
@@ -33,28 +33,28 @@ Apply flux boundary conditions to a field `c` by adding the associated flux dive
3333
the source term `Gc` at the left and right.
3434
"""
3535
apply_x_bcs!(Gc, grid::AbstractGrid, c, west_bc, east_bc, arch::AbstractArchitecture, dep, args...) =
36-
launch!(arch, grid, :yz, _apply_x_bcs!, Gc, instantiated_location(Gc), grid, west_bc, east_bc, Tuple(args), dependencies=dep)
36+
launch!(arch, grid, :yz, _apply_x_bcs!, Gc, instantiated_location(Gc), grid, west_bc, east_bc, args..., dependencies=dep)
3737

3838
"""
3939
Apply flux boundary conditions to a field `c` by adding the associated flux divergence to
4040
the source term `Gc` at the left and right.
4141
"""
4242
apply_y_bcs!(Gc, grid::AbstractGrid, c, south_bc, north_bc, arch::AbstractArchitecture, dep, args...) =
43-
launch!(arch, grid, :xz, _apply_y_bcs!, Gc, instantiated_location(Gc), grid, south_bc, north_bc, Tuple(args), dependencies=dep)
43+
launch!(arch, grid, :xz, _apply_y_bcs!, Gc, instantiated_location(Gc), grid, south_bc, north_bc, args..., dependencies=dep)
4444

4545
"""
4646
Apply flux boundary conditions to a field `c` by adding the associated flux divergence to
4747
the source term `Gc` at the top and bottom.
4848
"""
4949
apply_z_bcs!(Gc, grid::AbstractGrid, c, bottom_bc, top_bc, arch::AbstractArchitecture, dep, args...) =
50-
launch!(arch, grid, :xy, _apply_z_bcs!, Gc, instantiated_location(Gc), grid, bottom_bc, top_bc, Tuple(args), dependencies=dep)
50+
launch!(arch, grid, :xy, _apply_z_bcs!, Gc, instantiated_location(Gc), grid, bottom_bc, top_bc, args..., dependencies=dep)
5151

5252
"""
5353
_apply_x_bcs!(Gc, grid, west_bc, east_bc, args...)
5454
5555
Apply a west and/or east boundary condition to variable `c`.
5656
"""
57-
@kernel function _apply_x_bcs!(Gc, loc, grid, west_bc, east_bc, args)
57+
@kernel function _apply_x_bcs!(Gc, loc, grid, west_bc, east_bc, args...)
5858
j, k = @index(Global, NTuple)
5959
apply_x_west_bc!(Gc, loc, west_bc, j, k, grid, args...)
6060
apply_x_east_bc!(Gc, loc, east_bc, j, k, grid, args...)
@@ -65,7 +65,7 @@ end
6565
6666
Apply a south and/or north boundary condition to variable `c`.
6767
"""
68-
@kernel function _apply_y_bcs!(Gc, loc, grid, south_bc, north_bc, args)
68+
@kernel function _apply_y_bcs!(Gc, loc, grid, south_bc, north_bc, args...)
6969
i, k = @index(Global, NTuple)
7070
apply_y_south_bc!(Gc, loc, south_bc, i, k, grid, args...)
7171
apply_y_north_bc!(Gc, loc, north_bc, i, k, grid, args...)
@@ -76,7 +76,7 @@ end
7676
7777
Apply a top and/or bottom boundary condition to variable `c`.
7878
"""
79-
@kernel function _apply_z_bcs!(Gc, loc, grid, bottom_bc, top_bc, args)
79+
@kernel function _apply_z_bcs!(Gc, loc, grid, bottom_bc, top_bc, args...)
8080
i, j = @index(Global, NTuple)
8181
apply_z_bottom_bc!(Gc, loc, bottom_bc, i, j, grid, args...)
8282
apply_z_top_bc!(Gc, loc, top_bc, i, j, grid, args...)

‎src/BoundaryConditions/fill_halo_regions.jl

+9-9
Original file line numberDiff line numberDiff line change
@@ -170,23 +170,23 @@ fill_first(bc1, bc2) = true
170170
##### General fill_halo! kernels
171171
#####
172172

173-
@kernel function _fill_west_and_east_halo!(c, west_bc, east_bc, offset, loc, grid, args)
173+
@kernel function _fill_west_and_east_halo!(c, west_bc, east_bc, offset, loc, grid, args...)
174174
j, k = @index(Global, NTuple)
175175
j′ = j + offset[1]
176176
k′ = k + offset[2]
177177
_fill_west_halo!(j′, k′, grid, c, west_bc, loc, args...)
178178
_fill_east_halo!(j′, k′, grid, c, east_bc, loc, args...)
179179
end
180180

181-
@kernel function _fill_south_and_north_halo!(c, south_bc, north_bc, offset, loc, grid, args)
181+
@kernel function _fill_south_and_north_halo!(c, south_bc, north_bc, offset, loc, grid, args...)
182182
i, k = @index(Global, NTuple)
183183
i′ = i + offset[1]
184184
k′ = k + offset[2]
185185
_fill_south_halo!(i′, k′, grid, c, south_bc, loc, args...)
186186
_fill_north_halo!(i′, k′, grid, c, north_bc, loc, args...)
187187
end
188188

189-
@kernel function _fill_bottom_and_top_halo!(c, bottom_bc, top_bc, offset, loc, grid, args)
189+
@kernel function _fill_bottom_and_top_halo!(c, bottom_bc, top_bc, offset, loc, grid, args...)
190190
i, j = @index(Global, NTuple)
191191
i′ = i + offset[1]
192192
j′ = j + offset[2]
@@ -200,7 +200,7 @@ end
200200

201201
import Oceananigans.Utils: @constprop
202202

203-
@kernel function _fill_west_and_east_halo!(c::NTuple, west_bc, east_bc, offset, loc, grid, args)
203+
@kernel function _fill_west_and_east_halo!(c::NTuple, west_bc, east_bc, offset, loc, grid, args...)
204204
j, k = @index(Global, NTuple)
205205
ntuple(Val(length(west_bc))) do n
206206
Base.@_inline_meta
@@ -212,7 +212,7 @@ import Oceananigans.Utils: @constprop
212212
end
213213
end
214214

215-
@kernel function _fill_south_and_north_halo!(c::NTuple, south_bc, north_bc, offset, loc, grid, args)
215+
@kernel function _fill_south_and_north_halo!(c::NTuple, south_bc, north_bc, offset, loc, grid, args...)
216216
i, k = @index(Global, NTuple)
217217
ntuple(Val(length(south_bc))) do n
218218
Base.@_inline_meta
@@ -224,7 +224,7 @@ end
224224
end
225225
end
226226

227-
@kernel function _fill_bottom_and_top_halo!(c::NTuple, bottom_bc, top_bc, offset, loc, grid, args)
227+
@kernel function _fill_bottom_and_top_halo!(c::NTuple, bottom_bc, top_bc, offset, loc, grid, args...)
228228
i, j = @index(Global, NTuple)
229229
ntuple(Val(length(bottom_bc))) do n
230230
Base.@_inline_meta
@@ -237,13 +237,13 @@ end
237237
end
238238

239239
fill_west_and_east_halo!(c, west_bc, east_bc, size, offset, loc, arch, dep, grid, args...; kwargs...) =
240-
launch!(arch, grid, size, _fill_west_and_east_halo!, c, west_bc, east_bc, offset, loc, grid, Tuple(args); dependencies=dep, kwargs...)
240+
launch!(arch, grid, size, _fill_west_and_east_halo!, c, west_bc, east_bc, offset, loc, grid, args...; dependencies=dep, kwargs...)
241241

242242
fill_south_and_north_halo!(c, south_bc, north_bc, size, offset, loc, arch, dep, grid, args...; kwargs...) =
243-
launch!(arch, grid, size, _fill_south_and_north_halo!, c, south_bc, north_bc, offset, loc, grid, Tuple(args); dependencies=dep, kwargs...)
243+
launch!(arch, grid, size, _fill_south_and_north_halo!, c, south_bc, north_bc, offset, loc, grid, args...; dependencies=dep, kwargs...)
244244

245245
fill_bottom_and_top_halo!(c, bottom_bc, top_bc, size, offset, loc, arch, dep, grid, args...; kwargs...) =
246-
launch!(arch, grid, size, _fill_bottom_and_top_halo!, c, bottom_bc, top_bc, offset, loc, grid, Tuple(args); dependencies=dep, kwargs...)
246+
launch!(arch, grid, size, _fill_bottom_and_top_halo!, c, bottom_bc, top_bc, offset, loc, grid, args...; dependencies=dep, kwargs...)
247247

248248
#####
249249
##### Calculate kernel size and offset for Windowed and Sliced Fields

‎src/BoundaryConditions/fill_halo_regions_open.jl

+9-9
Original file line numberDiff line numberDiff line change
@@ -9,33 +9,33 @@
99
# because the boundary-normal index can vary (and array boundary conditions need to be
1010
# 3D in general).
1111

12-
@kernel function set_west_or_east_u!(u, offset, i_boundary, bc, grid, args)
12+
@kernel function set_west_or_east_u!(u, offset, i_boundary, bc, grid, args...)
1313
j, k = @index(Global, NTuple)
1414
j′ = j + offset[1]
1515
k′ = k + offset[2]
1616
@inbounds u[i_boundary, j′, k′] = getbc(bc, j′, k′, grid, args...)
1717
end
1818

19-
@kernel function set_south_or_north_v!(v, offset, j_boundary, bc, grid, args)
19+
@kernel function set_south_or_north_v!(v, offset, j_boundary, bc, grid, args...)
2020
i, k = @index(Global, NTuple)
2121
i′ = i + offset[1]
2222
k′ = k + offset[2]
2323
@inbounds v[i′, j_boundary, k′] = getbc(bc, i′, k′, grid, args...)
2424
end
2525

26-
@kernel function set_bottom_or_top_w!(w, offset, k_boundary, bc, grid, args)
26+
@kernel function set_bottom_or_top_w!(w, offset, k_boundary, bc, grid, args...)
2727
i, j = @index(Global, NTuple)
2828
i′ = i + offset[1]
2929
j′ = j + offset[2]
3030
@inbounds w[i′, j′, k_boundary] = getbc(bc, i′, j′, grid, args...)
3131
end
3232

33-
@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...)
34-
@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...)
35-
@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...)
36-
@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...)
37-
@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...)
38-
@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...)
33+
@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...)
34+
@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...)
35+
@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...)
36+
@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...)
37+
@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...)
38+
@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...)
3939

4040
@inline _fill_west_halo!(j, k, grid, c, bc::OBC, loc, args...) = @inbounds c[1, j, k] = getbc(bc, j, k, grid, args...)
4141
@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...)

‎src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_tendencies.jl

+32-35
Original file line numberDiff line numberDiff line change
@@ -104,28 +104,26 @@ function calculate_hydrostatic_free_surface_interior_tendency_contributions!(mod
104104
model.closure,
105105
model.diffusivity_fields)
106106

107-
args = tuple(Val(tracer_index),
108-
Val(tracer_name),
109-
c_advection,
110-
closure,
111-
c_immersed_bc,
112-
model.buoyancy,
113-
model.biogeochemistry,
114-
model.velocities,
115-
model.free_surface,
116-
model.tracers,
117-
top_tracer_bcs,
118-
diffusivity_fields,
119-
model.auxiliary_fields,
120-
c_forcing,
121-
model.clock)
122-
123107
Gc_event = launch!(arch, grid, :xyz,
124108
calculate_hydrostatic_free_surface_Gc!,
125109
c_tendency,
126110
c_kernel_function,
127111
grid,
128-
args;
112+
Val(tracer_index),
113+
Val(tracer_name),
114+
c_advection,
115+
closure,
116+
c_immersed_bc,
117+
model.buoyancy,
118+
model.biogeochemistry,
119+
model.velocities,
120+
model.free_surface,
121+
model.tracers,
122+
top_tracer_bcs,
123+
diffusivity_fields,
124+
model.auxiliary_fields,
125+
c_forcing,
126+
model.clock;
129127
dependencies = barrier,
130128
only_active_cells)
131129

@@ -155,17 +153,15 @@ function calculate_free_surface_tendency!(grid, model, dependencies)
155153

156154
arch = architecture(grid)
157155

158-
args = tuple(model.velocities,
159-
model.free_surface,
160-
model.tracers,
161-
model.auxiliary_fields,
162-
model.forcing,
163-
model.clock)
164-
165156
Gη_event = launch!(arch, grid, :xy,
166157
calculate_hydrostatic_free_surface_Gη!, model.timestepper.Gⁿ.η,
167158
grid,
168-
args;
159+
model.velocities,
160+
model.free_surface,
161+
model.tracers,
162+
model.auxiliary_fields,
163+
model.forcing,
164+
model.clock;
169165
dependencies = dependencies)
170166

171167
return Gη_event
@@ -181,7 +177,8 @@ function calculate_hydrostatic_momentum_tendencies!(model, velocities; dependenc
181177
u_immersed_bc = immersed_boundary_condition(velocities.u)
182178
v_immersed_bc = immersed_boundary_condition(velocities.v)
183179

184-
start_momentum_kernel_args = (model.advection.momentum,
180+
start_momentum_kernel_args = (grid,
181+
model.advection.momentum,
185182
model.coriolis,
186183
model.closure)
187184

@@ -201,11 +198,11 @@ function calculate_hydrostatic_momentum_tendencies!(model, velocities; dependenc
201198
only_active_cells = use_only_active_cells(grid)
202199

203200
Gu_event = launch!(arch, grid, :xyz,
204-
calculate_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid, u_kernel_args;
201+
calculate_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, u_kernel_args...;
205202
dependencies = dependencies, only_active_cells)
206203

207204
Gv_event = launch!(arch, grid, :xyz,
208-
calculate_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid, v_kernel_args;
205+
calculate_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, v_kernel_args...;
209206
dependencies = dependencies, only_active_cells)
210207

211208
Gη_event = calculate_free_surface_tendency!(grid, model, dependencies)
@@ -246,24 +243,24 @@ end
246243
#####
247244

248245
""" Calculate the right-hand-side of the u-velocity equation. """
249-
@kernel function calculate_hydrostatic_free_surface_Gu!(Gu, grid, args)
246+
@kernel function calculate_hydrostatic_free_surface_Gu!(Gu, grid, args...)
250247
i, j, k = @index(Global, NTuple)
251248
@inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...)
252249
end
253250

254-
@kernel function calculate_hydrostatic_free_surface_Gu!(Gu, grid::ActiveCellsIBG, args)
251+
@kernel function calculate_hydrostatic_free_surface_Gu!(Gu, grid::ActiveCellsIBG, args...)
255252
idx = @index(Global, Linear)
256253
i, j, k = active_linear_index_to_ntuple(idx, grid)
257254
@inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...)
258255
end
259256

260257
""" Calculate the right-hand-side of the v-velocity equation. """
261-
@kernel function calculate_hydrostatic_free_surface_Gv!(Gv, grid, args)
258+
@kernel function calculate_hydrostatic_free_surface_Gv!(Gv, grid, args...)
262259
i, j, k = @index(Global, NTuple)
263260
@inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...)
264261
end
265262

266-
@kernel function calculate_hydrostatic_free_surface_Gv!(Gv, grid::ActiveCellsIBG, args)
263+
@kernel function calculate_hydrostatic_free_surface_Gv!(Gv, grid::ActiveCellsIBG, args...)
267264
idx = @index(Global, Linear)
268265
i, j, k = active_linear_index_to_ntuple(idx, grid)
269266
@inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...)
@@ -274,12 +271,12 @@ end
274271
#####
275272

276273
""" Calculate the right-hand-side of the tracer advection-diffusion equation. """
277-
@kernel function calculate_hydrostatic_free_surface_Gc!(Gc, tendency_kernel_function, grid, args)
274+
@kernel function calculate_hydrostatic_free_surface_Gc!(Gc, tendency_kernel_function, grid, args...)
278275
i, j, k = @index(Global, NTuple)
279276
@inbounds Gc[i, j, k] = tendency_kernel_function(i, j, k, grid, args...)
280277
end
281278

282-
@kernel function calculate_hydrostatic_free_surface_Gc!(Gc, tendency_kernel_function, grid::ActiveCellsIBG, args)
279+
@kernel function calculate_hydrostatic_free_surface_Gc!(Gc, tendency_kernel_function, grid::ActiveCellsIBG, args...)
283280
idx = @index(Global, Linear)
284281
i, j, k = active_linear_index_to_ntuple(idx, grid)
285282
@inbounds Gc[i, j, k] = tendency_kernel_function(i, j, k, grid, args...)
@@ -290,7 +287,7 @@ end
290287
#####
291288

292289
""" Calculate the right-hand-side of the free surface displacement (``η``) equation. """
293-
@kernel function calculate_hydrostatic_free_surface_Gη!(Gη, grid, args)
290+
@kernel function calculate_hydrostatic_free_surface_Gη!(Gη, grid, args...)
294291
i, j = @index(Global, NTuple)
295292
@inbounds Gη[i, j, grid.Nz+1] = free_surface_tendency(i, j, grid, args...)
296293
end

‎src/Models/NonhydrostaticModels/calculate_nonhydrostatic_tendencies.jl

+23-25
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,8 @@ function calculate_interior_tendency_contributions!(model; dependencies = device
6666
v_immersed_bc = velocities.v.boundary_conditions.immersed
6767
w_immersed_bc = velocities.w.boundary_conditions.immersed
6868

69-
start_momentum_kernel_args = (advection,
69+
start_momentum_kernel_args = (grid,
70+
advection,
7071
coriolis,
7172
stokes_drift,
7273
closure)
@@ -85,15 +86,15 @@ function calculate_interior_tendency_contributions!(model; dependencies = device
8586
only_active_cells = use_only_active_cells(grid)
8687

8788
Gu_event = launch!(arch, grid, :xyz, calculate_Gu!,
88-
tendencies.u, grid, u_kernel_args;
89+
tendencies.u, u_kernel_args...;
8990
dependencies, only_active_cells)
9091

9192
Gv_event = launch!(arch, grid, :xyz, calculate_Gv!,
92-
tendencies.v, grid, v_kernel_args;
93+
tendencies.v, v_kernel_args...;
9394
dependencies, only_active_cells)
9495

9596
Gw_event = launch!(arch, grid, :xyz, calculate_Gw!,
96-
tendencies.w, grid, w_kernel_args;
97+
tendencies.w, w_kernel_args...;
9798
dependencies, only_active_cells)
9899

99100
events = [Gu_event, Gv_event, Gw_event]
@@ -107,15 +108,12 @@ function calculate_interior_tendency_contributions!(model; dependencies = device
107108
@inbounds c_immersed_bc = tracers[tracer_index].boundary_conditions.immersed
108109
@inbounds tracer_name = keys(tracers)[tracer_index]
109110

110-
args = tuple(Val(tracer_index), Val(tracer_name),
111-
start_tracer_kernel_args...,
112-
c_immersed_bc,
113-
end_tracer_kernel_args...,
114-
forcing, clock)
115-
116-
117-
Gc_event = launch!(arch, grid, :xyz, calculate_Gc!,
118-
c_tendency, grid, args;
111+
Gc_event = launch!(arch, grid, :xyz, calculate_Gc!,
112+
c_tendency, grid, Val(tracer_index), Val(tracer_name),
113+
start_tracer_kernel_args...,
114+
c_immersed_bc,
115+
end_tracer_kernel_args...,
116+
forcing, clock;
119117
dependencies, only_active_cells)
120118

121119
push!(events, Gc_event)
@@ -131,36 +129,36 @@ end
131129
#####
132130

133131
""" Calculate the right-hand-side of the u-velocity equation. """
134-
@kernel function calculate_Gu!(Gu, grid, args)
132+
@kernel function calculate_Gu!(Gu, args...)
135133
i, j, k = @index(Global, NTuple)
136-
@inbounds Gu[i, j, k] = u_velocity_tendency(i, j, k, grid, args...)
134+
@inbounds Gu[i, j, k] = u_velocity_tendency(i, j, k, args...)
137135
end
138136

139-
@kernel function calculate_Gu!(Gu, grid::ActiveCellsIBG, args)
137+
@kernel function calculate_Gu!(Gu, grid::ActiveCellsIBG, args...)
140138
idx = @index(Global, Linear)
141139
i, j, k = active_linear_index_to_ntuple(idx, grid)
142140
@inbounds Gu[i, j, k] = u_velocity_tendency(i, j, k, grid, args...)
143141
end
144142

145143
""" Calculate the right-hand-side of the v-velocity equation. """
146-
@kernel function calculate_Gv!(Gv, grid, args)
144+
@kernel function calculate_Gv!(Gv, args...)
147145
i, j, k = @index(Global, NTuple)
148-
@inbounds Gv[i, j, k] = v_velocity_tendency(i, j, k, grid, args...)
146+
@inbounds Gv[i, j, k] = v_velocity_tendency(i, j, k, args...)
149147
end
150148

151-
@kernel function calculate_Gv!(Gv, grid::ActiveCellsIBG, args)
149+
@kernel function calculate_Gv!(Gv, grid::ActiveCellsIBG, args...)
152150
idx = @index(Global, Linear)
153151
i, j, k = active_linear_index_to_ntuple(idx, grid)
154152
@inbounds Gv[i, j, k] = v_velocity_tendency(i, j, k, grid, args...)
155153
end
156154

157155
""" Calculate the right-hand-side of the w-velocity equation. """
158-
@kernel function calculate_Gw!(Gw, grid, args)
156+
@kernel function calculate_Gw!(Gw, args...)
159157
i, j, k = @index(Global, NTuple)
160-
@inbounds Gw[i, j, k] = w_velocity_tendency(i, j, k, grid, args...)
158+
@inbounds Gw[i, j, k] = w_velocity_tendency(i, j, k, args...)
161159
end
162160

163-
@kernel function calculate_Gw!(Gw, grid::ActiveCellsIBG, args)
161+
@kernel function calculate_Gw!(Gw, grid::ActiveCellsIBG, args...)
164162
idx = @index(Global, Linear)
165163
i, j, k = active_linear_index_to_ntuple(idx, grid)
166164
@inbounds Gw[i, j, k] = w_velocity_tendency(i, j, k, grid, args...)
@@ -171,12 +169,12 @@ end
171169
#####
172170

173171
""" Calculate the right-hand-side of the tracer advection-diffusion equation. """
174-
@kernel function calculate_Gc!(Gc, grid, args)
172+
@kernel function calculate_Gc!(Gc, args...)
175173
i, j, k = @index(Global, NTuple)
176-
@inbounds Gc[i, j, k] = tracer_tendency(i, j, k, grid, args...)
174+
@inbounds Gc[i, j, k] = tracer_tendency(i, j, k, args...)
177175
end
178176

179-
@kernel function calculate_Gc!(Gc, grid::ActiveCellsIBG, args)
177+
@kernel function calculate_Gc!(Gc, grid::ActiveCellsIBG, args...)
180178
idx = @index(Global, Linear)
181179
i, j, k = active_linear_index_to_ntuple(idx, grid)
182180
@inbounds Gc[i, j, k] = tracer_tendency(i, j, k, grid, args...)

‎src/Models/ShallowWaterModels/shallow_water_diffusion_operators.jl

+3-5
Original file line numberDiff line numberDiff line change
@@ -55,10 +55,8 @@ Adapt.adapt_structure(to, closure::ShallowWaterScalarDiffusivity) =
5555
# The diffusivity for the shallow water model is calculated as h*ν in order to have a viscous term in the form
5656
# h⁻¹ ∇ ⋅ (hν t) where t is the 2D stress tensor plus a trace => t = ∇u + (∇u)ᵀ - ξI⋅(∇⋅u)
5757

58-
@kernel function _calculate_shallow_water_viscosity!(νₑ, grid, closure, clock, fields)
59-
i, j, k = @index(Global, NTuple)
60-
νₑ[i, j, k] = fields.h[i, j, k] * νᶜᶜᶜ(i, j, k, grid, viscosity_location(closure), closure.ν, clock, fields)
61-
end
58+
@inline calc_nonlinear_νᶜᶜᶜ(i, j, k, grid, closure::ShallowWaterScalarDiffusivity, clock, fields) =
59+
fields.h[i, j, k] * νᶜᶜᶜ(i, j, k, grid, viscosity_location(closure), closure.ν, clock, fields)
6260

6361
function calculate_diffusivities!(diffusivity_fields, closure::ShallowWaterScalarDiffusivity, model)
6462

@@ -69,7 +67,7 @@ function calculate_diffusivities!(diffusivity_fields, closure::ShallowWaterScala
6967
model_fields = shallow_water_fields(model.velocities, model.tracers, model.solution, formulation(model))
7068

7169
event = launch!(arch, grid, :xyz,
72-
_calculate_shallow_water_viscosity!,
70+
calculate_nonlinear_viscosity!,
7371
diffusivity_fields.νₑ, grid, closure, clock, model_fields,
7472
dependencies = device_event(arch))
7573

‎src/Operators/derivative_operators.jl

+13-13
Original file line numberDiff line numberDiff line change
@@ -60,21 +60,21 @@ for dir in (:x, :y, :z), L1 in (:ᶜ, :ᶠ), L2 in (:ᶜ, :ᶠ)
6060
fourth_order_center = Symbol(:∂⁴, dir, insert_symbol(dir, :ᶜ, L1, L2)...)
6161

6262
@eval begin
63-
@inline $second_order_face(i, j, k, grid, c) = $first_order_face(i, j, k, grid, $first_order_center, c)
64-
@inline $third_order_face(i, j, k, grid, c) = $first_order_face(i, j, k, grid, $second_order_center, c)
65-
@inline $fourth_order_face(i, j, k, grid, c) = $second_order_face(i, j, k, grid, $second_order_face, c)
63+
$second_order_face(i, j, k, grid, c) = $first_order_face(i, j, k, grid, $first_order_center, c)
64+
$third_order_face(i, j, k, grid, c) = $first_order_face(i, j, k, grid, $second_order_center, c)
65+
$fourth_order_face(i, j, k, grid, c) = $second_order_face(i, j, k, grid, $second_order_face, c)
6666

67-
@inline $second_order_center(i, j, k, grid, c) = $first_order_center(i, j, k, grid, $first_order_face, c)
68-
@inline $third_order_center(i, j, k, grid, c) = $first_order_center(i, j, k, grid, $second_order_face, c)
69-
@inline $fourth_order_center(i, j, k, grid, c) = $second_order_center(i, j, k, grid, $second_order_center, c)
67+
$second_order_center(i, j, k, grid, c) = $first_order_center(i, j, k, grid, $first_order_face, c)
68+
$third_order_center(i, j, k, grid, c) = $first_order_center(i, j, k, grid, $second_order_face, c)
69+
$fourth_order_center(i, j, k, grid, c) = $second_order_center(i, j, k, grid, $second_order_center, c)
7070

71-
@inline $second_order_face(i, j, k, grid, f::Function, args...) = $first_order_face(i, j, k, grid, $first_order_center, f::Function, args...)
72-
@inline $third_order_face(i, j, k, grid, f::Function, args...) = $first_order_face(i, j, k, grid, $second_order_center, f::Function, args...)
73-
@inline $fourth_order_face(i, j, k, grid, f::Function, args...) = $second_order_face(i, j, k, grid, $second_order_face, f::Function, args...)
71+
$second_order_face(i, j, k, grid, f::Function, args...) = $first_order_face(i, j, k, grid, $first_order_center, f::Function, args...)
72+
$third_order_face(i, j, k, grid, f::Function, args...) = $first_order_face(i, j, k, grid, $second_order_center, f::Function, args...)
73+
$fourth_order_face(i, j, k, grid, f::Function, args...) = $second_order_face(i, j, k, grid, $second_order_face, f::Function, args...)
7474

75-
@inline $second_order_center(i, j, k, grid, f::Function, args...) = $first_order_center(i, j, k, grid, $first_order_face, f::Function, args...)
76-
@inline $third_order_center(i, j, k, grid, f::Function, args...) = $first_order_center(i, j, k, grid, $second_order_face, f::Function, args...)
77-
@inline $fourth_order_center(i, j, k, grid, f::Function, args...) = $second_order_center(i, j, k, grid, $second_order_center, f::Function, args...)
75+
$second_order_center(i, j, k, grid, f::Function, args...) = $first_order_center(i, j, k, grid, $first_order_face, f::Function, args...)
76+
$third_order_center(i, j, k, grid, f::Function, args...) = $first_order_center(i, j, k, grid, $second_order_face, f::Function, args...)
77+
$fourth_order_center(i, j, k, grid, f::Function, args...) = $second_order_center(i, j, k, grid, $second_order_center, f::Function, args...)
7878
end
7979
end
8080

@@ -89,6 +89,6 @@ for dir in (:x, :y, :z), LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ), LZ in (:ᶜ, :
8989
derivative = Symbol(:∂, dir, LX, LY, LZ)
9090

9191
@eval begin
92-
@inline $operator(i, j, k, grid, c) = $area(i, j, k, grid) * $derivative(i, j, k, grid, c)
92+
$operator(i, j, k, grid, c) = $area(i, j, k, grid) * $derivative(i, j, k, grid, c)
9393
end
9494
end

‎src/Solvers/batched_tridiagonal_solver.jl

+2-2
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ function solve!(ϕ, solver::BatchedTridiagonalSolver, rhs, args...; dependencies
7777
grid = solver.grid
7878

7979
event = launch!(architecture(solver), grid, :xy,
80-
solve_batched_tridiagonal_system_kernel!, ϕ, a, b, c, rhs, t, grid, parameters, Tuple(args);
80+
solve_batched_tridiagonal_system_kernel!, ϕ, a, b, c, rhs, t, grid, parameters, args...,
8181
dependencies = dependencies)
8282

8383
wait(device(architecture(solver)), event)
@@ -88,7 +88,7 @@ end
8888
@inline float_eltype(ϕ::AbstractArray{T}) where T <: AbstractFloat = T
8989
@inline float_eltype(ϕ::AbstractArray{<:Complex{T}}) where T <: AbstractFloat = T
9090

91-
@kernel function solve_batched_tridiagonal_system_kernel!(ϕ, a, b, c, f, t, grid, p, args)
91+
@kernel function solve_batched_tridiagonal_system_kernel!(ϕ, a, b, c, f, t, grid, p, args...)
9292
Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz
9393

9494
i, j = @index(Global, NTuple)

‎src/TimeSteppers/quasi_adams_bashforth_2.jl

+1-1
Original file line numberDiff line numberDiff line change
@@ -165,4 +165,4 @@ Time step velocity fields via the 2nd-order quasi Adams-Bashforth method
165165
@inbounds u[i, j, k] += Δt * ((one_point_five + χ) * Gⁿ[i, j, k] - (oh_point_five + χ) * G⁻[i, j, k])
166166
end
167167

168-
@kernel ab2_step_field!(::FunctionField, Δt, χ, Gⁿ, G⁻) = nothing
168+
@kernel ab2_step_field!(::FunctionField, args...) = nothing

‎src/TurbulenceClosures/turbulence_closure_implementations/CATKEVerticalDiffusivities/CATKEVerticalDiffusivities.jl

+5-5
Original file line numberDiff line numberDiff line change
@@ -246,7 +246,7 @@ function calculate_diffusivities!(diffusivities, closure::FlavorOfCATKE, model)
246246
return nothing
247247
end
248248

249-
@kernel function calculate_CATKE_diffusivities!(diffusivities, grid, closure::FlavorOfCATKE, velocities, tracers, buoyancy, clock, top_tracer_bcs)
249+
@kernel function calculate_CATKE_diffusivities!(diffusivities, grid, closure::FlavorOfCATKE, velocities, tracers, buoyancy, args...)
250250
i, j, k, = @index(Global, NTuple)
251251

252252
# Ensure this works with "ensembles" of closures, in addition to ordinary single closures
@@ -255,9 +255,9 @@ end
255255
max_K = closure_ij.maximum_diffusivity
256256

257257
@inbounds begin
258-
diffusivities.Kᵘ[i, j, k] = min(max_K, Kuᶜᶜᶠ(i, j, k, grid, closure_ij, velocities, tracers, buoyancy, clock, top_tracer_bcs))
259-
diffusivities.Kᶜ[i, j, k] = min(max_K, Kcᶜᶜᶠ(i, j, k, grid, closure_ij, velocities, tracers, buoyancy, clock, top_tracer_bcs))
260-
diffusivities.Kᵉ[i, j, k] = min(max_K, Keᶜᶜᶠ(i, j, k, grid, closure_ij, velocities, tracers, buoyancy, clock, top_tracer_bcs))
258+
diffusivities.Kᵘ[i, j, k] = min(max_K, Kuᶜᶜᶠ(i, j, k, grid, closure_ij, velocities, tracers, buoyancy, args...))
259+
diffusivities.Kᶜ[i, j, k] = min(max_K, Kcᶜᶜᶠ(i, j, k, grid, closure_ij, velocities, tracers, buoyancy, args...))
260+
diffusivities.Kᵉ[i, j, k] = min(max_K, Keᶜᶜᶠ(i, j, k, grid, closure_ij, velocities, tracers, buoyancy, args...))
261261

262262
# "Patankar trick" for buoyancy production (cf Patankar 1980 or Burchard et al. 2003)
263263
# If buoyancy flux is a _sink_ of TKE, we treat it implicitly.
@@ -268,7 +268,7 @@ end
268268
dissipative_buoyancy_flux = sign(wb) * sign(eⁱʲᵏ) < 0
269269
wb_e = ifelse(dissipative_buoyancy_flux, wb / eⁱʲᵏ, zero(grid))
270270

271-
diffusivities.Lᵉ[i, j, k] = - wb_e + implicit_dissipation_coefficient(i, j, k, grid, closure_ij, velocities, tracers, buoyancy, clock, top_tracer_bcs)
271+
diffusivities.Lᵉ[i, j, k] = - wb_e + implicit_dissipation_coefficient(i, j, k, grid, closure_ij, velocities, tracers, buoyancy, args...)
272272
end
273273
end
274274

‎src/TurbulenceClosures/turbulence_closure_implementations/anisotropic_minimum_dissipation.jl

+1-1
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,7 @@ function calculate_diffusivities!(diffusivity_fields, closure::AnisotropicMinimu
190190
diffusivity_kernel! = calculate_nonlinear_tracer_diffusivity!(device(arch), workgroup, worksize)
191191

192192
barrier = device_event(arch)
193-
viscosity_event = viscosity_kernel!(diffusivity_fields.νₑ, grid, closure, buoyancy, velocities, tracers; dependencies=barrier)
193+
viscosity_event = viscosity_kernel!(diffusivity_fields.νₑ, grid, closure, buoyancy, velocities, tracers, dependencies=barrier)
194194

195195
events = [viscosity_event]
196196

‎src/TurbulenceClosures/turbulence_closure_utils.jl

+2-2
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,9 @@ end
2929
@inline geo_mean_Δᶠ(i, j, k, grid::AbstractGrid) =
3030
cbrt(Δxᶜᶜᶜ(i, j, k, grid) * Δyᶜᶜᶜ(i, j, k, grid) * Δzᶜᶜᶜ(i, j, k, grid))
3131

32-
@kernel function calculate_nonlinear_viscosity!(νₑ, grid, closure, buoyancy, velocities, tracers)
32+
@kernel function calculate_nonlinear_viscosity!(νₑ, grid, closure, args...)
3333
i, j, k = @index(Global, NTuple)
34-
@inbounds νₑ[i, j, k] = calc_nonlinear_νᶜᶜᶜ(i, j, k, grid, closure, buoyancy, velocities, tracers)
34+
@inbounds νₑ[i, j, k] = calc_nonlinear_νᶜᶜᶜ(i, j, k, grid, closure, args...)
3535
end
3636

3737
@kernel function calculate_nonlinear_tracer_diffusivity!(κₑ, grid, closure, tracer, tracer_index, U)

0 commit comments

Comments
 (0)
Please sign in to comment.