diff --git a/ext/OceananigansCUDAExt.jl b/ext/OceananigansCUDAExt.jl index 030bdb79907..1482863ae6a 100644 --- a/ext/OceananigansCUDAExt.jl +++ b/ext/OceananigansCUDAExt.jl @@ -56,6 +56,8 @@ Base.summary(::CUDAGPU) = "CUDAGPU" AC.architecture(::CuArray) = CUDAGPU() AC.architecture(::Type{CuArray}) = CUDAGPU() +AC.architecture(::CuDeviceArray) = CUDAGPU() +AC.architecture(::Type{CuDeviceArray}) = CUDAGPU() AC.architecture(::CuSparseMatrixCSC) = CUDAGPU() AC.array_type(::AC.GPU{CUDABackend}) = CuArray diff --git a/ext/OceananigansEnzymeExt.jl b/ext/OceananigansEnzymeExt.jl index 9b7ed4a5454..0439860fbc8 100644 --- a/ext/OceananigansEnzymeExt.jl +++ b/ext/OceananigansEnzymeExt.jl @@ -19,7 +19,7 @@ EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.AbstractOperations.m EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.Utils.flatten_reduced_dimensions), x...) = nothing EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.Utils.prettytime), x...) = nothing EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.Grids.total_size), x...) = nothing -EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.BoundaryConditions.parent_size_and_offset), x...) = nothing +EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.BoundaryConditions.periodic_size_and_offset), x...) = nothing @inline EnzymeCore.EnzymeRules.inactive_type(v::Type{Oceananigans.Utils.KernelParameters}) = true @inline batch(::Val{1}, ::Type{T}) where T = T diff --git a/src/BoundaryConditions/BoundaryConditions.jl b/src/BoundaryConditions/BoundaryConditions.jl index 1d00223a66f..b68a0ebf49c 100644 --- a/src/BoundaryConditions/BoundaryConditions.jl +++ b/src/BoundaryConditions/BoundaryConditions.jl @@ -9,7 +9,10 @@ export validate_boundary_condition_topology, validate_boundary_condition_architecture, FieldBoundaryConditions, compute_x_bcs!, compute_y_bcs!, compute_z_bcs!, - fill_halo_regions! + fill_halo_regions!, + WestAndEast, SouthAndNorth, BottomAndTop, + West, East, South, North, Bottom, Top, + DistributedFillHalo using Adapt using KernelAbstractions: @index, @kernel @@ -21,10 +24,22 @@ using Oceananigans.Grids import Adapt: adapt_structure +# All possible fill_halo! kernels +struct WestAndEast end +struct SouthAndNorth end +struct BottomAndTop end +struct West end +struct East end +struct South end +struct North end +struct Bottom end +struct Top end + include("boundary_condition_classifications.jl") include("boundary_condition.jl") include("discrete_boundary_function.jl") include("continuous_boundary_function.jl") +include("boundary_condition_ordering.jl") include("field_boundary_conditions.jl") include("show_boundary_conditions.jl") @@ -33,8 +48,8 @@ include("fill_halo_regions_value_gradient.jl") include("fill_halo_regions_open.jl") include("fill_halo_regions_periodic.jl") include("fill_halo_regions_flux.jl") -include("fill_halo_regions_nothing.jl") include("fill_halo_regions_zipper.jl") +include("fill_halo_kernels.jl") include("compute_flux_bcs.jl") diff --git a/src/BoundaryConditions/boundary_condition_ordering.jl b/src/BoundaryConditions/boundary_condition_ordering.jl new file mode 100644 index 00000000000..a99939ee430 --- /dev/null +++ b/src/BoundaryConditions/boundary_condition_ordering.jl @@ -0,0 +1,128 @@ +extract_bc(bcs, ::West) = tuple(bcs.west) +extract_bc(bcs, ::East) = tuple(bcs.east) +extract_bc(bcs, ::South) = tuple(bcs.south) +extract_bc(bcs, ::North) = tuple(bcs.north) +extract_bc(bcs, ::Bottom) = tuple(bcs.bottom) +extract_bc(bcs, ::Top) = tuple(bcs.top) + +extract_bc(bcs, ::BottomAndTop) = (bcs.bottom, bcs.top) +extract_bc(bcs, ::WestAndEast) = (bcs.west, bcs.east) +extract_bc(bcs, ::SouthAndNorth) = (bcs.south, bcs.north) + +# In case of a DistributedCommunication paired with a +# Flux, Value or Gradient boundary condition, we split the direction in two single-sided +# fill_halo! events (see issue #3342) +# `permute_boundary_conditions` returns a 2-tuple containing the ordered operations to execute in +# position [1] and the associated boundary conditions in position [2] +function permute_boundary_conditions(bcs) + + split_x_halo_filling = split_halo_filling(bcs.west, bcs.east) + split_y_halo_filling = split_halo_filling(bcs.south, bcs.north) + + if split_x_halo_filling + if split_y_halo_filling + sides = [West(), East(), South(), North(), BottomAndTop()] + bcs_array = [bcs.west, bcs.east, bcs.south, bcs.north, bcs.bottom] + else + sides = [West(), East(), SouthAndNorth(), BottomAndTop()] + bcs_array = [bcs.west, bcs.east, bcs.south, bcs.bottom] + end + else + if split_y_halo_filling + sides = [WestAndEast(), South(), North(), BottomAndTop()] + bcs_array = [bcs.west, bcs.south, bcs.north, bcs.bottom] + else + sides = [WestAndEast(), SouthAndNorth(), BottomAndTop()] + bcs_array = [bcs.west, bcs.south, bcs.bottom] + end + end + + perm = sortperm(bcs_array, lt=fill_first) + sides = tuple(sides[perm]...) + + boundary_conditions = Tuple(extract_bc(bcs, side) for side in sides) + + return sides, boundary_conditions +end + +side_name(::West) = :west +side_name(::East) = :east +side_name(::South) = :south +side_name(::North) = :north +side_name(::Bottom) = :bottom +side_name(::Top) = :top +side_name(::WestAndEast) = :west_and_east +side_name(::SouthAndNorth) = :south_and_north +side_name(::BottomAndTop) = :bottom_and_top + +# Split direction in two distinct fill_halo! events in case of a communication boundary condition +# (distributed DCBC), paired with a Flux, Value or Gradient boundary condition +split_halo_filling(bcs1, bcs2) = false +split_halo_filling(::DCBC, ::DCBC) = false +split_halo_filling(bcs1, ::DCBC) = true +split_halo_filling(::DCBC, bcs2) = true + +# Same thing for MultiRegion boundary conditions +split_halo_filling(::MCBC, ::MCBC) = false +split_halo_filling(bcs1, ::MCBC) = true +split_halo_filling(::MCBC, bcs2) = true + +# heterogenous distribute-shared communication is not supported +# TODO: support heterogeneous distributed-shared communication +split_halo_filling(::MCBC, ::DCBC) = throw("Cannot split MultiRegion and Distributed boundary conditions.") +split_halo_filling(::DCBC, ::MCBC) = throw("Cannot split MultiRegion and Distributed boundary conditions.") + +##### +##### Halo filling order +##### + +const PBCT = Union{PBC, Tuple{Vararg{PBC}}} +const MCBCT = Union{MCBC, Tuple{Vararg{MCBC}}} +const DCBCT = Union{DCBC, Tuple{Vararg{DCBC}}} +const OBCTC = Union{OBC, Tuple{Vararg{OBC}}} + +# Distributed halos have to be filled last to allow the +# possibility of asynchronous communication: +# If other halos are filled after we initiate the distributed communication, +# (but before communication is completed) the halos will be overwritten. +# For this reason we always want to perform local halo filling first and then +# initiate communication + +# Periodic is handled after Flux, Value, Gradient because +# Periodic fills also corners while Flux, Value, Gradient do not +# TODO: remove this ordering requirement (see issue https://github.com/CliMA/Oceananigans.jl/issues/3342) + +# Order of halo filling +# 1) Flux, Value, Gradient (TODO: remove these BC and apply them as fluxes) +# 2) Periodic (PBCT) +# 3) Shared Communication (MCBCT) +# 4) Distributed Communication (DCBCT) + +# We define "greater than" `>` and "lower than", for boundary conditions +# following the rules outlined in `fill_first` +# i.e. if `bc1 > bc2` then `bc2` precedes `bc1` in filling order +@inline Base.isless(bc1::BoundaryCondition, bc2::BoundaryCondition) = fill_first(bc1, bc2) + +# fallback for `Nothing` BC. +@inline Base.isless(::Nothing, ::Nothing) = true +@inline Base.isless(::BoundaryCondition, ::Nothing) = false +@inline Base.isless(::Nothing, ::BoundaryCondition) = true +@inline Base.isless(::BoundaryCondition, ::Missing) = false +@inline Base.isless(::Missing, ::BoundaryCondition) = true + +fill_first(bc1::DCBCT, bc2) = false +fill_first(bc1::PBCT, bc2::DCBCT) = true +fill_first(bc1::DCBCT, bc2::PBCT) = false +fill_first(bc1::MCBCT, bc2::DCBCT) = true +fill_first(bc1::DCBCT, bc2::MCBCT) = false +fill_first(bc1, bc2::DCBCT) = true +fill_first(bc1::DCBCT, bc2::DCBCT) = true +fill_first(bc1::PBCT, bc2) = false +fill_first(bc1::MCBCT, bc2) = false +fill_first(bc1::PBCT, bc2::MCBCT) = true +fill_first(bc1::MCBCT, bc2::PBCT) = false +fill_first(bc1, bc2::PBCT) = true +fill_first(bc1, bc2::MCBCT) = true +fill_first(bc1::PBCT, bc2::PBCT) = true +fill_first(bc1::MCBCT, bc2::MCBCT) = true +fill_first(bc1, bc2) = true diff --git a/src/BoundaryConditions/field_boundary_conditions.jl b/src/BoundaryConditions/field_boundary_conditions.jl index 2cf49a5a2dd..e8f3284645e 100644 --- a/src/BoundaryConditions/field_boundary_conditions.jl +++ b/src/BoundaryConditions/field_boundary_conditions.jl @@ -47,7 +47,7 @@ default_auxiliary_bc(grid, ::Val{:top}, loc) = _default_auxiliary_bc(topology ##### Field boundary conditions ##### -mutable struct FieldBoundaryConditions{W, E, S, N, B, T, I} +mutable struct FieldBoundaryConditions{W, E, S, N, B, T, I, K, O} west :: W east :: E south :: S @@ -55,6 +55,17 @@ mutable struct FieldBoundaryConditions{W, E, S, N, B, T, I} bottom :: B top :: T immersed :: I + kernels :: K # kernels used to fill halo regions + ordered_bcs :: O +end + +const boundarynames = (:west, :east, :south, :north, :bottom, :top, :immersed) + +const NoKernelFBC = FieldBoundaryConditions{W, E, S, N, B, T, I, Nothing, Nothing} where {W, E, S, N, B, T, I} + +# Internal constructor that fills up computational details in the "auxiliaries" spot. +function FieldBoundaryConditions(west, east, south, north, bottom, top, immersed) + return FieldBoundaryConditions(west, east, south, north, bottom, top, immersed, nothing, nothing) end function FieldBoundaryConditions(indices::Tuple, west, east, south, north, bottom, top, immersed) @@ -66,7 +77,7 @@ function FieldBoundaryConditions(indices::Tuple, west, east, south, north, botto end FieldBoundaryConditions(indices::Tuple, bcs::FieldBoundaryConditions) = - FieldBoundaryConditions(indices, (getproperty(bcs, side) for side in propertynames(bcs))...) + FieldBoundaryConditions(indices, (getproperty(bcs, side) for side in boundarynames)...) FieldBoundaryConditions(indices::Tuple, ::Nothing) = nothing FieldBoundaryConditions(indices::Tuple, ::Missing) = nothing @@ -76,6 +87,9 @@ window_boundary_conditions(::UnitRange, left, right) = nothing, nothing window_boundary_conditions(::Base.OneTo, left, right) = nothing, nothing window_boundary_conditions(::Colon, left, right) = left, right +# The only thing we need +Adapt.adapt_structure(to, fbcs::FieldBoundaryConditions) = (kernels = fbcs.kernels, ordered_bcs = Adapt.adapt(to, fbcs.ordered_bcs)) + on_architecture(arch, fbcs::FieldBoundaryConditions) = FieldBoundaryConditions(on_architecture(arch, fbcs.west), on_architecture(arch, fbcs.east), @@ -83,7 +97,9 @@ on_architecture(arch, fbcs::FieldBoundaryConditions) = on_architecture(arch, fbcs.north), on_architecture(arch, fbcs.bottom), on_architecture(arch, fbcs.top), - on_architecture(arch, fbcs.immersed)) + on_architecture(arch, fbcs.immersed), + fbcs.kernels, + on_architecture(arch, fbcs.ordered_bcs)) """ FieldBoundaryConditions(; kwargs...) diff --git a/src/BoundaryConditions/fill_halo_kernels.jl b/src/BoundaryConditions/fill_halo_kernels.jl new file mode 100644 index 00000000000..fc535004fd2 --- /dev/null +++ b/src/BoundaryConditions/fill_halo_kernels.jl @@ -0,0 +1,163 @@ +using Oceananigans.Utils: configure_kernel + +""" + construct_boundary_conditions_kernels(bcs::FieldBoundaryConditions, + data::OffsetArray, + grid::AbstractGrid, + loc, indices) + +Construct preconfigured boundary condition kernels for a given `data` array, `grid`, +and the provided `bcs` (a FieldBoundaryConditions` object). +Return a new `FieldBoundaryConditions` object with the preconfigured kernels and +ordered boundary conditions. +""" +function construct_boundary_conditions_kernels(bcs::FieldBoundaryConditions, + data::OffsetArray, + grid::AbstractGrid, + loc, indices) + + kernels!, ordered_bcs = fill_halo_kernels(bcs, data, grid, loc, indices) + regularized_bcs = FieldBoundaryConditions(bcs.west, bcs.east, bcs.south, bcs.north, + bcs.bottom, bcs.top, bcs.immersed, + kernels!, ordered_bcs) + return regularized_bcs +end + +# If the bcs are nothing or missing... they remain nothing or missing +construct_boundary_conditions_kernels(::Nothing, data, grid, loc, indices) = nothing +construct_boundary_conditions_kernels(::Missing, data, grid, loc, indices) = missing + +# Select the valid BC out of a tuple to configure the kernel +@inline select_bc(bcs::Tuple) = @inbounds bcs[1] +@inline select_bc(bcs::Tuple{<:Nothing, <:BoundaryCondition}) = @inbounds bcs[2] +@inline select_bc(bcs::Tuple{<:BoundaryCondition, <:Nothing}) = @inbounds bcs[1] +@inline select_bc(bcs::BoundaryCondition) = bc + +@inline function fill_halo_kernels(bcs::FieldBoundaryConditions, + data::OffsetArray, + grid::AbstractGrid, + loc, indices) + + arch = architecture(grid) + sides, ordered_bcs = permute_boundary_conditions(bcs) + reduced_dimensions = findall(x -> x isa Nothing, loc) + reduced_dimensions = tuple(reduced_dimensions...) + names = Tuple(side_name(side) for side in sides) + kernels! = [] + + for task in 1:length(sides) + side = sides[task] + bc = select_bc(ordered_bcs[task]) + + size = fill_halo_size(data, side, indices, bc, loc, grid) + offset = fill_halo_offset(size, side, indices) + kernel! = fill_halo_kernel!(side, bc, grid, size, offset, data, reduced_dimensions) + + push!(kernels!, kernel!) + end + + kernels! = tuple(kernels!...) + + return NamedTuple{names}(kernels!), NamedTuple{names}(ordered_bcs) +end + +@inline get_boundary_kernels(bcs::NoKernelFBC, data, grid, loc, indices) = fill_halo_kernels(bcs, data, grid, loc, indices) +@inline get_boundary_kernels(bcs, args...) = bcs.kernels, bcs.ordered_bcs + +@inline fix_halo_offsets(o, co) = co > 0 ? o - co : o # Windowed fields have only positive offsets to correct + +@inline periodic_size_and_offset(c, dim1, dim2, size, offset) = (size, fix_halo_offsets.(offset, c.offsets[[dim1, dim2]])) +@inline periodic_size_and_offset(c, dim1, dim2, ::Symbol, offset) = (size(parent(c))[[dim1, dim2]], (0, 0)) + +#### +#### Fill halo configured kernels +#### + +const NoBC = Union{Nothing, Missing} + +fill_halo_kernel!(value, bc::NoBC, args...) = nothing + +@inline kernel_parameters(size, offset) = KernelParameters(size, offset) +@inline kernel_parameters(size::Symbol, offset) = size + +##### +##### Two-sided fill halo kernels +##### + +fill_halo_kernel!(::WestAndEast, bc::BoundaryCondition, grid, size, offset, data, reduced_dimensions) = + configure_kernel(architecture(grid), grid, kernel_parameters(size, offset), _fill_west_and_east_halo!; reduced_dimensions)[1] + +fill_halo_kernel!(::SouthAndNorth, bc::BoundaryCondition, grid, size, offset, data, reduced_dimensions) = + configure_kernel(architecture(grid), grid, kernel_parameters(size, offset), _fill_south_and_north_halo!; reduced_dimensions)[1] + +fill_halo_kernel!(::BottomAndTop, bc::BoundaryCondition, grid, size, offset, data, reduced_dimensions) = + configure_kernel(architecture(grid), grid, kernel_parameters(size, offset), _fill_bottom_and_top_halo!; reduced_dimensions)[1] + +##### +##### One-sided fill halo kernels +##### + +fill_halo_kernel!(::West, bc::BoundaryCondition, grid, size, offset, data, reduced_dimensions) = + configure_kernel(architecture(grid), grid, kernel_parameters(size, offset), _fill_only_west_halo!; reduced_dimensions)[1] + +fill_halo_kernel!(::East, bc::BoundaryCondition, grid, size, offset, data, reduced_dimensions) = + configure_kernel(architecture(grid), grid, kernel_parameters(size, offset), _fill_only_east_halo!; reduced_dimensions)[1] + +fill_halo_kernel!(::South, bc::BoundaryCondition, grid, size, offset, data, reduced_dimensions) = + configure_kernel(architecture(grid), grid, kernel_parameters(size, offset), _fill_only_south_halo!; reduced_dimensions)[1] + +fill_halo_kernel!(::North, bc::BoundaryCondition, grid, size, offset, data, reduced_dimensions) = + configure_kernel(architecture(grid), grid, kernel_parameters(size, offset), _fill_only_north_halo!; reduced_dimensions)[1] + +fill_halo_kernel!(::Bottom, bc::BoundaryCondition, grid, size, offset, data, reduced_dimensions) = + configure_kernel(architecture(grid), grid, kernel_parameters(size, offset), _fill_only_bottom_halo!; reduced_dimensions)[1] + +fill_halo_kernel!(::Top, bc::BoundaryCondition, grid, size, offset, data, reduced_dimensions) = + configure_kernel(architecture(grid), grid, kernel_parameters(size, offset), _fill_only_top_halo!; reduced_dimensions)[1] + +##### +##### Periodic fill halo kernels (Always two-sided) +##### + +function fill_halo_kernel!(::WestAndEast, bc::PBC, grid, size, offset, data, reduced_dimensions) + yz_size, offset = periodic_size_and_offset(data, 2, 3, size, offset) + return configure_kernel(architecture(grid), grid, kernel_parameters(yz_size, offset), _fill_periodic_west_and_east_halo!)[1] +end + +function fill_halo_kernel!(::SouthAndNorth, bc::PBC, grid, size, offset, data, reduced_dimensions) + xz_size, offset = periodic_size_and_offset(data, 1, 3, size, offset) + return configure_kernel(architecture(grid), grid, kernel_parameters(xz_size, offset), _fill_periodic_south_and_north_halo!)[1] +end + +function fill_halo_kernel!(::BottomAndTop, bc::PBC, grid, size, offset, data, reduced_dimensions) + xy_size, offset = periodic_size_and_offset(data, 1, 2, size, offset) + return configure_kernel(architecture(grid), grid, kernel_parameters(xy_size, offset), _fill_periodic_bottom_and_top_halo!)[1] +end + +##### +##### Distributed Boundary Conditions +##### + +# A struct to hold the side of the fill_halo kernel +# These are defined in `src/DistributedComputations/halo_communication.jl` +struct DistributedFillHalo{S} + side :: S +end + +for Side in (:WestAndEast, :SouthAndNorth, :BottomAndTop, :West, :East, :South, :North, :Bottom, :Top) + @eval fill_halo_kernel!(::$Side, bc::DCBC, grid, size, offset, data, reduced_dimensions) = DistributedFillHalo($Side()) +end + +##### +##### MultiRegion Boundary Conditions +##### + +# A struct to hold the side of the fill_halo kernel +# These are defined in `src/MultiRegion/multi_region_boundary_conditions.jl` +struct MultiRegionFillHalo{S} + side :: S +end + +for Side in (:WestAndEast, :SouthAndNorth, :BottomAndTop, :West, :East, :South, :North, :Bottom, :Top) + @eval fill_halo_kernel!(::$Side, bc::MCBC, grid, size, offset, data, reduced_dimensions) = MultiRegionFillHalo($Side()) +end diff --git a/src/BoundaryConditions/fill_halo_regions.jl b/src/BoundaryConditions/fill_halo_regions.jl index 2aa41231494..a07def3b0bd 100644 --- a/src/BoundaryConditions/fill_halo_regions.jl +++ b/src/BoundaryConditions/fill_halo_regions.jl @@ -8,8 +8,8 @@ import Base ##### General halo filling functions ##### +fill_halo_regions!(::Ref, args...; kwargs...) = nothing # a lot of Refs are passed around, so we need this fill_halo_regions!(::Nothing, args...; kwargs...) = nothing -fill_halo_regions!(::NamedTuple{(), Tuple{}}, args...; kwargs...) = nothing """ fill_halo_regions!(fields::Union{Tuple, NamedTuple}, arch, args...) @@ -20,180 +20,36 @@ conditions, possibly recursing into `fields` if it is a nested tuple-of-tuples. # Some fields have `nothing` boundary conditions, such as `FunctionField` and `ZeroField`. fill_halo_regions!(c::OffsetArray, ::Nothing, args...; kwargs...) = nothing -retrieve_bc(bc) = bc - -# Returns the boundary conditions a specific side for `FieldBoundaryConditions` inputs and -# a tuple of boundary conditions for `NTuple{N, <:FieldBoundaryConditions}` inputs -for dir in (:west, :east, :south, :north, :bottom, :top) - extract_side_bc = Symbol(:extract_, dir, :_bc) - @eval begin - @inline $extract_side_bc(bc) = retrieve_bc(bc.$dir) - @inline $extract_side_bc(bc::Tuple) = map($extract_side_bc, bc) - end -end -# -@inline extract_bc(bc, ::Val{:west}) = tuple(extract_west_bc(bc)) -@inline extract_bc(bc, ::Val{:east}) = tuple(extract_east_bc(bc)) -@inline extract_bc(bc, ::Val{:south}) = tuple(extract_south_bc(bc)) -@inline extract_bc(bc, ::Val{:north}) = tuple(extract_north_bc(bc)) -@inline extract_bc(bc, ::Val{:bottom}) = tuple(extract_bottom_bc(bc)) -@inline extract_bc(bc, ::Val{:top}) = tuple(extract_top_bc(bc)) - -@inline extract_bc(bc, ::Val{:west_and_east}) = (extract_west_bc(bc), extract_east_bc(bc)) -@inline extract_bc(bc, ::Val{:south_and_north}) = (extract_south_bc(bc), extract_north_bc(bc)) -@inline extract_bc(bc, ::Val{:bottom_and_top}) = (extract_bottom_bc(bc), extract_top_bc(bc)) - -# Finally, the true fill_halo! -const MaybeTupledData = Union{OffsetArray, Tuple{Vararg{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...) - arch = architecture(grid) - - if fill_boundary_normal_velocities - fill_open_boundary_regions!(c, boundary_conditions, indices, loc, grid, args...; kwargs...) - end - - fill_halos!, bcs = permute_boundary_conditions(boundary_conditions) - number_of_tasks = length(fill_halos!) +function fill_halo_regions!(c::OffsetArray, boundary_conditions, indices, loc, grid, args...; kwargs...) + + kernels!, bcs = get_boundary_kernels(boundary_conditions, c, grid, loc, indices) + number_of_tasks = length(kernels!) # Fill halo in the three permuted directions (1, 2, and 3), making sure dependencies are fulfilled for task = 1:number_of_tasks - fill_halo_event!(c, fill_halos![task], bcs[task], indices, loc, arch, grid, args...; kwargs...) + @inbounds fill_halo_event!(c, kernels![task], bcs[task], loc, grid, args...; kwargs...) end return nothing end -function fill_halo_event!(c, fill_halos!, bcs, indices, loc, arch, grid, args...; - async = false, # This kwargs is specific to DistributedGrids, here is does nothing - kwargs...) - - # Calculate size and offset of the fill_halo kernel - # We assume that the kernel size is the same for west and east boundaries, - # south and north boundaries, and bottom and top boundaries - size = fill_halo_size(c, fill_halos!, indices, bcs[1], loc, grid) - offset = fill_halo_offset(size, fill_halos!, indices) +const NoBCs = Union{Nothing, Tuple{Vararg{Nothing}}} - fill_halos!(c, bcs..., size, offset, loc, arch, grid, args...; kwargs...) - - return nothing -end - -# In case of a DistributedCommunication paired with a -# Flux, Value or Gradient boundary condition, we split the direction in two single-sided -# fill_halo! events (see issue #3342) -# `permute_boundary_conditions` returns a 2-tuple containing the ordered operations to execute in -# position [1] and the associated boundary conditions in position [2] -function permute_boundary_conditions(boundary_conditions) - - split_x_halo_filling = split_halo_filling(extract_west_bc(boundary_conditions), extract_east_bc(boundary_conditions)) - split_y_halo_filling = split_halo_filling(extract_south_bc(boundary_conditions), extract_north_bc(boundary_conditions)) - - west_bc = extract_west_bc(boundary_conditions) - east_bc = extract_east_bc(boundary_conditions) - south_bc = extract_south_bc(boundary_conditions) - north_bc = extract_north_bc(boundary_conditions) - - if split_x_halo_filling - if split_y_halo_filling - fill_halos! = [fill_west_halo!, fill_east_halo!, fill_south_halo!, fill_north_halo!, fill_bottom_and_top_halo!] - sides = [:west, :east, :south, :north, :bottom_and_top] - bcs_array = [west_bc, east_bc, south_bc, north_bc, extract_bottom_bc(boundary_conditions)] - else - fill_halos! = [fill_west_halo!, fill_east_halo!, fill_south_and_north_halo!, fill_bottom_and_top_halo!] - sides = [:west, :east, :south_and_north, :bottom_and_top] - bcs_array = [west_bc, east_bc, south_bc, extract_bottom_bc(boundary_conditions)] - end - else - if split_y_halo_filling - fill_halos! = [fill_west_and_east_halo!, fill_south_halo!, fill_north_halo!, fill_bottom_and_top_halo!] - sides = [:west_and_east, :south, :north, :bottom_and_top] - bcs_array = [west_bc, south_bc, north_bc, extract_bottom_bc(boundary_conditions)] - else - fill_halos! = [fill_west_and_east_halo!, fill_south_and_north_halo!, fill_bottom_and_top_halo!] - sides = [:west_and_east, :south_and_north, :bottom_and_top] - bcs_array = [west_bc, south_bc, extract_bottom_bc(boundary_conditions)] - end - end - - perm = sortperm(bcs_array, lt=fill_first) - fill_halos! = fill_halos![perm] - sides = sides[perm] - - boundary_conditions = Tuple(extract_bc(boundary_conditions, Val(side)) for side in sides) - - return fill_halos!, boundary_conditions -end - -# Split direction in two distinct fill_halo! events in case of a communication boundary condition -# (distributed DCBC), paired with a Flux, Value or Gradient boundary condition -split_halo_filling(bcs1, bcs2) = false -split_halo_filling(::DCBC, ::DCBC) = false -split_halo_filling(bcs1, ::DCBC) = true -split_halo_filling(::DCBC, bcs2) = true - -# TODO: support heterogeneous distributed-shared communication -# split_halo_filling(::MCBC, ::DCBC) = false -# split_halo_filling(::DCBC, ::MCBC) = false -# split_halo_filling(::MCBC, ::MCBC) = false -# split_halo_filling(bcs1, ::MCBC) = true -# split_halo_filling(::MCBC, bcs2) = true +@inline fill_halo_event!(c, kernel!, bcs, loc, grid, args...; kwargs...) = kernel!(c, bcs..., loc, grid, Tuple(args)) +@inline fill_halo_event!(c, ::Nothing, ::NoBCs, loc, grid, args...; kwargs...) = nothing ##### -##### Halo filling order +##### Nothing BCs ##### -const PBCT = Union{PBC, Tuple{Vararg{PBC}}} -const MCBCT = Union{MCBC, Tuple{Vararg{MCBC}}} -const DCBCT = Union{DCBC, Tuple{Vararg{DCBC}}} - -# Distributed halos have to be filled last to allow the -# possibility of asynchronous communication: -# If other halos are filled after we initiate the distributed communication, -# (but before communication is completed) the halos will be overwritten. -# For this reason we always want to perform local halo filling first and then -# initiate communication - -# Periodic is handled after Flux, Value, Gradient because -# Periodic fills also corners while Flux, Value, Gradient do not -# TODO: remove this ordering requirement (see issue https://github.com/CliMA/Oceananigans.jl/issues/3342) - -# Order of halo filling -# 1) Flux, Value, Gradient (TODO: remove these BC and apply them as fluxes) -# 2) Periodic (PBCT) -# 3) Shared Communication (MCBCT) -# 4) Distributed Communication (DCBCT) - -# We define "greater than" `>` and "lower than", for boundary conditions -# following the rules outlined in `fill_first` -# i.e. if `bc1 > bc2` then `bc2` precedes `bc1` in filling order -@inline Base.isless(bc1::BoundaryCondition, bc2::BoundaryCondition) = fill_first(bc1, bc2) - -# fallback for `Nothing` BC. -@inline Base.isless(::Nothing, ::Nothing) = true -@inline Base.isless(::BoundaryCondition, ::Nothing) = false -@inline Base.isless(::Nothing, ::BoundaryCondition) = true -@inline Base.isless(::BoundaryCondition, ::Missing) = false -@inline Base.isless(::Missing, ::BoundaryCondition) = true - -fill_first(bc1::DCBCT, bc2) = false -fill_first(bc1::PBCT, bc2::DCBCT) = true -fill_first(bc1::DCBCT, bc2::PBCT) = false -fill_first(bc1::MCBCT, bc2::DCBCT) = true -fill_first(bc1::DCBCT, bc2::MCBCT) = false -fill_first(bc1, bc2::DCBCT) = true -fill_first(bc1::DCBCT, bc2::DCBCT) = true -fill_first(bc1::PBCT, bc2) = false -fill_first(bc1::MCBCT, bc2) = false -fill_first(bc1::PBCT, bc2::MCBCT) = true -fill_first(bc1::MCBCT, bc2::PBCT) = false -fill_first(bc1, bc2::PBCT) = true -fill_first(bc1, bc2::MCBCT) = true -fill_first(bc1::PBCT, bc2::PBCT) = true -fill_first(bc1::MCBCT, bc2::MCBCT) = true -fill_first(bc1, bc2) = true +@inline _fill_west_halo!(j, k, grid, c, ::Nothing, args...) = nothing +@inline _fill_east_halo!(j, k, grid, c, ::Nothing, args...) = nothing +@inline _fill_south_halo!(i, k, grid, c, ::Nothing, args...) = nothing +@inline _fill_north_halo!(i, k, grid, c, ::Nothing, args...) = nothing +@inline _fill_bottom_halo!(i, j, grid, c, ::Nothing, args...) = nothing +@inline _fill_top_halo!(i, j, grid, c, ::Nothing, args...) = nothing ##### ##### Double-sided fill_halo! kernels @@ -251,104 +107,13 @@ end _fill_top_halo!(i, j, grid, c, bc, loc, args...) end -##### -##### Tupled double-sided fill_halo! kernels -##### - -# Note, we do not need tupled single-sided fill_halo! kernels since `DCBC` do not -# support tupled halo filling -import Oceananigans.Utils: @constprop - -@kernel function _fill_west_and_east_halo!(c::Tuple, west_bc, east_bc, loc, grid, args) - j, k = @index(Global, NTuple) - ntuple(Val(length(west_bc))) do n - Base.@_inline_meta - @constprop(:aggressive) # TODO constprop failure on `loc[n]` - @inbounds begin - _fill_west_halo!(j, k, grid, c[n], west_bc[n], loc[n], args...) - _fill_east_halo!(j, k, grid, c[n], east_bc[n], loc[n], args...) - end - end -end - -@kernel function _fill_south_and_north_halo!(c::Tuple, south_bc, north_bc, loc, grid, args) - i, k = @index(Global, NTuple) - ntuple(Val(length(south_bc))) do n - Base.@_inline_meta - @constprop(:aggressive) # TODO constprop failure on `loc[n]` - @inbounds begin - _fill_south_halo!(i, k, grid, c[n], south_bc[n], loc[n], args...) - _fill_north_halo!(i, k, grid, c[n], north_bc[n], loc[n], args...) - end - end -end - -@kernel function _fill_bottom_and_top_halo!(c::Tuple, bottom_bc, top_bc, loc, grid, args) - i, j = @index(Global, NTuple) - ntuple(Val(length(bottom_bc))) do n - Base.@_inline_meta - @constprop(:aggressive) # TODO constprop failure on `loc[n]` - @inbounds begin - _fill_bottom_halo!(i, j, grid, c[n], bottom_bc[n], loc[n], args...) - _fill_top_halo!(i, j, grid, c[n], top_bc[n], loc[n], args...) - end - end -end - -##### -##### Kernel launchers for single-sided fill_halos -##### - -fill_west_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) = - launch!(arch, grid, KernelParameters(size, offset), - _fill_only_west_halo!, c, bc, loc, grid, Tuple(args); kwargs...) - -fill_east_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) = - launch!(arch, grid, KernelParameters(size, offset), - _fill_only_east_halo!, c, bc, loc, grid, Tuple(args); kwargs...) - -fill_south_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) = - launch!(arch, grid, KernelParameters(size, offset), - _fill_only_south_halo!, c, bc, loc, grid, Tuple(args); kwargs...) - -fill_north_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) = - launch!(arch, grid, KernelParameters(size, offset), - _fill_only_north_halo!, c, bc, loc, grid, Tuple(args); kwargs...) - -fill_bottom_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) = - launch!(arch, grid, KernelParameters(size, offset), - _fill_only_bottom_halo!, c, bc, loc, grid, Tuple(args); kwargs...) - -fill_top_halo!(c, bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) = - launch!(arch, grid, KernelParameters(size, offset), - _fill_only_top_halo!, c, bc, loc, grid, Tuple(args); kwargs...) - -##### -##### Kernel launchers for double-sided fill_halos -##### - -function fill_west_and_east_halo!(c, west_bc, east_bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) - return launch!(arch, grid, KernelParameters(size, offset), - _fill_west_and_east_halo!, c, west_bc, east_bc, loc, grid, Tuple(args); kwargs...) -end - -function fill_south_and_north_halo!(c, south_bc, north_bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) - return launch!(arch, grid, KernelParameters(size, offset), - _fill_south_and_north_halo!, c, south_bc, north_bc, loc, grid, Tuple(args); kwargs...) -end - -function fill_bottom_and_top_halo!(c, bottom_bc, top_bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) - return launch!(arch, grid, KernelParameters(size, offset), - _fill_bottom_and_top_halo!, c, bottom_bc, top_bc, loc, grid, Tuple(args); kwargs...) -end - ##### ##### Calculate kernel size and offset for Windowed and Sliced Fields ##### -const WEB = Union{typeof(fill_west_and_east_halo!), typeof(fill_west_halo!), typeof(fill_east_halo!)} -const SNB = Union{typeof(fill_south_and_north_halo!), typeof(fill_south_halo!), typeof(fill_north_halo!)} -const TBB = Union{typeof(fill_bottom_and_top_halo!), typeof(fill_bottom_halo!), typeof(fill_top_halo!)} +const WEB = Union{WestAndEast, West, East} +const SNB = Union{SouthAndNorth, South, North} +const TBB = Union{BottomAndTop, Bottom, Top} # Tupled halo filling _only_ deals with full fields! @inline fill_halo_size(::Tuple, ::WEB, args...) = :yz diff --git a/src/BoundaryConditions/fill_halo_regions_flux.jl b/src/BoundaryConditions/fill_halo_regions_flux.jl index 64cbaf9fb21..dd5a9612850 100644 --- a/src/BoundaryConditions/fill_halo_regions_flux.jl +++ b/src/BoundaryConditions/fill_halo_regions_flux.jl @@ -1,11 +1,6 @@ -##### Kernels that ensure 'no-flux' from second- and fourth-order diffusion operators. -##### Kernel functions that ensure 'no-flux' from second- and fourth-order diffusion operators. +##### Kernel functions that ensure 'no-flux' from second-order diffusion operators. ##### Note that flux divergence associated with a flux boundary condition is added ##### in a separate step. -##### -##### We implement two sets of kernel functions: one for filling one boundary at a time, and -##### a second that fills both boundaries at the same as a performance optimization. -##### ##### ##### Low-level functions that set data @@ -20,7 +15,6 @@ @inline _fill_flux_bottom_halo!(i, j, k, grid, c) = @inbounds c[i, j, 1-k] = c[i, j, k] @inline _fill_flux_top_halo!(i, j, k, grid, c) = @inbounds c[i, j, grid.Nz+k] = c[i, j, grid.Nz+1-k] -##### ##### ##### Combined halo filling functions ##### diff --git a/src/BoundaryConditions/fill_halo_regions_nothing.jl b/src/BoundaryConditions/fill_halo_regions_nothing.jl deleted file mode 100644 index 5721e275ca9..00000000000 --- a/src/BoundaryConditions/fill_halo_regions_nothing.jl +++ /dev/null @@ -1,18 +0,0 @@ -##### -##### Nothing happens when your boundary condition is nothing -##### - -fill_west_and_east_halo!(c, ::Nothing, ::Nothing, args...; kwargs...) = nothing -fill_south_and_north_halo!(c,::Nothing, ::Nothing, args...; kwargs...) = nothing -fill_bottom_and_top_halo!(c, ::Nothing, ::Nothing, args...; kwargs...) = nothing - -for dir in (:west, :east, :south, :north, :bottom, :top) - fill_nothing! = Symbol( :fill_, dir, :_halo!) - alt_fill_nothing! = Symbol(:_fill_, dir, :_halo!) - @eval begin - @inline $fill_nothing!(c, ::Nothing, args...; kwargs...) = nothing - @inline $alt_fill_nothing!(i, j, grid, c, ::Nothing, args...) = nothing - @inline $alt_fill_nothing!(i, j, grid, ::Nothing, ::Nothing, args...) = nothing - @inline $alt_fill_nothing!(i, j, grid, ::Nothing, args...) = nothing - end -end diff --git a/src/BoundaryConditions/fill_halo_regions_open.jl b/src/BoundaryConditions/fill_halo_regions_open.jl index 807ba4a345e..293f2fe6ede 100644 --- a/src/BoundaryConditions/fill_halo_regions_open.jl +++ b/src/BoundaryConditions/fill_halo_regions_open.jl @@ -1,66 +1,3 @@ -@inline fill_open_boundary_regions!(field, args...) = - fill_open_boundary_regions!(field, field.boundary_conditions, field.indices, instantiated_location(field), field.grid) - -""" - fill_open_boundary_regions!(fields, boundary_conditions, indices, loc, grid, args...; kwargs...) - -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...; only_local_halos = false, kwargs...) - arch = architecture(grid) - - # 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) - - left_bc = left_open_boundary_condition(boundary_conditions, loc) - right_bc = right_open_boundary_condition(boundary_conditions, loc) - - bcs_tuple = (left_bc, right_bc) - - if !isnothing(fill_halo!) && any(!isnothing, bcs_tuple) - - # Overwrite the `only_local_halos` keyword argument, because open boundaries - # are always local boundaries that do not require communication - only_local_halos = true - - fill_halo_event!(field, fill_halo!, bcs_tuple, indices, loc, arch, grid, args...; only_local_halos, kwargs...) - end - - return nothing -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! - -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 - -@inline retrieve_open_bc(bc::OBC) = bc -@inline retrieve_open_bc(bc) = nothing - -@inline retrieve_bc(bc::OBC) = nothing - -# for regular halo fills, return nothing if the BC is not an OBC -@inline left_open_boundary_condition(boundary_conditions, 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) - # Open 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...) @@ -68,3 +5,10 @@ end @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...) + +@inline function fill_halo_event!(c, kernel!, bcs::OBCTC, loc, grid, args...; fill_open_bcs=true, kwargs...) + if fill_open_bcs + return kernel!(c, bcs..., loc, grid, Tuple(args)) + end + return nothing +end diff --git a/src/BoundaryConditions/fill_halo_regions_periodic.jl b/src/BoundaryConditions/fill_halo_regions_periodic.jl index 3fe1d1f5597..ba51a40cbe9 100644 --- a/src/BoundaryConditions/fill_halo_regions_periodic.jl +++ b/src/BoundaryConditions/fill_halo_regions_periodic.jl @@ -1,122 +1,33 @@ -using KernelAbstractions.Extras.LoopInfo: @unroll - -##### -##### Periodic boundary conditions -##### - -@inline parent_size_and_offset(c, dim1, dim2, size, offset) = (parent(c), size, fix_halo_offsets.(offset, c.offsets[[dim1, dim2]])) -@inline parent_size_and_offset(c, dim1, dim2, ::Symbol, offset) = (parent(c), size(parent(c))[[dim1, dim2]], (0, 0)) - -@inline function parent_size_and_offset(c::Tuple, dim1, dim2, ::Symbol, offset) - p = parent.(c) - p_size = (minimum([size(t, dim1) for t in p]), minimum([size(t, dim2) for t in p])) - return p, p_size, (0, 0) -end - -@inline fix_halo_offsets(o, co) = co > 0 ? o - co : o # Windowed fields have only positive offsets to correct - -function fill_west_and_east_halo!(c, ::PBCT, ::PBCT, size, offset, loc, arch, grid, args...; only_local_halos = false, kw...) - c_parent, yz_size, offset = parent_size_and_offset(c, 2, 3, size, offset) - launch!(arch, grid, KernelParameters(yz_size, offset), fill_periodic_west_and_east_halo!, c_parent, Val(grid.Hx), grid.Nx; kw...) - return nothing -end - -function fill_south_and_north_halo!(c, ::PBCT, ::PBCT, size, offset, loc, arch, grid, args...; only_local_halos = false, kw...) - c_parent, xz_size, offset = parent_size_and_offset(c, 1, 3, size, offset) - launch!(arch, grid, KernelParameters(xz_size, offset), fill_periodic_south_and_north_halo!, c_parent, Val(grid.Hy), grid.Ny; kw...) - return nothing -end - -function fill_bottom_and_top_halo!(c, ::PBCT, ::PBCT, size, offset, loc, arch, grid, args...; only_local_halos = false, kw...) - c_parent, xy_size, offset = parent_size_and_offset(c, 1, 2, size, offset) - launch!(arch, grid, KernelParameters(xy_size, offset), fill_periodic_bottom_and_top_halo!, c_parent, Val(grid.Hz), grid.Nz; kw...) - return nothing -end - ##### ##### Periodic boundary condition kernels ##### -@kernel function fill_periodic_west_and_east_halo!(c, ::Val{H}, N) where H - j, k = @index(Global, NTuple) - @unroll for i = 1:H - @inbounds begin - c[i, j, k] = c[N+i, j, k] # west - c[N+H+i, j, k] = c[H+i, j, k] # east - end - end -end - -@kernel function fill_periodic_south_and_north_halo!(c, ::Val{H}, N) where H - i, k = @index(Global, NTuple) - @unroll for j = 1:H - @inbounds begin - c[i, j, k] = c[i, N+j, k] # south - c[i, N+H+j, k] = c[i, H+j, k] # north - end - end -end - -@kernel function fill_periodic_bottom_and_top_halo!(c, ::Val{H}, N) where H - i, j = @index(Global, NTuple) - @unroll for k = 1:H - @inbounds begin - c[i, j, k] = c[i, j, N+k] # top - c[i, j, N+H+k] = c[i, j, H+k] # bottom - end - end -end - -#### -#### Tupled periodic boundary condition -#### - -@kernel function fill_periodic_west_and_east_halo!(c::Tuple, ::Val{H}, N) where {H} +@kernel function _fill_periodic_west_and_east_halo!(c, west_bc, east_bc, loc, grid, args) j, k = @index(Global, NTuple) - ntuple(Val(length(c))) do n - Base.@_inline_meta - @unroll for i = 1:H - @inbounds begin - c[n][i, j, k] = c[n][N+i, j, k] # west - c[n][N+H+i, j, k] = c[n][H+i, j, k] # east - end - end + H = grid.Hx + N = grid.Nx + @inbounds for i = 1:H + parent(c)[i, j, k] = parent(c)[N+i, j, k] # west + parent(c)[N+H+i, j, k] = parent(c)[H+i, j, k] # east end end -@kernel function fill_periodic_south_and_north_halo!(c::Tuple, ::Val{H}, N) where {H} +@kernel function _fill_periodic_south_and_north_halo!(c, south_bc, north_bc, loc, grid, args) i, k = @index(Global, NTuple) - ntuple(Val(length(c))) do n - Base.@_inline_meta - @unroll for j = 1:H - @inbounds begin - c[n][i, j, k] = c[n][i, N+j, k] # south - c[n][i, N+H+j, k] = c[n][i, H+j, k] # north - end - end + H = grid.Hy + N = grid.Ny + @inbounds for j = 1:H + parent(c)[i, j, k] = parent(c)[i, N+j, k] # south + parent(c)[i, N+H+j, k] = parent(c)[i, H+j, k] # north end end -@kernel function fill_periodic_bottom_and_top_halo!(c::Tuple, ::Val{H}, N) where {H} +@kernel function _fill_periodic_bottom_and_top_halo!(c, bottom_bc, top_bc, loc, grid, args) i, j = @index(Global, NTuple) - ntuple(Val(length(c))) do n - Base.@_inline_meta - @unroll for k = 1:H - @inbounds begin - c[n][i, j, k] = c[n][i, j, N+k] # top - c[n][i, j, N+H+k] = c[n][i, j, H+k] # bottom - end - end + H = grid.Hz + N = grid.Nz + @inbounds for k = 1:H + parent(c)[i, j, k] = parent(c)[i, j, N+k] # top + parent(c)[i, j, N+H+k] = parent(c)[i, j, H+k] # bottom end -end - -##### -##### Throw error if single-sided periodic boundary conditions are used -##### - -fill_west_halo!(c, ::PBCT, args...; kwargs...) = throw(ArgumentError("Periodic boundary conditions must be applied to both sides")) -fill_east_halo!(c, ::PBCT, args...; kwargs...) = throw(ArgumentError("Periodic boundary conditions must be applied to both sides")) -fill_south_halo!(c, ::PBCT, args...; kwargs...) = throw(ArgumentError("Periodic boundary conditions must be applied to both sides")) -fill_north_halo!(c, ::PBCT, args...; kwargs...) = throw(ArgumentError("Periodic boundary conditions must be applied to both sides")) -fill_bottom_halo!(c, ::PBCT, args...; kwargs...) = throw(ArgumentError("Periodic boundary conditions must be applied to both sides")) -fill_top_halo!(c, ::PBCT, args...; kwargs...) = throw(ArgumentError("Periodic boundary conditions must be applied to both sides")) +end \ No newline at end of file diff --git a/src/BoundaryConditions/polar_boundary_condition.jl b/src/BoundaryConditions/polar_boundary_condition.jl index 4cdd5a7425e..a74bee285fc 100644 --- a/src/BoundaryConditions/polar_boundary_condition.jl +++ b/src/BoundaryConditions/polar_boundary_condition.jl @@ -50,39 +50,33 @@ function update_pole_value!(bc::PolarValue, c, grid, loc) Nz = size(c, 3) Oz = c.offsets[3] params = KernelParameters(1:1, 1:1, 1+Oz:Nz+Oz) - launch!(architecture(grid), grid, params, _average_pole_value!, bc.data, c, j, grid, loc) + launch!(architecture(bc.data), grid, params, _average_pole_value!, bc.data, c, j, grid, loc) return nothing end -function fill_south_halo!(c, bc::PolarBoundaryCondition, size, offset, loc, arch, grid, args...; only_local_halos=false, kwargs...) - update_pole_value!(bc.condition, c, grid, loc) - return launch!(arch, grid, KernelParameters(size, offset), - _fill_only_south_halo!, c, bc, loc, grid, Tuple(args); kwargs...) -end +const SouthPolarBC = Tuple{<:PolarBoundaryCondition, <:BoundaryCondition} +const NorthPolarBC = Tuple{<:BoundaryCondition, <:PolarBoundaryCondition} +const SouthAndNorthPolarBC = Tuple{<:PolarBoundaryCondition, <:PolarBoundaryCondition} -function fill_north_halo!(c, bc::PolarBoundaryCondition, size, offset, loc, arch, grid, args...; only_local_halos=false, kwargs...) +# fill_halo_event!(c, kernels![task], bcs[task], loc, grid, args...; kwargs...) +function fill_halo_event!(c, kernel!, bc::PolarBoundaryCondition, loc, grid, args...; kwargs...) update_pole_value!(bc.condition, c, grid, loc) - return launch!(arch, grid, KernelParameters(size, offset), - _fill_only_north_halo!, c, bc, loc, grid, Tuple(args); kwargs...) + return kernel!(c, bc, loc, grid, Tuple(args)) end -function fill_south_and_north_halo!(c, south_bc::PolarBoundaryCondition, north_bc, size, offset, loc, arch, grid, args...; only_local_halos=false, kwargs...) - update_pole_value!(south_bc.condition, c, grid, loc) - return launch!(arch, grid, KernelParameters(size, offset), - _fill_south_and_north_halo!, c, south_bc, north_bc, loc, grid, Tuple(args); kwargs...) +function fill_halo_event!(c, kernel!, bcs::SouthPolarBC, loc, grid, args...; kwargs...) + update_pole_value!(bcs[1].condition, c, grid, loc) + return kernel!(c, bcs..., loc, grid, Tuple(args)) end -function fill_south_and_north_halo!(c, south_bc, north_bc::PolarBoundaryCondition, size, offset, loc, arch, grid, args...; only_local_halos=false, kwargs...) - update_pole_value!(north_bc.condition, c, grid, loc) - return launch!(arch, grid, KernelParameters(size, offset), - _fill_south_and_north_halo!, c, south_bc, north_bc, loc, grid, Tuple(args); kwargs...) +function fill_halo_event!(c, kernel!, bcs::NorthPolarBC, loc, grid, args...; kwargs...) + update_pole_value!(bcs[2].condition, c, grid, loc) + return kernel!(c, bcs..., loc, grid, Tuple(args)) end -function fill_south_and_north_halo!(c, south_bc::PolarBoundaryCondition, north_bc::PolarBoundaryCondition, size, offset, loc, arch, grid, args...; only_local_halos=false, kwargs...) - update_pole_value!(south_bc.condition, c, grid, loc) - update_pole_value!(north_bc.condition, c, grid, loc) - return launch!(arch, grid, KernelParameters(size, offset), - _fill_south_and_north_halo!, c, south_bc, north_bc, loc, grid, Tuple(args); kwargs...) +function fill_halo_event!(c, kernel!, bcs::SouthAndNorthPolarBC, loc, grid, args...; kwargs...) + update_pole_value!(bcs[1].condition, c, grid, loc) + update_pole_value!(bcs[2].condition, c, grid, loc) + return kernel!(c, bcs..., loc, grid, Tuple(args)) end - diff --git a/src/DistributedComputations/communication_buffers.jl b/src/DistributedComputations/communication_buffers.jl index bb2200c7efb..fd26eb0b2ed 100644 --- a/src/DistributedComputations/communication_buffers.jl +++ b/src/DistributedComputations/communication_buffers.jl @@ -1,3 +1,4 @@ +using Oceananigans.BoundaryConditions using Oceananigans.BoundaryConditions: FieldBoundaryConditions, BoundaryCondition using Oceananigans.BoundaryConditions: MultiRegionCommunication, DistributedCommunication using Oceananigans.Grids: halo_size, size @@ -151,22 +152,22 @@ end ##### Single sided fill_send_buffers! ##### -fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:west}) = +fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::West) = _fill_west_send_buffer!(parent(c), buff, buff.west, halo_size(grid)[[1, 2]]..., size(grid)[[1, 2]]...) -fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:east}) = +fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::East) = _fill_east_send_buffer!(parent(c), buff, buff.east, halo_size(grid)[[1, 2]]..., size(grid)[[1, 2]]...) -fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:south}) = +fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::South) = _fill_south_send_buffer!(parent(c), buff, buff.south, halo_size(grid)[[1, 2]]..., size(grid)[[1, 2]]...) -fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:north}) = +fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::North) = _fill_north_send_buffer!(parent(c), buff, buff.north, halo_size(grid)[[1, 2]]..., size(grid)[[1, 2]]...) -fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:bottom}) = nothing -fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:top}) = nothing +fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Bottom) = nothing +fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Top) = nothing ##### ##### Double sided fill_send_buffers! ##### -function fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:west_and_east}) +function fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::WestAndEast) Hx, Hy, _ = halo_size(grid) Nx, Ny, _ = size(grid) @@ -174,7 +175,7 @@ function fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, :: _fill_east_send_buffer!(parent(c), buff, buff.east, Hx, Hy, Nx, Ny) end -function fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:south_and_north}) +function fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::SouthAndNorth) Hx, Hy, _ = halo_size(grid) Nx, Ny, _ = size(grid) @@ -182,7 +183,7 @@ function fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, :: _fill_north_send_buffer!(parent(c), buff, buff.north, Hx, Hy, Nx, Ny) end -fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:bottom_and_top}) = nothing +fill_send_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::BottomAndTop) = nothing """ recv_from_buffers!(c::OffsetArray, buffers::CommunicationBuffers, grid) @@ -222,22 +223,22 @@ end ##### Single sided recv_from_buffers! ##### -recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:west}) = +recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::West) = _recv_from_west_buffer!(parent(c), buff, buff.west, halo_size(grid)[[1, 2]]..., size(grid)[[1, 2]]...) -recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:east}) = +recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::East) = _recv_from_east_buffer!(parent(c), buff, buff.east, halo_size(grid)[[1, 2]]..., size(grid)[[1, 2]]...) -recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:south}) = +recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::South) = _recv_from_south_buffer!(parent(c), buff, buff.south, halo_size(grid)[[1, 2]]..., size(grid)[[1, 2]]...) -recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:north}) = +recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::North) = _recv_from_north_buffer!(parent(c), buff, buff.north, halo_size(grid)[[1, 2]]..., size(grid)[[1, 2]]...) -recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:bottom}) = nothing -recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:top}) = nothing +recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Bottom) = nothing +recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Top) = nothing ##### ##### Double sided recv_from_buffers! ##### -function recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:west_and_east}) +function recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::WestAndEast) Hx, Hy, _ = halo_size(grid) Nx, Ny, _ = size(grid) @@ -247,7 +248,7 @@ function recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, :: return nothing end -function recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:south_and_north}) +function recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::SouthAndNorth) Hx, Hy, _ = halo_size(grid) Nx, Ny, _ = size(grid) @@ -257,7 +258,7 @@ function recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, :: return nothing end -recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::Val{:bottom_and_top}) = nothing +recv_from_buffers!(c::OffsetArray, buff::CommunicationBuffers, grid, ::BottomAndTop) = nothing ##### ##### Individual _fill_send_buffers and _recv_from_buffer kernels diff --git a/src/DistributedComputations/halo_communication.jl b/src/DistributedComputations/halo_communication.jl index b15e98b3fdb..f14222dda05 100644 --- a/src/DistributedComputations/halo_communication.jl +++ b/src/DistributedComputations/halo_communication.jl @@ -1,25 +1,13 @@ using KernelAbstractions: @kernel, @index -using Oceananigans.Fields: reduced_dimensions, - instantiated_location, - fill_reduced_field_halos! - -import Oceananigans.Fields: tupled_fill_halo_regions! +using Oceananigans.Fields: instantiated_location +using Oceananigans.BoundaryConditions using Oceananigans.BoundaryConditions: - fill_halo_size, - fill_halo_offset, - permute_boundary_conditions, - fill_open_boundary_regions!, - PBCT, DCBCT # tuples - -import Oceananigans.BoundaryConditions: - fill_halo_regions!, fill_first, fill_halo_event!, - fill_west_halo!, fill_east_halo!, fill_south_halo!, - fill_north_halo!, fill_bottom_halo!, fill_top_halo!, - fill_west_and_east_halo!, - fill_south_and_north_halo!, - fill_bottom_and_top_halo! + DistributedFillHalo, + PBCT, DCBCT, get_boundary_kernels + +import Oceananigans.BoundaryConditions: fill_halo_event!, fill_halo_regions! ##### ##### MPI tags for halo communication BCs @@ -86,43 +74,27 @@ end ##### Filling halos for halo communication boundary conditions ##### -function tupled_fill_halo_regions!(fields, grid::DistributedGrid, args...; kwargs...) - not_reduced_fields = fill_reduced_field_halos!(fields, args...; kwargs) - - for field in not_reduced_fields - # Make sure we are filling a `Field` type. - field isa Field && fill_halo_regions!(field, args...; kwargs...) - end -end +fill_halo_regions!(field::DistributedField, args...; kwargs...) = + fill_halo_regions!(field.data, + field.boundary_conditions, + field.indices, + instantiated_location(field), + field.grid, + field.communication_buffers, + args...; + kwargs...) -function fill_halo_regions!(field::DistributedField, args...; kwargs...) - reduced_dims = reduced_dimensions(field) - - return fill_halo_regions!(field.data, - field.boundary_conditions, - field.indices, - instantiated_location(field), - field.grid, - field.communication_buffers, - args...; - reduced_dimensions = reduced_dims, - kwargs...) -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, buffers, args...; kwargs...) - end +function fill_halo_regions!(c::OffsetArray, boundary_conditions, indices, loc, grid::DistributedGrid, buffers, args...; + fill_open_bcs=true, kwargs...) arch = architecture(grid) - fill_halos!, bcs = permute_boundary_conditions(bcs) - number_of_tasks = length(fill_halos!) + kernels!, bcs = get_boundary_kernels(boundary_conditions, c, grid, loc, indices) + + number_of_tasks = length(kernels!) outstanding_requests = length(arch.mpi_requests) for task = 1:number_of_tasks - fill_halo_event!(c, fill_halos![task], bcs[task], indices, loc, arch, grid, buffers, args...; kwargs...) + fill_halo_event!(c, kernels![task], bcs[task], loc, grid, buffers, args...; kwargs...) end fill_corners!(c, arch.connectivity, indices, loc, arch, grid, buffers, args...; kwargs...) @@ -156,7 +128,7 @@ end # Reset MPI tag arch.mpi_tag[] -= arch.mpi_tag[] - recv_from_buffers!(c, buffers, grid, Val(side)) + recv_from_buffers!(c, buffers, grid, side) return nothing end @@ -184,43 +156,31 @@ function fill_corners!(c, connectivity, indices, loc, arch, grid, buffers, args. !isnothing(reqnw) && push!(requests, reqnw...) !isnothing(reqne) && push!(requests, reqne...) - pool_requests_or_complete_comm!(c, arch, grid, buffers, requests, async, :corners) + pool_requests_or_complete_comm!(c, arch, grid, buffers, requests, async, Val(:corners)) return nothing end -@inline communication_side(::Val{fill_west_and_east_halo!}) = :west_and_east -@inline communication_side(::Val{fill_south_and_north_halo!}) = :south_and_north -@inline communication_side(::Val{fill_bottom_and_top_halo!}) = :bottom_and_top -@inline communication_side(::Val{fill_west_halo!}) = :west -@inline communication_side(::Val{fill_east_halo!}) = :east -@inline communication_side(::Val{fill_south_halo!}) = :south -@inline communication_side(::Val{fill_north_halo!}) = :north -@inline communication_side(::Val{fill_bottom_halo!}) = :bottom -@inline communication_side(::Val{fill_top_halo!}) = :top - cooperative_wait(req::MPI.Request) = MPI.Waitall(req) cooperative_waitall!(req::Array{MPI.Request}) = MPI.Waitall(req) # There are two additional keyword arguments (with respect to serial `fill_halo_event!`s) that take an effect on `DistributedGrids`: # - only_local_halos: if true, only the local halos are filled, i.e. corresponding to non-communicating boundary conditions # - async: if true, ansynchronous MPI communication is enabled -function fill_halo_event!(c, fill_halos!, bcs, indices, loc, arch, grid::DistributedGrid, buffers, args...; +function fill_halo_event!(c, kernel!::DistributedFillHalo, bcs, loc, grid::DistributedGrid, buffers, args...; async = false, only_local_halos = false, kwargs...) - buffer_side = communication_side(Val(fill_halos!)) - - if !only_local_halos # Then we need to fill the `send` buffers - fill_send_buffers!(c, buffers, grid, Val(buffer_side)) + if only_local_halos # No need to do anything here + return nothing end - # Calculate size and offset of the fill_halo kernel - # We assume that the kernel size is the same for west and east boundaries, - # south and north boundaries and bottom and top boundaries - size = fill_halo_size(c, fill_halos!, indices, bcs[1], loc, grid) - offset = fill_halo_offset(size, fill_halos!, indices) + buffer_side = kernel!.side + arch = architecture(grid) + + fill_send_buffers!(c, buffers, grid, buffer_side) + sync_device!(arch) # We need to synchronize the device before we start the communication - requests = fill_halos!(c, bcs..., size, offset, loc, arch, grid, buffers, args...; only_local_halos, kwargs...) + requests = kernel!(c, bcs..., loc, grid, arch, buffers) pool_requests_or_complete_comm!(c, arch, grid, buffers, requests, async, buffer_side) return nothing @@ -233,7 +193,7 @@ end for side in [:southwest, :southeast, :northwest, :northeast] fill_corner_halo! = Symbol("fill_$(side)_halo!") send_side_halo = Symbol("send_$(side)_halo") - recv_and_fill_side_halo! = Symbol("recv_and_fill_$(side)_halo!") + recv_side_halo! = Symbol("recv_$(side)_halo!") @eval begin $fill_corner_halo!(c, corner, indices, loc, arch, grid, buffers, ::Nothing, args...; kwargs...) = nothing @@ -242,7 +202,7 @@ for side in [:southwest, :southeast, :northwest, :northeast] child_arch = child_architecture(arch) local_rank = arch.local_rank - recv_req = $recv_and_fill_side_halo!(c, grid, arch, loc, local_rank, corner, buffers) + recv_req = $recv_side_halo!(c, grid, arch, loc, local_rank, corner, buffers) send_req = $send_side_halo(c, grid, arch, loc, local_rank, corner, buffers) return [send_req, recv_req] @@ -254,62 +214,71 @@ end ##### Double-sided Distributed fill_halo!s ##### -for (side, opposite_side) in zip([:west, :south], [:east, :north]) - fill_both_halo! = Symbol("fill_$(side)_and_$(opposite_side)_halo!") - send_side_halo = Symbol("send_$(side)_halo") - send_opposite_side_halo = Symbol("send_$(opposite_side)_halo") - recv_and_fill_side_halo! = Symbol("recv_and_fill_$(side)_halo!") - recv_and_fill_opposite_side_halo! = Symbol("recv_and_fill_$(opposite_side)_halo!") +function (::DistributedFillHalo{<:WestAndEast})(c, west_bc, east_bc, loc, grid, arch, buffers) + @assert west_bc.condition.from == east_bc.condition.from # Extra protection in case of bugs + local_rank = west_bc.condition.from - @eval begin - function $fill_both_halo!(c, bc_side::DCBCT, bc_opposite_side::DCBCT, size, offset, loc, arch::Distributed, - grid::DistributedGrid, buffers, args...; only_local_halos = false, kwargs...) + recv_req1 = recv_west_halo!(c, grid, arch, loc, local_rank, west_bc.condition.to, buffers) + recv_req2 = recv_east_halo!(c, grid, arch, loc, local_rank, east_bc.condition.to, buffers) - only_local_halos && return nothing + send_req1 = send_west_halo(c, grid, arch, loc, local_rank, west_bc.condition.to, buffers) + send_req2 = send_east_halo(c, grid, arch, loc, local_rank, east_bc.condition.to, buffers) - sync_device!(arch) + return [send_req1, send_req2, recv_req1, recv_req2] +end - @assert bc_side.condition.from == bc_opposite_side.condition.from # Extra protection in case of bugs - local_rank = bc_side.condition.from +function (::DistributedFillHalo{<:SouthAndNorth})(c, south_bc, north_bc, loc, grid, arch, buffers) + @assert south_bc.condition.from == north_bc.condition.from # Extra protection in case of bugs + local_rank = south_bc.condition.from - recv_req1 = $recv_and_fill_side_halo!(c, grid, arch, loc, local_rank, bc_side.condition.to, buffers) - recv_req2 = $recv_and_fill_opposite_side_halo!(c, grid, arch, loc, local_rank, bc_opposite_side.condition.to, buffers) + recv_req1 = recv_south_halo!(c, grid, arch, loc, local_rank, south_bc.condition.to, buffers) + recv_req2 = recv_north_halo!(c, grid, arch, loc, local_rank, north_bc.condition.to, buffers) - send_req1 = $send_side_halo(c, grid, arch, loc, local_rank, bc_side.condition.to, buffers) - send_req2 = $send_opposite_side_halo(c, grid, arch, loc, local_rank, bc_opposite_side.condition.to, buffers) + send_req1 = send_south_halo(c, grid, arch, loc, local_rank, south_bc.condition.to, buffers) + send_req2 = send_north_halo(c, grid, arch, loc, local_rank, north_bc.condition.to, buffers) - return [send_req1, send_req2, recv_req1, recv_req2] - end - end + return [send_req1, send_req2, recv_req1, recv_req2] end ##### ##### Single-sided Distributed fill_halo!s ##### -for side in [:west, :east, :south, :north] - fill_side_halo! = Symbol("fill_$(side)_halo!") - send_side_halo = Symbol("send_$(side)_halo") - recv_and_fill_side_halo! = Symbol("recv_and_fill_$(side)_halo!") - - @eval begin - function $fill_side_halo!(c, bc_side::DCBCT, size, offset, loc, arch::Distributed, grid::DistributedGrid, - buffers, args...; only_local_halos = false, kwargs...) +function (::DistributedFillHalo{<:West})(c, bc, loc, grid, arch, buffers) + local_rank = bc.condition.from + recv_req = recv_west_halo!(c, grid, arch, loc, local_rank, bc.condition.to, buffers) + send_req = send_west_halo(c, grid, arch, loc, local_rank, bc.condition.to, buffers) + return [send_req, recv_req] +end - only_local_halos && return nothing +function (::DistributedFillHalo{<:East})(c, bc, loc, grid, arch, buffers) + local_rank = bc.condition.from + recv_req = recv_east_halo!(c, grid, arch, loc, local_rank, bc.condition.to, buffers) + send_req = send_east_halo(c, grid, arch, loc, local_rank, bc.condition.to, buffers) + return [send_req, recv_req] +end - sync_device!(arch) +function (::DistributedFillHalo{<:South})(c, bc, loc, grid, arch, buffers) + local_rank = bc.condition.from + recv_req = recv_south_halo!(c, grid, arch, loc, local_rank, bc.condition.to, buffers) + send_req = send_south_halo(c, grid, arch, loc, local_rank, bc.condition.to, buffers) + return [send_req, recv_req] +end - child_arch = child_architecture(arch) - local_rank = bc_side.condition.from +function (::DistributedFillHalo{<:North})(c, bc, loc, grid, arch, buffers) + local_rank = bc.condition.from + recv_req = recv_north_halo!(c, grid, arch, loc, local_rank, bc.condition.to, buffers) + send_req = send_north_halo(c, grid, arch, loc, local_rank, bc.condition.to, buffers) + return [send_req, recv_req] +end - recv_req = $recv_and_fill_side_halo!(c, grid, arch, loc, local_rank, bc_side.condition.to, buffers) - send_req = $send_side_halo(c, grid, arch, loc, local_rank, bc_side.condition.to, buffers) +##### +##### No communication in the vertical direction +##### - return [send_req, recv_req] - end - end -end +(::DistributedFillHalo{<:BottomAndTop})(args...) = nothing +(::DistributedFillHalo{<:Bottom})(args...) = nothing +(::DistributedFillHalo{<:Top})(args...) = nothing ##### ##### Sending halos @@ -343,13 +312,13 @@ end for side in sides side_str = string(side) - recv_and_fill_side_halo! = Symbol("recv_and_fill_$(side)_halo!") + recv_side_halo! = Symbol("recv_$(side)_halo!") underlying_side_halo = Symbol("underlying_$(side)_halo") side_recv_tag = Symbol("$(side)_recv_tag") get_side_recv_buffer = Symbol("get_$(side)_recv_buffer") @eval begin - function $recv_and_fill_side_halo!(c, grid, arch, location, local_rank, rank_to_recv_from, buffers) + function $recv_side_halo!(c, grid, arch, location, local_rank, rank_to_recv_from, buffers) recv_buffer = $get_side_recv_buffer(c, grid, buffers, arch) recv_tag = $side_recv_tag(arch, grid, location) diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 250cf4e789c..eeb245d7b13 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -1,6 +1,7 @@ -using Oceananigans.BoundaryConditions: OBC, MCBC, BoundaryCondition, Zipper +using Oceananigans.BoundaryConditions: OBC, MCBC, BoundaryCondition, Zipper, construct_boundary_conditions_kernels using Oceananigans.Grids: parent_index_range, index_range_offset, default_indices, all_indices, validate_indices using Oceananigans.Grids: index_range_contains +using Oceananigans.Architectures: convert_to_device using Adapt using LinearAlgebra @@ -31,7 +32,8 @@ struct Field{LX, LY, LZ, O, G, I, D, T, B, S, F} <: AbstractField{LX, LY, LZ, G, # Inner constructor that does not validate _anything_! function Field{LX, LY, LZ}(grid::G, data::D, bcs::B, indices::I, op::O, status::S, buffers::F) where {LX, LY, LZ, G, D, B, O, S, I, F} T = eltype(data) - return new{LX, LY, LZ, O, G, I, D, T, B, S, F}(grid, data, bcs, indices, op, status, buffers) + @apply_regionally new_bcs = construct_boundary_conditions_kernels(bcs, data, grid, (LX(), LY(), LZ()), indices) # Adding the kernels to the bcs + return new{LX, LY, LZ, O, G, I, D, T, typeof(new_bcs), S, F}(grid, data, new_bcs, indices, op, status, buffers) end end @@ -804,18 +806,22 @@ end ##### fill_halo_regions! ##### -function fill_halo_regions!(field::Field, args...; kwargs...) - reduced_dims = reduced_dimensions(field) +function fill_halo_regions!(field::Field, positional_args...; kwargs...) - fill_halo_regions!(field.data, - field.boundary_conditions, - field.indices, - instantiated_location(field), - field.grid, - args...; - reduced_dimensions = reduced_dims, - kwargs...) + arch = architecture(field.grid) + args = (field.data, + field.boundary_conditions, + field.indices, + instantiated_location(field), + field.grid, + positional_args...) + + # Manually convert args... to be + # passed to the fill_halo_regions! function. + GC.@preserve args begin + converted_args = convert_to_device(arch, args) + fill_halo_regions!(converted_args...; kwargs...) + end return nothing end - diff --git a/src/Fields/field_tuples.jl b/src/Fields/field_tuples.jl index be8cf71dcbc..af47099cdee 100644 --- a/src/Fields/field_tuples.jl +++ b/src/Fields/field_tuples.jl @@ -1,3 +1,4 @@ +using Oceananigans.Grids: AbstractGrid using Oceananigans.BoundaryConditions: FieldBoundaryConditions, regularize_field_boundary_conditions ##### @@ -53,74 +54,18 @@ Fill halo regions for all `fields`. The algorithm: 4. In every direction, the halo regions in each of the remaining `Field` tuple are filled simultaneously. """ -function fill_halo_regions!(maybe_nested_tuple::Union{NamedTuple, Tuple}, args...; - signed = true, # This kwarg is active only for a `ConformalCubedSphereGrid`, here we discard it. - kwargs...) - - flattened = flattened_unique_values(maybe_nested_tuple) - - # Look for grid within the flattened field tuple: - for f in flattened - if isdefined(f, :grid) - grid = f.grid - return tupled_fill_halo_regions!(flattened, grid, args...; kwargs...) - end - end - - return tupled_fill_halo_regions!(flattened, args...; kwargs...) -end - -# Version where we find grid amongst ordinary fields: -function tupled_fill_halo_regions!(fields, args...; kwargs...) +function fill_halo_regions!(fields::Union{NamedTuple, Tuple}, args...; kwargs...) - not_reduced_fields = fill_reduced_field_halos!(fields, args...; kwargs) - - if !isempty(not_reduced_fields) # ie not reduced, and with default_indices - grid = first(not_reduced_fields).grid - fill_halo_regions!(map(data, not_reduced_fields), - map(boundary_conditions, not_reduced_fields), - default_indices(3), - map(instantiated_location, not_reduced_fields), - grid, args...; kwargs...) + for field in fields + fill_halo_regions!(field, args...; kwargs...) end return nothing end -# Version where grid is provided: -function tupled_fill_halo_regions!(fields, grid::AbstractGrid, args...; kwargs...) - - not_reduced_fields = fill_reduced_field_halos!(fields, args...; kwargs) - - if !isempty(not_reduced_fields) # ie not reduced, and with default_indices - fill_halo_regions!(map(data, not_reduced_fields), - map(boundary_conditions, not_reduced_fields), - default_indices(3), - map(instantiated_location, not_reduced_fields), - grid, args...; kwargs...) - end - - return nothing -end - -# Helper function to create the tuple of ordinary fields: -function fill_reduced_field_halos!(fields, args...; kwargs) - - not_reduced_fields = Field[] - for f in fields - bcs = boundary_conditions(f) - if !isnothing(bcs) - if f isa ReducedField || !(f isa FullField) - # Windowed and reduced fields - fill_halo_regions!(f, args...; kwargs...) - else - push!(not_reduced_fields, f) - end - end - end - - return tuple(not_reduced_fields...) -end +# This is a convenience function that allows `fill_halo_regions!` to be dispatched on the grid type. +fill_halo_regions!(fields::NamedTuple, grid::AbstractGrid, args...; signed=true, kwargs...) = fill_halo_regions!(fields, args...; kwargs...) +fill_halo_regions!(fields::Tuple, grid::AbstractGrid, args...; signed=true, kwargs...) = fill_halo_regions!(fields, args...; kwargs...) ##### ##### Tracer names diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl index a1a3dba62e9..b42a3132307 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl @@ -131,7 +131,7 @@ function compute_free_surface_tendency!(grid, model, ::SplitExplicitFreeSurface) @apply_regionally compute_split_explicit_forcing!(GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ, baroclinic_timestepper, Val(stage)) fields_to_fill = (GUⁿ, GVⁿ) - fill_halo_regions!(fields_to_fill; async = true) + fill_halo_regions!(fields_to_fill; async=true) return nothing end diff --git a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl index 5b912262f92..16587f9d77d 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl @@ -8,7 +8,6 @@ using Oceananigans.TurbulenceClosures: compute_diffusivities! using Oceananigans.ImmersedBoundaries: mask_immersed_field!, mask_immersed_field_xy!, inactive_node using Oceananigans.Models: update_model_field_time_series! using Oceananigans.Models.NonhydrostaticModels: update_hydrostatic_pressure!, p_kernel_parameters -using Oceananigans.Fields: tupled_fill_halo_regions! import Oceananigans.Models.NonhydrostaticModels: compute_auxiliaries! import Oceananigans.TimeSteppers: update_state! @@ -39,11 +38,12 @@ function update_state!(model::HydrostaticFreeSurfaceModel, grid, callbacks; comp # Update the boundary conditions @apply_regionally update_boundary_conditions!(fields(model), model) - tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model); async=true) + # Fill the halos + fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model); async=true) @apply_regionally compute_auxiliaries!(model) - fill_halo_regions!(model.diffusivity_fields; only_local_halos = true) + fill_halo_regions!(model.diffusivity_fields; only_local_halos=true) [callback(model) for callback in callbacks if callback.callsite isa UpdateStateCallsite] diff --git a/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl b/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl index 25054fa5a04..ac16bcad60a 100644 --- a/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl +++ b/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl @@ -31,8 +31,7 @@ function update_state!(model::NonhydrostaticModel, callbacks=[]; compute_tendenc update_boundary_conditions!(fields(model), model) # Fill halos for velocities and tracers - fill_halo_regions!(merge(model.velocities, model.tracers), model.clock, fields(model); - fill_boundary_normal_velocities = false, async = true) + fill_halo_regions!(merge(model.velocities, model.tracers), model.grid, model.clock, fields(model); fill_open_bcs=false, async=true) # Compute auxiliary fields for aux_field in model.auxiliary_fields diff --git a/src/MultiRegion/MultiRegion.jl b/src/MultiRegion/MultiRegion.jl index 1cff639c6de..0b1c5e19396 100644 --- a/src/MultiRegion/MultiRegion.jl +++ b/src/MultiRegion/MultiRegion.jl @@ -2,7 +2,6 @@ module MultiRegion export MultiRegionGrid, MultiRegionField export XPartition, YPartition, Connectivity -export AbstractRegionSide, East, West, North, South export CubedSpherePartition, ConformalCubedSphereGrid, CubedSphereField using Oceananigans @@ -41,13 +40,6 @@ abstract type AbstractPartition end abstract type AbstractConnectivity end -abstract type AbstractRegionSide end - -struct West <: AbstractRegionSide end -struct East <: AbstractRegionSide end -struct North <: AbstractRegionSide end -struct South <: AbstractRegionSide end - struct XPartition{N} <: AbstractPartition div :: N diff --git a/src/MultiRegion/cubed_sphere_boundary_conditions.jl b/src/MultiRegion/cubed_sphere_boundary_conditions.jl index 1050d954a86..7c0f54cc0e6 100644 --- a/src/MultiRegion/cubed_sphere_boundary_conditions.jl +++ b/src/MultiRegion/cubed_sphere_boundary_conditions.jl @@ -1,5 +1,5 @@ -using Oceananigans.BoundaryConditions: fill_halo_size, fill_halo_offset, fill_west_and_east_halo!, - fill_south_and_north_halo! +using Oceananigans.BoundaryConditions +using Oceananigans.BoundaryConditions: fill_halo_size, fill_halo_offset using Oceananigans.Fields: reduced_dimensions using Oceananigans.MultiRegion: number_of_regions @@ -12,19 +12,20 @@ function fill_halo_regions!(field::CubedSphereField{<:Center, <:Center}; kwargs. regions = Iterate(1:6) @apply_regionally fill_cubed_sphere_field_halo_event!(grid, field, multiregion_field, regions, - grid.connectivity.connections, fill_west_and_east_halo!, + grid.connectivity.connections, WestAndEast(), _fill_cubed_sphere_center_center_field_east_west_halo_regions!) @apply_regionally fill_cubed_sphere_field_halo_event!(grid, field, multiregion_field, regions, - grid.connectivity.connections, fill_south_and_north_halo!, + grid.connectivity.connections, SouthAndNorth(), _fill_cubed_sphere_center_center_field_north_south_halo_regions!) return nothing end @inline function fill_cubed_sphere_field_halo_event!(grid, field, multiregion_field, region, connections, - fill_halo_function!, _fill_halo_kernel!) - sz = fill_halo_size(field.data, fill_halo_function!, field.indices, FullyConnected, location(field), grid) - of = fill_halo_offset(sz, fill_halo_function!, field.indices) + side, _fill_halo_kernel!) + + sz = fill_halo_size(field.data, side, field.indices, FullyConnected, location(field), grid) + of = fill_halo_offset(sz, side, field.indices) kernel_parameters = KernelParameters(sz, of) reduced_dims = reduced_dimensions(field) @@ -105,10 +106,10 @@ function fill_halo_regions!(field::CubedSphereField{<:Face, <:Face}; kwargs...) regions = Iterate(1:6) @apply_regionally fill_cubed_sphere_field_halo_event!(grid, field, multiregion_field, regions, - grid.connectivity.connections, fill_west_and_east_halo!, + grid.connectivity.connections, WestAndEast(), _fill_cubed_sphere_face_face_field_east_west_halo_regions!) @apply_regionally fill_cubed_sphere_field_halo_event!(grid, field, multiregion_field, regions, - grid.connectivity.connections, fill_south_and_north_halo!, + grid.connectivity.connections, SouthAndNorth(), _fill_cubed_sphere_face_face_field_north_south_halo_regions!) return nothing @@ -220,10 +221,10 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Center, <:Center}, regions = Iterate(1:6) @apply_regionally fill_cubed_sphere_field_pairs_halo_event!(grid, field_1, multiregion_field_1, field_2, - multiregion_field_2, regions, grid.connectivity.connections, plmn, fill_west_and_east_halo!, + multiregion_field_2, regions, grid.connectivity.connections, plmn, WestAndEast(), _fill_cubed_sphere_center_center_center_center_field_pairs_east_west_halo_regions!) @apply_regionally fill_cubed_sphere_field_pairs_halo_event!(grid, field_1, multiregion_field_1, field_2, - multiregion_field_2, regions, grid.connectivity.connections, plmn, fill_south_and_north_halo!, + multiregion_field_2, regions, grid.connectivity.connections, plmn, SouthAndNorth(), _fill_cubed_sphere_center_center_center_center_field_pairs_north_south_halo_regions!) return nothing @@ -231,9 +232,10 @@ end @inline function fill_cubed_sphere_field_pairs_halo_event!(grid, field_1, multiregion_field_1, field_2, multiregion_field_2, region, connections, plmn, - fill_halo_function!, _fill_halo_kernel!) - sz = fill_halo_size(field_1.data, fill_halo_function!, field_1.indices, FullyConnected, location(field_1), grid) - of = fill_halo_offset(sz, fill_halo_function!, field_1.indices) + side, _fill_halo_kernel!) + + sz = fill_halo_size(field_1.data, side, field_1.indices, FullyConnected, location(field_1), grid) + of = fill_halo_offset(sz, side, field_1.indices) kernel_parameters = KernelParameters(sz, of) reduced_dims = reduced_dimensions(field_1) @@ -351,10 +353,10 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Face, <:Center}, regions = Iterate(1:6) @apply_regionally fill_cubed_sphere_field_pairs_halo_event!(grid, field_1, multiregion_field_1, field_2, - multiregion_field_2, regions, grid.connectivity.connections, plmn, fill_west_and_east_halo!, + multiregion_field_2, regions, grid.connectivity.connections, plmn, WestAndEast(), _fill_cubed_sphere_face_center_center_face_field_pairs_east_west_halo_regions!) @apply_regionally fill_cubed_sphere_field_pairs_halo_event!(grid, field_1, multiregion_field_1, field_2, - multiregion_field_2, regions, grid.connectivity.connections, plmn, fill_south_and_north_halo!, + multiregion_field_2, regions, grid.connectivity.connections, plmn, SouthAndNorth(), _fill_cubed_sphere_face_center_center_face_field_pairs_north_south_halo_regions!) #= @@ -580,10 +582,10 @@ function fill_halo_regions!(field_1::CubedSphereField{<:Face, <:Face}, regions = Iterate(1:6) @apply_regionally fill_cubed_sphere_field_pairs_halo_event!(grid, field_1, multiregion_field_1, field_2, - multiregion_field_2, regions, grid.connectivity.connections, plmn, fill_west_and_east_halo!, + multiregion_field_2, regions, grid.connectivity.connections, plmn, WestAndEast(), _fill_cubed_sphere_face_face_face_face_field_pairs_east_west_halo_regions!) @apply_regionally fill_cubed_sphere_field_pairs_halo_event!(grid, field_1, multiregion_field_1, field_2, - multiregion_field_2, regions, grid.connectivity.connections, plmn, fill_south_and_north_halo!, + multiregion_field_2, regions, grid.connectivity.connections, plmn, SouthAndNorth(), _fill_cubed_sphere_face_face_face_face_field_pairs_north_south_halo_regions!) return nothing diff --git a/src/MultiRegion/cubed_sphere_connectivity.jl b/src/MultiRegion/cubed_sphere_connectivity.jl index 434374a9ddc..5fe791b85bc 100644 --- a/src/MultiRegion/cubed_sphere_connectivity.jl +++ b/src/MultiRegion/cubed_sphere_connectivity.jl @@ -56,7 +56,7 @@ struct CubedSphereRegionalConnectivity{S, FS, R} <: AbstractConnectivity CubedSphereRegionalConnectivity(rank, from_rank, side, from_side, rotation=nothing) Return a `CubedSphereRegionalConnectivity`: `from_rank :: Int` → `rank :: Int` and - `from_side :: AbstractRegionSide` → `side :: AbstractRegionSide`. The rotation of + `from_side` → `side`. The rotation of the adjacent region relative to the host region is prescribed via `rotation` argument (default `rotation=nothing`). @@ -73,8 +73,8 @@ struct CubedSphereRegionalConnectivity{S, FS, R} <: AbstractConnectivity julia> CubedSphereRegionalConnectivity(1, 2, East(), West()) CubedSphereRegionalConnectivity - ├── from: Oceananigans.MultiRegion.West side, region 2 - ├── to: Oceananigans.MultiRegion.East side, region 1 + ├── from: Oceananigans.BoundaryConditions.West side, region 2 + ├── to: Oceananigans.BoundaryConditions.East side, region 1 └── no rotation ``` @@ -84,8 +84,8 @@ struct CubedSphereRegionalConnectivity{S, FS, R} <: AbstractConnectivity ```jldoctest cubedsphereconnectivity julia> CubedSphereRegionalConnectivity(1, 3, North(), East(), ↺()) CubedSphereRegionalConnectivity - ├── from: Oceananigans.MultiRegion.East side, region 3 - ├── to: Oceananigans.MultiRegion.North side, region 1 + ├── from: Oceananigans.BoundaryConditions.East side, region 3 + ├── to: Oceananigans.BoundaryConditions.North side, region 1 └── counter-clockwise rotation ↺ ``` """ diff --git a/src/MultiRegion/cubed_sphere_grid.jl b/src/MultiRegion/cubed_sphere_grid.jl index a77105c811b..d763eca430b 100644 --- a/src/MultiRegion/cubed_sphere_grid.jl +++ b/src/MultiRegion/cubed_sphere_grid.jl @@ -162,28 +162,28 @@ julia> using Oceananigans.MultiRegion: East, North, West, South julia> for region in 1:length(grid); println(grid.connectivity.connections[region].south); end CubedSphereRegionalConnectivity -├── from: Oceananigans.MultiRegion.North side, region 6 -├── to: Oceananigans.MultiRegion.South side, region 1 +├── from: Oceananigans.BoundaryConditions.North side, region 6 +├── to: Oceananigans.BoundaryConditions.South side, region 1 └── no rotation CubedSphereRegionalConnectivity -├── from: Oceananigans.MultiRegion.East side, region 6 -├── to: Oceananigans.MultiRegion.South side, region 2 +├── from: Oceananigans.BoundaryConditions.East side, region 6 +├── to: Oceananigans.BoundaryConditions.South side, region 2 └── counter-clockwise rotation ↺ CubedSphereRegionalConnectivity -├── from: Oceananigans.MultiRegion.North side, region 2 -├── to: Oceananigans.MultiRegion.South side, region 3 +├── from: Oceananigans.BoundaryConditions.North side, region 2 +├── to: Oceananigans.BoundaryConditions.South side, region 3 └── no rotation CubedSphereRegionalConnectivity -├── from: Oceananigans.MultiRegion.East side, region 2 -├── to: Oceananigans.MultiRegion.South side, region 4 +├── from: Oceananigans.BoundaryConditions.East side, region 2 +├── to: Oceananigans.BoundaryConditions.South side, region 4 └── counter-clockwise rotation ↺ CubedSphereRegionalConnectivity -├── from: Oceananigans.MultiRegion.North side, region 4 -├── to: Oceananigans.MultiRegion.South side, region 5 +├── from: Oceananigans.BoundaryConditions.North side, region 4 +├── to: Oceananigans.BoundaryConditions.South side, region 5 └── no rotation CubedSphereRegionalConnectivity -├── from: Oceananigans.MultiRegion.East side, region 4 -├── to: Oceananigans.MultiRegion.South side, region 6 +├── from: Oceananigans.BoundaryConditions.East side, region 4 +├── to: Oceananigans.BoundaryConditions.South side, region 6 └── counter-clockwise rotation ↺ ``` """ diff --git a/src/MultiRegion/multi_region_boundary_conditions.jl b/src/MultiRegion/multi_region_boundary_conditions.jl index 193df81eba1..2caeafa37fb 100644 --- a/src/MultiRegion/multi_region_boundary_conditions.jl +++ b/src/MultiRegion/multi_region_boundary_conditions.jl @@ -2,36 +2,20 @@ using Oceananigans: instantiated_location using Oceananigans.Architectures: on_architecture, device_copy_to! using Oceananigans.Operators: assumed_field_location using Oceananigans.Fields: reduced_dimensions -using Oceananigans.DistributedComputations: communication_side, fill_send_buffers! +using Oceananigans.DistributedComputations: fill_send_buffers! using Oceananigans.BoundaryConditions: ContinuousBoundaryFunction, DiscreteBoundaryFunction, - permute_boundary_conditions, - extract_west_bc, extract_east_bc, extract_south_bc, - extract_north_bc, extract_top_bc, extract_bottom_bc, - fill_halo_event!, - MCBCT, - MCBC, - fill_open_boundary_regions! - -import Oceananigans.Fields: tupled_fill_halo_regions!, boundary_conditions, data - -import Oceananigans.BoundaryConditions: - fill_halo_regions!, - fill_west_and_east_halo!, - fill_south_and_north_halo!, - fill_bottom_and_top_halo!, - fill_west_halo!, - fill_east_halo!, - fill_south_halo!, - fill_north_halo! + fill_halo_event!, get_boundary_kernels, + MultiRegionFillHalo, + MCBCT, MCBC + +import Oceananigans.BoundaryConditions: fill_halo_regions!, fill_halo_event! @inline bc_str(::MultiRegionObject) = "MultiRegion Boundary Conditions" -@inline extract_field_buffers(field::Field) = field.communication_buffers -@inline boundary_conditions(field::MultiRegionField) = field.boundary_conditions -@inline function tupled_fill_halo_regions!(fields::NamedTuple, grid::ConformalCubedSphereGridOfSomeKind, args...; kwargs...) +@inline function fill_halo_regions!(fields::NamedTuple, grid::ConformalCubedSphereGridOfSomeKind, args...; kwargs...) u = haskey(fields, :u) ? fields.u : nothing v = haskey(fields, :v) ? fields.v : nothing @@ -56,19 +40,15 @@ import Oceananigans.BoundaryConditions: return nothing end -function fill_halo_regions!(field::MultiRegionField, args...; kwargs...) - reduced_dims = reduced_dimensions(field) - - return fill_halo_regions!(field.data, - field.boundary_conditions, - field.indices, - instantiated_location(field), - field.grid, - field.communication_buffers, - args...; - reduced_dimensions = reduced_dims, - kwargs...) -end +fill_halo_regions!(field::MultiRegionField, args...; kwargs...) = + fill_halo_regions!(field.data, + field.boundary_conditions, + field.indices, + instantiated_location(field), + field.grid, + field.communication_buffers, + args...; + kwargs...) fill_halo_regions!(c::MultiRegionObject, ::Nothing, args...; kwargs...) = nothing @@ -76,124 +56,38 @@ fill_halo_regions!(c::MultiRegionObject, ::Nothing, args...; kwargs...) = nothin ##### fill_halo_regions! for a MultiRegionObject ##### -# this can be used once we don't need to fill Value, Flux and Gradient anymore! -# fill_halo_regions!(c::MultiRegionObject, bcs, loc, mrg::MultiRegionGrid, buffers, args...; kwargs...) = -# apply_regionally!(fill_halo_regions!, c, bcs, loc, mrg, Reference(c.regional_objects), Reference(buffers.regional_objects), args...; kwargs...) - -# TODO: Adapt MultiRegion boundary conditions to the Distributed logic: aka split the -# halo sides in two different sides. Requires synchronizing all workers. -# Might be difficult for `CubedSphereGrids` where all buffers must be filled prior -# communicating. -# For the moment we keep the old version because asynchronous communication, -# which requires splitting, is not yet implemented. -extract_west_or_east_bc(bc) = max(bc.west, bc.east) -extract_south_or_north_bc(bc) = max(bc.south, bc.north) -extract_bottom_or_top_bc(bc) = max(bc.bottom, bc.top) - -function multi_region_permute_boundary_conditions(bcs) - fill_halos! = [ - fill_west_and_east_halo!, - fill_south_and_north_halo!, - fill_bottom_and_top_halo!, - ] - - boundary_conditions_array = [ - extract_west_or_east_bc(bcs), - extract_south_or_north_bc(bcs), - extract_bottom_or_top_bc(bcs) - ] - - boundary_conditions = [ - (extract_west_bc(bcs), extract_east_bc(bcs)), - (extract_south_bc(bcs), extract_north_bc(bcs)), - (extract_bottom_bc(bcs), extract_top_bc(bcs)) - ] - - perm = sortperm(boundary_conditions_array) - fill_halos! = fill_halos![perm] - boundary_conditions = boundary_conditions[perm] - - return (fill_halos!, boundary_conditions) -end - -function fill_halo_regions!(c::MultiRegionObject, bcs, indices, loc, mrg::MultiRegionGrid, buffers, args...; fill_boundary_normal_velocities = true, kwargs...) - arch = architecture(mrg) - - if fill_boundary_normal_velocities - apply_regionally!(fill_open_boundary_regions!, c, bcs, indices, loc, mrg, args...) - end +# Fill halo regions with a double pass for the moment. +# TODO: Optimize this some way (needs infra-regional-function synchronization between regions). +# The complication here is the possibility of different regions having different number of tasks, +# Which might happen, for example, for a grid that partitioned in a Bounded direction. +function fill_halo_regions!(c::MultiRegionObject, bcs, indices, loc, mrg::MultiRegionGrid, buffers, args...; fill_open_bcs=true, kwargs...) + arch = architecture(mrg) + buff_ref = Reference(buffers.regional_objects) - @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(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) - - apply_regionally!(fill_halo_event!, c, fill_halo_side!, bcs_side, - indices, loc, arch, mrg, buff, - args...; kwargs...) - end - - return nothing -end - -# Find a better way to do this (this will not work for corners!!) -function fill_multiregion_send_buffers!(c, buffers, grid, bcs) - - if !isempty(filter(x -> x isa MCBCT, bcs)) - fill_send_buffers!(c, buffers, grid) - end + apply_regionally!(fill_send_buffers!, c, buffers, mrg) + apply_regionally!(fill_halo_regions!, c, bcs, indices, loc, mrg, buff_ref, args...; fill_open_bcs, kwargs...) + apply_regionally!(fill_send_buffers!, c, buffers, mrg) + apply_regionally!(fill_halo_regions!, c, bcs, indices, loc, mrg, buff_ref, args...; fill_open_bcs, kwargs...) return nothing end ##### -##### fill_halo! for Communicating boundary condition +##### fill halo event, splat the args... ##### -## Fill communicating boundary condition halos -for (lside, rside) in zip([:west, :south, :bottom], [:east, :north, :top]) - fill_both_halo! = Symbol(:fill_, lside, :_and_, rside, :_halo!) - fill_left_halo! = Symbol(:fill_, lside, :_halo!) - fill_right_halo! = Symbol(:fill_, rside, :_halo!) - - @eval begin - function $fill_both_halo!(c, left_bc::MCBC, right_bc, kernel_size, offset, loc, arch, grid, buffers, args...; kwargs...) - $fill_right_halo!(c, right_bc, kernel_size, offset, loc, arch, grid, args...; kwargs...) - $fill_left_halo!(c, left_bc, kernel_size, offset, loc, arch, grid, buffers, args...; kwargs...) - return nothing - end - - function $fill_both_halo!(c, left_bc, right_bc::MCBC, kernel_size, offset, loc, arch, grid, buffers, args...; kwargs...) - $fill_left_halo!(c, left_bc, kernel_size, offset, loc, arch, grid, args...; kwargs...) - $fill_right_halo!(c, right_bc, kernel_size, offset, loc, arch, grid, buffers, args...; kwargs...) - return nothing - end - end -end +fill_halo_event!(c, kernel!::MultiRegionFillHalo, bcs, loc, grid, buffers, args...; kwargs...) = kernel!(c, bcs..., loc, grid, buffers) getside(x, ::North) = x.north getside(x, ::South) = x.south getside(x, ::West) = x.west getside(x, ::East) = x.east -fill_west_and_east_halo!(c, westbc::MCBC, eastbc::MCBC, kernel_size, offset, loc, arch, grid, ::Nothing, args...; kwargs...) = nothing -fill_south_and_north_halo!(c, southbc::MCBC, northbc::MCBC, kernel_size, offset, loc, arch, grid, ::Nothing, args...; kwargs...) = nothing - -fill_west_halo!(c, bc::MCBC, kernel_size, offset, loc, arch, grid, ::Nothing, args...; kwargs...) = nothing -fill_east_halo!(c, bc::MCBC, kernel_size, offset, loc, arch, grid, ::Nothing, args...; kwargs...) = nothing -fill_south_halo!(c, bc::MCBC, kernel_size, offset, loc, arch, grid, ::Nothing, args...; kwargs...) = nothing -fill_north_halo!(c, bc::MCBC, kernel_size, offset, loc, arch, grid, ::Nothing, args...; kwargs...) = nothing +##### +##### Double-sided MultiRegion filling kernels +##### -function fill_west_and_east_halo!(c, westbc::MCBC, eastbc::MCBC, kernel_size, offset, loc, arch, grid, buffers, args...; kwargs...) +function (::MultiRegionFillHalo{<:WestAndEast})(c, westbc, eastbc, loc, grid, buffers) H = halo_size(grid)[1] N = size(grid)[1] @@ -233,7 +127,7 @@ function fill_west_and_east_halo!(c, westbc::MCBC, eastbc::MCBC, kernel_size, of return nothing end -function fill_south_and_north_halo!(c, southbc::MCBC, northbc::MCBC, kernel_size, offset, loc, arch, grid, buffers, args...; kwargs...) +function (::MultiRegionFillHalo{<:SouthAndNorth})(c, southbc, northbc, loc, grid, buffers) H = halo_size(grid)[2] N = size(grid)[2] @@ -277,13 +171,10 @@ function fill_south_and_north_halo!(c, southbc::MCBC, northbc::MCBC, kernel_size end ##### -##### Single fill_halo! for Communicating boundary condition +##### Single-sided MultiRegion filling kernels ##### -# TODO: Allow support for `Bounded` Cubed sphere grids -# (i.e. with the same shift implemented in the double-sided fill halo) - -function fill_west_halo!(c, bc::MCBC, kernel_size, offset, loc, arch, grid, buffers, args...; kwargs...) +function (::MultiRegionFillHalo{<:West})(c, bc, loc, grid, buffers) H = halo_size(grid)[1] N = size(grid)[1] @@ -303,7 +194,7 @@ function fill_west_halo!(c, bc::MCBC, kernel_size, offset, loc, arch, grid, buff return nothing end -function fill_east_halo!(c, bc::MCBC, kernel_size, offset, loc, arch, grid, buffers, args...; kwargs...) +function (::MultiRegionFillHalo{<:East})(c, bc, loc, grid, buffers) H = halo_size(grid)[1] N = size(grid)[1] @@ -323,7 +214,7 @@ function fill_east_halo!(c, bc::MCBC, kernel_size, offset, loc, arch, grid, buff return nothing end -function fill_south_halo!(c, bc::MCBC, kernel_size, offset, loc, arch, grid, buffers, args...; kwargs...) +function (::MultiRegionFillHalo{<:South})(c, bc, loc, grid, buffers) H = halo_size(grid)[2] N = size(grid)[2] @@ -343,7 +234,7 @@ function fill_south_halo!(c, bc::MCBC, kernel_size, offset, loc, arch, grid, buf return nothing end -function fill_north_halo!(c, bc::MCBC, kernel_size, offset, loc, arch, grid, buffers, args...; kwargs...) +function (::MultiRegionFillHalo{<:North})(c, bc, loc, grid, buffers) H = halo_size(grid)[2] N = size(grid)[2] @@ -380,7 +271,9 @@ end _getregion(fc.north, i), _getregion(fc.bottom, i), _getregion(fc.top, i), - fc.immersed) + fc.immersed, + fc.kernels, + _getregion(fc.ordered_bcs, i)) @inline getregion(bc::BoundaryCondition, i) = BoundaryCondition(bc.classification, _getregion(bc.condition, i)) @@ -401,7 +294,9 @@ end getregion(fc.north, i), getregion(fc.bottom, i), getregion(fc.top, i), - fc.immersed) + fc.immersed, + fc.kernels, + getregion(fc.ordered_bcs, i)) @inline _getregion(bc::BoundaryCondition, i) = BoundaryCondition(bc.classification, getregion(bc.condition, i)) diff --git a/src/MultiRegion/multi_region_connectivity.jl b/src/MultiRegion/multi_region_connectivity.jl index d1b8a89fe26..8ebb0368a1c 100644 --- a/src/MultiRegion/multi_region_connectivity.jl +++ b/src/MultiRegion/multi_region_connectivity.jl @@ -1,13 +1,13 @@ using Oceananigans.Grids: topology """ - struct RegionalConnectivity{S <: AbstractRegionSide, FS <: AbstractRegionSide} <: AbstractConnectivity + struct RegionalConnectivity{S, FS} <: AbstractConnectivity The connectivity among various regions in a multi-region partition. $(TYPEDFIELDS) """ -struct RegionalConnectivity{S <: AbstractRegionSide, FS <: AbstractRegionSide} <: AbstractConnectivity +struct RegionalConnectivity{S, FS} <: AbstractConnectivity "the current region rank" rank :: Int "the region from which boundary condition comes from" diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 2f5dc1c9cc7..1d5b83566bc 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -1,6 +1,5 @@ -using Oceananigans.BoundaryConditions: fill_open_boundary_regions!, - permute_boundary_conditions, - fill_halo_event!, +using Oceananigans.BoundaryConditions: permute_boundary_conditions, + fill_halo_event!, get_boundary_kernels, DistributedCommunication using Oceananigans.DistributedComputations: cooperative_waitall!, @@ -60,21 +59,18 @@ end end function fill_halo_regions!(c::OffsetArray, bcs, indices, loc, grid::DistributedTripolarGridOfSomeKind, buffers, args...; - only_local_halos=false, fill_boundary_normal_velocities=true, kwargs...) - - if fill_boundary_normal_velocities - fill_open_boundary_regions!(c, bcs, indices, loc, grid, args...; kwargs...) - end + only_local_halos=false, fill_open_bcs=true, kwargs...) north_bc = bcs.north arch = architecture(grid) - fill_halos!, bcs = permute_boundary_conditions(bcs) - number_of_tasks = length(fill_halos!) + kernels!, ordered_bcs = get_boundary_kernels(bcs, c, grid, loc, indices) + + number_of_tasks = length(kernels!) outstanding_requests = length(arch.mpi_requests) for task = 1:number_of_tasks - fill_halo_event!(c, fill_halos![task], bcs[task], indices, loc, arch, grid, buffers, args...; only_local_halos, kwargs...) + @inbounds fill_halo_event!(c, kernels![task], ordered_bcs[task], loc, grid, buffers, args...; kwargs...) end fill_corners!(c, arch.connectivity, indices, loc, arch, grid, buffers, args...; only_local_halos, kwargs...) @@ -113,4 +109,4 @@ function synchronize_communication!(field::Field{<:Any, <:Any, <:Any, <:Any, <:D switch_north_halos!(field, north_bc, field.grid, instantiated_location) return nothing -end \ No newline at end of file +end diff --git a/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl b/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl index d7157f438ec..b28ef56feac 100644 --- a/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl +++ b/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl @@ -9,7 +9,6 @@ using Oceananigans.ImmersedBoundaries import Oceananigans.BoundaryConditions: default_auxiliary_bc, regularize_field_boundary_conditions import Oceananigans.Grids: x_domain, y_domain -import Oceananigans.Fields: tupled_fill_halo_regions! # A tripolar grid is always between 0 and 360 in longitude # and always caps at the north pole (90°N) @@ -49,10 +48,3 @@ function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, end default_auxiliary_bc(grid::TripolarGridOfSomeKind, ::Val{:north}, loc) = ZipperBoundaryCondition(1) - -# Not sure this is needed, but it is here for now -function tupled_fill_halo_regions!(full_fields, grid::TripolarGridOfSomeKind, args...; kwargs...) - for field in full_fields - fill_halo_regions!(field, args...; kwargs...) - end -end \ No newline at end of file diff --git a/src/OrthogonalSphericalShellGrids/tripolar_grid.jl b/src/OrthogonalSphericalShellGrids/tripolar_grid.jl index 6cd01bda1cb..9792f012352 100644 --- a/src/OrthogonalSphericalShellGrids/tripolar_grid.jl +++ b/src/OrthogonalSphericalShellGrids/tripolar_grid.jl @@ -153,7 +153,7 @@ function TripolarGrid(arch = CPU(), FT::DataType = Float64; # ZipperBoundaryCondition that edge fields need to switch sign (which we definitely do not # want for coordinates and metrics) default_boundary_conditions = FieldBoundaryConditions(north = ZipperBoundaryCondition(), - south = nothing, # The south should be `continued` + south = NoFluxBoundaryCondition(), # The south should be `continued` west = Oceananigans.PeriodicBoundaryCondition(), east = Oceananigans.PeriodicBoundaryCondition(), top = nothing, diff --git a/src/OutputWriters/output_writer_utils.jl b/src/OutputWriters/output_writer_utils.jl index 0873de4e264..b118b3f2699 100644 --- a/src/OutputWriters/output_writer_utils.jl +++ b/src/OutputWriters/output_writer_utils.jl @@ -95,13 +95,15 @@ end # Special saveproperty! so boundary conditions are easily readable outside julia. function saveproperty!(file, address, bcs::FieldBoundaryConditions) for boundary in propertynames(bcs) - bc = getproperty(bcs, boundary) - file[address * "/$boundary/type"] = bc_str(bc) - - if bc === nothing || bc.condition isa Function || bc.condition isa ContinuousBoundaryFunction - file[address * "/$boundary/condition"] = missing - else - file[address * "/$boundary/condition"] = on_architecture(CPU(), bc.condition) + if !(boundary == :kernels || boundary == :ordered_bcs) # Skip kernels and ordered_bcs. + bc = getproperty(bcs, boundary) + file[address * "/$boundary/type"] = bc_str(bc) + + if bc === nothing || bc.condition isa Function || bc.condition isa ContinuousBoundaryFunction + file[address * "/$boundary/condition"] = missing + else + file[address * "/$boundary/condition"] = on_architecture(CPU(), bc.condition) + end end end end @@ -133,14 +135,21 @@ function serializeproperty!(file, address, grid::DistributedGrid) file[address] = on_architecture(cpu_arch, grid) end +function remove_function_bcs(fbcs::FieldBoundaryConditions) + west = has_reference(Function, fbcs.west) ? missing : fbcs.west + east = has_reference(Function, fbcs.east) ? missing : fbcs.east + south = has_reference(Function, fbcs.south) ? missing : fbcs.south + north = has_reference(Function, fbcs.north) ? missing : fbcs.north + bottom = has_reference(Function, fbcs.bottom) ? missing : fbcs.bottom + top = has_reference(Function, fbcs.top) ? missing : fbcs.top + immersed = has_reference(Function, fbcs.immersed) ? missing : fbcs.immersed + new_fbcs = FieldBoundaryConditions(west, east, south, north, bottom, top, immersed) + return new_fbcs +end + function serializeproperty!(file, address, fbcs::FieldBoundaryConditions) - # TODO: it'd be better to "filter" `FieldBoundaryCondition` and then serialize - # rather than punting with `missing` instead. - if has_reference(Function, fbcs) - file[address] = missing - else - file[address] = on_architecture(CPU(), fbcs) - end + new_fbcs = remove_function_bcs(fbcs) + file[address] = on_architecture(CPU(), new_fbcs) end function serializeproperty!(file, address, f::Field) @@ -183,6 +192,13 @@ has_reference(::Type{T}, ::NTuple{N, <:T}) where {N, T} = true has_reference(T::Type{Function}, f::Field) = has_reference(T, f.data) || has_reference(T, f.boundary_conditions) +# Short circuit on boundary conditions. +has_reference(T::Type{Function}, bcs::FieldBoundaryConditions) = + has_reference(T, bcs.west) || has_reference(T, bcs.east) || + has_reference(T, bcs.south) || has_reference(T, bcs.north) || + has_reference(T, bcs.bottom) || has_reference(T, bcs.top) || + has_reference(T, bcs.immersed) + """ has_reference(has_type, obj) diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/dynamic_coefficient.jl b/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/dynamic_coefficient.jl index 4f974b5dddb..8f37b7a4328 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/dynamic_coefficient.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/dynamic_coefficient.jl @@ -343,5 +343,3 @@ function allocate_coefficient_fields(closure::LagrangianAveragedDynamicSmagorins return (; Σ, Σ̄, 𝒥ᴸᴹ, 𝒥ᴹᴹ, 𝒥ᴸᴹ⁻, 𝒥ᴹᴹ⁻, previous_compute_time) end - - diff --git a/src/Utils/kernel_launching.jl b/src/Utils/kernel_launching.jl index 29ed2f56032..0385b29023d 100644 --- a/src/Utils/kernel_launching.jl +++ b/src/Utils/kernel_launching.jl @@ -269,7 +269,7 @@ Keyword Arguments # Transform keyword arguments into arguments to be able to dispatch correctly return configure_kernel(arch, grid, workspec, kernel!, active_cells_map, exclude_periphery; - reduced_dimensions = (), + reduced_dimensions, location = nothing) end diff --git a/src/Utils/multi_region_transformation.jl b/src/Utils/multi_region_transformation.jl index 80df52329dd..c4faa85a89b 100644 --- a/src/Utils/multi_region_transformation.jl +++ b/src/Utils/multi_region_transformation.jl @@ -29,6 +29,7 @@ struct MultiRegionObject{R, D, B} return new{R, D, B}(regional_objects, devices, backend) end end + MultiRegionObject(arch::AbstractArchitecture, regional_objects...; devices=Tuple(CPU() for _ in regional_objects)) = MultiRegionObject(device(arch), regional_objects...; devices=devices) MultiRegionObject(arch::AbstractArchitecture, regional_objects::Tuple, devices::Tuple) = @@ -136,7 +137,6 @@ on_architecture(arch::GPU, mo::MultiRegionObject) = devs = isnothing(multi_region_args) ? multi_region_kwargs : multi_region_args devs = devices(devs) - for (r, dev) in enumerate(devs) switch_device!(dev) regional_func!((getregion(arg, r) for arg in args)...; (getregion(kwarg, r) for kwarg in kwargs)...) @@ -161,7 +161,6 @@ end devs = isnothing(multi_region_args) ? multi_region_kwargs : multi_region_args devs = devices(devs) - # Dig out the backend since we don't have access to arch. backend = nothing for arg in args @@ -170,9 +169,11 @@ end break end end + if backend isa Nothing backend = devs[1] end + # Evaluate regional_func on the device of that region and collect # return values regional_return_values = Vector(undef, length(devs)) diff --git a/test/test_boundary_conditions_integration.jl b/test/test_boundary_conditions_integration.jl index 7567406ffb6..13b194da7a0 100644 --- a/test/test_boundary_conditions_integration.jl +++ b/test/test_boundary_conditions_integration.jl @@ -111,8 +111,6 @@ end_position(::Val{1}, grid) = (grid.Nx+1, 1, 1) end_position(::Val{2}, grid) = (1, grid.Ny+1, 1) end_position(::Val{3}, grid) = (1, 1, grid.Nz+1) - - wall_normal_boundary_condition(::Val{1}, obc) = (; u = FieldBoundaryConditions(east = obc, west = obc)) wall_normal_boundary_condition(::Val{2}, obc) = (; v = FieldBoundaryConditions(south = obc, north = obc)) wall_normal_boundary_condition(::Val{3}, obc) = (; w = FieldBoundaryConditions(bottom = obc, top = obc)) @@ -180,7 +178,7 @@ function test_open_boundary_condition_mass_conservation(arch, FT, boundary_condi run!(simulation) compute!(∫∇u) - @test (@allowscalar ∫∇u[]) ≈ 0 atol=3*eps(FT) + @test (@allowscalar ∫∇u[]) ≈ 0 atol=3.5*eps(FT) end test_boundary_conditions(C, FT, ArrayType) = (integer_bc(C, FT, ArrayType), diff --git a/test/test_output_readers.jl b/test/test_output_readers.jl index ed2de960a9c..dc1da63dc47 100644 --- a/test/test_output_readers.jl +++ b/test/test_output_readers.jl @@ -335,38 +335,37 @@ end end end - if arch isa GPU - @testset "FieldTimeSeries with CuArray boundary conditions [$(typeof(arch))]" begin - @info " Testing FieldTimeSeries with CuArray boundary conditions..." - - x = y = z = (0, 1) - grid = RectilinearGrid(GPU(); size=(1, 1, 1), x, y, z) - - τx = CuArray(zeros(size(grid)...)) - τy = Field{Center, Face, Nothing}(grid) - u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(τx)) - v_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(τy)) - model = NonhydrostaticModel(; grid, boundary_conditions = (; u=u_bcs, v=v_bcs)) - simulation = Simulation(model; Δt=1, stop_iteration=1) - - simulation.output_writers[:jld2] = JLD2Writer(model, model.velocities, - filename = "test_cuarray_bc.jld2", - schedule=IterationInterval(1), - overwrite_existing = true) - - run!(simulation) - - ut = FieldTimeSeries("test_cuarray_bc.jld2", "u") - vt = FieldTimeSeries("test_cuarray_bc.jld2", "v") - @test ut.boundary_conditions.top.classification isa Flux - @test ut.boundary_conditions.top.condition isa Array - - τy_ow = vt.boundary_conditions.top.condition - @test τy_ow isa Field{Center, Face, Nothing} - @test architecture(τy_ow) isa CPU - @test parent(τy_ow) isa Array - rm("test_cuarray_bc.jld2") - end + @testset "FieldTimeSeries with Array boundary conditions [$(typeof(arch))]" begin + @info " Testing FieldTimeSeries with Array boundary conditions..." + x = y = z = (0, 1) + grid = RectilinearGrid(arch; size=(1, 1, 1), x, y, z) + + τx = on_architecture(arch, zeros(size(grid)...)) + τy = Field{Center, Face, Nothing}(grid) + u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(τx)) + v_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(τy)) + model = NonhydrostaticModel(; grid, boundary_conditions = (; u=u_bcs, v=v_bcs)) + simulation = Simulation(model; Δt=1, stop_iteration=1) + + filename = arch isa GPU ? "test_cuarray_bc.jld2" : "test_array_bc.jld2" + + simulation.output_writers[:jld2] = JLD2Writer(model, model.velocities; + filename, + schedule=IterationInterval(1), + overwrite_existing = true) + + run!(simulation) + + ut = FieldTimeSeries(filename, "u") + vt = FieldTimeSeries(filename, "v") + @test ut.boundary_conditions.top.classification isa Flux + @test ut.boundary_conditions.top.condition isa Array + + τy_ow = vt.boundary_conditions.top.condition + @test τy_ow isa Field{Center, Face, Nothing} + @test architecture(τy_ow) isa CPU + @test parent(τy_ow) isa Array + rm(filename) end end