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

support for active_cells_map in kernels #3920

Draft
wants to merge 19 commits into
base: main
Choose a base branch
from
Draft

support for active_cells_map in kernels #3920

Show file tree
Hide file tree
Changes from 12 commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
279e4c4
support for MappedFunctions
simone-silvestri Nov 12, 2024
6182e34
remove double kernels everywhere
simone-silvestri Nov 12, 2024
5545605
make sure also GPU works
simone-silvestri Nov 12, 2024
e239dbe
make sure there is no return in the function
simone-silvestri Nov 12, 2024
32e0c52
return nothing
simone-silvestri Nov 12, 2024
0b56304
adapt for GPU usage
simone-silvestri Nov 12, 2024
052f56e
not sure about the GPU
simone-silvestri Nov 12, 2024
d0d56ac
hmmm
simone-silvestri Nov 12, 2024
7e41737
probably like this it will work on the GPU?
simone-silvestri Nov 12, 2024
42d4989
Merge branch 'ss/active-index-macro' of github.com:CliMA/Oceananigans…
simone-silvestri Nov 12, 2024
130a357
some cleanup plus testing a test?
simone-silvestri Nov 12, 2024
3be13dd
works on gpu, to fix the dynamic check
simone-silvestri Nov 12, 2024
64cad18
move lwargs
simone-silvestri Nov 13, 2024
333bfd7
kwarg management
simone-silvestri Nov 13, 2024
e069051
reduce time for tridiagonal solve
simone-silvestri Nov 13, 2024
f0a887a
see if tests pass (especially distributed where we have active_cells)
simone-silvestri Nov 13, 2024
7b33e7a
bugfix
simone-silvestri Nov 13, 2024
26194f3
grid from solver
simone-silvestri Nov 13, 2024
db84ef2
Merge branch 'main' into ss/active-index-macro
simone-silvestri Nov 14, 2024
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
Empty file added kernel_maps.jl
View file
Open in desktop
Empty file.
6 changes: 5 additions & 1 deletion src/BoundaryConditions/fill_halo_regions.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,11 @@ const MaybeTupledData = Union{OffsetArray, NTuple{<:Any, OffsetArray}}

"Fill halo regions in ``x``, ``y``, and ``z`` for a given field's data."
function fill_halo_regions!(c::MaybeTupledData, boundary_conditions, indices, loc, grid, args...;
fill_boundary_normal_velocities = true, kwargs...)
fill_boundary_normal_velocities = true,
only_local_halos = false, # Only valid for `DistributedGrids`, we throw it away here
async = false, # Only valid for `DistributedGrids`, we throw it away here
kwargs...)

arch = architecture(grid)

if fill_boundary_normal_velocities
Expand Down
29 changes: 5 additions & 24 deletions src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,6 @@ function compute_hydrostatic_free_surface_tendency_contributions!(model, kernel_
compute_hydrostatic_free_surface_Gc!,
c_tendency,
grid,
active_cells_map,
args;
active_cells_map)
end
Expand Down Expand Up @@ -161,12 +160,12 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para

launch!(arch, grid, kernel_parameters,
compute_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid,
active_cells_map, u_kernel_args;
u_kernel_args;
active_cells_map)

launch!(arch, grid, kernel_parameters,
compute_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid,
active_cells_map, v_kernel_args;
v_kernel_args;
active_cells_map)

compute_free_surface_tendency!(grid, model, :xy)
Expand Down Expand Up @@ -200,45 +199,27 @@ end
#####

""" Calculate the right-hand-side of the u-velocity equation. """
@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, ::Nothing, args)
@kernel function compute_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 compute_hydrostatic_free_surface_Gu!(Gu, grid, active_cells_map, args)
idx = @index(Global, Linear)
i, j, k = active_linear_index_to_tuple(idx, active_cells_map)
@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 compute_hydrostatic_free_surface_Gv!(Gv, grid, ::Nothing, args)
@kernel function compute_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 compute_hydrostatic_free_surface_Gv!(Gv, grid, active_cells_map, args)
idx = @index(Global, Linear)
i, j, k = active_linear_index_to_tuple(idx, active_cells_map)
@inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...)
end

#####
##### Tendency calculators for tracers
#####

""" Calculate the right-hand-side of the tracer advection-diffusion equation. """
@kernel function compute_hydrostatic_free_surface_Gc!(Gc, grid, ::Nothing, args)
@kernel function compute_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 compute_hydrostatic_free_surface_Gc!(Gc, grid, active_cells_map, args)
idx = @index(Global, Linear)
i, j, k = active_linear_index_to_tuple(idx, active_cells_map)
@inbounds Gc[i, j, k] = hydrostatic_free_surface_tracer_tendency(i, j, k, grid, args...)
end

#####
##### Tendency calculators for an explicit free surface
#####
Expand Down
1 change: 0 additions & 1 deletion src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
using Oceananigans.Fields: location
using Oceananigans.TimeSteppers: ab2_step_field!
using Oceananigans.TurbulenceClosures: implicit_step!
using Oceananigans.ImmersedBoundaries: retrieve_interior_active_cells_map, retrieve_surface_active_cells_map

import Oceananigans.TimeSteppers: ab2_step!

Expand Down
21 changes: 2 additions & 19 deletions src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,7 @@ end

