Skip to content

Commit ce8c49e

Browse files
simone-silvestrisimone-silvestri
authored
Merge branch 'main' into ss/tripolar-grid-in-oceananigans
2 parents e2a7157 + 9111a8f commit ce8c49e

File tree

3 files changed

+56
-97
lines changed

3 files changed

+56
-97
lines changed

‎src/BoundaryConditions/fill_halo_regions_open.jl

Lines changed: 47 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -6,104 +6,63 @@
66
77
Fill open boundary halo regions by filling boundary conditions on field faces with `open_fill`.
88
"""
9-
function fill_open_boundary_regions!(field, boundary_conditions, indices, loc, grid, args...; kwargs...)
9+
function fill_open_boundary_regions!(field, boundary_conditions, indices, loc, grid, args...; only_local_halos = false, kwargs...)
1010
arch = architecture(grid)
1111

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

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

20-
launch!(arch, grid, fill_size, open_fill, field, left_bc, right_bc, loc, grid, args)
22+
bcs_tuple = (left_bc, right_bc)
2123

22-
return nothing
23-
end
24-
25-
fill_open_boundary_regions!(fields::NTuple, boundary_conditions, indices, loc, grid, args...; kwargs...) =
26-
[fill_open_boundary_regions!(field, boundary_conditions[n], indices, loc[n], grid, args...; kwargs...) for (n, field) in enumerate(fields)]
27-
28-
# for regular halo fills
29-
@inline left_velocity_open_boundary_condition(boundary_condition, loc) = nothing
30-
@inline left_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Face, Center, Center}) = boundary_conditions.west
31-
@inline left_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Center, Face, Center}) = boundary_conditions.south
32-
@inline left_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Center, Center, Face}) = boundary_conditions.bottom
33-
34-
@inline right_velocity_open_boundary_condition(boundary_conditions, loc) = nothing
35-
@inline right_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Face, Center, Center}) = boundary_conditions.east
36-
@inline right_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Center, Face, Center}) = boundary_conditions.north
37-
@inline right_velocity_open_boundary_condition(boundary_conditions, ::Tuple{Center, Center, Face}) = boundary_conditions.top
24+
if !isnothing(fill_halo!) && any(!isnothing, bcs_tuple)
3825

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

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

48-
@inline get_open_halo_filling_functions(loc) = _no_fill!, _no_fill!
49-
@inline get_open_halo_filling_functions(::Tuple{Face, Center, Center}) = _fill_west_and_east_open_halo!, fill_west_and_east_halo!
50-
@inline get_open_halo_filling_functions(::Tuple{Center, Face, Center}) = _fill_south_and_north_open_halo!, fill_south_and_north_halo!
51-
@inline get_open_halo_filling_functions(::Tuple{Center, Center, Face}) = _fill_bottom_and_top_open_halo!, fill_bottom_and_top_halo!
52-
53-
@kernel _no_fill!(args...) = nothing
54-
55-
@inline fill_halo_size(field, ::typeof(_no_fill!), args...) = (0, 0)
56-
57-
@kernel function _fill_west_and_east_open_halo!(c, west_bc, east_bc, loc, grid, args)
58-
j, k = @index(Global, NTuple)
59-
_fill_west_open_halo!(j, k, grid, c, west_bc, loc, args...)
60-
_fill_east_open_halo!(j, k, grid, c, east_bc, loc, args...)
33+
return nothing
6134
end
6235

63-
@kernel function _fill_south_and_north_open_halo!(c, south_bc, north_bc, loc, grid, args)
64-
i, k = @index(Global, NTuple)
65-
_fill_south_open_halo!(i, k, grid, c, south_bc, loc, args...)
66-
_fill_north_open_halo!(i, k, grid, c, north_bc, loc, args...)
67-
end
36+
@inline get_open_halo_filling_functions(loc) = nothing
37+
@inline get_open_halo_filling_functions(::Tuple{Face, Center, Center}) = fill_west_and_east_halo!
38+
@inline get_open_halo_filling_functions(::Tuple{Center, Face, Center}) = fill_south_and_north_halo!
39+
@inline get_open_halo_filling_functions(::Tuple{Center, Center, Face}) = fill_bottom_and_top_halo!
6840

69-
@kernel function _fill_bottom_and_top_open_halo!(c, bottom_bc, top_bc, loc, grid, args)
70-
i, j = @index(Global, NTuple)
71-
_fill_bottom_open_halo!(i, j, grid, c, bottom_bc, loc, args...)
72-
_fill_top_open_halo!(i, j, grid, c, top_bc, loc, args...)
41+
function fill_open_boundary_regions!(fields::Tuple, boundary_conditions, indices, loc, grid, args...; kwargs...)
42+
for n in eachindex(fields)
43+
fill_open_boundary_regions!(fields[n], boundary_conditions[n], indices, loc[n], grid, args...; kwargs...)
44+
end
45+
return nothing
7346
end
7447

75-
# Generic fallback
76-
77-
@inline _fill_west_open_halo!(j, k, grid, c, bc, loc, args...) = nothing
78-
@inline _fill_east_open_halo!(j, k, grid, c, bc, loc, args...) = nothing
79-
@inline _fill_south_open_halo!(i, k, grid, c, bc, loc, args...) = nothing
80-
@inline _fill_north_open_halo!(i, k, grid, c, bc, loc, args...) = nothing
81-
@inline _fill_bottom_open_halo!(i, j, grid, c, bc, loc, args...) = nothing
82-
@inline _fill_top_open_halo!(i, j, grid, c, bc, loc, args...) = nothing
83-
84-
# Open boundary condition fallback
85-
86-
@inline _fill_west_open_halo!(j, k, grid, c, bc::OBC, loc, args...) = @inbounds c[1, j, k] = getbc(bc, j, k, grid, args...)
87-
@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...)
88-
@inline _fill_south_open_halo!(i, k, grid, c, bc::OBC, loc, args...) = @inbounds c[i, 1, k] = getbc(bc, i, k, grid, args...)
89-
@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...)
90-
@inline _fill_bottom_open_halo!(i, j, grid, c, bc::OBC, loc, args...) = @inbounds c[i, j, 1] = getbc(bc, i, j, grid, args...)
91-
@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...)
92-
93-
# Regular boundary fill defaults
94-
95-
@inline _fill_west_halo!(j, k, grid, c, bc::OBC, loc, args...) = _fill_west_open_halo!(j, k, grid, c, bc, loc, args...)
96-
@inline _fill_east_halo!(j, k, grid, c, bc::OBC, loc, args...) = _fill_east_open_halo!(j, k, grid, c, bc, loc, args...)
97-
@inline _fill_south_halo!(i, k, grid, c, bc::OBC, loc, args...) = _fill_south_open_halo!(i, k, grid, c, bc, loc, args...)
98-
@inline _fill_north_halo!(i, k, grid, c, bc::OBC, loc, args...) = _fill_north_open_halo!(i, k, grid, c, bc, loc, args...)
99-
@inline _fill_bottom_halo!(i, j, grid, c, bc::OBC, loc, args...) = _fill_bottom_open_halo!(i, j, grid, c, bc, loc, args...)
100-
@inline _fill_top_halo!(i, j, grid, c, bc::OBC, loc, args...) = _fill_top_open_halo!(i, j, grid, c, bc, loc, args...)
101-
102-
# Regular boundary fill for wall normal velocities
103-
104-
@inline _fill_west_halo!(j, k, grid, c, bc::OBC, ::Tuple{Face, <:Any, <:Any}, args...) = nothing
105-
@inline _fill_east_halo!(j, k, grid, c, bc::OBC, ::Tuple{Face, <:Any, <:Any}, args...) = nothing
106-
@inline _fill_south_halo!(i, k, grid, c, bc::OBC, ::Tuple{<:Any, Face, <:Any}, args...) = nothing
107-
@inline _fill_north_halo!(i, k, grid, c, bc::OBC, ::Tuple{<:Any, Face, <:Any}, args...) = nothing
108-
@inline _fill_bottom_halo!(i, j, grid, c, bc::OBC, ::Tuple{<:Any, <:Any, Face}, args...) = nothing
109-
@inline _fill_top_halo!(i, j, grid, c, bc::OBC, ::Tuple{<:Any, <:Any, Face}, args...) = nothing
48+
@inline retrieve_open_bc(bc::OBC) = bc
49+
@inline retrieve_open_bc(bc) = nothing
50+
51+
# for regular halo fills, return nothing if the BC is not an OBC
52+
@inline left_open_boundary_condition(boundary_condition, loc) = nothing
53+
@inline left_open_boundary_condition(boundary_conditions, ::Tuple{Face, Center, Center}) = retrieve_open_bc(boundary_conditions.west)
54+
@inline left_open_boundary_condition(boundary_conditions, ::Tuple{Center, Face, Center}) = retrieve_open_bc(boundary_conditions.south)
55+
@inline left_open_boundary_condition(boundary_conditions, ::Tuple{Center, Center, Face}) = retrieve_open_bc(boundary_conditions.bottom)
56+
57+
@inline right_open_boundary_condition(boundary_conditions, loc) = nothing
58+
@inline right_open_boundary_condition(boundary_conditions, ::Tuple{Face, Center, Center}) = retrieve_open_bc(boundary_conditions.east)
59+
@inline right_open_boundary_condition(boundary_conditions, ::Tuple{Center, Face, Center}) = retrieve_open_bc(boundary_conditions.north)
60+
@inline right_open_boundary_condition(boundary_conditions, ::Tuple{Center, Center, Face}) = retrieve_open_bc(boundary_conditions.top)
61+
62+
# Opern boundary fill
63+
@inline _fill_west_halo!(j, k, grid, c, bc::OBC, loc, args...) = @inbounds c[1, j, k] = getbc(bc, j, k, grid, args...)
64+
@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...)
65+
@inline _fill_south_halo!(i, k, grid, c, bc::OBC, loc, args...) = @inbounds c[i, 1, k] = getbc(bc, i, k, grid, args...)
66+
@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...)
67+
@inline _fill_bottom_halo!(i, j, grid, c, bc::OBC, loc, args...) = @inbounds c[i, j, 1] = getbc(bc, i, j, grid, args...)
68+
@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...)

‎src/DistributedComputations/halo_communication.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ end
100100

101101
function fill_halo_regions!(c::OffsetArray, bcs, indices, loc, grid::DistributedGrid, buffers, args...; fill_boundary_normal_velocities = true, kwargs...)
102102
if fill_boundary_normal_velocities
103-
fill_open_boundary_regions!(c, bcs, indices, loc, grid, args...; kwargs...)
103+
fill_open_boundary_regions!(c, bcs, indices, loc, grid, buffers, args...; kwargs...)
104104
end
105105

106106
arch = architecture(grid)

‎src/MultiRegion/multi_region_boundary_conditions.jl

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -101,28 +101,28 @@ end
101101

102102
function fill_halo_regions!(c::MultiRegionObject, bcs, indices, loc, mrg::MultiRegionGrid, buffers, args...; fill_boundary_normal_velocities = true, kwargs...)
103103
arch = architecture(mrg)
104-
@apply_regionally fill_halos!, bcs = multi_region_permute_boundary_conditions(bcs)
105-
104+
105+
if fill_boundary_normal_velocities
106+
apply_regionally!(fill_open_boundary_regions!, c, bcs, indices, loc, mrg, args...)
107+
end
108+
109+
@apply_regionally fill_halos!, permuted_bcs = multi_region_permute_boundary_conditions(bcs)
110+
106111
# The number of tasks is fixed to 3 (see `multi_region_permute_boundary_conditions`).
107112
# When we want to allow asynchronous communication, we will might need to split the halos sides
108113
# and the number of tasks might increase.
109114
for task in 1:3
110115
@apply_regionally begin
111-
bcs_side = getindex(bcs, task)
116+
bcs_side = getindex(permuted_bcs, task)
112117
fill_halo_side! = getindex(fill_halos!, task)
113118
fill_multiregion_send_buffers!(c, buffers, mrg, bcs_side)
114119
end
115120

116121
buff = Reference(buffers.regional_objects)
117122

118-
if fill_boundary_normal_velocities
119-
apply_regionally!(fill_open_boundary_regions!, c, bcs_side, indices, loc, mrg, args...)
120-
end
121-
122123
apply_regionally!(fill_halo_event!, c, fill_halo_side!, bcs_side,
123124
indices, loc, arch, mrg, buff,
124125
args...; kwargs...)
125-
126126
end
127127

128128
return nothing

0 commit comments

Comments
 (0)