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

0 commit comments

Comments
 (0)
Please sign in to comment.