Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fixes CPU slowdown with KernelAbstractions >= v0.8 #3026

Merged
merged 7 commits into from
Mar 24, 2023
Merged

Fixes CPU slowdown with KernelAbstractions >= v0.8 #3026

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
9fc0bc6
fixes
simone-silvestri Mar 23, 2023
33cec55
new way
simone-silvestri Mar 23, 2023
f4d40f5
Merge branch 'main' into ss/fixes_performance_issues
glwagner Mar 23, 2023
c641bc5
Add inline to derivative ops
glwagner Mar 23, 2023
87f6ec5
Add missing grid to tracer-tendency
glwagner Mar 23, 2023
5f21327
bugfixes
simone-silvestri Mar 23, 2023
3bcb433
Merge branch 'ss/fixes_performance_issues' of github.com:CliMA/Oceana…
simone-silvestri Mar 23, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions src/BoundaryConditions/apply_flux_bcs.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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, args..., dependencies=dep)
launch!(arch, grid, :yz, _apply_x_bcs!, Gc, instantiated_location(Gc), grid, west_bc, east_bc, Tuple(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, args..., dependencies=dep)
launch!(arch, grid, :xz, _apply_y_bcs!, Gc, instantiated_location(Gc), grid, south_bc, north_bc, Tuple(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, args..., dependencies=dep)
launch!(arch, grid, :xy, _apply_z_bcs!, Gc, instantiated_location(Gc), grid, bottom_bc, top_bc, Tuple(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...)
Expand All @@ -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...)
Expand All @@ -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...)
Expand Down
18 changes: 9 additions & 9 deletions src/BoundaryConditions/fill_halo_regions.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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, args...; dependencies=dep, kwargs...)
launch!(arch, grid, size, _fill_west_and_east_halo!, c, west_bc, east_bc, offset, loc, grid, Tuple(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, args...; dependencies=dep, kwargs...)
launch!(arch, grid, size, _fill_south_and_north_halo!, c, south_bc, north_bc, offset, loc, grid, Tuple(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, args...; dependencies=dep, kwargs...)
launch!(arch, grid, size, _fill_bottom_and_top_halo!, c, bottom_bc, top_bc, offset, loc, grid, Tuple(args); dependencies=dep, kwargs...)

#####
##### Calculate kernel size and offset for Windowed and Sliced Fields
Expand Down
18 changes: 9 additions & 9 deletions src/BoundaryConditions/fill_halo_regions_open.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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, 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!(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!(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...)
Expand Down
67 changes: 35 additions & 32 deletions src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_tendencies.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,26 +104,28 @@ 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,
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;
args;
dependencies = barrier,
only_active_cells)

Expand Down Expand Up @@ -153,15 +155,17 @@ 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,
model.velocities,
model.free_surface,
model.tracers,
model.auxiliary_fields,
model.forcing,
model.clock;
args;
dependencies = dependencies)

return Gη_event
Expand All @@ -177,8 +181,7 @@ 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 = (grid,
model.advection.momentum,
start_momentum_kernel_args = (model.advection.momentum,
model.coriolis,
model.closure)

Expand All @@ -198,11 +201,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, u_kernel_args...;
calculate_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid, u_kernel_args;
dependencies = dependencies, only_active_cells)

Gv_event = launch!(arch, grid, :xyz,
calculate_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, v_kernel_args...;
calculate_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid, v_kernel_args;
dependencies = dependencies, only_active_cells)

Gη_event = calculate_free_surface_tendency!(grid, model, dependencies)
Expand Down Expand Up @@ -243,24 +246,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...)
Expand All @@ -271,12 +274,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...)
Expand All @@ -287,7 +290,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
Expand Down
Loading