Add Y-partition and XY-partition tests

.buildkite/distributed/pipeline.yml

@@ -27,7 +27,7 @@ steps:
       - "julia --project -e 'using Pkg; Pkg.status()'"
 
     agents:
-      slurm_mem: 80G
+      slurm_mem: 120G
       slurm_gpus: 1
       slurm_cpus_per_task: 8
     env:
@@ -52,6 +52,7 @@ steps:
     commands:
       - "srun julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'"
     agents:
+      slurm_mem: 120G
       slurm_ntasks: 4
       slurm_gpus_per_task: 1
 
@@ -80,6 +81,7 @@ steps:
     commands:
       - "srun julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'"
     agents:
+      slurm_mem: 120G
       slurm_ntasks: 4
       slurm_gpus_per_task: 1
 

src/BoundaryConditions/fill_halo_regions.jl

@@ -21,30 +21,27 @@ 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
 
+# 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_bc = Symbol(:extract_, dir, :_bc)
+    extract_side_bc = Symbol(:extract_, dir, :_bc)
     @eval begin
-        @inline $extract_bc(bc) = bc.$dir
-        @inline $extract_bc(bc::Tuple) = map($extract_bc, bc)
+        @inline $extract_side_bc(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))
 
-# For inhomogeneous BC we extract the _last_ one 
-# example 
-# `bc.west <: DCBC`
-# `bc.east <: PBC`
-# `extract_west_or_east_bc(bc) == bc.west`
-# NOTE that `isless` follows order of execution, 
-# so `max(bcs...)` returns the last BC to execute
-
-  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)
-
-  extract_west_or_east_bc(bc::Tuple) =   map(extract_west_or_east_bc, bc)
-extract_south_or_north_bc(bc::Tuple) = map(extract_south_or_north_bc, bc)
- extract_bottom_or_top_bc(bc::Tuple) =  map(extract_bottom_or_top_bc, bc)
-
 # Finally, the true fill_halo!
 const MaybeTupledData = Union{OffsetArray, NTuple{<:Any, OffsetArray}}
 
