Skip to content

Commit

Permalink
Fix performance on CPU after upgrade to KernelAbstractions 0.8; issue #…
Browse files Browse the repository at this point in the history 
…2996 (#3030)

* test it out

* grammar

* maybe better solution?

* comment

* fixes implicit solver

* comment

* quick bugfix

* go back to 1 closure/diffusivity

* bugfix

---------

Co-authored-by: Navid C. Constantinou <navidcy@users.noreply.github.com>
  • Loading branch information
@simone-silvestri @navidcy
simone-silvestri and navidcy committed Mar 28, 2023
1 parent f09234d commit f4bbce1
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 24 deletions.
47 changes: 26 additions & 21 deletions src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_tendencies.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,34 +46,31 @@ end
using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: FlavorOfCATKE
using Oceananigans.TurbulenceClosures.MEWSVerticalDiffusivities: MEWS

const HFSM = HydrostaticFreeSurfaceModel

# Fallback
@inline tracer_tendency_kernel_function(model::HFSM, name, c, K) = hydrostatic_free_surface_tracer_tendency, c, K
@inline tracer_tendency_kernel_function(model::HFSM, ::Val{:K}, c::MEWS, K) = hydrostatic_turbulent_kinetic_energy_tendency, c, K
@inline tracer_tendency_kernel_function(model::HFSM, ::Val{:e}, c::FlavorOfCATKE, K) = hydrostatic_turbulent_kinetic_energy_tendency, c, K
@inline tracer_tendency_kernel_function(model::HFSM, name, c, K) = calculate_hydrostatic_free_surface_Gc!, c, K
@inline tracer_tendency_kernel_function(model::HFSM, ::Val{:K}, c::MEWS, K) = calculate_hydrostatic_free_surface_Ge!, c, K
@inline tracer_tendency_kernel_function(model::HFSM, ::Val{:e}, c::FlavorOfCATKE, K) = calculate_hydrostatic_free_surface_Ge!, c, K

function tracer_tendency_kernel_function(model::HFSM, ::Val{:e}, closures::Tuple, diffusivity_fields::Tuple)
catke_index = findfirst(c -> c isa FlavorOfCATKE, closures)

if isnothing(catke_index)
return hydrostatic_free_surface_tracer_tendency, closures, diffusivity_fields
return calculate_hydrostatic_free_surface_Gc!, closures, diffusivity_fields
else
catke_closure = closures[catke_index]
catke_diffusivity_fields = diffusivity_fields[catke_index]
return hydrostatic_turbulent_kinetic_energy_tendency, catke_closure, catke_diffusivity_fields
return calculate_hydrostatic_free_surface_Ge!, catke_closure, catke_diffusivity_fields
end
end

function tracer_tendency_kernel_function(model::HFSM, ::Val{:K}, closures::Tuple, diffusivity_fields::Tuple)
mews_index = findfirst(c -> c isa MEWS, closures)

if isnothing(mews_index)
return hydrostatic_free_surface_tracer_tendency, closures, diffusivity_fields
return calculate_hydrostatic_free_surface_Gc!, closures, diffusivity_fields
else
mews_closure = closures[mews_index]
mews_diffusivity_fields = diffusivity_fields[mews_index]
return hydrostatic_turbulent_kinetic_energy_tendency, mews_closure, mews_diffusivity_fields
return calculate_hydrostatic_free_surface_Ge!, mews_closure, mews_diffusivity_fields
end
end

Expand All @@ -99,10 +96,7 @@ function calculate_hydrostatic_free_surface_interior_tendency_contributions!(mod
c_forcing = model.forcing[tracer_name]
c_immersed_bc = immersed_boundary_condition(model.tracers[tracer_name])

c_kernel_function, closure, diffusivity_fields = tracer_tendency_kernel_function(model,
Val(tracer_name),
model.closure,
model.diffusivity_fields)
tendency_kernel!, closure, diffusivity = tracer_tendency_kernel_function(model, Val(tracer_name), model.closure, model.diffusivity_fields)

args = tuple(Val(tracer_index),
Val(tracer_name),
Expand All @@ -115,15 +109,14 @@ function calculate_hydrostatic_free_surface_interior_tendency_contributions!(mod
model.free_surface,
model.tracers,
top_tracer_bcs,
diffusivity_fields,
diffusivity,
model.auxiliary_fields,
c_forcing,
model.clock)

Gc_event = launch!(arch, grid, :xyz,
calculate_hydrostatic_free_surface_Gc!,
tendency_kernel!,
c_tendency,
c_kernel_function,
grid,
args;
dependencies = barrier,
Expand Down Expand Up @@ -274,15 +267,27 @@ 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, grid, 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 calculate_hydrostatic_free_surface_Gc!(Gc, grid::ActiveCellsIBG, args)
idx = @index(Global, Linear)
i, j, k = active_linear_index_to_ntuple(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 calculate_hydrostatic_free_surface_Ge!(Ge, grid, args)
i, j, k = @index(Global, NTuple)
@inbounds Gc[i, j, k] = tendency_kernel_function(i, j, k, grid, args...)
@inbounds Ge[i, j, k] = hydrostatic_turbulent_kinetic_energy_tendency(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_Ge!(Ge, 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...)
@inbounds Ge[i, j, k] = hydrostatic_turbulent_kinetic_energy_tendency(i, j, k, grid, args...)
end

#####
Expand Down
14 changes: 11 additions & 3 deletions src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@ using Oceananigans.Operators: Δzᵃᵃᶜ, Δzᵃᵃᶠ
using Oceananigans.AbstractOperations: flip
using Oceananigans.Solvers: BatchedTridiagonalSolver, solve!

import Oceananigans.Solvers: get_coefficient

#####
##### implicit_step! interface
#####
Expand Down Expand Up @@ -123,13 +125,19 @@ function implicit_diffusion_solver(::VerticallyImplicitTimeDiscretization, grid)
"grids that are Bounded in the z-direction.")

z_solver = BatchedTridiagonalSolver(grid;
lower_diagonal = maybe_tupled_ivd_lower_diagonal,
diagonal = ivd_diagonal,
upper_diagonal = maybe_tupled_ivd_upper_diagonal)
lower_diagonal = Val(:maybe_tupled_ivd_lower_diagonal),
diagonal = Val(:ivd_diagonal),
upper_diagonal = Val(:maybe_tupled_ivd_upper_diagonal))

return z_solver
end

# Extend the `get_coefficient` function to retrieve the correct `ivd_diagonal`, `ivd_lower_diagonal` and `ivd_upper_diagonal` functions
# REMEMBER: `get_coefficient(f::Function, args...)` leads to massive performance decrease on the CPU (https://github.com/CliMA/Oceananigans.jl/issues/2996)
@inline get_coefficient(::Val{:maybe_tupled_ivd_lower_diagonal}, i, j, k, grid, p, args...) = maybe_tupled_ivd_lower_diagonal(i, j, k, grid, args...)
@inline get_coefficient(::Val{:maybe_tupled_ivd_upper_diagonal}, i, j, k, grid, p, args...) = maybe_tupled_ivd_upper_diagonal(i, j, k, grid, args...)
@inline get_coefficient(::Val{:ivd_diagonal}, i, j, k, grid, p, args...) = ivd_diagonal(i, j, k, grid, args...)

#####
##### Implicit step functions
#####
Expand Down

2 comments on commit f4bbce1

@simone-silvestri
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request updated: JuliaRegistries/General/79587

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.80.0 -m "<description of version>" f4bbce10c3834ed505e1a46b3169a3edfcd609c9
git push origin v0.80.0

Please sign in to comment.