Skip to content

Commit

Permalink
Open halo filling on Windowed fields (#3834)
Browse files Browse the repository at this point in the history
* do this now

* fix open halo for windowed fields

* correct function that will eventually be removed

* some comment

* remove the launching of the kernels

* bugfix

* correct

* some reworking here

* fix a bit

* comment

* bugfix

* fix all tests

* some more comments plus remove not used code

---------

Co-authored-by: Gregory L. Wagner <[email protected]>
  • Loading branch information
simone-silvestri and glwagner authored Nov 12, 2024
1 parent 37d0e29 commit 9111a8f
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 97 deletions.
135 changes: 47 additions & 88 deletions src/BoundaryConditions/fill_halo_regions_open.jl
Original file line number Diff line number Diff line change
Expand Up @@ -6,104 +6,63 @@
Fill open boundary halo regions by filling boundary conditions on field faces with `open_fill`.
"""
function fill_open_boundary_regions!(field, boundary_conditions, indices, loc, grid, args...; kwargs...)
function fill_open_boundary_regions!(field, boundary_conditions, indices, loc, grid, args...; only_local_halos = false, kwargs...)
arch = architecture(grid)

left_bc = left_velocity_open_boundary_condition(boundary_conditions, loc)
right_bc = right_velocity_open_boundary_condition(boundary_conditions, loc)
# gets `fill_halo!`, the function which fills open boundaries at `loc`
# The underlying assumption is that open boundaries are uniquely defined by the location `loc`:
# (Face, Center, Center) -> fill west and east
# (Center, Face, Center) -> fill south and north
# (Center, Center, Face) -> fill bottom and top
fill_halo! = get_open_halo_filling_functions(loc)

# gets `open_fill`, the function which fills open boundaries at `loc`, as well as `regular_fill`
# which is the function which fills non-open boundaries at `loc` which informs `fill_halo_size`
open_fill, regular_fill = get_open_halo_filling_functions(loc)
fill_size = fill_halo_size(field, regular_fill, indices, boundary_conditions, loc, grid)
left_bc = left_open_boundary_condition(boundary_conditions, loc)
right_bc = right_open_boundary_condition(boundary_conditions, loc)

launch!(arch, grid, fill_size, open_fill, field, left_bc, right_bc, loc, grid, args)
bcs_tuple = (left_bc, right_bc)

return nothing
end

fill_open_boundary_regions!(fields::NTuple, boundary_conditions, indices, loc, grid, args...; kwargs...) =
[fill_open_boundary_regions!(field, boundary_conditions[n], indices, loc[n], grid, args...; kwargs...) for (n, field) in enumerate(fields)]

# for regular halo fills
@inline left_velocity_open_boundary_condition(boundary_condition, loc) = nothing
@inline left_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Face, Center, Center}) = boundary_conditions.west
@inline left_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Center, Face, Center}) = boundary_conditions.south
@inline left_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Center, Center, Face}) = boundary_conditions.bottom

@inline right_velocity_open_boundary_condition(boundary_conditions, loc) = nothing
@inline right_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Face, Center, Center}) = boundary_conditions.east
@inline right_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Center, Face, Center}) = boundary_conditions.north
@inline right_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Center, Center, Face}) = boundary_conditions.top
if !isnothing(fill_halo!) && any(!isnothing, bcs_tuple)

# for multi region halo fills
@inline left_velocity_open_boundary_condition(boundary_conditions::Tuple, ::Tuple{Face, Center, Center}) = @inbounds boundary_conditions[1]
@inline left_velocity_open_boundary_condition(boundary_conditions::Tuple, ::Tuple{Center, Face, Center}) = @inbounds boundary_conditions[1]
@inline left_velocity_open_boundary_condition(boundary_conditions::Tuple, ::Tuple{Center, Center, Face}) = @inbounds boundary_conditions[1]
# Overwrite the `only_local_halos` keyword argument, because open boundaries
# are always local boundaries that do not require communication
only_local_halos = true

@inline right_velocity_open_boundary_condition(boundary_conditions::Tuple, ::Tuple{Face, Center, Center}) = @inbounds boundary_conditions[2]
@inline right_velocity_open_boundary_condition(boundary_conditions::Tuple, ::Tuple{Center, Face, Center}) = @inbounds boundary_conditions[2]
@inline right_velocity_open_boundary_condition(boundary_conditions::Tuple, ::Tuple{Center, Center, Face}) = @inbounds boundary_conditions[2]
fill_halo_event!(field, fill_halo!, bcs_tuple, indices, loc, arch, grid, args...; only_local_halos, kwargs...)
end

@inline get_open_halo_filling_functions(loc) = _no_fill!, _no_fill!
@inline get_open_halo_filling_functions(::Tuple{Face, Center, Center}) = _fill_west_and_east_open_halo!, fill_west_and_east_halo!
@inline get_open_halo_filling_functions(::Tuple{Center, Face, Center}) = _fill_south_and_north_open_halo!, fill_south_and_north_halo!
@inline get_open_halo_filling_functions(::Tuple{Center, Center, Face}) = _fill_bottom_and_top_open_halo!, fill_bottom_and_top_halo!

@kernel _no_fill!(args...) = nothing

@inline fill_halo_size(field, ::typeof(_no_fill!), args...) = (0, 0)

@kernel function _fill_west_and_east_open_halo!(c, west_bc, east_bc, loc, grid, args)
j, k = @index(Global, NTuple)
_fill_west_open_halo!(j, k, grid, c, west_bc, loc, args...)
_fill_east_open_halo!(j, k, grid, c, east_bc, loc, args...)
return nothing
end

@kernel function _fill_south_and_north_open_halo!(c, south_bc, north_bc, loc, grid, args)
i, k = @index(Global, NTuple)
_fill_south_open_halo!(i, k, grid, c, south_bc, loc, args...)
_fill_north_open_halo!(i, k, grid, c, north_bc, loc, args...)
end
@inline get_open_halo_filling_functions(loc) = nothing
@inline get_open_halo_filling_functions(::Tuple{Face, Center, Center}) = fill_west_and_east_halo!
@inline get_open_halo_filling_functions(::Tuple{Center, Face, Center}) = fill_south_and_north_halo!
@inline get_open_halo_filling_functions(::Tuple{Center, Center, Face}) = fill_bottom_and_top_halo!

@kernel function _fill_bottom_and_top_open_halo!(c, bottom_bc, top_bc, loc, grid, args)
i, j = @index(Global, NTuple)
_fill_bottom_open_halo!(i, j, grid, c, bottom_bc, loc, args...)
_fill_top_open_halo!(i, j, grid, c, top_bc, loc, args...)
function fill_open_boundary_regions!(fields::Tuple, boundary_conditions, indices, loc, grid, args...; kwargs...)
for n in eachindex(fields)
fill_open_boundary_regions!(fields[n], boundary_conditions[n], indices, loc[n], grid, args...; kwargs...)
end
return nothing
end

# Generic fallback

@inline _fill_west_open_halo!(j, k, grid, c, bc, loc, args...) = nothing
@inline _fill_east_open_halo!(j, k, grid, c, bc, loc, args...) = nothing
@inline _fill_south_open_halo!(i, k, grid, c, bc, loc, args...) = nothing
@inline _fill_north_open_halo!(i, k, grid, c, bc, loc, args...) = nothing
@inline _fill_bottom_open_halo!(i, j, grid, c, bc, loc, args...) = nothing
@inline _fill_top_open_halo!(i, j, grid, c, bc, loc, args...) = nothing

# Open boundary condition fallback

@inline _fill_west_open_halo!(j, k, grid, c, bc::OBC, loc, args...) = @inbounds c[1, j, k] = getbc(bc, j, k, grid, args...)
@inline _fill_east_open_halo!(j, k, grid, c, bc::OBC, loc, args...) = @inbounds c[grid.Nx + 1, j, k] = getbc(bc, j, k, grid, args...)
@inline _fill_south_open_halo!(i, k, grid, c, bc::OBC, loc, args...) = @inbounds c[i, 1, k] = getbc(bc, i, k, grid, args...)
@inline _fill_north_open_halo!(i, k, grid, c, bc::OBC, loc, args...) = @inbounds c[i, grid.Ny + 1, k] = getbc(bc, i, k, grid, args...)
@inline _fill_bottom_open_halo!(i, j, grid, c, bc::OBC, loc, args...) = @inbounds c[i, j, 1] = getbc(bc, i, j, grid, args...)
@inline _fill_top_open_halo!(i, j, grid, c, bc::OBC, loc, args...) = @inbounds c[i, j, grid.Nz + 1] = getbc(bc, i, j, grid, args...)

# Regular boundary fill defaults

@inline _fill_west_halo!(j, k, grid, c, bc::OBC, loc, args...) = _fill_west_open_halo!(j, k, grid, c, bc, loc, args...)
@inline _fill_east_halo!(j, k, grid, c, bc::OBC, loc, args...) = _fill_east_open_halo!(j, k, grid, c, bc, loc, args...)
@inline _fill_south_halo!(i, k, grid, c, bc::OBC, loc, args...) = _fill_south_open_halo!(i, k, grid, c, bc, loc, args...)
@inline _fill_north_halo!(i, k, grid, c, bc::OBC, loc, args...) = _fill_north_open_halo!(i, k, grid, c, bc, loc, args...)
@inline _fill_bottom_halo!(i, j, grid, c, bc::OBC, loc, args...) = _fill_bottom_open_halo!(i, j, grid, c, bc, loc, args...)
@inline _fill_top_halo!(i, j, grid, c, bc::OBC, loc, args...) = _fill_top_open_halo!(i, j, grid, c, bc, loc, args...)

# Regular boundary fill for wall normal velocities

@inline _fill_west_halo!(j, k, grid, c, bc::OBC, ::Tuple{Face, <:Any, <:Any}, args...) = nothing
@inline _fill_east_halo!(j, k, grid, c, bc::OBC, ::Tuple{Face, <:Any, <:Any}, args...) = nothing
@inline _fill_south_halo!(i, k, grid, c, bc::OBC, ::Tuple{<:Any, Face, <:Any}, args...) = nothing
@inline _fill_north_halo!(i, k, grid, c, bc::OBC, ::Tuple{<:Any, Face, <:Any}, args...) = nothing
@inline _fill_bottom_halo!(i, j, grid, c, bc::OBC, ::Tuple{<:Any, <:Any, Face}, args...) = nothing
@inline _fill_top_halo!(i, j, grid, c, bc::OBC, ::Tuple{<:Any, <:Any, Face}, args...) = nothing
@inline retrieve_open_bc(bc::OBC) = bc
@inline retrieve_open_bc(bc) = nothing

# for regular halo fills, return nothing if the BC is not an OBC
@inline left_open_boundary_condition(boundary_condition, loc) = nothing
@inline left_open_boundary_condition(boundary_conditions, ::Tuple{Face, Center, Center}) = retrieve_open_bc(boundary_conditions.west)
@inline left_open_boundary_condition(boundary_conditions, ::Tuple{Center, Face, Center}) = retrieve_open_bc(boundary_conditions.south)
@inline left_open_boundary_condition(boundary_conditions, ::Tuple{Center, Center, Face}) = retrieve_open_bc(boundary_conditions.bottom)

@inline right_open_boundary_condition(boundary_conditions, loc) = nothing
@inline right_open_boundary_condition(boundary_conditions, ::Tuple{Face, Center, Center}) = retrieve_open_bc(boundary_conditions.east)
@inline right_open_boundary_condition(boundary_conditions, ::Tuple{Center, Face, Center}) = retrieve_open_bc(boundary_conditions.north)
@inline right_open_boundary_condition(boundary_conditions, ::Tuple{Center, Center, Face}) = retrieve_open_bc(boundary_conditions.top)

# Opern boundary fill
@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...)
@inline _fill_south_halo!(i, k, grid, c, bc::OBC, loc, args...) = @inbounds c[i, 1, k] = getbc(bc, i, k, grid, args...)
@inline _fill_north_halo!(i, k, grid, c, bc::OBC, loc, args...) = @inbounds c[i, grid.Ny + 1, k] = getbc(bc, i, k, grid, args...)
@inline _fill_bottom_halo!(i, j, grid, c, bc::OBC, loc, args...) = @inbounds c[i, j, 1] = getbc(bc, i, j, grid, args...)
@inline _fill_top_halo!(i, j, grid, c, bc::OBC, loc, args...) = @inbounds c[i, j, grid.Nz + 1] = getbc(bc, i, j, grid, args...)
2 changes: 1 addition & 1 deletion src/DistributedComputations/halo_communication.jl
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ end

function fill_halo_regions!(c::OffsetArray, bcs, indices, loc, grid::DistributedGrid, buffers, args...; fill_boundary_normal_velocities = true, kwargs...)
if fill_boundary_normal_velocities
fill_open_boundary_regions!(c, bcs, indices, loc, grid, args...; kwargs...)
fill_open_boundary_regions!(c, bcs, indices, loc, grid, buffers, args...; kwargs...)
end

arch = architecture(grid)
Expand Down
16 changes: 8 additions & 8 deletions src/MultiRegion/multi_region_boundary_conditions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -101,28 +101,28 @@ end

function fill_halo_regions!(c::MultiRegionObject, bcs, indices, loc, mrg::MultiRegionGrid, buffers, args...; fill_boundary_normal_velocities = true, kwargs...)
arch = architecture(mrg)
@apply_regionally fill_halos!, bcs = multi_region_permute_boundary_conditions(bcs)


if fill_boundary_normal_velocities
apply_regionally!(fill_open_boundary_regions!, c, bcs, indices, loc, mrg, args...)
end

@apply_regionally fill_halos!, permuted_bcs = multi_region_permute_boundary_conditions(bcs)

# The number of tasks is fixed to 3 (see `multi_region_permute_boundary_conditions`).
# When we want to allow asynchronous communication, we will might need to split the halos sides
# and the number of tasks might increase.
for task in 1:3
@apply_regionally begin
bcs_side = getindex(bcs, task)
bcs_side = getindex(permuted_bcs, task)
fill_halo_side! = getindex(fill_halos!, task)
fill_multiregion_send_buffers!(c, buffers, mrg, bcs_side)
end

buff = Reference(buffers.regional_objects)

if fill_boundary_normal_velocities
apply_regionally!(fill_open_boundary_regions!, c, bcs_side, indices, loc, mrg, args...)
end

apply_regionally!(fill_halo_event!, c, fill_halo_side!, bcs_side,
indices, loc, arch, mrg, buff,
args...; kwargs...)

end

return nothing
Expand Down

0 comments on commit 9111a8f

Please sign in to comment.