# Barotropic Model Kernels
# u_Δz = u * Δz
@kernel function _barotropic_mode_kernel!(U, V, grid, ::Nothing, u, v)
@kernel function _barotropic_mode_kernel!(U, V, grid, u, v)
i, j = @index(Global, NTuple)
k_top = grid.Nz+1

Expand All @@ -146,26 +146,9 @@ end
end
end

# Barotropic Model Kernels
# u_Δz = u * Δz
@kernel function _barotropic_mode_kernel!(U, V, grid, active_cells_map, u, v)
idx = @index(Global, Linear)
i, j = active_linear_index_to_tuple(idx, active_cells_map)
k_top = grid.Nz+1

@inbounds U[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * u[i, j, 1]
@inbounds V[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * v[i, j, 1]

for k in 2:grid.Nz
@inbounds U[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * u[i, j, k]
@inbounds V[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * v[i, j, k]
end
end

@inline function compute_barotropic_mode!(U, V, grid, u, v)
active_cells_map = retrieve_surface_active_cells_map(grid)

launch!(architecture(grid), grid, :xy, _barotropic_mode_kernel!, U, V, grid, active_cells_map, u, v; active_cells_map)
launch!(architecture(grid), grid, :xy, _barotropic_mode_kernel!, U, V, grid, u, v; active_cells_map)

return nothing
end
Expand Down
1 change: 0 additions & 1 deletion src/Models/HydrostaticFreeSurfaceModels/store_hydrostatic_free_surface_tendencies.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ using Oceananigans.TimeSteppers: store_field_tendencies!

using Oceananigans: prognostic_fields
using Oceananigans.Grids: AbstractGrid
using Oceananigans.ImmersedBoundaries: retrieve_interior_active_cells_map

using Oceananigans.Utils: launch!

Expand Down
39 changes: 8 additions & 31 deletions src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,15 +103,15 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti

exclude_periphery = true
launch!(arch, grid, kernel_parameters, compute_Gu!,
tendencies.u, grid, active_cells_map, u_kernel_args;
tendencies.u, grid, u_kernel_args;
active_cells_map, exclude_periphery)

launch!(arch, grid, kernel_parameters, compute_Gv!,
tendencies.v, grid, active_cells_map, v_kernel_args;
tendencies.v, grid, v_kernel_args;
active_cells_map, exclude_periphery)

launch!(arch, grid, kernel_parameters, compute_Gw!,
tendencies.w, grid, active_cells_map, w_kernel_args;
tendencies.w, grid, w_kernel_args;
active_cells_map, exclude_periphery)

start_tracer_kernel_args = (advection, closure)
Expand All @@ -131,7 +131,7 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti
forcing, clock)

launch!(arch, grid, kernel_parameters, compute_Gc!,
c_tendency, grid, active_cells_map, args;
c_tendency, grid, args;
active_cells_map)
end

Expand All @@ -143,57 +143,34 @@ end
#####

""" Calculate the right-hand-side of the u-velocity equation. """
@kernel function compute_Gu!(Gu, grid, ::Nothing, args)
@kernel function compute_Gu!(Gu, grid, args)
i, j, k = @index(Global, NTuple)
@inbounds Gu[i, j, k] = u_velocity_tendency(i, j, k, grid, args...)
end

@kernel function compute_Gu!(Gu, grid, interior_map, args)
idx = @index(Global, Linear)
i, j, k = active_linear_index_to_tuple(idx, interior_map)
@inbounds Gu[i, j, k] = u_velocity_tendency(i, j, k, grid, args...)
end

""" Calculate the right-hand-side of the v-velocity equation. """
@kernel function compute_Gv!(Gv, grid, ::Nothing, args)
@kernel function compute_Gv!(Gv, grid, args)
i, j, k = @index(Global, NTuple)
@inbounds Gv[i, j, k] = v_velocity_tendency(i, j, k, grid, args...)
end

@kernel function compute_Gv!(Gv, grid, interior_map, args)
idx = @index(Global, Linear)
i, j, k = active_linear_index_to_tuple(idx, interior_map)
@inbounds Gv[i, j, k] = v_velocity_tendency(i, j, k, grid, args...)
end

""" Calculate the right-hand-side of the w-velocity equation. """
@kernel function compute_Gw!(Gw, grid, ::Nothing, args)
@kernel function compute_Gw!(Gw, grid, args)
i, j, k = @index(Global, NTuple)
@inbounds Gw[i, j, k] = w_velocity_tendency(i, j, k, grid, args...)
end

@kernel function compute_Gw!(Gw, grid, interior_map, args)
idx = @index(Global, Linear)
i, j, k = active_linear_index_to_tuple(idx, interior_map)
@inbounds Gw[i, j, k] = w_velocity_tendency(i, j, k, grid, args...)
end

#####
##### Tracer(s)
#####

""" Calculate the right-hand-side of the tracer advection-diffusion equation. """
@kernel function compute_Gc!(Gc, grid, ::Nothing, args)
@kernel function compute_Gc!(Gc, grid, args)
i, j, k = @index(Global, NTuple)
@inbounds Gc[i, j, k] = tracer_tendency(i, j, k, grid, args...)
end

@kernel function compute_Gc!(Gc, grid, interior_map, args)
idx = @index(Global, Linear)
i, j, k = active_linear_index_to_tuple(idx, interior_map)
@inbounds Gc[i, j, k] = tracer_tendency(i, j, k, grid, args...)
end

#####
##### Boundary contributions to tendencies due to user-prescribed fluxes
#####
Expand Down
Loading