@@ -53,63 +50,95 @@ function fill_halo_regions!(c::MaybeTupledData, boundary_conditions, indices, lo
 
     arch = architecture(grid)
 
-    halo_tuple  = permute_boundary_conditions(boundary_conditions)
+    fill_halos!, bcs  = permute_boundary_conditions(boundary_conditions)
+    number_of_tasks = length(fill_halos!)
 
     # Fill halo in the three permuted directions (1, 2, and 3), making sure dependencies are fulfilled
-    for task in 1:3
-        fill_halo_event!(task, halo_tuple, c, indices, loc, arch, grid, args...; kwargs...)
+    for task = 1:number_of_tasks
+        fill_halo_event!(c, fill_halos![task], bcs[task], indices, loc, arch, grid, args...; kwargs...)
     end
 
     return nothing
 end
 
-function fill_halo_event!(task, halo_tuple, c, indices, loc, arch, grid, args...; kwargs...)
-    fill_halo!  = halo_tuple[1][task]
-    bc_left     = halo_tuple[2][task]
-    bc_right    = halo_tuple[3][task]
+function fill_halo_event!(c, fill_halos!, bcs, indices, loc, arch, grid, args...; kwargs...)
 
     # Calculate size and offset of the fill_halo kernel
-    size   = fill_halo_size(c, fill_halo!, indices, bc_left, loc, grid)
-    offset = fill_halo_offset(size, fill_halo!, indices)
+    # 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)
+
+    fill_halos!(c, bcs..., size, offset, loc, arch, grid, args...; kwargs...)
 
-    fill_halo!(c, bc_left, bc_right, 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_boundaries = split_boundary(extract_west_bc(boundary_conditions),  extract_east_bc(boundary_conditions))
+    split_y_boundaries = split_boundary(extract_south_bc(boundary_conditions), extract_north_bc(boundary_conditions))
+
+    fill_x_side_halo! = split_x_boundaries ? (fill_west_halo!,  fill_east_halo!)  : tuple(fill_west_and_east_halo!)
+    fill_y_side_halo! = split_y_boundaries ? (fill_south_halo!, fill_north_halo!) : tuple(fill_south_and_north_halo!)
+
+    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)
+
+    x_side_bcs = split_x_boundaries ? (west_bc, east_bc)   : tuple(west_bc)
+    y_side_bcs = split_y_boundaries ? (south_bc, north_bc) : tuple(south_bc)
+
+    x_sides = split_x_boundaries ? (:west, :east)   : tuple(:west_and_east)
+    y_sides = split_y_boundaries ? (:south, :north) : tuple(:south_and_north)
+
     fill_halos! = [
-        fill_west_and_east_halo!,
-        fill_south_and_north_halo!,
-        fill_bottom_and_top_halo!,
+        fill_x_side_halo!...,
+        fill_y_side_halo!...,
+        fill_bottom_and_top_halo! # bottom and top cannot be Distributed so never split them
     ]
-
-    boundary_conditions_array = [
-        extract_west_or_east_bc(boundary_conditions),
-        extract_south_or_north_bc(boundary_conditions),
-        extract_bottom_or_top_bc(boundary_conditions)
+    
+    sides = [
+        x_sides...,
+        y_sides...,
+        :bottom_and_top
     ]
 
-    boundary_conditions_array_left = [
-        extract_west_bc(boundary_conditions),
-        extract_south_bc(boundary_conditions),
+    bcs_array = [
+        x_side_bcs...,
+        y_side_bcs...,
         extract_bottom_bc(boundary_conditions)
     ]
 
-    boundary_conditions_array_right = [
-        extract_east_bc(boundary_conditions),
-        extract_north_bc(boundary_conditions),
-        extract_top_bc(boundary_conditions),
-    ]
-
-    perm = sortperm(boundary_conditions_array, lt=fill_first)
+    perm = sortperm(bcs_array, lt=fill_first)
     fill_halos! = fill_halos![perm]
-    boundary_conditions_array_left  = boundary_conditions_array_left[perm]
-    boundary_conditions_array_right = boundary_conditions_array_right[perm]
+    sides = sides[perm]
+
+    boundary_conditions = Tuple(extract_bc(boundary_conditions, Val(side)) for side in sides)
 
-    return (fill_halos!, boundary_conditions_array_left, boundary_conditions_array_right)
+    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_boundary(bcs1, bcs2)     = false
+split_boundary(::DCBC, ::DCBC) = false
+split_boundary(bcs1, ::DCBC)   = true
+split_boundary(::DCBC, bcs2)   = true
+
+# TODO: support heterogeneous distributed-shared communication
+# split_boundary(::MCBC, ::DCBC) = false
+# split_boundary(::DCBC, ::MCBC) = false
+# split_boundary(::MCBC, ::MCBC) = false
+# split_boundary(bcs1, ::MCBC)   = true
+# split_boundary(::MCBC, bcs2)   = true
+
 #####
 ##### Halo filling order
 #####
@@ -118,11 +147,16 @@ const PBCT  = Union{PBC,  NTuple{<:Any, <:PBC}}
 const MCBCT = Union{MCBC, NTuple{<:Any, <:MCBC}}
 const DCBCT = Union{DCBC, NTuple{<:Any, <:DCBC}}
 
-# Distributed halos have to be filled for last in case of 
-# buffered communication. Hence, we always fill them last
+# 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
 
-# The reasoning for filling Periodic after Flux, Value, Gradient 
+# 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)
@@ -160,7 +194,7 @@ fill_first(bc1::MCBCT, bc2::MCBCT) = true
 fill_first(bc1, bc2)               = true
 
 #####
-##### General fill_halo! kernels
+##### Double-sided fill_halo! kernels
 #####
 
 @kernel function _fill_west_and_east_halo!(c, west_bc, east_bc, loc, grid, args) 
@@ -182,9 +216,45 @@ end
 end
 
 #####
-##### Tuple fill_halo! kernels
+##### Single-sided fill_halo! kernels
+#####
+
+@kernel function _fill_only_west_halo!(c, bc, loc, grid, args)
+    j, k = @index(Global, NTuple)
+    _fill_west_halo!(j, k, grid, c, bc, loc, args...)
+end
+
+@kernel function _fill_only_south_halo!(c, bc, loc, grid, args)
+    i, k = @index(Global, NTuple)
+    _fill_south_halo!(i, k, grid, c, bc, loc, args...)
+end
+
+@kernel function _fill_only_bottom_halo!(c, bc, loc, grid, args)
+    i, j = @index(Global, NTuple)
+    _fill_bottom_halo!(i, j, grid, c, bc, loc, args...)
+end
+
+@kernel function _fill_only_east_halo!(c, bc, loc, grid, args)
+    j, k = @index(Global, NTuple)
+    _fill_east_halo!(j, k, grid, c, bc, loc, args...)
+end
+
+@kernel function _fill_only_north_halo!(c, bc, loc, grid, args)
+    i, k = @index(Global, NTuple)
+    _fill_north_halo!(i, k, grid, c, bc, loc, args...)
+end
+
+@kernel function _fill_only_top_halo!(c, bc, loc, grid, args)
+    i, j = @index(Global, NTuple)
+    _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::NTuple, west_bc, east_bc, loc, grid, args)
@@ -223,6 +293,27 @@ end
     end
 end
 
+#####
+##### Kernel launchers for single-sided fill_halos
+#####
+
+fill_west_halo!(c, bc, size, offset, loc, arch, grid, args...; 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...; 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...; 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...; 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...; 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...; 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
+#####
+
 fill_west_and_east_halo!(c, west_bc, east_bc, size, offset, loc, arch, grid, args...; kwargs...) =
     launch!(arch, grid, KernelParameters(size, offset), _fill_west_and_east_halo!, c, west_bc, east_bc, loc, grid, Tuple(args); kwargs...)
 
@@ -236,9 +327,9 @@ fill_bottom_and_top_halo!(c, bottom_bc, top_bc, size, offset, loc, arch, grid, a
 ##### Calculate kernel size and offset for Windowed and Sliced Fields
 #####
 
-const WEB = typeof(fill_west_and_east_halo!)
-const SNB = typeof(fill_south_and_north_halo!)
-const TBB = typeof(fill_bottom_and_top_halo!)
+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!)}
 
 # Tupled halo filling _only_ deals with full fields!
 @inline fill_halo_size(::Tuple, ::WEB, args...) = :yz
@@ -278,3 +369,4 @@ const TBB = typeof(fill_bottom_and_top_halo!)
 @inline fill_halo_offset(::Tuple, ::WEB, idx)  = (idx[2] == Colon() ? 0 : first(idx[2])-1, idx[3] == Colon() ? 0 : first(idx[3])-1)
 @inline fill_halo_offset(::Tuple, ::SNB, idx)  = (idx[1] == Colon() ? 0 : first(idx[1])-1, idx[3] == Colon() ? 0 : first(idx[3])-1)
 @inline fill_halo_offset(::Tuple, ::TBB, idx)  = (idx[1] == Colon() ? 0 : first(idx[1])-1, idx[2] == Colon() ? 0 : first(idx[2])-1)
+

src/BoundaryConditions/fill_halo_regions_flux.jl

@@ -33,55 +33,3 @@ using KernelAbstractions.Extras.LoopInfo: @unroll
 @inline _fill_north_halo!(i, k, grid, c, ::FBC, args...)  =  _fill_flux_north_halo!(i, 1, k, grid, c)
 @inline _fill_bottom_halo!(i, j, grid, c, ::FBC, args...) = _fill_flux_bottom_halo!(i, j, 1, grid, c)
 @inline _fill_top_halo!(i, j, grid, c, ::FBC, args...)    =    _fill_flux_top_halo!(i, j, 1, grid, c)
-
-#####
-##### Single halo filling kernels
-#####
-
-@kernel function fill_flux_west_halo!(c, grid)
-    j, k = @index(Global, NTuple)
-    _fill_flux_west_halo!(1, j, k, grid, c)
-end
-
-@kernel function fill_flux_south_halo!(c, grid)
-    i, k = @index(Global, NTuple)
-    _fill_flux_south_halo!(i, 1, k, grid, c)
-end
-
-@kernel function fill_flux_bottom_halo!(c, grid)
-    i, j = @index(Global, NTuple)
-    _fill_flux_bottom_halo!(i, j, 1, grid, c)
-end
-
-@kernel function fill_flux_east_halo!(c, grid)
-    j, k = @index(Global, NTuple)
-    _fill_flux_east_halo!(1, j, k, grid, c)
-end
-
-@kernel function fill_flux_north_halo!(c, grid)
-    i, k = @index(Global, NTuple)
-    _fill_flux_north_halo!(i, 1, k, grid, c)
-end
-
-@kernel function fill_flux_top_halo!(c, grid)
-    i, j = @index(Global, NTuple)
-    _fill_flux_top_halo!(i, j, 1, grid, c)
-end
-
-#####
-##### Kernel launchers for flux boundary conditions
-#####
-
-fill_west_halo!(c, bc::FBC, kernel_size, offset, loc, arch, grid, args...; kwargs...) = 
-            launch!(arch, grid, KernelParameters(kernel_size, offset), fill_flux_west_halo!, c,grid; kwargs...)
-fill_east_halo!(c, bc::FBC, kernel_size, offset, loc, arch, grid, args...; kwargs...) = 
-            launch!(arch, grid, KernelParameters(kernel_size, offset), fill_flux_east_halo!, c, grid; kwargs...)
-fill_south_halo!(c, bc::FBC, kernel_size, offset, loc, arch, grid, args...; kwargs...) = 
-            launch!(arch, grid, KernelParameters(kernel_size, offset), fill_flux_south_halo!, c, grid; kwargs...)
-fill_north_halo!(c, bc::FBC, kernel_size, offset, loc, arch, grid, args...; kwargs...) = 
-            launch!(arch, grid, KernelParameters(kernel_size, offset), fill_flux_north_halo!, c, grid; kwargs...)
-fill_bottom_halo!(c, bc::FBC, kernel_size, offset, loc, arch, grid, args...; kwargs...) = 
-            launch!(arch, grid, KernelParameters(kernel_size, offset), fill_flux_bottom_halo!, c, grid; kwargs...)
-fill_top_halo!(c, bc::FBC, kernel_size, offset, loc, arch, grid, args...; kwargs...) = 
-            launch!(arch, grid, KernelParameters(kernel_size, offset), fill_flux_top_halo!, c, grid; kwargs...)
-

src/BoundaryConditions/fill_halo_regions_periodic.jl

@@ -106,3 +106,14 @@ end
         end
     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"))