From 6ca6280e698ee1ce1173595e95138696130f0ced Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Fri, 27 Mar 2026 21:28:02 +1100 Subject: [PATCH 01/28] Add distributed fold topology types and fix tripolar halo fill Introduce 4 new topology types for distributed tripolar grid fold boundaries: - LeftConnectedRightCenterFolded (1xN UPivot) - LeftConnectedRightFaceFolded (1xN FPivot, Face-extended) - LeftConnectedRightCenterConnected{WestOfPivot/EastOfPivot} (MxN UPivot) - LeftConnectedRightFaceConnected{WestOfPivot/EastOfPivot} (MxN FPivot, Face-extended) These replace the previous FullyConnected y-topology on northernmost distributed ranks, which lost fold information and caused FPivot CF/FF fields to crash (OOB) because they need Ny+1 Face points. Key changes: - Grids.jl: new types, WestOfPivot/EastOfPivot, global_fold_topology() - grid_utils.jl: BoundedTopology union includes Face-extended fold types - distributed_grids.jl: insert_connected_topology 5-arg methods for fold topologies - distributed_zipper.jl: complete rewrite with topology-based buffer dispatch, corrected y-ranges for Face-extended grids, fold-line WoP/EoP handling - Tripolar struct simplified to 3 type params (fold info now in grid topology) - Split-explicit, advection, and BC unions updated for new types Co-Authored-By: Claude Opus 4.6 (1M context) --- ...topologically_conditional_interpolation.jl | 4 +- .../field_boundary_conditions.jl | 7 + .../distributed_grids.jl | 22 + src/Grids/Grids.jl | 56 +++ src/Grids/grid_utils.jl | 4 +- .../SplitExplicitFreeSurfaces.jl | 6 +- .../split_explicit_free_surface.jl | 11 +- .../distributed_tripolar_grid.jl | 35 +- .../distributed_zipper.jl | 426 ++++++++++++++++-- .../tripolar_field_extensions.jl | 21 +- .../halo_fill_viz.jl | 7 +- 11 files changed, 524 insertions(+), 75 deletions(-) diff --git a/src/Advection/topologically_conditional_interpolation.jl b/src/Advection/topologically_conditional_interpolation.jl index c7d1068ede4..4aad0044db4 100644 --- a/src/Advection/topologically_conditional_interpolation.jl +++ b/src/Advection/topologically_conditional_interpolation.jl @@ -20,7 +20,9 @@ using Oceananigans.Grids: AbstractGrid, const AG = AbstractGrid # topologies bounded at least on one side -const BT = Union{Bounded, RightConnected, LeftConnected} +const BT = Union{Bounded, RightConnected, LeftConnected, + LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded, + LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected} # Bounded underlying Grids const AGX = AG{<:Any, <:BT} diff --git a/src/BoundaryConditions/field_boundary_conditions.jl b/src/BoundaryConditions/field_boundary_conditions.jl index a7da506c111..194babc0498 100644 --- a/src/BoundaryConditions/field_boundary_conditions.jl +++ b/src/BoundaryConditions/field_boundary_conditions.jl @@ -21,6 +21,10 @@ default_prognostic_bc(::RightConnected, ::Center, default) = default.boundary_c default_prognostic_bc(::RightFaceFolded, ::Center, default) = default.boundary_condition default_prognostic_bc(::RightCenterFolded, ::Center, default) = default.boundary_condition +const DistributedFoldTopology = Union{LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded, + LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected} + +default_prognostic_bc(::DistributedFoldTopology, ::Center, default) = default.boundary_condition # TODO: make model constructors enforce impenetrability on velocity components to simplify this code default_prognostic_bc(::Bounded, ::Face, default) = ImpenetrableBoundaryCondition() @@ -28,6 +32,7 @@ default_prognostic_bc(::LeftConnected, ::Face, default) = ImpenetrableBoundaryC default_prognostic_bc(::RightConnected, ::Face, default) = ImpenetrableBoundaryCondition() default_prognostic_bc(::RightFaceFolded, ::Face, default) = ImpenetrableBoundaryCondition() default_prognostic_bc(::RightCenterFolded, ::Face, default) = ImpenetrableBoundaryCondition() +default_prognostic_bc(::DistributedFoldTopology, ::Face, default) = ImpenetrableBoundaryCondition() default_prognostic_bc(::Bounded, ::Nothing, default) = nothing default_prognostic_bc(::Flat, ::Nothing, default) = nothing @@ -37,6 +42,7 @@ default_prognostic_bc(::LeftConnected, ::Nothing, default) = nothing default_prognostic_bc(::RightConnected, ::Nothing, default) = nothing default_prognostic_bc(::RightFaceFolded, ::Nothing, default) = nothing default_prognostic_bc(::RightCenterFolded, ::Nothing, default) = nothing +default_prognostic_bc(::DistributedFoldTopology, ::Nothing, default) = nothing _default_auxiliary_bc(topo, loc) = default_prognostic_bc(topo, loc, DefaultBoundaryCondition()) _default_auxiliary_bc(::Bounded, ::Face) = nothing @@ -44,6 +50,7 @@ _default_auxiliary_bc(::RightConnected, ::Face) = nothing _default_auxiliary_bc(::LeftConnected, ::Face) = nothing _default_auxiliary_bc(::RightFaceFolded, ::Face) = nothing _default_auxiliary_bc(::RightCenterFolded, ::Face) = nothing +_default_auxiliary_bc(::DistributedFoldTopology, ::Face) = nothing default_auxiliary_bc(grid, ::Val{:east}, loc) = _default_auxiliary_bc(topology(grid, 1)(), loc[1]) default_auxiliary_bc(grid, ::Val{:west}, loc) = _default_auxiliary_bc(topology(grid, 1)(), loc[1]) diff --git a/src/DistributedComputations/distributed_grids.jl b/src/DistributedComputations/distributed_grids.jl index 7eded1e9d73..f274bb2be7e 100644 --- a/src/DistributedComputations/distributed_grids.jl +++ b/src/DistributedComputations/distributed_grids.jl @@ -349,6 +349,28 @@ insert_connected_topology(::Type{Bounded}, R, r) = ifelse(R == 1, Bounded, insert_connected_topology(::Type{Periodic}, R, r) = ifelse(R == 1, Periodic, FullyConnected) +# Fold-aware topology insertion for distributed tripolar grids. +# These take 5 arguments: (global_y_topology, Ry, ry, Rx, rx) where +# Ry/ry are y-rank count/index and Rx/rx are x-rank count/index (all 1-based). + +pivot_side(Rx, rx) = rx ≤ Rx ÷ 2 ? WestOfPivot : EastOfPivot + +local_fold_topology(::Type{RightCenterFolded}, Rx, rx) = + Rx == 1 ? LeftConnectedRightCenterFolded : LeftConnectedRightCenterConnected{pivot_side(Rx, rx)} + +local_fold_topology(::Type{RightFaceFolded}, Rx, rx) = + Rx == 1 ? LeftConnectedRightFaceFolded : LeftConnectedRightFaceConnected{pivot_side(Rx, rx)} + +function insert_connected_topology(T::Type{<:Union{RightCenterFolded, RightFaceFolded}}, Ry, ry, Rx, rx) + if ry == 1 + return RightConnected + elseif ry == Ry + return local_fold_topology(T, Rx, rx) + else + return FullyConnected + end +end + """ reconstruct_global_topology(T, R, r, r1, r2, arch) diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index 7843ee1826b..2e08d2db7de 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -4,6 +4,10 @@ export Center, Face export AbstractTopology, topology export Periodic, Bounded, Flat, FullyConnected, LeftConnected, RightConnected export RightFaceFolded, RightCenterFolded +export LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded +export LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected +export WestOfPivot, EastOfPivot +export global_fold_topology export AbstractGrid, AbstractUnderlyingGrid, halo_size, total_size export RectilinearGrid export AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid @@ -117,6 +121,58 @@ Grid topology for tripolar U-point pivot connection. """ struct RightCenterFolded <: AbstractTopology end +# Pivot side indicators for distributed fold topologies +struct WestOfPivot end +struct EastOfPivot end + +""" + LeftConnectedRightCenterFolded + +Local grid topology for the northernmost y-rank of a 1×N distributed tripolar grid +with U-point pivot (serial fold). Connected to the south neighbor on the left, +center-folded on the right (north). +""" +struct LeftConnectedRightCenterFolded <: AbstractTopology end + +""" + LeftConnectedRightFaceFolded + +Local grid y topology for the northernmost y-rank of a 1×N distributed tripolar grid +with F-point pivot (serial fold). Connected to the south neighbor on the left, +face-folded on the right (north). Face-extended (Ny+1 Face points in y). +""" +struct LeftConnectedRightFaceFolded <: AbstractTopology end + +""" + LeftConnectedRightCenterConnected{P} + +Local grid y topology for the northernmost y-rank of an M×N distributed tripolar grid +with U-point pivot (distributed zipper). `P` is `WestOfPivot` or `EastOfPivot`. +""" +struct LeftConnectedRightCenterConnected{P} <: AbstractTopology end + +""" + LeftConnectedRightFaceConnected{P} + +Local grid y topology for the northernmost y-rank of an M×N distributed tripolar grid +with F-point pivot (distributed zipper). `P` is `WestOfPivot` or `EastOfPivot`. +Face-extended (Ny+1 Face points in y). +""" +struct LeftConnectedRightFaceConnected{P} <: AbstractTopology end + +""" + global_fold_topology(T) + +Return the global grid y fold topology (`RightCenterFolded` or `RightFaceFolded`) from +the local grid y topology type. +""" +global_fold_topology(::Type{RightCenterFolded}) = RightCenterFolded +global_fold_topology(::Type{RightFaceFolded}) = RightFaceFolded +global_fold_topology(::Type{LeftConnectedRightCenterFolded}) = RightCenterFolded +global_fold_topology(::Type{LeftConnectedRightFaceFolded}) = RightFaceFolded +global_fold_topology(::Type{<:LeftConnectedRightCenterConnected}) = RightCenterFolded +global_fold_topology(::Type{<:LeftConnectedRightFaceConnected}) = RightFaceFolded + ##### ##### Directions (for tilted domains) ##### diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index 4c17c7efbd5..d5f7112117d 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -40,7 +40,9 @@ end end end -const BoundedTopology = Union{Bounded, LeftConnected, RightFaceFolded} +const BoundedTopology = Union{Bounded, LeftConnected, RightFaceFolded, + LeftConnectedRightFaceFolded, + LeftConnectedRightFaceConnected} const AT = AbstractTopology Base.length(::Face, ::BoundedTopology, N) = N + 1 diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl index 1ac4e176980..e6adc8480db 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl @@ -10,7 +10,11 @@ using Oceananigans.ImmersedBoundaries: column_depthTᶠᶜᵃ, column_depthTᶜ using Oceananigans.Operators: ∂xᵣTᶠᶜᶠ, ∂xᵣᶠᶜᶠ, ∂yᵣTᶜᶠᶠ, ∂yᵣᶜᶠᶠ, δxTᶜᵃᵃ, δxᶜᵃᵃ, δyTᵃᶜᵃ, δyᵃᶜᵃ using Oceananigans.BoundaryConditions: FieldBoundaryConditions, fill_halo_regions! using Oceananigans.Fields: Field -using Oceananigans.Grids: Center, Face, get_active_column_map, topology +using Oceananigans.Grids: Center, Face, get_active_column_map, topology, + LeftConnected, RightConnected, FullyConnected, + RightCenterFolded, RightFaceFolded, + LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded, + LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected using Oceananigans.ImmersedBoundaries: mask_immersed_field! using Oceananigans.Models.HydrostaticFreeSurfaceModels: AbstractFreeSurface, free_surface_displacement_field, diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl index 9edbcf39afa..e567e666478 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl @@ -225,7 +225,10 @@ function hydrostatic_tendency_fields(velocities, free_surface::SplitExplicitFree return merge((u=u, v=v, U=U, V=V), tracers) end -const ConnectedTopology = Union{LeftConnected, RightConnected, FullyConnected, RightCenterFolded, RightFaceFolded} +const ConnectedTopology = Union{LeftConnected, RightConnected, FullyConnected, + RightCenterFolded, RightFaceFolded, + LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded, + LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected} # Internal function for HydrostaticFreeSurfaceModel function materialize_free_surface(free_surface::SplitExplicitFreeSurface{extend_halos}, velocities, grid) where {extend_halos} @@ -403,6 +406,12 @@ split_explicit_kernel_size(::Type{LeftConnected}, N, H) = -H+2:N split_explicit_kernel_size(::Type{RightCenterFolded}, N, H) = 1:N+H-1 split_explicit_kernel_size(::Type{RightFaceFolded}, N, H) = 1:N+H-1 +# Distributed fold topologies: connected on both sides (left=MPI, right=fold/zipper) +split_explicit_kernel_size(::Type{LeftConnectedRightCenterFolded}, N, H) = -H+2:N+H-1 +split_explicit_kernel_size(::Type{LeftConnectedRightFaceFolded}, N, H) = -H+2:N+H-1 +split_explicit_kernel_size(::Type{<:LeftConnectedRightCenterConnected}, N, H) = -H+2:N+H-1 +split_explicit_kernel_size(::Type{<:LeftConnectedRightFaceConnected}, N, H) = -H+2:N+H-1 + # Adapt Adapt.adapt_structure(to, free_surface::SplitExplicitFreeSurface{extend_halos}) where {extend_halos} = SplitExplicitFreeSurface{extend_halos}(Adapt.adapt(to, free_surface.displacement), diff --git a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl index 788d6499573..5d744d03090 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl @@ -1,4 +1,4 @@ -using Oceananigans.BoundaryConditions: DistributedCommunicationBoundaryCondition +using Oceananigans.BoundaryConditions: DistributedCommunicationBoundaryCondition, FZBC using Oceananigans.Fields: validate_indices, validate_field_data using Oceananigans.DistributedComputations: DistributedComputations, @@ -11,7 +11,8 @@ using Oceananigans.DistributedComputations: concatenate_local_sizes, communication_buffers -using Oceananigans.Grids: topology, RightConnected, FullyConnected +using Oceananigans.Grids: topology, RightConnected, FullyConnected, global_fold_topology +using Oceananigans.DistributedComputations: insert_connected_topology import Oceananigans.Fields: Field, validate_indices, validate_boundary_conditions @@ -109,8 +110,14 @@ function TripolarGrid(arch::Distributed, FT::DataType=Float64; Azᶜᶠᵃ = partition_tripolar_metric(global_grid, :Azᶜᶠᵃ, irange, jrange) Azᶠᶠᵃ = partition_tripolar_metric(global_grid, :Azᶠᶠᵃ, irange, jrange) - LY = yrank == 0 ? RightConnected : FullyConnected LX = workers[1] == 1 ? Periodic : FullyConnected + + global_fold_topology = topology(global_grid, 2) + + # 1-based indices for insert_connected_topology + Rx, Ry = workers[1], workers[2] + rx, ry = xrank + 1, yrank + 1 + LY = insert_connected_topology(global_fold_topology, Ry, ry, Rx, rx) ny = nylocal[yrank+1] nx = nxlocal[xrank+1] @@ -228,9 +235,10 @@ function receiving_rank(arch; receive_idx_x = ranks(arch)[1] - arch.local_index[ return receive_rank end -# a distributed `TripolarGrid` needs a `UPivotZipperBoundaryCondition` for the north boundary -# only on the last rank -# TODO: generalize to any ZipperBoundaryCondition +zipper_bc(::Type{RightCenterFolded}, sign) = UPivotZipperBoundaryCondition(sign) +zipper_bc(::Type{RightFaceFolded}, sign) = FPivotZipperBoundaryCondition(sign) + +const ZBC = Union{<:UZBC, <:FZBC} function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, grid::MPITripolarGridOfSomeKind, field_name::Symbol, @@ -248,7 +256,7 @@ function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, south = regularize_boundary_condition(bcs.south, grid, loc, 2, LeftBoundary, prognostic_names) north = if yrank == processor_size[2] - 1 && processor_size[1] == 1 - UPivotZipperBoundaryCondition(sign) + zipper_bc(global_fold_topology(topology(grid, 2)), sign) elseif yrank == processor_size[2] - 1 && processor_size[1] != 1 from = arch.local_rank @@ -280,26 +288,21 @@ function Field(loc::Tuple{<:LX, <:LY, <:LZ}, grid::MPITripolarGridOfSomeKind, da validate_field_data(loc, data, grid, indices) validate_boundary_conditions(loc, grid, old_bcs) - default_zipper = UPivotZipperBoundaryCondition(sign(LX, LY)) - if isnothing(old_bcs) || ismissing(old_bcs) new_bcs = old_bcs else new_bcs = inject_halo_communication_boundary_conditions(old_bcs, arch.local_rank, arch.connectivity, topology(grid)) - # North boundary conditions are "special". If we are at the top of the domain, i.e. - # the last rank, then we need to substitute the BC only if the old one is not already - # a zipper boundary condition. Otherwise we always substitute because we need to - # inject the halo boundary conditions. if yrank == processor_size[2] - 1 && processor_size[1] == 1 - north_bc = if !(old_bcs.north isa UZBC) + default_zipper = zipper_bc(global_fold_topology(topology(grid, 2)), sign(LX, LY)) + north_bc = if !(old_bcs.north isa ZBC) default_zipper else old_bcs.north end elseif yrank == processor_size[2] - 1 && processor_size[1] != 1 - sgn = old_bcs.north isa UZBC ? old_bcs.north.condition : sign(LX, LY) + sgn = old_bcs.north isa ZBC ? old_bcs.north.condition : sign(LX, LY) from = arch.local_rank to = arch.connectivity.north halo_communication = ZipperHaloCommunicationRanks(sgn; from, to) @@ -317,7 +320,7 @@ function Field(loc::Tuple{<:LX, <:LY, <:LZ}, grid::MPITripolarGridOfSomeKind, da bottom=new_bcs.bottom) end - buffers = communication_buffers(grid, data, new_bcs) + buffers = communication_buffers(grid, data, new_bcs, (LX(), LY(), LZ())) return Field{LX, LY, LZ}(grid, data, new_bcs, indices, op, status, buffers) end diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 4c80de8be28..8278ed81eea 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -1,57 +1,409 @@ using Oceananigans.BoundaryConditions: get_boundary_kernels, DistributedCommunication using Oceananigans.DistributedComputations: cooperative_waitall!, recv_from_buffers!, distributed_fill_halo_event! using Oceananigans.DistributedComputations: CommunicationBuffers, fill_corners!, loc_id, AsynchronousDistributed +using Oceananigans.Grids: AbstractGrid, topology, + RightCenterFolded, RightFaceFolded, + LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded, + LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected, + WestOfPivot, EastOfPivot +using Oceananigans.DistributedComputations: Distributed, on_architecture, ranks, x_communication_buffer using Oceananigans.Fields: instantiated_location import Oceananigans.BoundaryConditions: fill_halo_regions! -import Oceananigans.DistributedComputations: synchronize_communication! +import Oceananigans.DistributedComputations: synchronize_communication!, + y_communication_buffer, corner_communication_buffer, + _fill_north_send_buffer!, _recv_from_north_buffer!, + _fill_northwest_send_buffer!, _fill_northeast_send_buffer!, + _recv_from_northwest_buffer!, _recv_from_northeast_buffer! +import Oceananigans.Fields: communication_buffers @inline instantiate(T::DataType) = T() @inline instantiate(T) = T const DistributedZipper = BoundaryCondition{<:DistributedCommunication, <:ZipperHaloCommunicationRanks} -switch_north_halos!(c, north_bc, grid, loc) = nothing +##### +##### Topology unions for dispatch +##### -function switch_north_halos!(c, north_bc::DistributedZipper, grid, loc) - sign = north_bc.condition.sign - hz = halo_size(grid) - sz = size(grid) +const UPivotTopology = Union{RightCenterFolded, + LeftConnectedRightCenterFolded, + LeftConnectedRightCenterConnected} +const FPivotTopology = Union{RightFaceFolded, + LeftConnectedRightFaceFolded, + LeftConnectedRightFaceConnected} - _switch_north_halos!(parent(c), loc, sign, sz, hz) +# 1D fold (1xN, y-partitioned only) vs 2D fold (MxN, x+y partitioned) +const OneDFoldTopology = Union{LeftConnectedRightCenterFolded, + LeftConnectedRightFaceFolded} +const TwoDFoldTopology = Union{LeftConnectedRightCenterConnected, + LeftConnectedRightFaceConnected} - return nothing +# UPivot fold line is at Center-y; FPivot fold line is at Face-y +has_fold_line(::Type{<:UPivotTopology}, ::Center) = true +has_fold_line(::Type{<:UPivotTopology}, ::Face) = false +has_fold_line(::Type{<:FPivotTopology}, ::Center) = false +has_fold_line(::Type{<:FPivotTopology}, ::Face) = true + +##### +##### Zipper communication buffers with fold-aware packing +##### + +struct OneDZipperBuffer{Loc, FoT, B, S} + send :: B + recv :: B + sign :: S +end + +struct TwoDZipperBuffer{Loc, FoT, B, S} + send :: B + recv :: B + sign :: S +end + +struct ZipperCornerBuffer{Loc, FoT, B, S} + send :: B + recv :: B + sign :: S end -@inline reversed_halos(::Tuple{<:Any, <:Center, <:Any}, Ny, Hy) = Ny+2Hy-1:-1:Ny+Hy+1 -@inline reversed_halos(::Tuple{<:Any, <:Face, <:Any}, Ny, Hy) = Ny+2Hy:-1:Ny+Hy+2 +# Value-argument constructors: all types inferred +OneDZipperBuffer(loc::Loc, fot::FoT, send::B, recv::B, sign::S) where {Loc, FoT, B, S} = OneDZipperBuffer{Loc, FoT, B, S}(send, recv, sign) +TwoDZipperBuffer(loc::Loc, fot::FoT, send::B, recv::B, sign::S) where {Loc, FoT, B, S} = TwoDZipperBuffer{Loc, FoT, B, S}(send, recv, sign) +ZipperCornerBuffer(loc::Loc, fot::FoT, send::B, recv::B, sign::S) where {Loc, FoT, B, S} = ZipperCornerBuffer{Loc, FoT, B, S}(send, recv, sign) -@inline adjust_x_face!(c, loc, north_halos, Px) = nothing -@inline adjust_x_face!(c, ::Tuple{<:Face, <:Any, <:Any}, north_halos, Px) = view(c, 2:Px, north_halos, :) .= view(c, 1:Px-1, north_halos, :) +Adapt.adapt_structure(to, buff::OneDZipperBuffer) = nothing +Adapt.adapt_structure(to, buff::TwoDZipperBuffer) = nothing +Adapt.adapt_structure(to, buff::ZipperCornerBuffer) = nothing -# We throw away the first point! -@inline function _switch_north_halos!(c, loc, sign, (Nx, Ny, Nz), (Hx, Hy, Hz)) +##### +##### Buffer construction for tripolar grids +##### - # Domain indices common for all locations - north_halos = Ny+Hy+1:Ny+2Hy-1 - west_corner = 1:Hx - east_corner = Nx+Hx+1:Nx+2Hx - interior = Hx+1:Nx+Hx +function communication_buffers(grid::MPITripolarGridOfSomeKind, data, bcs, loc) + Hx, Hy, Hz = halo_size(grid) + arch = architecture(grid) - # Location - dependent halo indices - reversed_north_halos = reversed_halos(loc, Ny, Hy) + west = x_communication_buffer(arch, grid, data, Hx, bcs.west) + east = x_communication_buffer(arch, grid, data, Hx, bcs.east) + south = y_communication_buffer(arch, grid, data, Hy, bcs.south) + north = y_tripolar_buffer(arch, grid, data, Hy, bcs.north, loc) - view(c, west_corner, north_halos, :) .= sign .* reverse(view(c, west_corner, reversed_north_halos, :), dims = 1) - view(c, east_corner, north_halos, :) .= sign .* reverse(view(c, east_corner, reversed_north_halos, :), dims = 1) - view(c, interior, north_halos, :) .= sign .* reverse(view(c, interior, reversed_north_halos, :), dims = 1) + sw = corner_communication_buffer(arch, grid, data, Hx, Hy, west, south) + se = corner_communication_buffer(arch, grid, data, Hx, Hy, east, south) + nw = northwest_tripolar_buffer(arch, grid, data, Hx, Hy, west, north) + ne = northeast_tripolar_buffer(arch, grid, data, Hx, Hy, east, north) - # throw out first point for the x - face locations - adjust_x_face!(c, loc, north_halos, size(c, 1)) + return CommunicationBuffers(west, east, south, north, sw, se, nw, ne) +end - return nothing +# Fallback: non-zipper north BC uses standard buffer +y_tripolar_buffer(arch, grid, data, H, bc, loc) = y_communication_buffer(arch, grid, data, H, bc) + +# 1D fold (1xN) → OneDZipperBuffer (full-width) +function y_tripolar_buffer(arch::Distributed, grid::AbstractGrid{<:Any, <:Any, Topo}, + data, H, bc::DistributedZipper, loc::Loc) where {Topo <: OneDFoldTopology, Loc} + Tx, _, Tz = size(parent(data)) + FT = eltype(data) + sgn = bc.condition.sign + send = on_architecture(arch, zeros(FT, Tx, H, Tz)) + recv = on_architecture(arch, zeros(FT, Tx, H, Tz)) + return OneDZipperBuffer(loc, Topo(), send, recv, sgn) +end + +# 2D fold (MxN) → TwoDZipperBuffer (interior-width, Hy or Hy+1 rows for fold line) +function y_tripolar_buffer(arch::Distributed, grid::AbstractGrid{<:Any, <:Any, Topo}, + data, H, bc::DistributedZipper, loc::Loc) where {Topo <: TwoDFoldTopology, Loc} + Nx = size(grid, 1) + _, _, Tz = size(parent(data)) + FT = eltype(data) + sgn = bc.condition.sign + Hbuf = has_fold_line(Topo, loc[2]) ? H + 1 : H + send = on_architecture(arch, zeros(FT, Nx, Hbuf, Tz)) + recv = on_architecture(arch, zeros(FT, Nx, Hbuf, Tz)) + return TwoDZipperBuffer(loc, Topo(), send, recv, sgn) +end + +# Fallbacks: non-zipper corners +northwest_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge) = corner_communication_buffer(arch, grid, data, Hx, Hy, xedge, yedge) +northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge) = corner_communication_buffer(arch, grid, data, Hx, Hy, xedge, yedge) + +# ── NW corner: Center-x (Hx columns) ── +function northwest_tripolar_buffer(arch::Distributed, grid, data, Hx, Hy, xedge, + yedge::Union{<:OneDZipperBuffer{Loc, FoT}, + <:TwoDZipperBuffer{Loc, FoT}}) where {Loc <: Tuple{<:Center, <:Any, <:Any}, FoT} + Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign + send = on_architecture(arch, zeros(FT, Hx, Hy, Tz)) + recv = on_architecture(arch, zeros(FT, Hx, Hy, Tz)) + return ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sgn)}(send, recv, sgn) +end + +# ── NW corner: Face-x (Hx+1 columns for the Face-x shift) ── +function northwest_tripolar_buffer(arch::Distributed, grid, data, Hx, Hy, xedge, + yedge::Union{<:OneDZipperBuffer{Loc, FoT}, + <:TwoDZipperBuffer{Loc, FoT}}) where {Loc <: Tuple{<:Face, <:Any, <:Any}, FoT} + Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign + send = on_architecture(arch, zeros(FT, Hx+1, Hy, Tz)) + recv = on_architecture(arch, zeros(FT, Hx+1, Hy, Tz)) + return ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sgn)}(send, recv, sgn) +end + +# ── NE corner: Center-x (Hx columns) ── +function northeast_tripolar_buffer(arch::Distributed, grid, data, Hx, Hy, xedge, + yedge::Union{<:OneDZipperBuffer{Loc, FoT}, + <:TwoDZipperBuffer{Loc, FoT}}) where {Loc <: Tuple{<:Center, <:Any, <:Any}, FoT} + Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign + send = on_architecture(arch, zeros(FT, Hx, Hy, Tz)) + recv = on_architecture(arch, zeros(FT, Hx, Hy, Tz)) + return ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sgn)}(send, recv, sgn) +end + +# ── NE corner: Face-x (Hx-1 columns, TwoDBuffer covers the extra) ── +function northeast_tripolar_buffer(arch::Distributed, grid, data, Hx, Hy, xedge, + yedge::Union{<:OneDZipperBuffer{Loc, FoT}, + <:TwoDZipperBuffer{Loc, FoT}}) where {Loc <: Tuple{<:Face, <:Any, <:Any}, FoT} + Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign + Hx_ne = max(Hx - 1, 0) + send = on_architecture(arch, zeros(FT, Hx_ne, Hy, Tz)) + recv = on_architecture(arch, zeros(FT, Hx_ne, Hy, Tz)) + return ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sgn)}(send, recv, sgn) +end + +##### +##### Location aliases for dispatch +##### + +const CC = Tuple{<:Center, <:Center, <:Any} +const FC = Tuple{<:Face, <:Center, <:Any} +const CF = Tuple{<:Center, <:Face, <:Any} +const FF = Tuple{<:Face, <:Face, <:Any} + +##### +##### Fold-aware send buffer packing: _fill_north_send_buffer! +##### +##### Y-range rules (parent-array coords, Ny = size(grid, 2)): +##### UPivot Center-y (CC/FC): skip fold at Ny → Hy rows from Ny-1: Ny+Hy-1:-1:Ny +##### All other cases: Hy rows from Ny: Ny+Hy:-1:Ny+1 +##### +##### For TwoD fold-line fields (Hy+1 buffer): row 1 = fold line, rows 2:Hy+1 = halo sources +##### + +# ── TwoDZipperBuffer: UPivot Center-y (CC/FC) — fold line + Hy halo rows ── + +function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) + view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:Ny+Hy, :) + view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy-1:-1:Ny, :) +end + +function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) + view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:Ny+Hy, :) + view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy-1:-1:Ny, :) +end + +# ── TwoDZipperBuffer: UPivot Face-y (CF/FF) — no fold line, Hy halo rows ── + +_fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +_fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) + +# ── TwoDZipperBuffer: FPivot Center-y (CC/FC) — no fold line, Hy halo rows ── + +_fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +_fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) + +# ── TwoDZipperBuffer: FPivot Face-y (CF/FF) — fold line + Hy halo rows ── + +function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) + view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+1+Hy:Ny+1+Hy, :) + view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end -# Disambiguation +function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) + view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+1+Hy:Ny+1+Hy, :) + view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +end + +# ── OneDZipperBuffer: UPivot Center-y (CC/FC) — simple reverse ── + +_fill_north_send_buffer!(c, b::OneDZipperBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, size(c,1):-1:1, Ny+Hy-1:-1:Ny, :) + +function _fill_north_send_buffer!(c, b::OneDZipperBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) + Tx = size(c, 1) + view(b.send, 2:Tx, :, :) .= b.sign .* view(c, Tx:-1:2, Ny+Hy-1:-1:Ny, :) + view(b.send, 1:1, :, :) .= b.sign .* view(c, 1:1, Ny+Hy-1:-1:Ny, :) +end + +# ── OneDZipperBuffer: UPivot Face-y (CF/FF) ── + +_fill_north_send_buffer!(c, b::OneDZipperBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, size(c,1):-1:1, Ny+Hy:-1:Ny+1, :) + +function _fill_north_send_buffer!(c, b::OneDZipperBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) + Tx = size(c, 1) + view(b.send, 2:Tx, :, :) .= b.sign .* view(c, Tx:-1:2, Ny+Hy:-1:Ny+1, :) + view(b.send, 1:1, :, :) .= b.sign .* view(c, 1:1, Ny+Hy:-1:Ny+1, :) +end + +# ── OneDZipperBuffer: FPivot Center-y (CC/FC) ── + +_fill_north_send_buffer!(c, b::OneDZipperBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, size(c,1):-1:1, Ny+Hy:-1:Ny+1, :) + +function _fill_north_send_buffer!(c, b::OneDZipperBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) + Tx = size(c, 1) + view(b.send, 2:Tx, :, :) .= b.sign .* view(c, Tx:-1:2, Ny+Hy:-1:Ny+1, :) + view(b.send, 1:1, :, :) .= b.sign .* view(c, 1:1, Ny+Hy:-1:Ny+1, :) +end + +# ── OneDZipperBuffer: FPivot Face-y (CF/FF) ── + +_fill_north_send_buffer!(c, b::OneDZipperBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, size(c,1):-1:1, Ny+Hy:-1:Ny+1, :) + +function _fill_north_send_buffer!(c, b::OneDZipperBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) + Tx = size(c, 1) + view(b.send, 2:Tx, :, :) .= b.sign .* view(c, Tx:-1:2, Ny+Hy:-1:Ny+1, :) + view(b.send, 1:1, :, :) .= b.sign .* view(c, 1:1, Ny+Hy:-1:Ny+1, :) +end + +##### +##### Fold-aware send buffer packing: NW corner (leftmost Hx interior columns) +##### + +# Center-x NW +_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy-1:-1:Ny, :) +_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) + +# Face-x NW: Hx+1 columns from leftmost Hx+1 interior columns +_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy-1:-1:Ny, :) +_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) + +##### +##### Fold-aware send buffer packing: NE corner (rightmost Hx interior columns) +##### + +# Center-x NE +_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy-1:-1:Ny, :) +_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) +_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) +_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) + +# Face-x NE: Hx-1 columns from rightmost Hx-1 interior columns +_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy-1:-1:Ny, :) +_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) +_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) +_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) + +##### +##### Recv methods: direct placement (data is already folded by sender) +##### +##### For TwoD fold-line fields (Hy+1 buffer): +##### EastOfPivot: write fold line + halos +##### WestOfPivot: write halos only (skip fold line in row 1) +##### +##### For FPivot Face-y: halos start at parent Ny+Hy+2 (field Ny+2, past fold line at Ny+1) +##### + +# ── TwoDZipperBuffer: UPivot CC/FC — fold line dispatch on WoP/EoP ── + +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:LeftConnectedRightCenterConnected{EastOfPivot}}, Hx, Hy, Nx, Ny) + view(c, 1+Hx:Nx+Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) + view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +end + +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:LeftConnectedRightCenterConnected{WestOfPivot}}, Hx, Hy, Nx, Ny) + view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +end + +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected{EastOfPivot}}, Hx, Hy, Nx, Ny) + view(c, 2+Hx:Nx+Hx+1, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) + view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +end + +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected{WestOfPivot}}, Hx, Hy, Nx, Ny) + view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +end + +# ── TwoDZipperBuffer: UPivot CF/FF — no fold line, standard placement ── + +_recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv + +# ── TwoDZipperBuffer: FPivot CC/FC — no fold line, standard placement ── + +_recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv + +# ── TwoDZipperBuffer: FPivot CF/FF — fold line dispatch, shifted +1 for Face-extended ── + +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CF, <:LeftConnectedRightFaceConnected{EastOfPivot}}, Hx, Hy, Nx, Ny) + view(c, 1+Hx:Nx+Hx, Ny+1+Hy:Ny+1+Hy, :) .= view(buff.recv, :, 1:1, :) + view(c, 1+Hx:Nx+Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +end + +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CF, <:LeftConnectedRightFaceConnected{WestOfPivot}}, Hx, Hy, Nx, Ny) + view(c, 1+Hx:Nx+Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +end + +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FF, <:LeftConnectedRightFaceConnected{EastOfPivot}}, Hx, Hy, Nx, Ny) + view(c, 2+Hx:Nx+Hx+1, Ny+1+Hy:Ny+1+Hy, :) .= view(buff.recv, :, 1:1, :) + view(c, 2+Hx:Nx+Hx+1, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +end + +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FF, <:LeftConnectedRightFaceConnected{WestOfPivot}}, Hx, Hy, Nx, Ny) + view(c, 2+Hx:Nx+Hx+1, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +end + +# ── OneDZipperBuffer: UPivot — standard placement ── + +_recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:Any, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, :, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv + +# ── OneDZipperBuffer: FPivot CC/FC — standard placement ── + +_recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, :, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, :, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv + +# ── OneDZipperBuffer: FPivot CF/FF — shifted +1 for Face-extended ── + +_recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, :, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv +_recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, :, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv + +# ── Corner recv: UPivot — standard placement ── + +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv + +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv + +# ── Corner recv: FPivot CC/FC — standard placement ── + +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv + +# ── Corner recv: FPivot CF/FF — shifted +1 for Face-extended ── + +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv + +##### +##### _switch_north_halos! — no-op (fold logic in send buffers) +##### + +switch_north_halos!(c, north_bc, grid, loc) = nothing + +##### +##### fill_halo_regions! for distributed tripolar grids +##### + fill_halo_regions!(c::OffsetArray, ::Nothing, indices, loc, ::MPITripolarGridOfSomeKind, args...; kwargs...) = nothing function fill_halo_regions!(c::OffsetArray, bcs, indices, loc, grid::MPITripolarGridOfSomeKind, buffers::CommunicationBuffers, args...; kwargs...) @@ -68,40 +420,24 @@ function fill_halo_regions!(c::OffsetArray, bcs, indices, loc, grid::MPITripolar fill_corners!(c, arch.connectivity, indices, loc, arch, grid, buffers, args...; kwargs...) - # We increment the request counter only if we have actually initiated the MPI communication. - # This is the case only if at least one of the boundary conditions is a distributed communication - # boundary condition (DCBCT) _and_ the `only_local_halos` keyword argument is false. if length(arch.mpi_requests) > outstanding_requests arch.mpi_tag[] += 1 end - if arch.mpi_tag[] == 0 # The communication has been reset, switch the north halos! - north_bc = bcs.north - switch_north_halos!(c, north_bc, grid, loc) - end - return nothing end function synchronize_communication!(field::Field{<:Any, <:Any, <:Any, <:Any, <:MPITripolarGridOfSomeKind}) arch = architecture(field.grid) - if arch isa AsynchronousDistributed # Otherwise no need to synchronize - # Wait for outstanding requests + if arch isa AsynchronousDistributed if !isempty(arch.mpi_requests) cooperative_waitall!(arch.mpi_requests) - - # Reset MPI tag arch.mpi_tag[] = 0 - - # Reset MPI requests empty!(arch.mpi_requests) end recv_from_buffers!(field.data, field.communication_buffers, field.grid) - - north_bc = field.boundary_conditions.north - switch_north_halos!(field, north_bc, field.grid, instantiated_location(field)) end return nothing diff --git a/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl b/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl index d619ee50fb5..b583f5f4256 100644 --- a/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl +++ b/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl @@ -4,7 +4,9 @@ using Oceananigans.BoundaryConditions: FieldBoundaryConditions, regularize_immersed_boundary_condition, LeftBoundary, RightBoundary -using Oceananigans.Grids: Grids, Center, Face +using Oceananigans.Grids: Grids, Center, Face, + LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded, + LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected using Oceananigans.BoundaryConditions: BoundaryConditions # A tripolar grid is always between 0 and 360 in longitude @@ -19,12 +21,17 @@ sign(::Type{Face}, ::Type{Center}) = - 1 # u-velocity type sign(::Type{Center}, ::Type{Face}) = - 1 # v-velocity type sign(::Type{Center}, ::Type{Center}) = 1 -# Determine the appropriate north fold boundary condition based on grid topology -# TODO: Implement proper topologies for distributed tripolar grids -# and remove the default fallback for AbstractTopology -north_fold_boundary_condition(::Type{<:AbstractTopology}) = UPivotZipperBoundaryCondition # Default fallback for distribtuted to work -north_fold_boundary_condition(::Type{RightCenterFolded}) = UPivotZipperBoundaryCondition -north_fold_boundary_condition(::Type{RightFaceFolded}) = FPivotZipperBoundaryCondition +# Determine the appropriate north fold boundary condition based on grid topology. +# Non-fold topologies (FullyConnected, RightConnected, etc.) default to UPivot — this +# value is only used as a placeholder; these ranks get their north BC overridden by +# inject_halo_communication_boundary_conditions or regularize_field_boundary_conditions. +north_fold_boundary_condition(::Type{<:AbstractTopology}) = UPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{RightCenterFolded}) = UPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{RightFaceFolded}) = FPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{LeftConnectedRightCenterFolded}) = UPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{LeftConnectedRightFaceFolded}) = FPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{<:LeftConnectedRightCenterConnected}) = UPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{<:LeftConnectedRightFaceConnected}) = FPivotZipperBoundaryCondition north_fold_boundary_condition(grid::TripolarGridOfSomeKind) = north_fold_boundary_condition(topology(grid, 2)) # a `TripolarGrid` needs a `UPivotZipperBoundaryCondition` for the north boundary diff --git a/validation/orthogonal_spherical_shell_grid/halo_fill_viz.jl b/validation/orthogonal_spherical_shell_grid/halo_fill_viz.jl index 29f92fddd2e..857e408cc5f 100644 --- a/validation/orthogonal_spherical_shell_grid/halo_fill_viz.jl +++ b/validation/orthogonal_spherical_shell_grid/halo_fill_viz.jl @@ -60,7 +60,7 @@ for fold_topology in fold_topologies color = :red, marker = :star5, markersize = 15, label = "Pivot points" ) # default grid underlay - offset = (fold_topology == RightFaceFolded) ? 1 : 1/2 # how much more grid north of pivots + offset = (fold_topology == RightFaceFolded) ? 0 : 1/2 # how much more grid north of pivots interior_poly_x = pivoti .+ [-Nx ÷ 2, Nx ÷ 2, Nx ÷ 2, -Nx ÷ 2] interior_poly_y = pivotj + offset .+ [0, 0, -Ny, -Ny] poly!(ax, interior_poly_x, interior_poly_y, color = (:black, 0.05), strokecolor = :black, strokewidth = 1) @@ -74,7 +74,8 @@ for fold_topology in fold_topologies radius = sqrt(sum((source .- destination) .^ 2)) / 2 start_angle = atan(destination[2] - origin[2], destination[1] - origin[1]) stop_angle = start_angle - π - scsrc = scatter!(ax, source...; color = Cycled(j), label = "Interior source points") + jcolor = Nyfield + 1 - source[2] + scsrc = scatter!(ax, source...; color = Cycled(jcolor), label = "Interior source points") # scatter!(ax, destination...; marker = :rect, color = Cycled(j), markersize = 20) scdest = scatter!(ax, destination...; marker = :rect, color = :white, markersize = 15, @@ -84,7 +85,7 @@ for fold_topology in fold_topologies # dtriangle = Polygon([Point(0, 0), Point(1, 1), Point(-1, 1)]) scutri = scatter!(ax, destination...; marker = :dtriangle, rotation = stop_angle, color = Cycled(j)) translate!(scutri, 0, 0, 2) # Hack more - ar = arc!(ax, origin, radius, start_angle, stop_angle, color = Cycled(j)) + ar = arc!(ax, origin, radius, start_angle, stop_angle, color = Cycled(jcolor)) translate!(ar, 0, 0, 2) # Hack more # add legend once From 221be456801f7d5d99be377ce734f476fe8d1129 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Mon, 30 Mar 2026 02:12:27 +1100 Subject: [PATCH 02/28] Store fold topology in Tripolar type parameter and fix distributed halo fill MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Make Tripolar{N,F,S,FT} carry the fold topology as a phantom type parameter (isbits for GPU). Add fold_topology() getter, remove global_fold_topology() which could not resolve the fold type from non-fold rank topologies (FullyConnected, RightConnected). - Fix reconstruct_global_grid and with_halo for distributed TripolarGrid: pass fold_topology from conformal_mapping so FPivot grids reconstruct correctly (was defaulting to RightCenterFolded). - Fix distributed zipper corner buffers: use has_fold_line to size corner buffers (Hy+1 when fold line present), dispatch on TwoDZipperBuffer only (corners not needed for 1D partitions), remove redundant arch::Distributed constraint, rename H→Hy. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/Grids/Grids.jl | 14 --- .../distributed_tripolar_grid.jl | 16 ++- .../distributed_zipper.jl | 103 +++++++++--------- .../tripolar_grid.jl | 25 ++++- 4 files changed, 82 insertions(+), 76 deletions(-) diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index 2e08d2db7de..995bbec9668 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -7,7 +7,6 @@ export RightFaceFolded, RightCenterFolded export LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded export LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected export WestOfPivot, EastOfPivot -export global_fold_topology export AbstractGrid, AbstractUnderlyingGrid, halo_size, total_size export RectilinearGrid export AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid @@ -160,19 +159,6 @@ Face-extended (Ny+1 Face points in y). """ struct LeftConnectedRightFaceConnected{P} <: AbstractTopology end -""" - global_fold_topology(T) - -Return the global grid y fold topology (`RightCenterFolded` or `RightFaceFolded`) from -the local grid y topology type. -""" -global_fold_topology(::Type{RightCenterFolded}) = RightCenterFolded -global_fold_topology(::Type{RightFaceFolded}) = RightFaceFolded -global_fold_topology(::Type{LeftConnectedRightCenterFolded}) = RightCenterFolded -global_fold_topology(::Type{LeftConnectedRightFaceFolded}) = RightFaceFolded -global_fold_topology(::Type{<:LeftConnectedRightCenterConnected}) = RightCenterFolded -global_fold_topology(::Type{<:LeftConnectedRightFaceConnected}) = RightFaceFolded - ##### ##### Directions (for tilted domains) ##### diff --git a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl index 5d744d03090..844e1bff9b4 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl @@ -11,7 +11,7 @@ using Oceananigans.DistributedComputations: concatenate_local_sizes, communication_buffers -using Oceananigans.Grids: topology, RightConnected, FullyConnected, global_fold_topology +using Oceananigans.Grids: topology, RightConnected, FullyConnected using Oceananigans.DistributedComputations: insert_connected_topology import Oceananigans.Fields: Field, validate_indices, validate_boundary_conditions @@ -256,7 +256,7 @@ function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, south = regularize_boundary_condition(bcs.south, grid, loc, 2, LeftBoundary, prognostic_names) north = if yrank == processor_size[2] - 1 && processor_size[1] == 1 - zipper_bc(global_fold_topology(topology(grid, 2)), sign) + zipper_bc(fold_topology(grid.conformal_mapping), sign) elseif yrank == processor_size[2] - 1 && processor_size[1] != 1 from = arch.local_rank @@ -294,9 +294,8 @@ function Field(loc::Tuple{<:LX, <:LY, <:LZ}, grid::MPITripolarGridOfSomeKind, da new_bcs = inject_halo_communication_boundary_conditions(old_bcs, arch.local_rank, arch.connectivity, topology(grid)) if yrank == processor_size[2] - 1 && processor_size[1] == 1 - default_zipper = zipper_bc(global_fold_topology(topology(grid, 2)), sign(LX, LY)) north_bc = if !(old_bcs.north isa ZBC) - default_zipper + zipper_bc(fold_topology(grid.conformal_mapping), sign(LX, LY)) else old_bcs.north end @@ -344,13 +343,16 @@ function DistributedComputations.reconstruct_global_grid(grid::MPITripolarGrid) first_pole_longitude = grid.conformal_mapping.first_pole_longitude southernmost_latitude = grid.conformal_mapping.southernmost_latitude + fold_topology = Oceananigans.OrthogonalSphericalShellGrids.fold_topology(grid.conformal_mapping) + return TripolarGrid(child_arch, FT; halo, size, north_poles_latitude, first_pole_longitude, southernmost_latitude, - z) + z, + fold_topology) end function Grids.with_halo(new_halo, old_grid::MPITripolarGrid) @@ -364,6 +366,7 @@ function Grids.with_halo(new_halo, old_grid::MPITripolarGrid) north_poles_latitude = old_grid.conformal_mapping.north_poles_latitude first_pole_longitude = old_grid.conformal_mapping.first_pole_longitude southernmost_latitude = old_grid.conformal_mapping.southernmost_latitude + fold_topology = Oceananigans.OrthogonalSphericalShellGrids.fold_topology(old_grid.conformal_mapping) return TripolarGrid(arch, eltype(old_grid); halo = new_halo, @@ -371,5 +374,6 @@ function Grids.with_halo(new_halo, old_grid::MPITripolarGrid) north_poles_latitude, first_pole_longitude, southernmost_latitude, - z) + z, + fold_topology) end diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 8278ed81eea..915e612d67b 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -98,29 +98,29 @@ function communication_buffers(grid::MPITripolarGridOfSomeKind, data, bcs, loc) end # Fallback: non-zipper north BC uses standard buffer -y_tripolar_buffer(arch, grid, data, H, bc, loc) = y_communication_buffer(arch, grid, data, H, bc) +y_tripolar_buffer(arch, grid, data, Hy, bc, loc) = y_communication_buffer(arch, grid, data, Hy, bc) # 1D fold (1xN) → OneDZipperBuffer (full-width) -function y_tripolar_buffer(arch::Distributed, grid::AbstractGrid{<:Any, <:Any, Topo}, - data, H, bc::DistributedZipper, loc::Loc) where {Topo <: OneDFoldTopology, Loc} +function y_tripolar_buffer(arch, grid::AbstractGrid{<:Any, <:Any, Topo}, + data, Hy, bc::DistributedZipper, loc::Loc) where {Topo <: OneDFoldTopology, Loc} Tx, _, Tz = size(parent(data)) FT = eltype(data) sgn = bc.condition.sign - send = on_architecture(arch, zeros(FT, Tx, H, Tz)) - recv = on_architecture(arch, zeros(FT, Tx, H, Tz)) + send = on_architecture(arch, zeros(FT, Tx, Hy, Tz)) + recv = on_architecture(arch, zeros(FT, Tx, Hy, Tz)) return OneDZipperBuffer(loc, Topo(), send, recv, sgn) end -# 2D fold (MxN) → TwoDZipperBuffer (interior-width, Hy or Hy+1 rows for fold line) -function y_tripolar_buffer(arch::Distributed, grid::AbstractGrid{<:Any, <:Any, Topo}, - data, H, bc::DistributedZipper, loc::Loc) where {Topo <: TwoDFoldTopology, Loc} +# 2D fold (MxN) → TwoDZipperBuffer (interior-width, Hy′ = Hy or Hy+1 rows for fold line) +function y_tripolar_buffer(arch, grid::AbstractGrid{<:Any, <:Any, Topo}, + data, Hy, bc::DistributedZipper, loc::Loc) where {Topo <: TwoDFoldTopology, Loc} Nx = size(grid, 1) _, _, Tz = size(parent(data)) FT = eltype(data) sgn = bc.condition.sign - Hbuf = has_fold_line(Topo, loc[2]) ? H + 1 : H - send = on_architecture(arch, zeros(FT, Nx, Hbuf, Tz)) - recv = on_architecture(arch, zeros(FT, Nx, Hbuf, Tz)) + Hy′ = has_fold_line(Topo, loc[2]) ? Hy + 1 : Hy + send = on_architecture(arch, zeros(FT, Nx, Hy′, Tz)) + recv = on_architecture(arch, zeros(FT, Nx, Hy′, Tz)) return TwoDZipperBuffer(loc, Topo(), send, recv, sgn) end @@ -128,44 +128,47 @@ end northwest_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge) = corner_communication_buffer(arch, grid, data, Hx, Hy, xedge, yedge) northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge) = corner_communication_buffer(arch, grid, data, Hx, Hy, xedge, yedge) +# Corner buffers are only needed for 2D (MxN) partitions. +# For FPivot CF/FF the fold line sits at Face-y, requiring Hy′ = Hy + 1 rows +# (mirroring the TwoDZipperBuffer). All other cases use Hy′ = Hy. + # ── NW corner: Center-x (Hx columns) ── -function northwest_tripolar_buffer(arch::Distributed, grid, data, Hx, Hy, xedge, - yedge::Union{<:OneDZipperBuffer{Loc, FoT}, - <:TwoDZipperBuffer{Loc, FoT}}) where {Loc <: Tuple{<:Center, <:Any, <:Any}, FoT} +function northwest_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, + yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Center, <:Any, <:Any}, FoT} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - send = on_architecture(arch, zeros(FT, Hx, Hy, Tz)) - recv = on_architecture(arch, zeros(FT, Hx, Hy, Tz)) + Hy′ = has_fold_line(FoT, Loc.parameters[2]()) ? Hy + 1 : Hy + send = on_architecture(arch, zeros(FT, Hx, Hy′, Tz)) + recv = on_architecture(arch, zeros(FT, Hx, Hy′, Tz)) return ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sgn)}(send, recv, sgn) end # ── NW corner: Face-x (Hx+1 columns for the Face-x shift) ── -function northwest_tripolar_buffer(arch::Distributed, grid, data, Hx, Hy, xedge, - yedge::Union{<:OneDZipperBuffer{Loc, FoT}, - <:TwoDZipperBuffer{Loc, FoT}}) where {Loc <: Tuple{<:Face, <:Any, <:Any}, FoT} +function northwest_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, + yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Face, <:Any, <:Any}, FoT} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - send = on_architecture(arch, zeros(FT, Hx+1, Hy, Tz)) - recv = on_architecture(arch, zeros(FT, Hx+1, Hy, Tz)) + Hy′ = has_fold_line(FoT, Loc.parameters[2]()) ? Hy + 1 : Hy + send = on_architecture(arch, zeros(FT, Hx+1, Hy′, Tz)) + recv = on_architecture(arch, zeros(FT, Hx+1, Hy′, Tz)) return ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sgn)}(send, recv, sgn) end # ── NE corner: Center-x (Hx columns) ── -function northeast_tripolar_buffer(arch::Distributed, grid, data, Hx, Hy, xedge, - yedge::Union{<:OneDZipperBuffer{Loc, FoT}, - <:TwoDZipperBuffer{Loc, FoT}}) where {Loc <: Tuple{<:Center, <:Any, <:Any}, FoT} +function northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, + yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Center, <:Any, <:Any}, FoT} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - send = on_architecture(arch, zeros(FT, Hx, Hy, Tz)) - recv = on_architecture(arch, zeros(FT, Hx, Hy, Tz)) + Hy′ = has_fold_line(FoT, Loc.parameters[2]()) ? Hy + 1 : Hy + send = on_architecture(arch, zeros(FT, Hx, Hy′, Tz)) + recv = on_architecture(arch, zeros(FT, Hx, Hy′, Tz)) return ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sgn)}(send, recv, sgn) end # ── NE corner: Face-x (Hx-1 columns, TwoDBuffer covers the extra) ── -function northeast_tripolar_buffer(arch::Distributed, grid, data, Hx, Hy, xedge, - yedge::Union{<:OneDZipperBuffer{Loc, FoT}, - <:TwoDZipperBuffer{Loc, FoT}}) where {Loc <: Tuple{<:Face, <:Any, <:Any}, FoT} +function northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, + yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Face, <:Any, <:Any}, FoT} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - Hx_ne = max(Hx - 1, 0) - send = on_architecture(arch, zeros(FT, Hx_ne, Hy, Tz)) - recv = on_architecture(arch, zeros(FT, Hx_ne, Hy, Tz)) + Hy′ = has_fold_line(FoT, Loc.parameters[2]()) ? Hy + 1 : Hy + send = on_architecture(arch, zeros(FT, Hx - 1, Hy′, Tz)) + recv = on_architecture(arch, zeros(FT, Hx - 1, Hy′, Tz)) return ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sgn)}(send, recv, sgn) end @@ -266,33 +269,33 @@ end ##### Fold-aware send buffer packing: NW corner (leftmost Hx interior columns) ##### -# Center-x NW -_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy-1:-1:Ny, :) +# Center-x NW: Hy rows, or Hy+1 when fold line is at this y-location +_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny, :) _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+1+Hy:-1:Ny+1, :) # Face-x NW: Hx+1 columns from leftmost Hx+1 interior columns -_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy-1:-1:Ny, :) +_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny, :) _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+1+Hy:-1:Ny+1, :) ##### ##### Fold-aware send buffer packing: NE corner (rightmost Hx interior columns) ##### -# Center-x NE -_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy-1:-1:Ny, :) +# Center-x NE: Hy rows, or Hy+1 when fold line is at this y-location +_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny, :) _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) -_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) +_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+1+Hy:-1:Ny+1, :) # Face-x NE: Hx-1 columns from rightmost Hx-1 interior columns -_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy-1:-1:Ny, :) +_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny, :) _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) -_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) +_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+1+Hy:-1:Ny+1, :) ##### ##### Recv methods: direct placement (data is already folded by sender) @@ -370,14 +373,14 @@ _recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:FF, <:FPivotTopology}, Hx, # ── Corner recv: UPivot — standard placement ── -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv # ── Corner recv: FPivot CC/FC — standard placement ── @@ -389,10 +392,10 @@ _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:FPivotTopology} # ── Corner recv: FPivot CF/FF — shifted +1 for Face-extended ── -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:1+Ny+2Hy, :) .= buff.recv +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:1+Ny+2Hy, :) .= buff.recv +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:1+Ny+2Hy, :) .= buff.recv +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:1+Ny+2Hy, :) .= buff.recv ##### ##### _switch_north_halos! — no-op (fold logic in send buffers) diff --git a/src/OrthogonalSphericalShellGrids/tripolar_grid.jl b/src/OrthogonalSphericalShellGrids/tripolar_grid.jl index 38dbc58dfa4..81b97edf952 100644 --- a/src/OrthogonalSphericalShellGrids/tripolar_grid.jl +++ b/src/OrthogonalSphericalShellGrids/tripolar_grid.jl @@ -1,24 +1,37 @@ using Oceananigans.BoundaryConditions: UPivotZipperBoundaryCondition, FPivotZipperBoundaryCondition, NoFluxBoundaryCondition using Oceananigans.Grids: Grids, Bounded, Flat, OrthogonalSphericalShellGrid, Periodic, RectilinearGrid, architecture, cpu_face_constructor_z, validate_dimension_specification, - RightCenterFolded, RightFaceFolded + AbstractTopology, RightCenterFolded, RightFaceFolded using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid """ - struct Tripolar{N, F, S} + struct Tripolar{N, F, S, FT<:AbstractTopology} A structure to represent a tripolar grid on an orthogonal spherical shell. +The fold topology `FT` (e.g., `RightCenterFolded` or `RightFaceFolded`) is stored +as a type parameter rather than a field, keeping the struct `isbits` for GPU kernels. """ -struct Tripolar{N, F, S} +struct Tripolar{N, F, S, FT<:AbstractTopology} north_poles_latitude :: N first_pole_longitude :: F southernmost_latitude :: S end -Adapt.adapt_structure(to, t::Tripolar) = +# Getter: returns the fold topology Type +fold_topology(::Tripolar{<:Any, <:Any, <:Any, FT}) where FT = FT + +# Constructor accepting fold topology as a Type argument +Tripolar(n, f, s, ::Type{FT}) where {FT<:AbstractTopology} = + Tripolar{typeof(n), typeof(f), typeof(s), FT}(n, f, s) + +# Backward-compatible constructor (defaults to UPivot) +Tripolar(n, f, s) = Tripolar(n, f, s, RightCenterFolded) + +Adapt.adapt_structure(to, t::Tripolar{<:Any, <:Any, <:Any, FT}) where FT = Tripolar(Adapt.adapt(to, t.north_poles_latitude), Adapt.adapt(to, t.first_pole_longitude), - Adapt.adapt(to, t.southernmost_latitude)) + Adapt.adapt(to, t.southernmost_latitude), + FT) const TripolarGrid{FT, TX, TY, TZ, CZ, CC, FC, CF, FF, Arch} = OrthogonalSphericalShellGrid{FT, TX, TY, TZ, CZ, <:Tripolar, CC, FC, CF, FF, Arch} const TripolarGridOfSomeKind{FT, TX, TY, TZ} = Union{TripolarGrid{FT, TX, TY, TZ}, ImmersedBoundaryGrid{FT, TX, TY, TZ, <:TripolarGrid}} @@ -336,7 +349,7 @@ function TripolarGrid(arch = CPU(), FT::DataType = Oceananigans.defaults.FloatTy on_architecture(arch, map(FT, Azᶜᶠᵃ)), on_architecture(arch, map(FT, Azᶠᶠᵃ)), convert(FT, radius), - Tripolar(north_poles_latitude, first_pole_longitude, southernmost_latitude)) + Tripolar(north_poles_latitude, first_pole_longitude, southernmost_latitude, fold_topology)) return grid end From ddd9b19d5e27d228b5a8d8682dcfc71b7ad0fda6 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Mon, 30 Mar 2026 10:52:27 +1100 Subject: [PATCH 03/28] Replace WestOfPivot/EastOfPivot with per-buffer FL/WFL fold-line dispatch Remove the WestOfPivot/EastOfPivot type parameter from topology types and replace with per-buffer fold-line control via two new type parameters on TwoDZipperBuffer and ZipperCornerBuffer: - FL (fold line in buffer): true when the buffer contains the fold-line row (Hy+1 rows). Set to has_fold_line() for ALL buffers to ensure MPI size matching between mirror partners. - WFL (writes fold line): true when the recv should write the fold-line row. Computed per-buffer from the rank's x-position relative to the pivot (Nx/2), accounting for periodic wrap at domain boundaries. Helper functions north_writes_fold_line, northwest_writes_fold_line, northeast_writes_fold_line determine WFL based on rx and Rx. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../distributed_grids.jl | 6 +- src/Grids/Grids.jl | 17 +- .../distributed_zipper.jl | 171 +++++++++++++----- 3 files changed, 134 insertions(+), 60 deletions(-) diff --git a/src/DistributedComputations/distributed_grids.jl b/src/DistributedComputations/distributed_grids.jl index f274bb2be7e..ea28729d51d 100644 --- a/src/DistributedComputations/distributed_grids.jl +++ b/src/DistributedComputations/distributed_grids.jl @@ -353,13 +353,11 @@ insert_connected_topology(::Type{Periodic}, R, r) = ifelse(R == 1, Periodic, Ful # These take 5 arguments: (global_y_topology, Ry, ry, Rx, rx) where # Ry/ry are y-rank count/index and Rx/rx are x-rank count/index (all 1-based). -pivot_side(Rx, rx) = rx ≤ Rx ÷ 2 ? WestOfPivot : EastOfPivot - local_fold_topology(::Type{RightCenterFolded}, Rx, rx) = - Rx == 1 ? LeftConnectedRightCenterFolded : LeftConnectedRightCenterConnected{pivot_side(Rx, rx)} + Rx == 1 ? LeftConnectedRightCenterFolded : LeftConnectedRightCenterConnected local_fold_topology(::Type{RightFaceFolded}, Rx, rx) = - Rx == 1 ? LeftConnectedRightFaceFolded : LeftConnectedRightFaceConnected{pivot_side(Rx, rx)} + Rx == 1 ? LeftConnectedRightFaceFolded : LeftConnectedRightFaceConnected function insert_connected_topology(T::Type{<:Union{RightCenterFolded, RightFaceFolded}}, Ry, ry, Rx, rx) if ry == 1 diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index 995bbec9668..11d171ecaa9 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -6,7 +6,6 @@ export Periodic, Bounded, Flat, FullyConnected, LeftConnected, RightConnected export RightFaceFolded, RightCenterFolded export LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded export LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected -export WestOfPivot, EastOfPivot export AbstractGrid, AbstractUnderlyingGrid, halo_size, total_size export RectilinearGrid export AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid @@ -120,10 +119,6 @@ Grid topology for tripolar U-point pivot connection. """ struct RightCenterFolded <: AbstractTopology end -# Pivot side indicators for distributed fold topologies -struct WestOfPivot end -struct EastOfPivot end - """ LeftConnectedRightCenterFolded @@ -143,21 +138,21 @@ face-folded on the right (north). Face-extended (Ny+1 Face points in y). struct LeftConnectedRightFaceFolded <: AbstractTopology end """ - LeftConnectedRightCenterConnected{P} + LeftConnectedRightCenterConnected Local grid y topology for the northernmost y-rank of an M×N distributed tripolar grid -with U-point pivot (distributed zipper). `P` is `WestOfPivot` or `EastOfPivot`. +with U-point pivot (distributed zipper). """ -struct LeftConnectedRightCenterConnected{P} <: AbstractTopology end +struct LeftConnectedRightCenterConnected <: AbstractTopology end """ - LeftConnectedRightFaceConnected{P} + LeftConnectedRightFaceConnected Local grid y topology for the northernmost y-rank of an M×N distributed tripolar grid -with F-point pivot (distributed zipper). `P` is `WestOfPivot` or `EastOfPivot`. +with F-point pivot (distributed zipper). Face-extended (Ny+1 Face points in y). """ -struct LeftConnectedRightFaceConnected{P} <: AbstractTopology end +struct LeftConnectedRightFaceConnected <: AbstractTopology end ##### ##### Directions (for tilted domains) diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 915e612d67b..4164e40f1ab 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -4,8 +4,7 @@ using Oceananigans.DistributedComputations: CommunicationBuffers, fill_corners!, using Oceananigans.Grids: AbstractGrid, topology, RightCenterFolded, RightFaceFolded, LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded, - LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected, - WestOfPivot, EastOfPivot + LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected using Oceananigans.DistributedComputations: Distributed, on_architecture, ranks, x_communication_buffer using Oceananigans.Fields: instantiated_location @@ -45,6 +44,33 @@ has_fold_line(::Type{<:UPivotTopology}, ::Face) = false has_fold_line(::Type{<:FPivotTopology}, ::Center) = false has_fold_line(::Type{<:FPivotTopology}, ::Face) = true +# The fold has two pivot points due to x-periodicity: +# 1st pivot: between rx = Rx÷2 and rx = Rx÷2+1 +# 2nd pivot: at the periodic boundary between rx = Rx and rx = 1 +# The serial ifelse(i > Nx÷2, ...) writes the east half only. +# Corner conditions account for periodic wrap at rx=1 (NW) and rx=Rx (NE). + +# North buffer: east-of-pivot ranks write fold line +function north_writes_fold_line(arch) + rx = arch.local_index[1] + Rx = ranks(arch)[1] + return rx > Rx ÷ 2 +end + +# NW corner: west edge past pivot, or rx=1 (periodic wrap to east half) +function northwest_writes_fold_line(arch) + rx = arch.local_index[1] + Rx = ranks(arch)[1] + return rx > Rx ÷ 2 + 1 || rx == 1 +end + +# NE corner: east edge past pivot, but not rx=Rx (periodic wrap to west half) +function northeast_writes_fold_line(arch) + rx = arch.local_index[1] + Rx = ranks(arch)[1] + return rx >= Rx ÷ 2 && rx < Rx +end + ##### ##### Zipper communication buffers with fold-aware packing ##### @@ -55,13 +81,13 @@ struct OneDZipperBuffer{Loc, FoT, B, S} sign :: S end -struct TwoDZipperBuffer{Loc, FoT, B, S} +struct TwoDZipperBuffer{Loc, FoT, B, S, FL, WFL} send :: B recv :: B sign :: S end -struct ZipperCornerBuffer{Loc, FoT, B, S} +struct ZipperCornerBuffer{Loc, FoT, B, S, FL, WFL} send :: B recv :: B sign :: S @@ -69,8 +95,12 @@ end # Value-argument constructors: all types inferred OneDZipperBuffer(loc::Loc, fot::FoT, send::B, recv::B, sign::S) where {Loc, FoT, B, S} = OneDZipperBuffer{Loc, FoT, B, S}(send, recv, sign) -TwoDZipperBuffer(loc::Loc, fot::FoT, send::B, recv::B, sign::S) where {Loc, FoT, B, S} = TwoDZipperBuffer{Loc, FoT, B, S}(send, recv, sign) -ZipperCornerBuffer(loc::Loc, fot::FoT, send::B, recv::B, sign::S) where {Loc, FoT, B, S} = ZipperCornerBuffer{Loc, FoT, B, S}(send, recv, sign) +TwoDZipperBuffer(loc, fot, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {FL, WFL} = + TwoDZipperBuffer{typeof(loc), typeof(fot), typeof(send), typeof(sign), FL, WFL}(send, recv, sign) +ZipperCornerBuffer(loc, fot, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {FL, WFL} = + ZipperCornerBuffer{typeof(loc), typeof(fot), typeof(send), typeof(sign), FL, WFL}(send, recv, sign) +ZipperCornerBuffer(::Type{Loc}, ::Type{FoT}, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {Loc, FoT, FL, WFL} = + ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sign), FL, WFL}(send, recv, sign) Adapt.adapt_structure(to, buff::OneDZipperBuffer) = nothing Adapt.adapt_structure(to, buff::TwoDZipperBuffer) = nothing @@ -118,10 +148,12 @@ function y_tripolar_buffer(arch, grid::AbstractGrid{<:Any, <:Any, Topo}, _, _, Tz = size(parent(data)) FT = eltype(data) sgn = bc.condition.sign - Hy′ = has_fold_line(Topo, loc[2]) ? Hy + 1 : Hy + fl = has_fold_line(Topo, loc[2]) + Hy′ = fl ? Hy + 1 : Hy + wfl = fl && north_writes_fold_line(arch) send = on_architecture(arch, zeros(FT, Nx, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Nx, Hy′, Tz)) - return TwoDZipperBuffer(loc, Topo(), send, recv, sgn) + return TwoDZipperBuffer(loc, Topo(), send, recv, sgn, Val(fl), Val(wfl)) end # Fallbacks: non-zipper corners @@ -129,47 +161,56 @@ northwest_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge) = corner_commu northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge) = corner_communication_buffer(arch, grid, data, Hx, Hy, xedge, yedge) # Corner buffers are only needed for 2D (MxN) partitions. -# For FPivot CF/FF the fold line sits at Face-y, requiring Hy′ = Hy + 1 rows -# (mirroring the TwoDZipperBuffer). All other cases use Hy′ = Hy. +# FL (fold line in buffer) is true for ALL corners when has_fold_line, to ensure +# MPI size matching between mirror partners. WFL (writes fold line) is per-corner, +# based on whether the corner's global x-position is past the pivot. # ── NW corner: Center-x (Hx columns) ── function northwest_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Center, <:Any, <:Any}, FoT} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - Hy′ = has_fold_line(FoT, Loc.parameters[2]()) ? Hy + 1 : Hy + fl = has_fold_line(FoT, Loc.parameters[2]()) + Hy′ = fl ? Hy + 1 : Hy + wfl = fl && northwest_writes_fold_line(arch) send = on_architecture(arch, zeros(FT, Hx, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Hx, Hy′, Tz)) - return ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sgn)}(send, recv, sgn) + return ZipperCornerBuffer(Loc, FoT, send, recv, sgn, Val(fl), Val(wfl)) end # ── NW corner: Face-x (Hx+1 columns for the Face-x shift) ── function northwest_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Face, <:Any, <:Any}, FoT} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - Hy′ = has_fold_line(FoT, Loc.parameters[2]()) ? Hy + 1 : Hy + fl = has_fold_line(FoT, Loc.parameters[2]()) + Hy′ = fl ? Hy + 1 : Hy + wfl = fl && northwest_writes_fold_line(arch) send = on_architecture(arch, zeros(FT, Hx+1, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Hx+1, Hy′, Tz)) - return ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sgn)}(send, recv, sgn) + return ZipperCornerBuffer(Loc, FoT, send, recv, sgn, Val(fl), Val(wfl)) end # ── NE corner: Center-x (Hx columns) ── function northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Center, <:Any, <:Any}, FoT} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - Hy′ = has_fold_line(FoT, Loc.parameters[2]()) ? Hy + 1 : Hy + fl = has_fold_line(FoT, Loc.parameters[2]()) + Hy′ = fl ? Hy + 1 : Hy + wfl = fl && northeast_writes_fold_line(arch) send = on_architecture(arch, zeros(FT, Hx, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Hx, Hy′, Tz)) - return ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sgn)}(send, recv, sgn) + return ZipperCornerBuffer(Loc, FoT, send, recv, sgn, Val(fl), Val(wfl)) end # ── NE corner: Face-x (Hx-1 columns, TwoDBuffer covers the extra) ── function northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Face, <:Any, <:Any}, FoT} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - Hy′ = has_fold_line(FoT, Loc.parameters[2]()) ? Hy + 1 : Hy + fl = has_fold_line(FoT, Loc.parameters[2]()) + Hy′ = fl ? Hy + 1 : Hy + wfl = fl && northeast_writes_fold_line(arch) send = on_architecture(arch, zeros(FT, Hx - 1, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Hx - 1, Hy′, Tz)) - return ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sgn)}(send, recv, sgn) + return ZipperCornerBuffer(Loc, FoT, send, recv, sgn, Val(fl), Val(wfl)) end ##### @@ -300,60 +341,60 @@ _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, H ##### ##### Recv methods: direct placement (data is already folded by sender) ##### -##### For TwoD fold-line fields (Hy+1 buffer): -##### EastOfPivot: write fold line + halos -##### WestOfPivot: write halos only (skip fold line in row 1) +##### For TwoD fold-line fields (FL=true, Hy+1 buffer): +##### WFL=true: write fold line (row 1) + halos (rows 2:Hy+1) +##### WFL=false: write halos only (rows 2:Hy+1, skip fold line) ##### ##### For FPivot Face-y: halos start at parent Ny+Hy+2 (field Ny+2, past fold line at Ny+1) ##### -# ── TwoDZipperBuffer: UPivot CC/FC — fold line dispatch on WoP/EoP ── +# ── TwoDZipperBuffer: UPivot CC/FC — FL=true, dispatch on WFL ── -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:LeftConnectedRightCenterConnected{EastOfPivot}}, Hx, Hy, Nx, Ny) +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) view(c, 1+Hx:Nx+Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:LeftConnectedRightCenterConnected{WestOfPivot}}, Hx, Hy, Nx, Ny) +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected{EastOfPivot}}, Hx, Hy, Nx, Ny) +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) view(c, 2+Hx:Nx+Hx+1, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected{WestOfPivot}}, Hx, Hy, Nx, Ny) +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end -# ── TwoDZipperBuffer: UPivot CF/FF — no fold line, standard placement ── +# ── TwoDZipperBuffer: UPivot CF/FF — FL=false, no fold line ── _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -# ── TwoDZipperBuffer: FPivot CC/FC — no fold line, standard placement ── +# ── TwoDZipperBuffer: FPivot CC/FC — FL=false, no fold line ── _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -# ── TwoDZipperBuffer: FPivot CF/FF — fold line dispatch, shifted +1 for Face-extended ── +# ── TwoDZipperBuffer: FPivot CF/FF — FL=true, dispatch on WFL, shifted +1 for Face-y ── -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CF, <:LeftConnectedRightFaceConnected{EastOfPivot}}, Hx, Hy, Nx, Ny) +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CF, <:LeftConnectedRightFaceConnected, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) view(c, 1+Hx:Nx+Hx, Ny+1+Hy:Ny+1+Hy, :) .= view(buff.recv, :, 1:1, :) view(c, 1+Hx:Nx+Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CF, <:LeftConnectedRightFaceConnected{WestOfPivot}}, Hx, Hy, Nx, Ny) +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CF, <:LeftConnectedRightFaceConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) view(c, 1+Hx:Nx+Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FF, <:LeftConnectedRightFaceConnected{EastOfPivot}}, Hx, Hy, Nx, Ny) +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FF, <:LeftConnectedRightFaceConnected, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) view(c, 2+Hx:Nx+Hx+1, Ny+1+Hy:Ny+1+Hy, :) .= view(buff.recv, :, 1:1, :) view(c, 2+Hx:Nx+Hx+1, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FF, <:LeftConnectedRightFaceConnected{WestOfPivot}}, Hx, Hy, Nx, Ny) +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FF, <:LeftConnectedRightFaceConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) view(c, 2+Hx:Nx+Hx+1, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end @@ -371,31 +412,71 @@ _recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:FC, <:FPivotTopology}, Hx, _recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, :, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv _recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, :, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv -# ── Corner recv: UPivot — standard placement ── +# ── Corner recv: UPivot CC/FC — FL=true, dispatch on WFL ── + +# WFL=true: write fold line + halos +function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) + view(c, 1:Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) + view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +end +function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) + view(c, 1:Hx+1, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) + view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +end +function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) + view(c, 1+Nx+Hx:Nx+2Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) + view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +end +function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) + view(c, 2+Nx+Hx:Nx+2Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) + view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +end + +# WFL=false: write halos only (skip fold line) +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) + +# ── Corner recv: UPivot CF/FF — FL=false, no fold line, Hy rows ── -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv - -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -# ── Corner recv: FPivot CC/FC — standard placement ── +# ── Corner recv: FPivot CC/FC — FL=false, no fold line, Hy rows ── _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -# ── Corner recv: FPivot CF/FF — shifted +1 for Face-extended ── +# ── Corner recv: FPivot CF/FF — FL=true, dispatch on WFL, shifted +1 for Face-y ── + +# WFL=true: write fold line + halos +function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) + view(c, 1:Hx, 1+Ny+Hy:1+Ny+Hy, :) .= view(buff.recv, :, 1:1, :) + view(c, 1:Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +end +function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) + view(c, 1:Hx+1, 1+Ny+Hy:1+Ny+Hy, :) .= view(buff.recv, :, 1:1, :) + view(c, 1:Hx+1, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +end +function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) + view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:1+Ny+Hy, :) .= view(buff.recv, :, 1:1, :) + view(c, 1+Nx+Hx:Nx+2Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +end +function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) + view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:1+Ny+Hy, :) .= view(buff.recv, :, 1:1, :) + view(c, 2+Nx+Hx:Nx+2Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +end -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:1+Ny+2Hy, :) .= buff.recv -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:1+Ny+2Hy, :) .= buff.recv -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:1+Ny+2Hy, :) .= buff.recv -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:1+Ny+2Hy, :) .= buff.recv +# WFL=false: write halos only (skip fold line) +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) ##### ##### _switch_north_halos! — no-op (fold logic in send buffers) From f788a07de9e3616f4a172a724aa43b7588e24af2 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Mon, 30 Mar 2026 20:40:08 +1100 Subject: [PATCH 04/28] Replace WestOfPivot/EastOfPivot with per-buffer FL/WFL fold-line dispatch - Remove WestOfPivot/EastOfPivot type parameters and pivot_side function - Add FL (fold-line in buffer) and WFL (writes fold line) type parameters to TwoDZipperBuffer and ZipperCornerBuffer - Per-buffer fold-line helpers: north/northwest/northeast_writes_fold_line accounting for both pivot points (x-periodicity) - TripolarXBuffer with location-aware y-size using length(loc_y, topo, Ny) instead of has_fold_line (only FPivot Face-y gets Ny+1, not UPivot) - Temporary @info diagnostics for send/recv tracing (to be removed) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../distributed_zipper.jl | 177 +++++++++++++++--- 1 file changed, 152 insertions(+), 25 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 4164e40f1ab..df0b237e2be 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -13,7 +13,9 @@ import Oceananigans.DistributedComputations: synchronize_communication!, y_communication_buffer, corner_communication_buffer, _fill_north_send_buffer!, _recv_from_north_buffer!, _fill_northwest_send_buffer!, _fill_northeast_send_buffer!, - _recv_from_northwest_buffer!, _recv_from_northeast_buffer! + _recv_from_northwest_buffer!, _recv_from_northeast_buffer!, + _fill_west_send_buffer!, _fill_east_send_buffer!, + _recv_from_west_buffer!, _recv_from_east_buffer! import Oceananigans.Fields: communication_buffers @inline instantiate(T::DataType) = T() @@ -106,6 +108,52 @@ Adapt.adapt_structure(to, buff::OneDZipperBuffer) = nothing Adapt.adapt_structure(to, buff::TwoDZipperBuffer) = nothing Adapt.adapt_structure(to, buff::ZipperCornerBuffer) = nothing +# X-direction buffer for tripolar grids: like TwoDBuffer but with location-aware y-size. +# For fold-line fields on the north row, the x-buffer needs Ny+1 y-rows so that the +# fold-line row is exchanged periodically. Corners then overwrite where needed (WFL=true). +struct TripolarXBuffer{B} + send :: B + recv :: B +end + +Adapt.adapt_structure(to, buff::TripolarXBuffer) = nothing + +# TripolarXBuffer send/recv: use the buffer's actual y-size +function _fill_west_send_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) + @info "TripolarX west send: c x=$(1+Hx):$(2Hx), y=$(1+Hy):$(size(buff.send,2)+Hy), buff size=$(size(buff.send))" + buff.send .= view(c, 1+Hx:2Hx, 1+Hy:size(buff.send,2)+Hy, :) +end +function _fill_east_send_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) + @info "TripolarX east send: c x=$(1+Nx):$(Nx+Hx), y=$(1+Hy):$(size(buff.send,2)+Hy), buff size=$(size(buff.send))" + buff.send .= view(c, 1+Nx:Nx+Hx, 1+Hy:size(buff.send,2)+Hy, :) +end +function _recv_from_west_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) + @info "TripolarX west recv: c x=1:$(Hx), y=$(1+Hy):$(size(buff.recv,2)+Hy), buff size=$(size(buff.recv))" + view(c, 1:Hx, 1+Hy:size(buff.recv,2)+Hy, :) .= buff.recv +end +function _recv_from_east_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) + tgt_x, tgt_y = Nx+Hx+1, size(buff.recv,2)+Hy + @info "TripolarX east recv: c x=$(1+Nx+Hx):$(Nx+2Hx), y=$(1+Hy):$(tgt_y), buff size=$(size(buff.recv)), c[$tgt_x,$tgt_y] BEFORE=$(c[tgt_x, tgt_y, (size(c,3)+1)÷2]), buff[1,end,1]=$(buff.recv[1, end, (size(buff.recv,3)+1)÷2])" + view(c, 1+Nx+Hx:Nx+2Hx, 1+Hy:size(buff.recv,2)+Hy, :) .= buff.recv + @info "TripolarX east recv: c[$tgt_x,$tgt_y] AFTER=$(c[tgt_x, tgt_y, (size(c,3)+1)÷2])" +end + +# Fold-aware x-buffer: uniform Ny+1 for fold-line fields on north row. +# Fallback for non-zipper north (south ranks): standard x-communication. +x_tripolar_buffer(arch, grid, data, Hx, bc, loc, north) = x_communication_buffer(arch, grid, data, Hx, bc) + +# When north IS a TwoDZipperBuffer (fold north row): use location-aware y-size +# (Ny+1 for BoundedTopology Face-y, i.e., FPivot CF/FF; Ny otherwise) +function x_tripolar_buffer(arch, grid, data, Hx, bc, loc, + north::TwoDZipperBuffer{Loc, FoT}) where {Loc, FoT} + Ny_buf = length(Loc.parameters[2](), FoT(), size(grid, 2)) + _, _, Tz = size(parent(data)) + FT = eltype(data) + send = on_architecture(arch, zeros(FT, Hx, Ny_buf, Tz)) + recv = on_architecture(arch, zeros(FT, Hx, Ny_buf, Tz)) + return TripolarXBuffer(send, recv) +end + ##### ##### Buffer construction for tripolar grids ##### @@ -114,11 +162,15 @@ function communication_buffers(grid::MPITripolarGridOfSomeKind, data, bcs, loc) Hx, Hy, Hz = halo_size(grid) arch = architecture(grid) - west = x_communication_buffer(arch, grid, data, Hx, bcs.west) - east = x_communication_buffer(arch, grid, data, Hx, bcs.east) south = y_communication_buffer(arch, grid, data, Hy, bcs.south) north = y_tripolar_buffer(arch, grid, data, Hy, bcs.north, loc) + # x-buffers dispatch on north: Ny+1 for fold-line fields on north row + west = x_tripolar_buffer(arch, grid, data, Hx, bcs.west, loc, north) + east = x_tripolar_buffer(arch, grid, data, Hx, bcs.east, loc, north) + + @info "communication_buffers: rank=$(arch.local_index), loc=$loc, west=$(typeof(west)), east=$(typeof(east)), north=$(typeof(north))" + sw = corner_communication_buffer(arch, grid, data, Hx, Hy, west, south) se = corner_communication_buffer(arch, grid, data, Hx, Hy, east, south) nw = northwest_tripolar_buffer(arch, grid, data, Hx, Hy, west, north) @@ -201,7 +253,7 @@ function northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, return ZipperCornerBuffer(Loc, FoT, send, recv, sgn, Val(fl), Val(wfl)) end -# ── NE corner: Face-x (Hx-1 columns, TwoDBuffer covers the extra) ── +# ── NE corner: Face-x (Hx-1 columns, north buffer covers the extra Face-x column) ── function northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Face, <:Any, <:Any}, FoT} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign @@ -235,33 +287,50 @@ const FF = Tuple{<:Face, <:Face, <:Any} # ── TwoDZipperBuffer: UPivot Center-y (CC/FC) — fold line + Hy halo rows ── function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) + @info "CC UPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y fold=$(Ny+Hy), y halo=$(Ny)..$(Ny+Hy-1), buff size=$(size(b.send))" view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:Ny+Hy, :) view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy-1:-1:Ny, :) end function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) + @info "FC UPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y fold=$(Ny+Hy), y halo=$(Ny)..$(Ny+Hy-1), buff size=$(size(b.send)), c[$(1+Hx),$(Ny+Hy-1),1]=$(c[1+Hx, Ny+Hy-1, (size(c,3)+1)÷2])" view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:Ny+Hy, :) view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy-1:-1:Ny, :) + @info "FC UPivot north send: buff[end,2,1]=$(b.send[end, 2, (size(b.send,3)+1)÷2]) (should go to mirror's last x at first halo row)" end # ── TwoDZipperBuffer: UPivot Face-y (CF/FF) — no fold line, Hy halo rows ── -_fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -_fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) + @info "CF UPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y=$(Ny+1)..$(Ny+Hy), buff size=$(size(b.send))" + b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +end +function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) + @info "FF UPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y=$(Ny+1)..$(Ny+Hy), buff size=$(size(b.send))" + b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +end # ── TwoDZipperBuffer: FPivot Center-y (CC/FC) — no fold line, Hy halo rows ── -_fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -_fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) + @info "CC FPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y=$(Ny+1)..$(Ny+Hy), buff size=$(size(b.send))" + b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +end +function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) + @info "FC FPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y=$(Ny+1)..$(Ny+Hy), buff size=$(size(b.send))" + b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +end # ── TwoDZipperBuffer: FPivot Face-y (CF/FF) — fold line + Hy halo rows ── function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) + @info "CF FPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y fold=$(Ny+1+Hy), y halo=$(Ny+1)..$(Ny+Hy), buff size=$(size(b.send))" view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+1+Hy:Ny+1+Hy, :) view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) + @info "FF FPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y fold=$(Ny+1+Hy), y halo=$(Ny+1)..$(Ny+Hy), buff size=$(size(b.send))" view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+1+Hy:Ny+1+Hy, :) view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end @@ -311,32 +380,80 @@ end ##### # Center-x NW: Hy rows, or Hy+1 when fold line is at this y-location -_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny, :) -_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+1+Hy:-1:Ny+1, :) +function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) + @info "CC UPivot NW send: x=$(1+Hx):$(2Hx) (rev), y=$(Ny)..$(Ny+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny, :) +end +function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) + @info "CC FPivot NW send: x=$(1+Hx):$(2Hx) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +end +function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) + @info "CF UPivot NW send: x=$(1+Hx):$(2Hx) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +end +function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) + @info "CF FPivot NW send: x=$(1+Hx):$(2Hx) (rev), y=$(Ny+1)..$(Ny+1+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+1+Hy:-1:Ny+1, :) +end # Face-x NW: Hx+1 columns from leftmost Hx+1 interior columns -_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny, :) -_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -_fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+1+Hy:-1:Ny+1, :) +function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) + @info "FC UPivot NW send: x=$(1+Hx):$(2Hx+1) (rev), y=$(Ny)..$(Ny+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny, :) +end +function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) + @info "FC FPivot NW send: x=$(1+Hx):$(2Hx+1) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +end +function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) + @info "FF UPivot NW send: x=$(1+Hx):$(2Hx+1) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) +end +function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) + @info "FF FPivot NW send: x=$(1+Hx):$(2Hx+1) (rev), y=$(Ny+1)..$(Ny+1+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+1+Hy:-1:Ny+1, :) +end ##### ##### Fold-aware send buffer packing: NE corner (rightmost Hx interior columns) ##### # Center-x NE: Hy rows, or Hy+1 when fold line is at this y-location -_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny, :) -_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) -_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) -_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+1+Hy:-1:Ny+1, :) +function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) + @info "CC UPivot NE send: x=$(1+Nx):$(Nx+Hx) (rev), y=$(Ny)..$(Ny+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny, :) +end +function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) + @info "CC FPivot NE send: x=$(1+Nx):$(Nx+Hx) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) +end +function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) + @info "CF UPivot NE send: x=$(1+Nx):$(Nx+Hx) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) +end +function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) + @info "CF FPivot NE send: x=$(1+Nx):$(Nx+Hx) (rev), y=$(Ny+1)..$(Ny+1+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+1+Hy:-1:Ny+1, :) +end # Face-x NE: Hx-1 columns from rightmost Hx-1 interior columns -_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny, :) -_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) -_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) -_fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+1+Hy:-1:Ny+1, :) +function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) + @info "FC UPivot NE send: x=$(Nx+2):$(Nx+Hx) (rev), y=$(Ny)..$(Ny+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny, :) +end +function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) + @info "FC FPivot NE send: x=$(Nx+2):$(Nx+Hx) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) +end +function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) + @info "FF UPivot NE send: x=$(Nx+2):$(Nx+Hx) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) +end +function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) + @info "FF FPivot NE send: x=$(Nx+2):$(Nx+Hx) (rev), y=$(Ny+1)..$(Ny+1+Hy), buff=$(size(b.send))" + b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+1+Hy:-1:Ny+1, :) +end ##### ##### Recv methods: direct placement (data is already folded by sender) @@ -351,21 +468,27 @@ _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, H # ── TwoDZipperBuffer: UPivot CC/FC — FL=true, dispatch on WFL ── function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) + @info "CC UPivot north recv WFL=true: x=$(1+Hx):$(Nx+Hx), y=$(Ny+Hy):$(Ny+2Hy), buff size=$(size(buff.recv))" view(c, 1+Hx:Nx+Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) + @info "CC UPivot north recv WFL=false: x=$(1+Hx):$(Nx+Hx), y=$(1+Ny+Hy):$(Ny+2Hy), buff size=$(size(buff.recv))" view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) + @info "FC UPivot north recv WFL=true: x=$(2+Hx):$(Nx+Hx+1), y=$(Ny+Hy):$(Ny+2Hy), buff size=$(size(buff.recv))" view(c, 2+Hx:Nx+Hx+1, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) + tgt_x, tgt_y, tgt_z = Nx+Hx+1, 1+Ny+Hy, (size(c,3)+1)÷2 + @info "FC UPivot north recv WFL=false: x=$(2+Hx):$(Nx+Hx+1), y=$(1+Ny+Hy):$(Ny+2Hy), c[$tgt_x,$tgt_y,$tgt_z] BEFORE=$(c[tgt_x, tgt_y, tgt_z]), buff[end,2,$tgt_z]=$(buff.recv[end, 2, tgt_z])" view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) + @info "FC UPivot north recv WFL=false: c[$tgt_x,$tgt_y,$tgt_z] AFTER=$(c[tgt_x, tgt_y, tgt_z])" end # ── TwoDZipperBuffer: UPivot CF/FF — FL=false, no fold line ── @@ -428,6 +551,7 @@ function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivot view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) end function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) + @info "FC UPivot NE corner recv WFL=true: x=$(2+Nx+Hx):$(Nx+2Hx), y=$(Ny+Hy):$(Ny+2Hy), buff size=$(size(buff.recv))" view(c, 2+Nx+Hx:Nx+2Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) end @@ -436,7 +560,10 @@ end _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) + @info "FC UPivot NE corner recv WFL=false: x=$(2+Nx+Hx):$(Nx+2Hx), y=$(1+Ny+Hy):$(Ny+2Hy), buff size=$(size(buff.recv))" + view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +end # ── Corner recv: UPivot CF/FF — FL=false, no fold line, Hy rows ── From 9f8a8cf4971fe2332ee65609275c1dd236286190 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Mon, 30 Mar 2026 20:44:27 +1100 Subject: [PATCH 05/28] Remove temporary @info diagnostics from distributed_zipper.jl Co-Authored-By: Claude Opus 4.6 (1M context) --- .../distributed_zipper.jl | 38 ------------------- 1 file changed, 38 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index df0b237e2be..25cb4e7c133 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -120,22 +120,17 @@ Adapt.adapt_structure(to, buff::TripolarXBuffer) = nothing # TripolarXBuffer send/recv: use the buffer's actual y-size function _fill_west_send_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) - @info "TripolarX west send: c x=$(1+Hx):$(2Hx), y=$(1+Hy):$(size(buff.send,2)+Hy), buff size=$(size(buff.send))" buff.send .= view(c, 1+Hx:2Hx, 1+Hy:size(buff.send,2)+Hy, :) end function _fill_east_send_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) - @info "TripolarX east send: c x=$(1+Nx):$(Nx+Hx), y=$(1+Hy):$(size(buff.send,2)+Hy), buff size=$(size(buff.send))" buff.send .= view(c, 1+Nx:Nx+Hx, 1+Hy:size(buff.send,2)+Hy, :) end function _recv_from_west_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) - @info "TripolarX west recv: c x=1:$(Hx), y=$(1+Hy):$(size(buff.recv,2)+Hy), buff size=$(size(buff.recv))" view(c, 1:Hx, 1+Hy:size(buff.recv,2)+Hy, :) .= buff.recv end function _recv_from_east_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) tgt_x, tgt_y = Nx+Hx+1, size(buff.recv,2)+Hy - @info "TripolarX east recv: c x=$(1+Nx+Hx):$(Nx+2Hx), y=$(1+Hy):$(tgt_y), buff size=$(size(buff.recv)), c[$tgt_x,$tgt_y] BEFORE=$(c[tgt_x, tgt_y, (size(c,3)+1)÷2]), buff[1,end,1]=$(buff.recv[1, end, (size(buff.recv,3)+1)÷2])" view(c, 1+Nx+Hx:Nx+2Hx, 1+Hy:size(buff.recv,2)+Hy, :) .= buff.recv - @info "TripolarX east recv: c[$tgt_x,$tgt_y] AFTER=$(c[tgt_x, tgt_y, (size(c,3)+1)÷2])" end # Fold-aware x-buffer: uniform Ny+1 for fold-line fields on north row. @@ -169,7 +164,6 @@ function communication_buffers(grid::MPITripolarGridOfSomeKind, data, bcs, loc) west = x_tripolar_buffer(arch, grid, data, Hx, bcs.west, loc, north) east = x_tripolar_buffer(arch, grid, data, Hx, bcs.east, loc, north) - @info "communication_buffers: rank=$(arch.local_index), loc=$loc, west=$(typeof(west)), east=$(typeof(east)), north=$(typeof(north))" sw = corner_communication_buffer(arch, grid, data, Hx, Hy, west, south) se = corner_communication_buffer(arch, grid, data, Hx, Hy, east, south) @@ -287,50 +281,41 @@ const FF = Tuple{<:Face, <:Face, <:Any} # ── TwoDZipperBuffer: UPivot Center-y (CC/FC) — fold line + Hy halo rows ── function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - @info "CC UPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y fold=$(Ny+Hy), y halo=$(Ny)..$(Ny+Hy-1), buff size=$(size(b.send))" view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:Ny+Hy, :) view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy-1:-1:Ny, :) end function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - @info "FC UPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y fold=$(Ny+Hy), y halo=$(Ny)..$(Ny+Hy-1), buff size=$(size(b.send)), c[$(1+Hx),$(Ny+Hy-1),1]=$(c[1+Hx, Ny+Hy-1, (size(c,3)+1)÷2])" view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:Ny+Hy, :) view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy-1:-1:Ny, :) - @info "FC UPivot north send: buff[end,2,1]=$(b.send[end, 2, (size(b.send,3)+1)÷2]) (should go to mirror's last x at first halo row)" end # ── TwoDZipperBuffer: UPivot Face-y (CF/FF) — no fold line, Hy halo rows ── function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - @info "CF UPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y=$(Ny+1)..$(Ny+Hy), buff size=$(size(b.send))" b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - @info "FF UPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y=$(Ny+1)..$(Ny+Hy), buff size=$(size(b.send))" b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end # ── TwoDZipperBuffer: FPivot Center-y (CC/FC) — no fold line, Hy halo rows ── function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - @info "CC FPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y=$(Ny+1)..$(Ny+Hy), buff size=$(size(b.send))" b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - @info "FC FPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y=$(Ny+1)..$(Ny+Hy), buff size=$(size(b.send))" b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end # ── TwoDZipperBuffer: FPivot Face-y (CF/FF) — fold line + Hy halo rows ── function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - @info "CF FPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y fold=$(Ny+1+Hy), y halo=$(Ny+1)..$(Ny+Hy), buff size=$(size(b.send))" view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+1+Hy:Ny+1+Hy, :) view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - @info "FF FPivot north send: x=$(1+Hx):$(Nx+Hx) (reversed), y fold=$(Ny+1+Hy), y halo=$(Ny+1)..$(Ny+Hy), buff size=$(size(b.send))" view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+1+Hy:Ny+1+Hy, :) view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end @@ -381,37 +366,29 @@ end # Center-x NW: Hy rows, or Hy+1 when fold line is at this y-location function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - @info "CC UPivot NW send: x=$(1+Hx):$(2Hx) (rev), y=$(Ny)..$(Ny+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny, :) end function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - @info "CC FPivot NW send: x=$(1+Hx):$(2Hx) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - @info "CF UPivot NW send: x=$(1+Hx):$(2Hx) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - @info "CF FPivot NW send: x=$(1+Hx):$(2Hx) (rev), y=$(Ny+1)..$(Ny+1+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+1+Hy:-1:Ny+1, :) end # Face-x NW: Hx+1 columns from leftmost Hx+1 interior columns function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - @info "FC UPivot NW send: x=$(1+Hx):$(2Hx+1) (rev), y=$(Ny)..$(Ny+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny, :) end function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - @info "FC FPivot NW send: x=$(1+Hx):$(2Hx+1) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - @info "FF UPivot NW send: x=$(1+Hx):$(2Hx+1) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) end function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - @info "FF FPivot NW send: x=$(1+Hx):$(2Hx+1) (rev), y=$(Ny+1)..$(Ny+1+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+1+Hy:-1:Ny+1, :) end @@ -421,37 +398,29 @@ end # Center-x NE: Hy rows, or Hy+1 when fold line is at this y-location function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - @info "CC UPivot NE send: x=$(1+Nx):$(Nx+Hx) (rev), y=$(Ny)..$(Ny+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny, :) end function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - @info "CC FPivot NE send: x=$(1+Nx):$(Nx+Hx) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) end function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - @info "CF UPivot NE send: x=$(1+Nx):$(Nx+Hx) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) end function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - @info "CF FPivot NE send: x=$(1+Nx):$(Nx+Hx) (rev), y=$(Ny+1)..$(Ny+1+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+1+Hy:-1:Ny+1, :) end # Face-x NE: Hx-1 columns from rightmost Hx-1 interior columns function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - @info "FC UPivot NE send: x=$(Nx+2):$(Nx+Hx) (rev), y=$(Ny)..$(Ny+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny, :) end function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - @info "FC FPivot NE send: x=$(Nx+2):$(Nx+Hx) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) end function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - @info "FF UPivot NE send: x=$(Nx+2):$(Nx+Hx) (rev), y=$(Ny+1)..$(Ny+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) end function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - @info "FF FPivot NE send: x=$(Nx+2):$(Nx+Hx) (rev), y=$(Ny+1)..$(Ny+1+Hy), buff=$(size(b.send))" b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+1+Hy:-1:Ny+1, :) end @@ -468,27 +437,22 @@ end # ── TwoDZipperBuffer: UPivot CC/FC — FL=true, dispatch on WFL ── function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - @info "CC UPivot north recv WFL=true: x=$(1+Hx):$(Nx+Hx), y=$(Ny+Hy):$(Ny+2Hy), buff size=$(size(buff.recv))" view(c, 1+Hx:Nx+Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) - @info "CC UPivot north recv WFL=false: x=$(1+Hx):$(Nx+Hx), y=$(1+Ny+Hy):$(Ny+2Hy), buff size=$(size(buff.recv))" view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - @info "FC UPivot north recv WFL=true: x=$(2+Hx):$(Nx+Hx+1), y=$(Ny+Hy):$(Ny+2Hy), buff size=$(size(buff.recv))" view(c, 2+Hx:Nx+Hx+1, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) tgt_x, tgt_y, tgt_z = Nx+Hx+1, 1+Ny+Hy, (size(c,3)+1)÷2 - @info "FC UPivot north recv WFL=false: x=$(2+Hx):$(Nx+Hx+1), y=$(1+Ny+Hy):$(Ny+2Hy), c[$tgt_x,$tgt_y,$tgt_z] BEFORE=$(c[tgt_x, tgt_y, tgt_z]), buff[end,2,$tgt_z]=$(buff.recv[end, 2, tgt_z])" view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) - @info "FC UPivot north recv WFL=false: c[$tgt_x,$tgt_y,$tgt_z] AFTER=$(c[tgt_x, tgt_y, tgt_z])" end # ── TwoDZipperBuffer: UPivot CF/FF — FL=false, no fold line ── @@ -551,7 +515,6 @@ function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivot view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) end function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - @info "FC UPivot NE corner recv WFL=true: x=$(2+Nx+Hx):$(Nx+2Hx), y=$(Ny+Hy):$(Ny+2Hy), buff size=$(size(buff.recv))" view(c, 2+Nx+Hx:Nx+2Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) end @@ -561,7 +524,6 @@ _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology, _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) - @info "FC UPivot NE corner recv WFL=false: x=$(2+Nx+Hx):$(Nx+2Hx), y=$(1+Ny+Hy):$(Ny+2Hy), buff size=$(size(buff.recv))" view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) end From 8fcb479ff6b6f05157b9c22c8282142fbbc26a0b Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Mon, 30 Mar 2026 21:58:46 +1100 Subject: [PATCH 06/28] Complementary x-buffer fold-line handling with FL/WFL on TripolarXBuffer Add FL/WFL type parameters to TripolarXBuffer, mirroring the north/corner buffer pattern. The west/east x-buffers complement the adjacent NW/NE corner: if the corner writes the fold line, the x-buffer skips it (and vice versa), ensuring no overlapping writes with async MPI. Separate west_tripolar_buffer and east_tripolar_buffer constructors since they complement different corners (NW vs NE). Remaining: Face-x pivot point (global x = Nx/2+1) at the fold line is unfilled in distributed because it falls in the gap between the NE corner (Hx-1 columns) and the east x-buffer (WFL=false). This only affects the exact pivot point which is under land in real ocean configurations. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../distributed_zipper.jl | 85 +++++++++++++------ 1 file changed, 60 insertions(+), 25 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 25cb4e7c133..831084258f9 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -108,47 +108,82 @@ Adapt.adapt_structure(to, buff::OneDZipperBuffer) = nothing Adapt.adapt_structure(to, buff::TwoDZipperBuffer) = nothing Adapt.adapt_structure(to, buff::ZipperCornerBuffer) = nothing -# X-direction buffer for tripolar grids: like TwoDBuffer but with location-aware y-size. -# For fold-line fields on the north row, the x-buffer needs Ny+1 y-rows so that the -# fold-line row is exchanged periodically. Corners then overwrite where needed (WFL=true). -struct TripolarXBuffer{B} +# X-direction buffer for tripolar grids: like TwoDBuffer but with location-aware y-size +# and fold-line awareness. FL/WFL mirror the corner buffer pattern: +# - FL: buffer has fold-line row (same as corners, for MPI size matching) +# - WFL: recv writes the fold-line row (complement of the adjacent corner) +# West WFL = !NW corner WFL, East WFL = !NE corner WFL. +struct TripolarXBuffer{B, FL, WFL} send :: B recv :: B end Adapt.adapt_structure(to, buff::TripolarXBuffer) = nothing -# TripolarXBuffer send/recv: use the buffer's actual y-size -function _fill_west_send_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) +TripolarXBuffer(send, recv, ::Val{FL}, ::Val{WFL}) where {FL, WFL} = + TripolarXBuffer{typeof(send), FL, WFL}(send, recv) + +# X-buffers WFL: complement of the ADJACENT corner. +west_writes_fold_line(arch) = !northwest_writes_fold_line(arch) +east_writes_fold_line(arch) = !northeast_writes_fold_line(arch) + +# TripolarXBuffer send: always pack full buffer (Ny_buf rows) +_fill_west_send_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) = buff.send .= view(c, 1+Hx:2Hx, 1+Hy:size(buff.send,2)+Hy, :) -end -function _fill_east_send_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) +_fill_east_send_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) = buff.send .= view(c, 1+Nx:Nx+Hx, 1+Hy:size(buff.send,2)+Hy, :) -end -function _recv_from_west_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) + +# TripolarXBuffer recv FL=false: no fold line, write all Ny_buf rows +_recv_from_west_buffer!(c, buff::TripolarXBuffer{<:Any, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Hy:size(buff.recv,2)+Hy, :) .= buff.recv -end -function _recv_from_east_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) - tgt_x, tgt_y = Nx+Hx+1, size(buff.recv,2)+Hy +_recv_from_east_buffer!(c, buff::TripolarXBuffer{<:Any, false}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Hy:size(buff.recv,2)+Hy, :) .= buff.recv -end -# Fold-aware x-buffer: uniform Ny+1 for fold-line fields on north row. +# TripolarXBuffer recv FL=true, WFL=true: write all Ny_buf rows (fold line included) +_recv_from_west_buffer!(c, buff::TripolarXBuffer{<:Any, true, true}, Hx, Hy, Nx, Ny) = + view(c, 1:Hx, 1+Hy:size(buff.recv,2)+Hy, :) .= buff.recv +_recv_from_east_buffer!(c, buff::TripolarXBuffer{<:Any, true, true}, Hx, Hy, Nx, Ny) = + view(c, 1+Nx+Hx:Nx+2Hx, 1+Hy:size(buff.recv,2)+Hy, :) .= buff.recv + +# TripolarXBuffer recv FL=true, WFL=false: skip last row (the fold line) +_recv_from_west_buffer!(c, buff::TripolarXBuffer{<:Any, true, false}, Hx, Hy, Nx, Ny) = + view(c, 1:Hx, 1+Hy:size(buff.recv,2)-1+Hy, :) .= view(buff.recv, :, 1:size(buff.recv,2)-1, :) +_recv_from_east_buffer!(c, buff::TripolarXBuffer{<:Any, true, false}, Hx, Hy, Nx, Ny) = + view(c, 1+Nx+Hx:Nx+2Hx, 1+Hy:size(buff.recv,2)-1+Hy, :) .= view(buff.recv, :, 1:size(buff.recv,2)-1, :) + +# Fold-aware x-buffer constructors. # Fallback for non-zipper north (south ranks): standard x-communication. -x_tripolar_buffer(arch, grid, data, Hx, bc, loc, north) = x_communication_buffer(arch, grid, data, Hx, bc) +# Separate west/east constructors since they complement different corners. +function west_tripolar_buffer(arch, grid, data, Hx, bc, loc, + north::TwoDZipperBuffer{Loc, FoT}) where {Loc, FoT} + loc_y = Loc.parameters[2]() + fl = has_fold_line(FoT, loc_y) + Ny_buf = length(loc_y, FoT(), size(grid, 2)) + wfl = fl && west_writes_fold_line(arch) + _, _, Tz = size(parent(data)) + FT = eltype(data) + send = on_architecture(arch, zeros(FT, Hx, Ny_buf, Tz)) + recv = on_architecture(arch, zeros(FT, Hx, Ny_buf, Tz)) + return TripolarXBuffer(send, recv, Val(fl), Val(wfl)) +end -# When north IS a TwoDZipperBuffer (fold north row): use location-aware y-size -# (Ny+1 for BoundedTopology Face-y, i.e., FPivot CF/FF; Ny otherwise) -function x_tripolar_buffer(arch, grid, data, Hx, bc, loc, - north::TwoDZipperBuffer{Loc, FoT}) where {Loc, FoT} - Ny_buf = length(Loc.parameters[2](), FoT(), size(grid, 2)) +function east_tripolar_buffer(arch, grid, data, Hx, bc, loc, + north::TwoDZipperBuffer{Loc, FoT}) where {Loc, FoT} + loc_y = Loc.parameters[2]() + fl = has_fold_line(FoT, loc_y) + Ny_buf = length(loc_y, FoT(), size(grid, 2)) + wfl = fl && east_writes_fold_line(arch) _, _, Tz = size(parent(data)) FT = eltype(data) send = on_architecture(arch, zeros(FT, Hx, Ny_buf, Tz)) recv = on_architecture(arch, zeros(FT, Hx, Ny_buf, Tz)) - return TripolarXBuffer(send, recv) + return TripolarXBuffer(send, recv, Val(fl), Val(wfl)) end +# Fallback when north is not a zipper (south ranks) +west_tripolar_buffer(arch, grid, data, Hx, bc, loc, north) = x_communication_buffer(arch, grid, data, Hx, bc) +east_tripolar_buffer(arch, grid, data, Hx, bc, loc, north) = x_communication_buffer(arch, grid, data, Hx, bc) + ##### ##### Buffer construction for tripolar grids ##### @@ -160,9 +195,9 @@ function communication_buffers(grid::MPITripolarGridOfSomeKind, data, bcs, loc) south = y_communication_buffer(arch, grid, data, Hy, bcs.south) north = y_tripolar_buffer(arch, grid, data, Hy, bcs.north, loc) - # x-buffers dispatch on north: Ny+1 for fold-line fields on north row - west = x_tripolar_buffer(arch, grid, data, Hx, bcs.west, loc, north) - east = x_tripolar_buffer(arch, grid, data, Hx, bcs.east, loc, north) + # x-buffers: separate west/east since they complement different corners (NW/NE) + west = west_tripolar_buffer(arch, grid, data, Hx, bcs.west, loc, north) + east = east_tripolar_buffer(arch, grid, data, Hx, bcs.east, loc, north) sw = corner_communication_buffer(arch, grid, data, Hx, Hy, west, south) From 21dae9472d8e8514bd5036076376f35c1b7ce6a7 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Mon, 30 Mar 2026 21:59:35 +1100 Subject: [PATCH 07/28] Remove leftover debug variable from FC UPivot north recv Co-Authored-By: Claude Opus 4.6 (1M context) --- src/OrthogonalSphericalShellGrids/distributed_zipper.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 831084258f9..040d2e667b5 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -486,7 +486,6 @@ function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnecte end function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) - tgt_x, tgt_y, tgt_z = Nx+Hx+1, 1+Ny+Hy, (size(c,3)+1)÷2 view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) end From 12b7ebaaa59f0ff23363952b1cf67837bcb29cb1 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Mon, 30 Mar 2026 22:10:18 +1100 Subject: [PATCH 08/28] Replace zipper_bc with north_fold_boundary_condition, remove redundant ZBC - Remove local zipper_bc function (duplicated north_fold_boundary_condition) - Use north_fold_boundary_condition(fold_topology(...))(sign) directly - Remove local const ZBC (already defined in BoundaryConditions) - Import ZBC from BoundaryConditions Co-Authored-By: Claude Opus 4.6 (1M context) --- .../distributed_tripolar_grid.jl | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl index 844e1bff9b4..5eeb384958c 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl @@ -1,4 +1,4 @@ -using Oceananigans.BoundaryConditions: DistributedCommunicationBoundaryCondition, FZBC +using Oceananigans.BoundaryConditions: DistributedCommunicationBoundaryCondition, FZBC, ZBC using Oceananigans.Fields: validate_indices, validate_field_data using Oceananigans.DistributedComputations: DistributedComputations, @@ -235,10 +235,6 @@ function receiving_rank(arch; receive_idx_x = ranks(arch)[1] - arch.local_index[ return receive_rank end -zipper_bc(::Type{RightCenterFolded}, sign) = UPivotZipperBoundaryCondition(sign) -zipper_bc(::Type{RightFaceFolded}, sign) = FPivotZipperBoundaryCondition(sign) - -const ZBC = Union{<:UZBC, <:FZBC} function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, grid::MPITripolarGridOfSomeKind, field_name::Symbol, @@ -256,7 +252,7 @@ function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, south = regularize_boundary_condition(bcs.south, grid, loc, 2, LeftBoundary, prognostic_names) north = if yrank == processor_size[2] - 1 && processor_size[1] == 1 - zipper_bc(fold_topology(grid.conformal_mapping), sign) + north_fold_boundary_condition(fold_topology(grid.conformal_mapping))(sign) elseif yrank == processor_size[2] - 1 && processor_size[1] != 1 from = arch.local_rank @@ -295,7 +291,7 @@ function Field(loc::Tuple{<:LX, <:LY, <:LZ}, grid::MPITripolarGridOfSomeKind, da if yrank == processor_size[2] - 1 && processor_size[1] == 1 north_bc = if !(old_bcs.north isa ZBC) - zipper_bc(fold_topology(grid.conformal_mapping), sign(LX, LY)) + north_fold_boundary_condition(fold_topology(grid.conformal_mapping))(sign(LX, LY)) else old_bcs.north end From c2e639b58854226ebc13d32fb3a46a4e1cfc9c0c Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Mon, 30 Mar 2026 22:10:47 +1100 Subject: [PATCH 09/28] Remove unnecessary <: from north_fold_boundary_condition for concrete types Co-Authored-By: Claude Opus 4.6 (1M context) --- .../tripolar_field_extensions.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl b/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl index b583f5f4256..4cf3dfc5015 100644 --- a/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl +++ b/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl @@ -30,8 +30,8 @@ north_fold_boundary_condition(::Type{RightCenterFolded}) = UPivot north_fold_boundary_condition(::Type{RightFaceFolded}) = FPivotZipperBoundaryCondition north_fold_boundary_condition(::Type{LeftConnectedRightCenterFolded}) = UPivotZipperBoundaryCondition north_fold_boundary_condition(::Type{LeftConnectedRightFaceFolded}) = FPivotZipperBoundaryCondition -north_fold_boundary_condition(::Type{<:LeftConnectedRightCenterConnected}) = UPivotZipperBoundaryCondition -north_fold_boundary_condition(::Type{<:LeftConnectedRightFaceConnected}) = FPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{LeftConnectedRightCenterConnected}) = UPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{LeftConnectedRightFaceConnected}) = FPivotZipperBoundaryCondition north_fold_boundary_condition(grid::TripolarGridOfSomeKind) = north_fold_boundary_condition(topology(grid, 2)) # a `TripolarGrid` needs a `UPivotZipperBoundaryCondition` for the north boundary From b37c868898720cfa7388cf570b09885badbdcd8c Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Mon, 30 Mar 2026 22:35:45 +1100 Subject: [PATCH 10/28] formatting --- .../split_explicit_free_surface.jl | 8 ++++---- .../tripolar_field_extensions.jl | 10 +++++----- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl index e567e666478..6d1083f5a74 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl @@ -407,10 +407,10 @@ split_explicit_kernel_size(::Type{RightCenterFolded}, N, H) = 1:N+H-1 split_explicit_kernel_size(::Type{RightFaceFolded}, N, H) = 1:N+H-1 # Distributed fold topologies: connected on both sides (left=MPI, right=fold/zipper) -split_explicit_kernel_size(::Type{LeftConnectedRightCenterFolded}, N, H) = -H+2:N+H-1 -split_explicit_kernel_size(::Type{LeftConnectedRightFaceFolded}, N, H) = -H+2:N+H-1 -split_explicit_kernel_size(::Type{<:LeftConnectedRightCenterConnected}, N, H) = -H+2:N+H-1 -split_explicit_kernel_size(::Type{<:LeftConnectedRightFaceConnected}, N, H) = -H+2:N+H-1 +split_explicit_kernel_size(::Type{LeftConnectedRightCenterFolded}, N, H) = -H+2:N+H-1 +split_explicit_kernel_size(::Type{LeftConnectedRightFaceFolded}, N, H) = -H+2:N+H-1 +split_explicit_kernel_size(::Type{LeftConnectedRightCenterConnected}, N, H) = -H+2:N+H-1 +split_explicit_kernel_size(::Type{LeftConnectedRightFaceConnected}, N, H) = -H+2:N+H-1 # Adapt Adapt.adapt_structure(to, free_surface::SplitExplicitFreeSurface{extend_halos}) where {extend_halos} = diff --git a/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl b/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl index 4cf3dfc5015..046731b0d31 100644 --- a/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl +++ b/src/OrthogonalSphericalShellGrids/tripolar_field_extensions.jl @@ -25,12 +25,12 @@ sign(::Type{Center}, ::Type{Center}) = 1 # Non-fold topologies (FullyConnected, RightConnected, etc.) default to UPivot — this # value is only used as a placeholder; these ranks get their north BC overridden by # inject_halo_communication_boundary_conditions or regularize_field_boundary_conditions. -north_fold_boundary_condition(::Type{<:AbstractTopology}) = UPivotZipperBoundaryCondition -north_fold_boundary_condition(::Type{RightCenterFolded}) = UPivotZipperBoundaryCondition -north_fold_boundary_condition(::Type{RightFaceFolded}) = FPivotZipperBoundaryCondition -north_fold_boundary_condition(::Type{LeftConnectedRightCenterFolded}) = UPivotZipperBoundaryCondition -north_fold_boundary_condition(::Type{LeftConnectedRightFaceFolded}) = FPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{<:AbstractTopology}) = UPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{RightCenterFolded}) = UPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{LeftConnectedRightCenterFolded}) = UPivotZipperBoundaryCondition north_fold_boundary_condition(::Type{LeftConnectedRightCenterConnected}) = UPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{RightFaceFolded}) = FPivotZipperBoundaryCondition +north_fold_boundary_condition(::Type{LeftConnectedRightFaceFolded}) = FPivotZipperBoundaryCondition north_fold_boundary_condition(::Type{LeftConnectedRightFaceConnected}) = FPivotZipperBoundaryCondition north_fold_boundary_condition(grid::TripolarGridOfSomeKind) = north_fold_boundary_condition(topology(grid, 2)) From 76a3148ac441bd968bccefa0ebc190ece722e735 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Pasquier?= <4486578+briochemc@users.noreply.github.com> Date: Mon, 30 Mar 2026 23:58:35 +1100 Subject: [PATCH 11/28] Update src/OrthogonalSphericalShellGrids/tripolar_grid.jl Co-authored-by: Simone Silvestri --- src/OrthogonalSphericalShellGrids/tripolar_grid.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/OrthogonalSphericalShellGrids/tripolar_grid.jl b/src/OrthogonalSphericalShellGrids/tripolar_grid.jl index 81b97edf952..10a4ff75d41 100644 --- a/src/OrthogonalSphericalShellGrids/tripolar_grid.jl +++ b/src/OrthogonalSphericalShellGrids/tripolar_grid.jl @@ -5,7 +5,7 @@ using Oceananigans.Grids: Grids, Bounded, Flat, OrthogonalSphericalShellGrid, Pe using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid """ - struct Tripolar{N, F, S, FT<:AbstractTopology} + struct Tripolar{N, F, S, TY<:AbstractTopology} A structure to represent a tripolar grid on an orthogonal spherical shell. The fold topology `FT` (e.g., `RightCenterFolded` or `RightFaceFolded`) is stored From 48585f606526d8172e6118a5f96b63621c532120 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Wed, 1 Apr 2026 14:06:35 +1100 Subject: [PATCH 12/28] Simplify distributed_zipper.jl and fix FPivot kernel coverage - Refactor distributed_zipper.jl: replace ~65 combinatorial send/recv methods with ~20 using helper functions for y-ranges, x-ranges, and type-parameter accessors. Move FL/WFL to front of type parameter list for clean dispatch. Consolidate 4 corner buffer constructors into 2. - Fix FPivot distributed simulation mismatch: extend async buffer tendency kernel parameters to use worksize(grid) instead of size(grid), ensuring the fold line at Ny+1 is covered by the buffer pass for RightFaceFolded grids. - Add worksize override for distributed FPivot grids (DRFTRG) so kernels correctly operate on the extra Face-y row at Ny+1. - Remove stale explicit imports (FZBC, UZBC, RightConnected, instantiated_location) and fix self-qualified fold_topology access. Co-Authored-By: Claude Opus 4.6 (1M context) --- ...ompute_nonhydrostatic_buffer_tendencies.jl | 35 +- .../OrthogonalSphericalShellGrids.jl | 4 +- .../distributed_tripolar_grid.jl | 20 +- .../distributed_zipper.jl | 419 ++++++------------ 4 files changed, 175 insertions(+), 303 deletions(-) diff --git a/src/Models/NonhydrostaticModels/compute_nonhydrostatic_buffer_tendencies.jl b/src/Models/NonhydrostaticModels/compute_nonhydrostatic_buffer_tendencies.jl index 0972d277412..fb34077bb0d 100644 --- a/src/Models/NonhydrostaticModels/compute_nonhydrostatic_buffer_tendencies.jl +++ b/src/Models/NonhydrostaticModels/compute_nonhydrostatic_buffer_tendencies.jl @@ -2,6 +2,7 @@ import Oceananigans.Models: compute_buffer_tendencies! using Oceananigans.TurbulenceClosures: required_halo_size_x, required_halo_size_y using Oceananigans.Grids: XFlatGrid, YFlatGrid +using Oceananigans.Utils: worksize # TODO: the code in this file is difficult to understand. # Rewriting it may be helpful. @@ -24,44 +25,48 @@ function compute_buffer_tendencies!(model::NonhydrostaticModel) return nothing end -# tendencies need computing in the range 1 : H and N - H + 1 : N +# Tendencies need computing in the range 1 : H and W - H + 1 : W where W = worksize. +# These buffer regions complement the halo-independent interior kernel (H+1:W-H). function buffer_tendency_kernel_parameters(grid, arch) Nx, Ny, Nz = size(grid) + Wx, Wy, _ = worksize(grid) Hx, Hy, _ = halo_size(grid) - param_west = (1:Hx, 1:Ny, 1:Nz) - param_east = (Nx-Hx+1:Nx, 1:Ny, 1:Nz) - param_south = (1:Nx, 1:Hy, 1:Nz) - param_north = (1:Nx, Ny-Hy+1:Ny, 1:Nz) + param_west = (1:Hx, 1:Wy, 1:Nz) + param_east = (Wx-Hx+1:Wx, 1:Wy, 1:Nz) + param_south = (1:Wx, 1:Hy, 1:Nz) + param_north = (1:Wx, Wy-Hy+1:Wy, 1:Nz) params = (param_west, param_east, param_south, param_north) return buffer_parameters(params, grid, arch) end -# p needs computing in the range 0 : 0 and N + 1 : N + 1 +# p needs computing in the range 0 : 0 and W + 1 : W + 1 where W = worksize function buffer_p_kernel_parameters(grid, arch) Nx, Ny, _ = size(grid) + Wx, Wy, _ = worksize(grid) - param_west = (0:0, 1:Ny) - param_east = (Nx+1:Nx+1, 1:Ny) - param_south = (1:Nx, 0:0) - param_north = (1:Nx, Ny+1:Ny+1) + param_west = (0:0, 1:Wy) + param_east = (Wx+1:Wx+1, 1:Wy) + param_south = (1:Wx, 0:0) + param_north = (1:Wx, Wy+1:Wy+1) params = (param_west, param_east, param_south, param_north) return buffer_parameters(params, grid, arch) end -# closure_fields need recomputing in the range 0 : B and N - B + 1 : N + 1 +# closure_fields need recomputing in the range 0 : B and W - B + 1 : W + 1 where W = worksize function buffer_κ_kernel_parameters(grid, closure, arch) Nx, Ny, Nz = size(grid) + Wx, Wy, _ = worksize(grid) Bx = required_halo_size_x(closure) By = required_halo_size_y(closure) - param_west = (0:Bx, 1:Ny, 1:Nz) - param_east = (Nx-Bx+1:Nx+1, 1:Ny, 1:Nz) - param_south = (1:Nx, 0:By, 1:Nz) - param_north = (1:Nx, Ny-By+1:Ny+1, 1:Nz) + param_west = (0:Bx, 1:Wy, 1:Nz) + param_east = (Wx-Bx+1:Wx+1, 1:Wy, 1:Nz) + param_south = (1:Wx, 0:By, 1:Nz) + param_north = (1:Wx, Wy-By+1:Wy+1, 1:Nz) params = (param_west, param_east, param_south, param_north) return buffer_parameters(params, grid, arch) diff --git a/src/OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl b/src/OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl index 65e6bca9e8c..10f52fa85a3 100644 --- a/src/OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl +++ b/src/OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl @@ -7,8 +7,8 @@ import Oceananigans import Oceananigans.Architectures: on_architecture using Oceananigans.Architectures: on_architecture, AbstractArchitecture, CPU, GPU -using Oceananigans.BoundaryConditions: BoundaryCondition, UZBC -using Oceananigans.Grids: AbstractTopology, RightConnected +using Oceananigans.BoundaryConditions: BoundaryCondition +using Oceananigans.Grids: AbstractTopology using Oceananigans.Grids: halo_size, generate_coordinate, topology using Oceananigans.Grids: total_length, add_halos, fill_metric_halo_regions! diff --git a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl index 4a8d1cfd09f..70c5ece4306 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl @@ -1,4 +1,4 @@ -using Oceananigans.BoundaryConditions: DistributedCommunicationBoundaryCondition, FZBC, ZBC +using Oceananigans.BoundaryConditions: DistributedCommunicationBoundaryCondition, ZBC using Oceananigans.Fields: validate_indices, validate_field_data using Oceananigans.DistributedComputations: DistributedComputations, @@ -11,8 +11,10 @@ using Oceananigans.DistributedComputations: concatenate_local_sizes, communication_buffers -using Oceananigans.Grids: topology, RightConnected, FullyConnected +using Oceananigans.Grids: topology, FullyConnected, + LeftConnectedRightFaceFolded, LeftConnectedRightFaceConnected using Oceananigans.DistributedComputations: insert_connected_topology +using Oceananigans.Utils: Utils import Oceananigans.Fields: Field, validate_indices, validate_boundary_conditions @@ -339,7 +341,7 @@ function DistributedComputations.reconstruct_global_grid(grid::MPITripolarGrid) first_pole_longitude = grid.conformal_mapping.first_pole_longitude southernmost_latitude = grid.conformal_mapping.southernmost_latitude - fold_topology = Oceananigans.OrthogonalSphericalShellGrids.fold_topology(grid.conformal_mapping) + fold_topology = OrthogonalSphericalShellGrids.fold_topology(grid.conformal_mapping) return TripolarGrid(child_arch, FT; halo, @@ -362,7 +364,7 @@ function Grids.with_halo(new_halo, old_grid::MPITripolarGrid) north_poles_latitude = old_grid.conformal_mapping.north_poles_latitude first_pole_longitude = old_grid.conformal_mapping.first_pole_longitude southernmost_latitude = old_grid.conformal_mapping.southernmost_latitude - fold_topology = Oceananigans.OrthogonalSphericalShellGrids.fold_topology(old_grid.conformal_mapping) + fold_topology = OrthogonalSphericalShellGrids.fold_topology(old_grid.conformal_mapping) return TripolarGrid(arch, eltype(old_grid); halo = new_halo, @@ -373,3 +375,13 @@ function Grids.with_halo(new_halo, old_grid::MPITripolarGrid) z, fold_topology) end + +##### +##### Extend worksize for distributed FPivot grids (matches RFTRG worksize for serial FPivot) +##### + +const DistributedFPivotTopology = Union{LeftConnectedRightFaceFolded, LeftConnectedRightFaceConnected} +const DRFTRG = Union{MPITripolarGrid{<:Any, <:Any, <:DistributedFPivotTopology}, + ImmersedBoundaryGrid{<:Any, <:Any, <:DistributedFPivotTopology, <:Any, <:MPITripolarGrid}} + +Utils.worksize(grid::DRFTRG) = grid.Nx, grid.Ny+1, grid.Nz diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 040d2e667b5..8fc96376e03 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -6,7 +6,6 @@ using Oceananigans.Grids: AbstractGrid, topology, LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded, LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected using Oceananigans.DistributedComputations: Distributed, on_architecture, ranks, x_communication_buffer -using Oceananigans.Fields: instantiated_location import Oceananigans.BoundaryConditions: fill_halo_regions! import Oceananigans.DistributedComputations: synchronize_communication!, @@ -83,13 +82,13 @@ struct OneDZipperBuffer{Loc, FoT, B, S} sign :: S end -struct TwoDZipperBuffer{Loc, FoT, B, S, FL, WFL} +struct TwoDZipperBuffer{FL, WFL, Loc, FoT, B, S} send :: B recv :: B sign :: S end -struct ZipperCornerBuffer{Loc, FoT, B, S, FL, WFL} +struct ZipperCornerBuffer{FL, WFL, Loc, FoT, B, S} send :: B recv :: B sign :: S @@ -98,11 +97,11 @@ end # Value-argument constructors: all types inferred OneDZipperBuffer(loc::Loc, fot::FoT, send::B, recv::B, sign::S) where {Loc, FoT, B, S} = OneDZipperBuffer{Loc, FoT, B, S}(send, recv, sign) TwoDZipperBuffer(loc, fot, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {FL, WFL} = - TwoDZipperBuffer{typeof(loc), typeof(fot), typeof(send), typeof(sign), FL, WFL}(send, recv, sign) + TwoDZipperBuffer{FL, WFL, typeof(loc), typeof(fot), typeof(send), typeof(sign)}(send, recv, sign) ZipperCornerBuffer(loc, fot, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {FL, WFL} = - ZipperCornerBuffer{typeof(loc), typeof(fot), typeof(send), typeof(sign), FL, WFL}(send, recv, sign) + ZipperCornerBuffer{FL, WFL, typeof(loc), typeof(fot), typeof(send), typeof(sign)}(send, recv, sign) ZipperCornerBuffer(::Type{Loc}, ::Type{FoT}, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {Loc, FoT, FL, WFL} = - ZipperCornerBuffer{Loc, FoT, typeof(send), typeof(sign), FL, WFL}(send, recv, sign) + ZipperCornerBuffer{FL, WFL, Loc, FoT, typeof(send), typeof(sign)}(send, recv, sign) Adapt.adapt_structure(to, buff::OneDZipperBuffer) = nothing Adapt.adapt_structure(to, buff::TwoDZipperBuffer) = nothing @@ -155,7 +154,7 @@ _recv_from_east_buffer!(c, buff::TripolarXBuffer{<:Any, true, false}, Hx, Hy, Nx # Fallback for non-zipper north (south ranks): standard x-communication. # Separate west/east constructors since they complement different corners. function west_tripolar_buffer(arch, grid, data, Hx, bc, loc, - north::TwoDZipperBuffer{Loc, FoT}) where {Loc, FoT} + north::TwoDZipperBuffer{<:Any, <:Any, Loc, FoT}) where {Loc, FoT} loc_y = Loc.parameters[2]() fl = has_fold_line(FoT, loc_y) Ny_buf = length(loc_y, FoT(), size(grid, 2)) @@ -168,7 +167,7 @@ function west_tripolar_buffer(arch, grid, data, Hx, bc, loc, end function east_tripolar_buffer(arch, grid, data, Hx, bc, loc, - north::TwoDZipperBuffer{Loc, FoT}) where {Loc, FoT} + north::TwoDZipperBuffer{<:Any, <:Any, Loc, FoT}) where {Loc, FoT} loc_y = Loc.parameters[2]() fl = has_fold_line(FoT, loc_y) Ny_buf = length(loc_y, FoT(), size(grid, 2)) @@ -246,51 +245,29 @@ northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge) = corner_commu # MPI size matching between mirror partners. WFL (writes fold line) is per-corner, # based on whether the corner's global x-position is past the pivot. -# ── NW corner: Center-x (Hx columns) ── function northwest_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, - yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Center, <:Any, <:Any}, FoT} + yedge::TwoDZipperBuffer{<:Any, <:Any, Loc, FoT}) where {Loc, FoT} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign + loc_x_inst = instantiate(Loc.parameters[1]) fl = has_fold_line(FoT, Loc.parameters[2]()) Hy′ = fl ? Hy + 1 : Hy wfl = fl && northwest_writes_fold_line(arch) - send = on_architecture(arch, zeros(FT, Hx, Hy′, Tz)) - recv = on_architecture(arch, zeros(FT, Hx, Hy′, Tz)) + Hx′ = nw_corner_nx(loc_x_inst, Hx) + send = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) + recv = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) return ZipperCornerBuffer(Loc, FoT, send, recv, sgn, Val(fl), Val(wfl)) end -# ── NW corner: Face-x (Hx+1 columns for the Face-x shift) ── -function northwest_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, - yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Face, <:Any, <:Any}, FoT} - Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - fl = has_fold_line(FoT, Loc.parameters[2]()) - Hy′ = fl ? Hy + 1 : Hy - wfl = fl && northwest_writes_fold_line(arch) - send = on_architecture(arch, zeros(FT, Hx+1, Hy′, Tz)) - recv = on_architecture(arch, zeros(FT, Hx+1, Hy′, Tz)) - return ZipperCornerBuffer(Loc, FoT, send, recv, sgn, Val(fl), Val(wfl)) -end - -# ── NE corner: Center-x (Hx columns) ── function northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, - yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Center, <:Any, <:Any}, FoT} + yedge::TwoDZipperBuffer{<:Any, <:Any, Loc, FoT}) where {Loc, FoT} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign + loc_x_inst = instantiate(Loc.parameters[1]) fl = has_fold_line(FoT, Loc.parameters[2]()) Hy′ = fl ? Hy + 1 : Hy wfl = fl && northeast_writes_fold_line(arch) - send = on_architecture(arch, zeros(FT, Hx, Hy′, Tz)) - recv = on_architecture(arch, zeros(FT, Hx, Hy′, Tz)) - return ZipperCornerBuffer(Loc, FoT, send, recv, sgn, Val(fl), Val(wfl)) -end - -# ── NE corner: Face-x (Hx-1 columns, north buffer covers the extra Face-x column) ── -function northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, - yedge::TwoDZipperBuffer{Loc, FoT}) where {Loc <: Tuple{<:Face, <:Any, <:Any}, FoT} - Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - fl = has_fold_line(FoT, Loc.parameters[2]()) - Hy′ = fl ? Hy + 1 : Hy - wfl = fl && northeast_writes_fold_line(arch) - send = on_architecture(arch, zeros(FT, Hx - 1, Hy′, Tz)) - recv = on_architecture(arch, zeros(FT, Hx - 1, Hy′, Tz)) + Hx′ = ne_corner_nx(loc_x_inst, Hx) + send = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) + recv = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) return ZipperCornerBuffer(Loc, FoT, send, recv, sgn, Val(fl), Val(wfl)) end @@ -304,302 +281,180 @@ const CF = Tuple{<:Center, <:Face, <:Any} const FF = Tuple{<:Face, <:Face, <:Any} ##### -##### Fold-aware send buffer packing: _fill_north_send_buffer! -##### -##### Y-range rules (parent-array coords, Ny = size(grid, 2)): -##### UPivot Center-y (CC/FC): skip fold at Ny → Hy rows from Ny-1: Ny+Hy-1:-1:Ny -##### All other cases: Hy rows from Ny: Ny+Hy:-1:Ny+1 -##### -##### For TwoD fold-line fields (Hy+1 buffer): row 1 = fold line, rows 2:Hy+1 = halo sources +##### Helper functions: compute ranges from location and topology ##### -# ── TwoDZipperBuffer: UPivot Center-y (CC/FC) — fold line + Hy halo rows ── +# Type-parameter accessors +@inline loc_x(::TwoDZipperBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[1]) +@inline loc_y(::TwoDZipperBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[2]) +@inline fold_topo(::TwoDZipperBuffer{<:Any, <:Any, <:Any, FoT}) where FoT = FoT() -function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:Ny+Hy, :) - view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy-1:-1:Ny, :) -end +@inline loc_x(::ZipperCornerBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[1]) +@inline loc_y(::ZipperCornerBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[2]) +@inline fold_topo(::ZipperCornerBuffer{<:Any, <:Any, <:Any, FoT}) where FoT = FoT() -function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:Ny+Hy, :) - view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy-1:-1:Ny, :) -end +@inline loc_x(::OneDZipperBuffer{Loc}) where Loc = instantiate(Loc.parameters[1]) +@inline loc_y(::OneDZipperBuffer{Loc}) where Loc = instantiate(Loc.parameters[2]) +@inline fold_topo(::OneDZipperBuffer{<:Any, FoT}) where FoT = FoT() -# ── TwoDZipperBuffer: UPivot Face-y (CF/FF) — no fold line, Hy halo rows ── +# Send y-ranges (source rows, reversed for fold) +@inline send_fold_y(::UPivotTopology, ::Center, Hy, Ny) = Ny + Hy +@inline send_fold_y(::FPivotTopology, ::Face, Hy, Ny) = Ny + 1 + Hy -function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -end -function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -end +@inline send_halo_y(::UPivotTopology, ::Center, Hy, Ny) = Ny+Hy-1:-1:Ny +@inline send_halo_y(topo, loc_y, Hy, Ny) = Ny+Hy:-1:Ny+1 -# ── TwoDZipperBuffer: FPivot Center-y (CC/FC) — no fold line, Hy halo rows ── +# Recv y-ranges (destination rows in parent) +@inline recv_fold_y(::UPivotTopology, ::Center, Hy, Ny) = Ny + Hy +@inline recv_fold_y(::FPivotTopology, ::Face, Hy, Ny) = Ny + 1 + Hy -function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -end -function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -end - -# ── TwoDZipperBuffer: FPivot Face-y (CF/FF) — fold line + Hy halo rows ── - -function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+1+Hy:Ny+1+Hy, :) - view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -end - -function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+1+Hy:Ny+1+Hy, :) - view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -end - -# ── OneDZipperBuffer: UPivot Center-y (CC/FC) — simple reverse ── - -_fill_north_send_buffer!(c, b::OneDZipperBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, size(c,1):-1:1, Ny+Hy-1:-1:Ny, :) - -function _fill_north_send_buffer!(c, b::OneDZipperBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - Tx = size(c, 1) - view(b.send, 2:Tx, :, :) .= b.sign .* view(c, Tx:-1:2, Ny+Hy-1:-1:Ny, :) - view(b.send, 1:1, :, :) .= b.sign .* view(c, 1:1, Ny+Hy-1:-1:Ny, :) -end +@inline recv_halo_y(::FPivotTopology, ::Face, Hy, Ny) = 2+Ny+Hy:1+Ny+2Hy +@inline recv_halo_y(topo, loc_y, Hy, Ny) = 1+Ny+Hy:Ny+2Hy -# ── OneDZipperBuffer: UPivot Face-y (CF/FF) ── +# Recv x-range for north buffer +@inline recv_x_range(::Center, Hx, Nx) = 1+Hx:Nx+Hx +@inline recv_x_range(::Face, Hx, Nx) = 2+Hx:Nx+Hx+1 -_fill_north_send_buffer!(c, b::OneDZipperBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, size(c,1):-1:1, Ny+Hy:-1:Ny+1, :) +# Corner send x-ranges (reversed for fold) +@inline nw_send_x(::Center, Hx, Nx) = 2Hx:-1:1+Hx +@inline nw_send_x(::Face, Hx, Nx) = 2Hx+1:-1:1+Hx +@inline ne_send_x(::Center, Hx, Nx) = Nx+Hx:-1:1+Nx +@inline ne_send_x(::Face, Hx, Nx) = Nx+Hx:-1:Nx+2 -function _fill_north_send_buffer!(c, b::OneDZipperBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - Tx = size(c, 1) - view(b.send, 2:Tx, :, :) .= b.sign .* view(c, Tx:-1:2, Ny+Hy:-1:Ny+1, :) - view(b.send, 1:1, :, :) .= b.sign .* view(c, 1:1, Ny+Hy:-1:Ny+1, :) -end +# Corner recv x-ranges +@inline nw_recv_x(::Center, Hx, Nx) = 1:Hx +@inline nw_recv_x(::Face, Hx, Nx) = 1:Hx+1 +@inline ne_recv_x(::Center, Hx, Nx) = 1+Nx+Hx:Nx+2Hx +@inline ne_recv_x(::Face, Hx, Nx) = 2+Nx+Hx:Nx+2Hx -# ── OneDZipperBuffer: FPivot Center-y (CC/FC) ── +# Corner buffer x-size (for constructors) +@inline nw_corner_nx(::Center, Hx) = Hx +@inline nw_corner_nx(::Face, Hx) = Hx + 1 +@inline ne_corner_nx(::Center, Hx) = Hx +@inline ne_corner_nx(::Face, Hx) = Hx - 1 -_fill_north_send_buffer!(c, b::OneDZipperBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, size(c,1):-1:1, Ny+Hy:-1:Ny+1, :) +##### +##### TwoDZipperBuffer: north send (FL=true has fold line, FL=false does not) +##### -function _fill_north_send_buffer!(c, b::OneDZipperBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - Tx = size(c, 1) - view(b.send, 2:Tx, :, :) .= b.sign .* view(c, Tx:-1:2, Ny+Hy:-1:Ny+1, :) - view(b.send, 1:1, :, :) .= b.sign .* view(c, 1:1, Ny+Hy:-1:Ny+1, :) +function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{true}, Hx, Hy, Nx, Ny) + topo, ly = fold_topo(b), loc_y(b) + fy = send_fold_y(topo, ly, Hy, Ny) + view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, fy:fy, :) + view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, send_halo_y(topo, ly, Hy, Ny), :) end -# ── OneDZipperBuffer: FPivot Face-y (CF/FF) ── - -_fill_north_send_buffer!(c, b::OneDZipperBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = b.send .= b.sign .* view(c, size(c,1):-1:1, Ny+Hy:-1:Ny+1, :) - -function _fill_north_send_buffer!(c, b::OneDZipperBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - Tx = size(c, 1) - view(b.send, 2:Tx, :, :) .= b.sign .* view(c, Tx:-1:2, Ny+Hy:-1:Ny+1, :) - view(b.send, 1:1, :, :) .= b.sign .* view(c, 1:1, Ny+Hy:-1:Ny+1, :) +function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{false}, Hx, Hy, Nx, Ny) + topo, ly = fold_topo(b), loc_y(b) + b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, send_halo_y(topo, ly, Hy, Ny), :) end ##### -##### Fold-aware send buffer packing: NW corner (leftmost Hx interior columns) +##### TwoDZipperBuffer: north recv (FL × WFL dispatch) ##### -# Center-x NW: Hy rows, or Hy+1 when fold line is at this y-location -function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny, :) -end -function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -end -function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -end -function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, 2Hx:-1:1+Hx, Ny+1+Hy:-1:Ny+1, :) +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{true, true}, Hx, Hy, Nx, Ny) + xr = recv_x_range(loc_x(buff), Hx, Nx) + topo, ly = fold_topo(buff), loc_y(buff) + fy = recv_fold_y(topo, ly, Hy, Ny) + view(c, xr, fy:fy , :) .= view(buff.recv, :, 1:1 , :) + view(c, xr, recv_halo_y(topo, ly, Hy, Ny), :) .= view(buff.recv, :, 2:Hy+1, :) end -# Face-x NW: Hx+1 columns from leftmost Hx+1 interior columns -function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny, :) +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{true, false}, Hx, Hy, Nx, Ny) + xr = recv_x_range(loc_x(buff), Hx, Nx) + view(c, xr, recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= view(buff.recv, :, 2:Hy+1, :) end -function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -end -function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+Hy:-1:Ny+1, :) -end -function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, 2Hx+1:-1:1+Hx, Ny+1+Hy:-1:Ny+1, :) + +function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{false}, Hx, Hy, Nx, Ny) + xr = recv_x_range(loc_x(buff), Hx, Nx) + view(c, xr, recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= buff.recv end ##### -##### Fold-aware send buffer packing: NE corner (rightmost Hx interior columns) +##### OneDZipperBuffer: north send (Center-x full reverse, Face-x partial reverse) ##### -# Center-x NE: Hy rows, or Hy+1 when fold line is at this y-location -function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny, :) -end -function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) -end -function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+Hy:-1:Ny+1, :) -end -function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, Nx+Hx:-1:1+Nx, Ny+1+Hy:-1:Ny+1, :) +function _fill_north_send_buffer!(c, b::OneDZipperBuffer{<:Union{CC,CF}}, Hx, Hy, Nx, Ny) + b.send .= b.sign .* view(c, size(c,1):-1:1, send_halo_y(fold_topo(b), loc_y(b), Hy, Ny), :) end -# Face-x NE: Hx-1 columns from rightmost Hx-1 interior columns -function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:UPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny, :) -end -function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) -end -function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+Hy:-1:Ny+1, :) -end -function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, Nx+Hx:-1:Nx+2, Ny+1+Hy:-1:Ny+1, :) +function _fill_north_send_buffer!(c, b::OneDZipperBuffer{<:Union{FC,FF}}, Hx, Hy, Nx, Ny) + Tx = size(c, 1) + hy = send_halo_y(fold_topo(b), loc_y(b), Hy, Ny) + view(b.send, 2:Tx, :, :) .= b.sign .* view(c, Tx:-1:2, hy, :) + view(b.send, 1:1, :, :) .= b.sign .* view(c, 1:1, hy, :) end ##### -##### Recv methods: direct placement (data is already folded by sender) -##### -##### For TwoD fold-line fields (FL=true, Hy+1 buffer): -##### WFL=true: write fold line (row 1) + halos (rows 2:Hy+1) -##### WFL=false: write halos only (rows 2:Hy+1, skip fold line) +##### OneDZipperBuffer: north recv ##### -##### For FPivot Face-y: halos start at parent Ny+Hy+2 (field Ny+2, past fold line at Ny+1) -##### - -# ── TwoDZipperBuffer: UPivot CC/FC — FL=true, dispatch on WFL ── -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - view(c, 1+Hx:Nx+Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) - view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) -end +_recv_from_north_buffer!(c, buff::OneDZipperBuffer, Hx, Hy, Nx, Ny) = view(c, :, recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= buff.recv -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) - view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) -end +##### +##### Corner send: NW and NE (FL=true includes fold line, FL=false halo only) +##### -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - view(c, 2+Hx:Nx+Hx+1, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) - view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{true}, Hx, Hy, Nx, Ny) + topo = fold_topo(b) + ly = loc_y(b) + fy = send_fold_y(topo, ly, Hy, Ny) + b.send .= b.sign .* view(c, nw_send_x(loc_x(b), Hx, Nx), fy:-1:fy-Hy, :) end -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:LeftConnectedRightCenterConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) - view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{false}, Hx, Hy, Nx, Ny) + topo = fold_topo(b) + fy = send_halo_y(topo, loc_y(b), Hy, Ny) + b.send .= b.sign .* view(c, nw_send_x(loc_x(b), Hx, Nx), fy, :) end -# ── TwoDZipperBuffer: UPivot CF/FF — FL=false, no fold line ── - -_recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv - -# ── TwoDZipperBuffer: FPivot CC/FC — FL=false, no fold line ── - -_recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Hx:Nx+Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Hx:Nx+Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv - -# ── TwoDZipperBuffer: FPivot CF/FF — FL=true, dispatch on WFL, shifted +1 for Face-y ── - -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CF, <:LeftConnectedRightFaceConnected, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - view(c, 1+Hx:Nx+Hx, Ny+1+Hy:Ny+1+Hy, :) .= view(buff.recv, :, 1:1, :) - view(c, 1+Hx:Nx+Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{true}, Hx, Hy, Nx, Ny) + topo = fold_topo(b) + ly = loc_y(b) + fy = send_fold_y(topo, ly, Hy, Ny) + b.send .= b.sign .* view(c, ne_send_x(loc_x(b), Hx, Nx), fy:-1:fy-Hy, :) end -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:CF, <:LeftConnectedRightFaceConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) - view(c, 1+Hx:Nx+Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{false}, Hx, Hy, Nx, Ny) + topo = fold_topo(b) + fy = send_halo_y(topo, loc_y(b), Hy, Ny) + b.send .= b.sign .* view(c, ne_send_x(loc_x(b), Hx, Nx), fy, :) end -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FF, <:LeftConnectedRightFaceConnected, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - view(c, 2+Hx:Nx+Hx+1, Ny+1+Hy:Ny+1+Hy, :) .= view(buff.recv, :, 1:1, :) - view(c, 2+Hx:Nx+Hx+1, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) -end +##### +##### Corner recv: NW and NE (FL × WFL dispatch) +##### -function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{<:FF, <:LeftConnectedRightFaceConnected, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) - view(c, 2+Hx:Nx+Hx+1, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:Hy+1, :) +function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{true, true}, Hx, Hy, Nx, Ny) + xr = nw_recv_x(loc_x(buff), Hx, Nx) + topo, ly = fold_topo(buff), loc_y(buff) + fy = recv_fold_y(topo, ly, Hy, Ny) + view(c, xr, fy:fy, :) .= view(buff.recv, :, 1:1, :) + view(c, xr, recv_halo_y(topo, ly, Hy, Ny), :) .= view(buff.recv, :, 2:size(buff.recv,2), :) end -# ── OneDZipperBuffer: UPivot — standard placement ── - -_recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:Any, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, :, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv - -# ── OneDZipperBuffer: FPivot CC/FC — standard placement ── - -_recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, :, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, :, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{true, false}, Hx, Hy, Nx, Ny) = + view(c, nw_recv_x(loc_x(buff), Hx, Nx), recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= + view(buff.recv, :, 2:size(buff.recv,2), :) -# ── OneDZipperBuffer: FPivot CF/FF — shifted +1 for Face-extended ── - -_recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:CF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, :, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv -_recv_from_north_buffer!(c, buff::OneDZipperBuffer{<:FF, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, :, 2+Ny+Hy:1+Ny+2Hy, :) .= buff.recv - -# ── Corner recv: UPivot CC/FC — FL=true, dispatch on WFL ── - -# WFL=true: write fold line + halos -function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - view(c, 1:Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) - view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -end -function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - view(c, 1:Hx+1, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) - view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -end -function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - view(c, 1+Nx+Hx:Nx+2Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) - view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -end -function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - view(c, 2+Nx+Hx:Nx+2Hx, Ny+Hy:Ny+Hy, :) .= view(buff.recv, :, 1:1, :) - view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -end +_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{false}, Hx, Hy, Nx, Ny) = + view(c, nw_recv_x(loc_x(buff), Hx, Nx), recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= buff.recv -# WFL=false: write halos only (skip fold line) -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:UPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) - view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{true, true}, Hx, Hy, Nx, Ny) + xr = ne_recv_x(loc_x(buff), Hx, Nx) + topo, ly = fold_topo(buff), loc_y(buff) + fy = recv_fold_y(topo, ly, Hy, Ny) + view(c, xr, fy:fy, :) .= view(buff.recv, :, 1:1, :) + view(c, xr, recv_halo_y(topo, ly, Hy, Ny), :) .= view(buff.recv, :, 2:size(buff.recv,2), :) end -# ── Corner recv: UPivot CF/FF — FL=false, no fold line, Hy rows ── - -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:UPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv - -# ── Corner recv: FPivot CC/FC — FL=false, no fold line, Hy rows ── - -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FC, <:FPivotTopology}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:Ny+2Hy, :) .= buff.recv - -# ── Corner recv: FPivot CF/FF — FL=true, dispatch on WFL, shifted +1 for Face-y ── - -# WFL=true: write fold line + halos -function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - view(c, 1:Hx, 1+Ny+Hy:1+Ny+Hy, :) .= view(buff.recv, :, 1:1, :) - view(c, 1:Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -end -function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - view(c, 1:Hx+1, 1+Ny+Hy:1+Ny+Hy, :) .= view(buff.recv, :, 1:1, :) - view(c, 1:Hx+1, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -end -function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - view(c, 1+Nx+Hx:Nx+2Hx, 1+Ny+Hy:1+Ny+Hy, :) .= view(buff.recv, :, 1:1, :) - view(c, 1+Nx+Hx:Nx+2Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -end -function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology, <:Any, <:Any, true, true}, Hx, Hy, Nx, Ny) - view(c, 2+Nx+Hx:Nx+2Hx, 1+Ny+Hy:1+Ny+Hy, :) .= view(buff.recv, :, 1:1, :) - view(c, 2+Nx+Hx:Nx+2Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -end +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{true, false}, Hx, Hy, Nx, Ny) = + view(c, ne_recv_x(loc_x(buff), Hx, Nx), recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= + view(buff.recv, :, 2:size(buff.recv,2), :) -# WFL=false: write halos only (skip fold line) -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx+1, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:CF, <:FPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{<:FF, <:FPivotTopology, <:Any, <:Any, true, false}, Hx, Hy, Nx, Ny) = view(c, 2+Nx+Hx:Nx+2Hx, 2+Ny+Hy:1+Ny+2Hy, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{false}, Hx, Hy, Nx, Ny) = + view(c, ne_recv_x(loc_x(buff), Hx, Nx), recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= buff.recv ##### ##### _switch_north_halos! — no-op (fold logic in send buffers) From bd8ce1fdd404e131900313cfe4cef15d610e8b57 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Wed, 1 Apr 2026 15:24:01 +1100 Subject: [PATCH 13/28] Remove stale UPivot fold BC overrides from test utils The overrides were modifying the serial fold BC to skip the fold-line consistency substitution, which is no longer needed now that the distributed fold handles this correctly. Co-Authored-By: Claude Opus 4.6 (1M context) --- test/distributed_tests_utils.jl | 41 --------------------------------- 1 file changed, 41 deletions(-) diff --git a/test/distributed_tests_utils.jl b/test/distributed_tests_utils.jl index c034abbc9ff..61dc75b1bed 100644 --- a/test/distributed_tests_utils.jl +++ b/test/distributed_tests_utils.jl @@ -4,8 +4,6 @@ using Oceananigans.DistributedComputations: reconstruct_global_field, reconstruc using Oceananigans.Units using Oceananigans.TimeSteppers: first_time_step! -import Oceananigans.BoundaryConditions: _fill_north_halo! -using Oceananigans.BoundaryConditions: UZBC, CCLocation, FCLocation include("dependencies_for_runtests.jl") @@ -30,45 +28,6 @@ function analytical_immersed_tripolar_grid(underlying_grid::TripolarGrid; radius return grid end -# The serial version of the TripolarGrid substitutes the second half of the last row of the grid. -# This is not done in the distributed version, so we need to undo this substitution if we want to -# compare the results. Otherwise very tiny differences caused by finite precision computations -# will appear in the last row of the grid. - -# tracers or similar fields -@inline _fill_north_halo!(i, k, grid, c, bc::UZBC, ::CCLocation, args...) = my_fold_north_center_center_upivot!(i, k, grid, bc.condition, c) -@inline _fill_north_halo!(i, k, grid, u, bc::UZBC, ::FCLocation, args...) = my_fold_north_face_center_upivot!(i, k, grid, bc.condition, u) - -@inline function my_fold_north_face_center_upivot!(i, k, grid, sign, c) - Nx, Ny, _ = size(grid) - - i′ = Nx - i + 2 # Remember! element Nx + 1 does not exist! - sign = ifelse(i′ > Nx , abs(sign), sign) # for periodic elements we change the sign - i′ = ifelse(i′ > Nx, i′ - Nx, i′) # Periodicity is hardcoded in the x-direction!! - Hy = grid.Hy - - for j = 1 : Hy - @inbounds begin - c[i, Ny + j, k] = sign * c[i′, Ny - j, k] # The Ny line is duplicated so we substitute starting Ny-1 - end - end - - return nothing -end - -@inline function my_fold_north_center_center_upivot!(i, k, grid, sign, c) - Nx, Ny, _ = size(grid) - - i′ = Nx - i + 1 - Hy = grid.Hy - - for j = 1 : Hy - @inbounds c[i, Ny + j, k] = sign * c[i′, Ny - j, k] # The Ny line is duplicated so we substitute starting Ny-1 - end - - return nothing -end - # Run the distributed grid simulation and save down reconstructed results function run_distributed_tripolar_grid(arch, filename) grid = TripolarGrid(arch; size = (40, 40, 1), z = (-1000, 0), halo = (5, 5, 5)) From c097b0c59837449eccad21ec3cbdc279faa6169b Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Wed, 1 Apr 2026 17:09:17 +1100 Subject: [PATCH 14/28] Fix self-qualified fold_topology access by inlining into keyword argument Co-Authored-By: Claude Opus 4.6 (1M context) --- .../distributed_tripolar_grid.jl | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl index 70c5ece4306..094d8a3b368 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl @@ -341,8 +341,6 @@ function DistributedComputations.reconstruct_global_grid(grid::MPITripolarGrid) first_pole_longitude = grid.conformal_mapping.first_pole_longitude southernmost_latitude = grid.conformal_mapping.southernmost_latitude - fold_topology = OrthogonalSphericalShellGrids.fold_topology(grid.conformal_mapping) - return TripolarGrid(child_arch, FT; halo, size, @@ -350,7 +348,7 @@ function DistributedComputations.reconstruct_global_grid(grid::MPITripolarGrid) first_pole_longitude, southernmost_latitude, z, - fold_topology) + fold_topology = fold_topology(grid.conformal_mapping)) end function Grids.with_halo(new_halo, old_grid::MPITripolarGrid) @@ -364,8 +362,6 @@ function Grids.with_halo(new_halo, old_grid::MPITripolarGrid) north_poles_latitude = old_grid.conformal_mapping.north_poles_latitude first_pole_longitude = old_grid.conformal_mapping.first_pole_longitude southernmost_latitude = old_grid.conformal_mapping.southernmost_latitude - fold_topology = OrthogonalSphericalShellGrids.fold_topology(old_grid.conformal_mapping) - return TripolarGrid(arch, eltype(old_grid); halo = new_halo, size = N, @@ -373,7 +369,7 @@ function Grids.with_halo(new_halo, old_grid::MPITripolarGrid) first_pole_longitude, southernmost_latitude, z, - fold_topology) + fold_topology = fold_topology(old_grid.conformal_mapping)) end ##### From 8ccbe750e5ae254e678c6f01f55dfbdffc201496 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Pasquier?= <4486578+briochemc@users.noreply.github.com> Date: Wed, 1 Apr 2026 19:09:06 +1100 Subject: [PATCH 15/28] Apply suggestions from code review Co-authored-by: Simone Silvestri --- src/OrthogonalSphericalShellGrids/distributed_zipper.jl | 5 ----- src/OrthogonalSphericalShellGrids/tripolar_grid.jl | 8 ++++---- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 8fc96376e03..02f84792560 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -456,11 +456,6 @@ _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{true, false}, Hx, Hy, N _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{false}, Hx, Hy, Nx, Ny) = view(c, ne_recv_x(loc_x(buff), Hx, Nx), recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= buff.recv -##### -##### _switch_north_halos! — no-op (fold logic in send buffers) -##### - -switch_north_halos!(c, north_bc, grid, loc) = nothing ##### ##### fill_halo_regions! for distributed tripolar grids diff --git a/src/OrthogonalSphericalShellGrids/tripolar_grid.jl b/src/OrthogonalSphericalShellGrids/tripolar_grid.jl index 10a4ff75d41..9593be7577c 100644 --- a/src/OrthogonalSphericalShellGrids/tripolar_grid.jl +++ b/src/OrthogonalSphericalShellGrids/tripolar_grid.jl @@ -11,18 +11,18 @@ A structure to represent a tripolar grid on an orthogonal spherical shell. The fold topology `FT` (e.g., `RightCenterFolded` or `RightFaceFolded`) is stored as a type parameter rather than a field, keeping the struct `isbits` for GPU kernels. """ -struct Tripolar{N, F, S, FT<:AbstractTopology} +struct Tripolar{N, F, S, TY<:AbstractTopology} north_poles_latitude :: N first_pole_longitude :: F southernmost_latitude :: S end # Getter: returns the fold topology Type -fold_topology(::Tripolar{<:Any, <:Any, <:Any, FT}) where FT = FT +fold_topology(::Tripolar{<:Any, <:Any, <:Any, TY}) where TY = TY # Constructor accepting fold topology as a Type argument -Tripolar(n, f, s, ::Type{FT}) where {FT<:AbstractTopology} = - Tripolar{typeof(n), typeof(f), typeof(s), FT}(n, f, s) +Tripolar(n, f, s, ::Type{TY}) where {TY<:AbstractTopology} = + Tripolar{typeof(n), typeof(f), typeof(s), TY}(n, f, s) # Backward-compatible constructor (defaults to UPivot) Tripolar(n, f, s) = Tripolar(n, f, s, RightCenterFolded) From 60869f67686ed8abc742d43b0bf26907a252aba2 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Wed, 1 Apr 2026 19:18:25 +1100 Subject: [PATCH 16/28] Fix fold topology type: FT -> TY --- src/OrthogonalSphericalShellGrids/tripolar_grid.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/tripolar_grid.jl b/src/OrthogonalSphericalShellGrids/tripolar_grid.jl index 9593be7577c..9afae712d18 100644 --- a/src/OrthogonalSphericalShellGrids/tripolar_grid.jl +++ b/src/OrthogonalSphericalShellGrids/tripolar_grid.jl @@ -27,11 +27,11 @@ Tripolar(n, f, s, ::Type{TY}) where {TY<:AbstractTopology} = # Backward-compatible constructor (defaults to UPivot) Tripolar(n, f, s) = Tripolar(n, f, s, RightCenterFolded) -Adapt.adapt_structure(to, t::Tripolar{<:Any, <:Any, <:Any, FT}) where FT = +Adapt.adapt_structure(to, t::Tripolar{<:Any, <:Any, <:Any, TY}) where TY = Tripolar(Adapt.adapt(to, t.north_poles_latitude), Adapt.adapt(to, t.first_pole_longitude), Adapt.adapt(to, t.southernmost_latitude), - FT) + TY) const TripolarGrid{FT, TX, TY, TZ, CZ, CC, FC, CF, FF, Arch} = OrthogonalSphericalShellGrid{FT, TX, TY, TZ, CZ, <:Tripolar, CC, FC, CF, FF, Arch} const TripolarGridOfSomeKind{FT, TX, TY, TZ} = Union{TripolarGrid{FT, TX, TY, TZ}, ImmersedBoundaryGrid{FT, TX, TY, TZ, <:TripolarGrid}} From 7695d8a7bf62420b22c2c6ceeed3fcabcb7821b2 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Wed, 1 Apr 2026 19:37:11 +1100 Subject: [PATCH 17/28] Remove unused methods --- .../distributed_zipper.jl | 52 ++----------------- 1 file changed, 3 insertions(+), 49 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 02f84792560..6748c19df8a 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -1,14 +1,12 @@ -using Oceananigans.BoundaryConditions: get_boundary_kernels, DistributedCommunication -using Oceananigans.DistributedComputations: cooperative_waitall!, recv_from_buffers!, distributed_fill_halo_event! -using Oceananigans.DistributedComputations: CommunicationBuffers, fill_corners!, loc_id, AsynchronousDistributed +using Oceananigans.BoundaryConditions: DistributedCommunication +using Oceananigans.DistributedComputations: CommunicationBuffers using Oceananigans.Grids: AbstractGrid, topology, RightCenterFolded, RightFaceFolded, LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded, LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected using Oceananigans.DistributedComputations: Distributed, on_architecture, ranks, x_communication_buffer -import Oceananigans.BoundaryConditions: fill_halo_regions! -import Oceananigans.DistributedComputations: synchronize_communication!, +import Oceananigans.DistributedComputations: y_communication_buffer, corner_communication_buffer, _fill_north_send_buffer!, _recv_from_north_buffer!, _fill_northwest_send_buffer!, _fill_northeast_send_buffer!, @@ -455,47 +453,3 @@ _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{true, false}, Hx, Hy, N _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{false}, Hx, Hy, Nx, Ny) = view(c, ne_recv_x(loc_x(buff), Hx, Nx), recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= buff.recv - - -##### -##### fill_halo_regions! for distributed tripolar grids -##### - -fill_halo_regions!(c::OffsetArray, ::Nothing, indices, loc, ::MPITripolarGridOfSomeKind, args...; kwargs...) = nothing - -function fill_halo_regions!(c::OffsetArray, bcs, indices, loc, grid::MPITripolarGridOfSomeKind, buffers::CommunicationBuffers, args...; kwargs...) - - arch = architecture(grid) - 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 - @inbounds distributed_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...; kwargs...) - - if length(arch.mpi_requests) > outstanding_requests - arch.mpi_tag[] += 1 - end - - return nothing -end - -function synchronize_communication!(field::Field{<:Any, <:Any, <:Any, <:Any, <:MPITripolarGridOfSomeKind}) - arch = architecture(field.grid) - - if arch isa AsynchronousDistributed - if !isempty(arch.mpi_requests) - cooperative_waitall!(arch.mpi_requests) - arch.mpi_tag[] = 0 - empty!(arch.mpi_requests) - end - - recv_from_buffers!(field.data, field.communication_buffers, field.grid) - end - - return nothing -end From 3731cef55429b5ceed3b31758bc2de7d38d0fbc6 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Wed, 1 Apr 2026 20:54:56 +1100 Subject: [PATCH 18/28] import fill_halo_regions! --- .../OrthogonalSphericalShellGrids.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl b/src/OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl index 10f52fa85a3..b7ff7eda1f8 100644 --- a/src/OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl +++ b/src/OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl @@ -11,6 +11,7 @@ using Oceananigans.BoundaryConditions: BoundaryCondition using Oceananigans.Grids: AbstractTopology using Oceananigans.Grids: halo_size, generate_coordinate, topology using Oceananigans.Grids: total_length, add_halos, fill_metric_halo_regions! +using Oceananigans.BoundaryConditions: fill_halo_regions! using Distances: haversine using Adapt: Adapt, adapt From e62d2b185822cb636010d5be150157b33a7f187c Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Wed, 1 Apr 2026 22:18:05 +1100 Subject: [PATCH 19/28] Remove dead OneDZipperBuffer code OneDZipperBuffer was never constructed because slab partitions (Rx=1) use the serial fold BC directly, not a DistributedZipper. The OneDFoldTopology + DistributedZipper combination was impossible. Re-add loc_id import needed by distributed_zipper_north_tags.jl. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../distributed_zipper.jl | 49 +------------------ 1 file changed, 1 insertion(+), 48 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 6748c19df8a..6e13930f4fa 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -1,5 +1,5 @@ using Oceananigans.BoundaryConditions: DistributedCommunication -using Oceananigans.DistributedComputations: CommunicationBuffers +using Oceananigans.DistributedComputations: CommunicationBuffers, loc_id using Oceananigans.Grids: AbstractGrid, topology, RightCenterFolded, RightFaceFolded, LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded, @@ -31,9 +31,6 @@ const FPivotTopology = Union{RightFaceFolded, LeftConnectedRightFaceFolded, LeftConnectedRightFaceConnected} -# 1D fold (1xN, y-partitioned only) vs 2D fold (MxN, x+y partitioned) -const OneDFoldTopology = Union{LeftConnectedRightCenterFolded, - LeftConnectedRightFaceFolded} const TwoDFoldTopology = Union{LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected} @@ -74,12 +71,6 @@ end ##### Zipper communication buffers with fold-aware packing ##### -struct OneDZipperBuffer{Loc, FoT, B, S} - send :: B - recv :: B - sign :: S -end - struct TwoDZipperBuffer{FL, WFL, Loc, FoT, B, S} send :: B recv :: B @@ -93,7 +84,6 @@ struct ZipperCornerBuffer{FL, WFL, Loc, FoT, B, S} end # Value-argument constructors: all types inferred -OneDZipperBuffer(loc::Loc, fot::FoT, send::B, recv::B, sign::S) where {Loc, FoT, B, S} = OneDZipperBuffer{Loc, FoT, B, S}(send, recv, sign) TwoDZipperBuffer(loc, fot, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {FL, WFL} = TwoDZipperBuffer{FL, WFL, typeof(loc), typeof(fot), typeof(send), typeof(sign)}(send, recv, sign) ZipperCornerBuffer(loc, fot, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {FL, WFL} = @@ -101,7 +91,6 @@ ZipperCornerBuffer(loc, fot, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {FL, ZipperCornerBuffer(::Type{Loc}, ::Type{FoT}, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {Loc, FoT, FL, WFL} = ZipperCornerBuffer{FL, WFL, Loc, FoT, typeof(send), typeof(sign)}(send, recv, sign) -Adapt.adapt_structure(to, buff::OneDZipperBuffer) = nothing Adapt.adapt_structure(to, buff::TwoDZipperBuffer) = nothing Adapt.adapt_structure(to, buff::ZipperCornerBuffer) = nothing @@ -208,17 +197,6 @@ end # Fallback: non-zipper north BC uses standard buffer y_tripolar_buffer(arch, grid, data, Hy, bc, loc) = y_communication_buffer(arch, grid, data, Hy, bc) -# 1D fold (1xN) → OneDZipperBuffer (full-width) -function y_tripolar_buffer(arch, grid::AbstractGrid{<:Any, <:Any, Topo}, - data, Hy, bc::DistributedZipper, loc::Loc) where {Topo <: OneDFoldTopology, Loc} - Tx, _, Tz = size(parent(data)) - FT = eltype(data) - sgn = bc.condition.sign - send = on_architecture(arch, zeros(FT, Tx, Hy, Tz)) - recv = on_architecture(arch, zeros(FT, Tx, Hy, Tz)) - return OneDZipperBuffer(loc, Topo(), send, recv, sgn) -end - # 2D fold (MxN) → TwoDZipperBuffer (interior-width, Hy′ = Hy or Hy+1 rows for fold line) function y_tripolar_buffer(arch, grid::AbstractGrid{<:Any, <:Any, Topo}, data, Hy, bc::DistributedZipper, loc::Loc) where {Topo <: TwoDFoldTopology, Loc} @@ -291,10 +269,6 @@ const FF = Tuple{<:Face, <:Face, <:Any} @inline loc_y(::ZipperCornerBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[2]) @inline fold_topo(::ZipperCornerBuffer{<:Any, <:Any, <:Any, FoT}) where FoT = FoT() -@inline loc_x(::OneDZipperBuffer{Loc}) where Loc = instantiate(Loc.parameters[1]) -@inline loc_y(::OneDZipperBuffer{Loc}) where Loc = instantiate(Loc.parameters[2]) -@inline fold_topo(::OneDZipperBuffer{<:Any, FoT}) where FoT = FoT() - # Send y-ranges (source rows, reversed for fold) @inline send_fold_y(::UPivotTopology, ::Center, Hy, Ny) = Ny + Hy @inline send_fold_y(::FPivotTopology, ::Face, Hy, Ny) = Ny + 1 + Hy @@ -369,27 +343,6 @@ function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{false}, Hx, Hy, Nx, view(c, xr, recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= buff.recv end -##### -##### OneDZipperBuffer: north send (Center-x full reverse, Face-x partial reverse) -##### - -function _fill_north_send_buffer!(c, b::OneDZipperBuffer{<:Union{CC,CF}}, Hx, Hy, Nx, Ny) - b.send .= b.sign .* view(c, size(c,1):-1:1, send_halo_y(fold_topo(b), loc_y(b), Hy, Ny), :) -end - -function _fill_north_send_buffer!(c, b::OneDZipperBuffer{<:Union{FC,FF}}, Hx, Hy, Nx, Ny) - Tx = size(c, 1) - hy = send_halo_y(fold_topo(b), loc_y(b), Hy, Ny) - view(b.send, 2:Tx, :, :) .= b.sign .* view(c, Tx:-1:2, hy, :) - view(b.send, 1:1, :, :) .= b.sign .* view(c, 1:1, hy, :) -end - -##### -##### OneDZipperBuffer: north recv -##### - -_recv_from_north_buffer!(c, buff::OneDZipperBuffer, Hx, Hy, Nx, Ny) = view(c, :, recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= buff.recv - ##### ##### Corner send: NW and NE (FL=true includes fold line, FL=false halo only) ##### From f22e1920cef19ffc3a3af28c2356e465aaa236e1 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Tue, 7 Apr 2026 09:59:08 +1000 Subject: [PATCH 20/28] Fix missing comma in using statement in SplitExplicitFreeSurfaces --- .../SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl index 941703b4899..da8a121532d 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl @@ -10,7 +10,7 @@ using Oceananigans.ImmersedBoundaries: column_depthTᶠᶜᵃ, column_depthTᶜ using Oceananigans.Operators: ∂xᵣTᶠᶜᶠ, ∂xᵣᶠᶜᶠ, ∂yᵣTᶜᶠᶠ, ∂yᵣᶜᶠᶠ, δxTᶜᵃᵃ, δxᶜᵃᵃ, δyTᵃᶜᵃ, δyᵃᶜᵃ using Oceananigans.BoundaryConditions: FieldBoundaryConditions, fill_halo_regions! using Oceananigans.Fields: Field -using Oceananigans.Grids: Center, Face, get_active_cells_map, topology +using Oceananigans.Grids: Center, Face, get_active_cells_map, topology, LeftConnected, RightConnected, FullyConnected, RightCenterFolded, RightFaceFolded, LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded, From 11e466d1f93d56707fc871e0fe2bef32e0e9f889 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Tue, 7 Apr 2026 22:44:26 +1000 Subject: [PATCH 21/28] Add RightFaceFolded (FPivot) coverage to distributed tripolar MPI tests Loop over both fold topologies (RightCenterFolded, RightFaceFolded) in grid reconstruction, field reconstruction, boundary condition, and simulation tests. The run_distributed_tripolar_grid helper now requires an explicit fold_topology keyword argument. Co-Authored-By: Claude Opus 4.6 (1M context) --- test/distributed_tests_utils.jl | 4 +- test/test_mpi_tripolar.jl | 134 +++++++++++++++++--------------- 2 files changed, 74 insertions(+), 64 deletions(-) diff --git a/test/distributed_tests_utils.jl b/test/distributed_tests_utils.jl index 61dc75b1bed..e56d256b655 100644 --- a/test/distributed_tests_utils.jl +++ b/test/distributed_tests_utils.jl @@ -29,8 +29,8 @@ function analytical_immersed_tripolar_grid(underlying_grid::TripolarGrid; radius end # Run the distributed grid simulation and save down reconstructed results -function run_distributed_tripolar_grid(arch, filename) - grid = TripolarGrid(arch; size = (40, 40, 1), z = (-1000, 0), halo = (5, 5, 5)) +function run_distributed_tripolar_grid(arch, filename; fold_topology) + grid = TripolarGrid(arch; size = (40, 40, 1), z = (-1000, 0), halo = (5, 5, 5), fold_topology) grid = analytical_immersed_tripolar_grid(grid) model = run_distributed_simulation(grid) diff --git a/test/test_mpi_tripolar.jl b/test/test_mpi_tripolar.jl index fa00af6d07a..daa66fbacaa 100644 --- a/test/test_mpi_tripolar.jl +++ b/test/test_mpi_tripolar.jl @@ -3,7 +3,9 @@ include("distributed_tests_utils.jl") using MPI -tripolar_reconstructed_grid = """ +fold_topologies = (RightCenterFolded, RightFaceFolded) + +tripolar_reconstructed_grid_script(fold_topology) = """ using MPI MPI.Init() using Test @@ -14,8 +16,8 @@ tripolar_reconstructed_grid = """ Distributed(CPU(), partition=Partition(2, 2))] for arch in archs - local_grid = TripolarGrid(arch; size = (12, 20, 1), z = (-1000, 0), halo = (2, 2, 2)) - global_grid = TripolarGrid(size = (12, 20, 1), z = (-1000, 0), halo = (2, 2, 2)) + local_grid = TripolarGrid(arch; size = (12, 20, 1), z = (-1000, 0), halo = (2, 2, 2), fold_topology = $fold_topology) + global_grid = TripolarGrid(size = (12, 20, 1), z = (-1000, 0), halo = (2, 2, 2), fold_topology = $fold_topology) reconstruct_grid = reconstruct_global_grid(local_grid) @@ -38,7 +40,7 @@ tripolar_reconstructed_grid = """ end """ -tripolar_reconstructed_field = """ +tripolar_reconstructed_field_script(fold_topology) = """ using MPI MPI.Init() using Test @@ -49,11 +51,11 @@ tripolar_reconstructed_field = """ Distributed(CPU(), partition=Partition(2, 2))] u = [i + 10 * j for i in 1:40, j in 1:40] - v = [i + 10 * j for i in 1:40, j in 1:40] + v = [i + 10 * j for i in 1:40, j in 1:$(fold_topology == RightCenterFolded ? 40 : 41)] c = [i + 10 * j for i in 1:40, j in 1:40] for arch in archs - local_grid = TripolarGrid(arch; size = (40, 40, 1), z = (-1000, 0), halo = (5, 5, 5)) + local_grid = TripolarGrid(arch; size = (40, 40, 1), z = (-1000, 0), halo = (5, 5, 5), fold_topology = $fold_topology) up = XFaceField(local_grid) vp = YFaceField(local_grid) @@ -63,7 +65,7 @@ tripolar_reconstructed_field = """ set!(vp, v) set!(cp, c) - global_grid = TripolarGrid(size = (40, 40, 1), z = (-1000, 0), halo = (5, 5, 5)) + global_grid = TripolarGrid(size = (40, 40, 1), z = (-1000, 0), halo = (5, 5, 5), fold_topology = $fold_topology) us = XFaceField(global_grid) vs = YFaceField(global_grid) @@ -79,35 +81,36 @@ tripolar_reconstructed_field = """ end """ -@testset "Test distributed TripolarGrid..." begin - write("distributed_tripolar_grid.jl", tripolar_reconstructed_grid) - run(`$(mpiexec()) -n 4 $(Base.julia_cmd()) -O0 distributed_tripolar_grid.jl`) - rm("distributed_tripolar_grid.jl") +@testset "Test distributed TripolarGrid $fold_topology..." for fold_topology in fold_topologies + write("distributed_tripolar_grid_$fold_topology.jl", tripolar_reconstructed_grid_script(fold_topology)) + run(`$(mpiexec()) -n 4 $(Base.julia_cmd()) -O0 $("distributed_tripolar_grid_$fold_topology.jl")`) + rm("distributed_tripolar_grid_$fold_topology.jl") - write("distributed_tripolar_field.jl", tripolar_reconstructed_field) - run(`$(mpiexec()) -n 4 $(Base.julia_cmd()) -O0 distributed_tripolar_field.jl`) - rm("distributed_tripolar_field.jl") + write("distributed_tripolar_field_$fold_topology.jl", tripolar_reconstructed_field_script(fold_topology)) + run(`$(mpiexec()) -n 4 $(Base.julia_cmd()) -O0 $("distributed_tripolar_field_$fold_topology.jl")`) + rm("distributed_tripolar_field_$fold_topology.jl") end -tripolar_boundary_conditions = """ +tripolar_boundary_conditions_script(fold_topology) = """ using MPI MPI.Init() include("distributed_tests_utils.jl") arch = Distributed(CPU(), partition = Partition(2, 2)) - grid = TripolarGrid(arch; size = (20, 20, 1), z = (-1000, 0)) + grid = TripolarGrid(arch; size = (20, 20, 1), z = (-1000, 0), fold_topology = $fold_topology) # Build initial condition - serial_grid = TripolarGrid(size = (20, 20, 1), z = (-1000, 0)) + serial_grid = TripolarGrid(size = (20, 20, 1), z = (-1000, 0), fold_topology = $fold_topology) vs = YFaceField(serial_grid) cs = CenterField(serial_grid) - I1 = [i + j * 100 for i in 1:20, j in 1:20] + I_v = [i + j * 100 for i in 1:20, j in 1:$(fold_topology == RightCenterFolded ? 20 : 21)] + I_c = [i + j * 100 for i in 1:20, j in 1:20] - set!(vs, I1) - set!(cs, I1) + set!(vs, I_v) + set!(cs, I_c) fill_halo_regions!((vs, cs)) @@ -118,7 +121,7 @@ tripolar_boundary_conditions = """ set!(c, cs) fill_halo_regions!((v, c)) - filename = "distributed_tripolar_boundary_conditions_" * string(arch.local_rank) * ".jld2" + filename = "distributed_$(fold_topology)_boundary_conditions_" * string(arch.local_rank) * ".jld2" jldopen(filename, "w") do file file["v"] = v.data @@ -129,68 +132,74 @@ tripolar_boundary_conditions = """ MPI.Finalize() """ -@testset "Test distributed TripolarGrid boundary conditions..." begin +@testset "Test distributed TripolarGrid boundary conditions $fold_topology..." for fold_topology in fold_topologies # Run the serial computation - grid = TripolarGrid(size = (20, 20, 1), z = (-1000, 0)) - - I1 = [i + j * 100 for i in 1:20, j in 1:20] + grid = TripolarGrid(size = (20, 20, 1), z = (-1000, 0); fold_topology) v = YFaceField(grid) c = CenterField(grid) - set!(v, I1) - set!(c, I1) + Nyv = fold_topology == RightCenterFolded ? 20 : 21 + I_v = [i + j * 100 for i in 1:20, j in 1:Nyv] + I_c = [i + j * 100 for i in 1:20, j in 1:20] + + set!(v, I_v) + set!(c, I_c) fill_halo_regions!((v, c)) - write("distributed_boundary_tests.jl", tripolar_boundary_conditions) - run(`$(mpiexec()) -n 4 $(Base.julia_cmd()) -O0 distributed_boundary_tests.jl`) - rm("distributed_boundary_tests.jl") + write("distributed_$(fold_topology)_boundary_tests.jl", tripolar_boundary_conditions_script(fold_topology)) + run(`$(mpiexec()) -n 4 $(Base.julia_cmd()) -O0 $("distributed_$(fold_topology)_boundary_tests.jl")`) + rm("distributed_$(fold_topology)_boundary_tests.jl") # Retrieve Parallel quantities from rank 1 (the north-west rank) - vp1 = jldopen("distributed_tripolar_boundary_conditions_1.jld2")["v"]; - cp1 = jldopen("distributed_tripolar_boundary_conditions_1.jld2")["c"]; + vp1 = jldopen("distributed_$(fold_topology)_boundary_conditions_1.jld2")["v"]; + cp1 = jldopen("distributed_$(fold_topology)_boundary_conditions_1.jld2")["c"]; # Retrieve Parallel quantities from rank 3 (the north-east rank) - vp3 = jldopen("distributed_tripolar_boundary_conditions_3.jld2")["v"]; - cp3 = jldopen("distributed_tripolar_boundary_conditions_3.jld2")["c"]; + vp3 = jldopen("distributed_$(fold_topology)_boundary_conditions_3.jld2")["v"]; + cp3 = jldopen("distributed_$(fold_topology)_boundary_conditions_3.jld2")["c"]; @test v.data[-3:14, end-3:end-1, 1] ≈ vp1.parent[:, end-3:end-1, 5] @test c.data[-3:14, end-3:end-1, 1] ≈ cp1.parent[:, end-3:end-1, 5] @test v.data[7:end, 7:end-1, 1] ≈ vp3.parent[:, 1:end-1, 5] @test c.data[7:end, 7:end-1, 1] ≈ cp3.parent[:, 1:end-1, 5] + + for rank in 0:3 + rm("distributed_$(fold_topology)_boundary_conditions_$(rank).jld2", force=true) + end end -run_slab_distributed_grid = """ +run_slab_distributed_grid(fold_topology) = """ using MPI MPI.Init() include("distributed_tests_utils.jl") arch = Distributed(CPU(), partition = Partition(1, 4)) - run_distributed_tripolar_grid(arch, "distributed_yslab_tripolar.jld2") + run_distributed_tripolar_grid(arch, "distributed_$(fold_topology)_yslab_tripolar.jld2"; fold_topology = $fold_topology) """ -run_pencil_distributed_grid = """ +run_pencil_distributed_grid(fold_topology) = """ using MPI MPI.Init() include("distributed_tests_utils.jl") arch = Distributed(CPU(), partition = Partition(2, 2)) - run_distributed_tripolar_grid(arch, "distributed_pencil_tripolar.jld2") + run_distributed_tripolar_grid(arch, "distributed_$(fold_topology)_pencil_tripolar.jld2"; fold_topology = $fold_topology) """ -run_large_pencil_distributed_grid = """ +run_large_pencil_distributed_grid(fold_topology) = """ using MPI MPI.Init() include("distributed_tests_utils.jl") arch = Distributed(CPU(), partition = Partition(4, 2)) - run_distributed_tripolar_grid(arch, "distributed_large_pencil_tripolar.jld2") + run_distributed_tripolar_grid(arch, "distributed_$(fold_topology)_large_pencil_tripolar.jld2"; fold_topology = $fold_topology) """ -@testset "Test distributed TripolarGrid simulations..." begin +@testset "Test distributed TripolarGrid simulations $fold_topology..." for fold_topology in fold_topologies # Run the serial computation - grid = TripolarGrid(size = (40, 40, 1), z = (-1000, 0), halo = (5, 5, 5)) + grid = TripolarGrid(size = (40, 40, 1), z = (-1000, 0), halo = (5, 5, 5); fold_topology) grid = analytical_immersed_tripolar_grid(grid) model = run_distributed_simulation(grid) @@ -202,18 +211,19 @@ run_large_pencil_distributed_grid = """ us = interior(us, :, :, 1) vs = interior(vs, :, :, 1) cs = interior(cs, :, :, 1) + # Run the distributed grid simulation with a slab configuration - write("distributed_slab_tests.jl", run_slab_distributed_grid) + write("distributed_slab_tests.jl", run_slab_distributed_grid(fold_topology)) run(`$(mpiexec()) -n 4 $(Base.julia_cmd()) -O0 distributed_slab_tests.jl`) rm("distributed_slab_tests.jl") # Retrieve Parallel quantities - up = jldopen("distributed_yslab_tripolar.jld2")["u"] - vp = jldopen("distributed_yslab_tripolar.jld2")["v"] - cp = jldopen("distributed_yslab_tripolar.jld2")["c"] - ηp = jldopen("distributed_yslab_tripolar.jld2")["η"] + up = jldopen("distributed_$(fold_topology)_yslab_tripolar.jld2")["u"] + vp = jldopen("distributed_$(fold_topology)_yslab_tripolar.jld2")["v"] + cp = jldopen("distributed_$(fold_topology)_yslab_tripolar.jld2")["c"] + ηp = jldopen("distributed_$(fold_topology)_yslab_tripolar.jld2")["η"] - rm("distributed_yslab_tripolar.jld2") + rm("distributed_$(fold_topology)_yslab_tripolar.jld2") # Test slab partitioning @test all(us .≈ up) @@ -222,17 +232,17 @@ run_large_pencil_distributed_grid = """ @test all(ηs .≈ ηp) # Run the distributed grid simulation with a pencil configuration - write("distributed_tests.jl", run_pencil_distributed_grid) - run(`$(mpiexec()) -n 4 $(Base.julia_cmd()) -O0 distributed_tests.jl`) - rm("distributed_tests.jl") + write("distributed_pencil_tests.jl", run_pencil_distributed_grid(fold_topology)) + run(`$(mpiexec()) -n 4 $(Base.julia_cmd()) -O0 distributed_pencil_tests.jl`) + rm("distributed_pencil_tests.jl") # Retrieve Parallel quantities - up = jldopen("distributed_pencil_tripolar.jld2")["u"] - vp = jldopen("distributed_pencil_tripolar.jld2")["v"] - ηp = jldopen("distributed_pencil_tripolar.jld2")["η"] - cp = jldopen("distributed_pencil_tripolar.jld2")["c"] + up = jldopen("distributed_$(fold_topology)_pencil_tripolar.jld2")["u"] + vp = jldopen("distributed_$(fold_topology)_pencil_tripolar.jld2")["v"] + ηp = jldopen("distributed_$(fold_topology)_pencil_tripolar.jld2")["η"] + cp = jldopen("distributed_$(fold_topology)_pencil_tripolar.jld2")["c"] - rm("distributed_pencil_tripolar.jld2") + rm("distributed_$(fold_topology)_pencil_tripolar.jld2") @test all(us .≈ up) @test all(vs .≈ vp) @@ -242,17 +252,17 @@ run_large_pencil_distributed_grid = """ # We try now with more ranks in the x-direction. This is not a trivial # test as we are now splitting, not only where the singularities are, but # also in the middle of the north fold. This is a more challenging test - write("distributed_large_pencil_tests.jl", run_large_pencil_distributed_grid) + write("distributed_large_pencil_tests.jl", run_large_pencil_distributed_grid(fold_topology)) run(`$(mpiexec()) -n 8 $(Base.julia_cmd()) -O0 distributed_large_pencil_tests.jl`) rm("distributed_large_pencil_tests.jl") # Retrieve Parallel quantities - up = jldopen("distributed_large_pencil_tripolar.jld2")["u"] - vp = jldopen("distributed_large_pencil_tripolar.jld2")["v"] - ηp = jldopen("distributed_large_pencil_tripolar.jld2")["η"] - cp = jldopen("distributed_large_pencil_tripolar.jld2")["c"] + up = jldopen("distributed_$(fold_topology)_large_pencil_tripolar.jld2")["u"] + vp = jldopen("distributed_$(fold_topology)_large_pencil_tripolar.jld2")["v"] + ηp = jldopen("distributed_$(fold_topology)_large_pencil_tripolar.jld2")["η"] + cp = jldopen("distributed_$(fold_topology)_large_pencil_tripolar.jld2")["c"] - rm("distributed_large_pencil_tripolar.jld2") + rm("distributed_$(fold_topology)_large_pencil_tripolar.jld2") @test all(us .≈ up) @test all(vs .≈ vp) From a5bd8e7751056408e528aae99295e1489628a282 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Wed, 8 Apr 2026 11:00:51 +1000 Subject: [PATCH 22/28] Fix analytical_immersed_tripolar_grid southern masking for FPivot grids MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Use φ < φm + radius instead of φ < φm to immerse the southernmost row on RightFaceFolded grids, where southernmost_latitude sits at the cell face so j=1 centers are slightly north of φm. Without this, the tracer field diverges between serial and distributed runs on FPivot grids. Co-Authored-By: Claude Opus 4.6 (1M context) --- test/distributed_tests_utils.jl | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/test/distributed_tests_utils.jl b/test/distributed_tests_utils.jl index e56d256b655..6cb85d8454f 100644 --- a/test/distributed_tests_utils.jl +++ b/test/distributed_tests_utils.jl @@ -18,10 +18,16 @@ function analytical_immersed_tripolar_grid(underlying_grid::TripolarGrid; radius Lz = underlying_grid.Lz - # We need a bottom height field that ``masks'' the singularities + # We need a bottom height field that ``masks'' the singularities. + # Note: For the test to pass we also need to mask the southernmost row near φm. + # We use φ < φm + radius (not just φ < φm) to also immerse the southernmost row on + # RightFaceFolded grids, for which φm sits at the cell face and j=1 centers + # are slightly north of φm. + # TODO: Investigate why masking near φm is necessary. Shouldn't the no-flux south + # BC already make the south halo behave as if everything south was immersed? bottom_height(λ, φ) = ((abs(λ - λp) < radius) & (abs(φp - φ) < radius)) | ((abs(λ - λp - 180) < radius) & (abs(φp - φ) < radius)) | - ((abs(λ - λp - 360) < radius) & (abs(φp - φ) < radius)) | (φ < φm) ? 0 : - Lz + ((abs(λ - λp - 360) < radius) & (abs(φp - φ) < radius)) | (φ < φm + radius) ? 0 : - Lz grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height)) From 3ebc99410451069e99e2988711028d04bb2f0098 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Pasquier?= <4486578+briochemc@users.noreply.github.com> Date: Wed, 8 Apr 2026 17:16:20 +1000 Subject: [PATCH 23/28] Apply suggestions from code review Co-authored-by: Simone Silvestri --- .../distributed_tripolar_grid.jl | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl index 094d8a3b368..ac16b0a35a0 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl @@ -120,8 +120,8 @@ function TripolarGrid(arch::Distributed, FT::DataType=Float64; Rx, Ry = workers[1], workers[2] rx, ry = xrank + 1, yrank + 1 LY = insert_connected_topology(global_fold_topology, Ry, ry, Rx, rx) - ny = nylocal[yrank+1] - nx = nxlocal[xrank+1] + ny = nylocal[ry] + nx = nxlocal[rx] z = on_architecture(arch, global_grid.z) radius = global_grid.radius @@ -254,7 +254,8 @@ function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions, south = regularize_boundary_condition(bcs.south, grid, loc, 2, LeftBoundary, prognostic_names) north = if yrank == processor_size[2] - 1 && processor_size[1] == 1 - north_fold_boundary_condition(fold_topology(grid.conformal_mapping))(sign) + TY = fold_topology(grid.conformal_mapping) + north_fold_boundary_condition(TY)(sign) elseif yrank == processor_size[2] - 1 && processor_size[1] != 1 from = arch.local_rank From ca61b8e47a06388194ee88f08afe89f06f5f6d50 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Wed, 8 Apr 2026 17:21:39 +1000 Subject: [PATCH 24/28] declutter line for north BC --- .../distributed_tripolar_grid.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl index ac16b0a35a0..16965328233 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl @@ -294,7 +294,9 @@ function Field(loc::Tuple{<:LX, <:LY, <:LZ}, grid::MPITripolarGridOfSomeKind, da if yrank == processor_size[2] - 1 && processor_size[1] == 1 north_bc = if !(old_bcs.north isa ZBC) - north_fold_boundary_condition(fold_topology(grid.conformal_mapping))(sign(LX, LY)) + TY = fold_topology(grid.conformal_mapping) + north_fold_boundary_condition(TY)(sign(LX, LY)) + else old_bcs.north end From 0d98e06684b66d4ae70a1d972b575ca3c83616fc Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Wed, 8 Apr 2026 17:41:36 +1000 Subject: [PATCH 25/28] Use TY throughout for fold_topology type --- .../distributed_zipper.jl | 48 +++++++++---------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 6e13930f4fa..8a72b43fbab 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -71,25 +71,25 @@ end ##### Zipper communication buffers with fold-aware packing ##### -struct TwoDZipperBuffer{FL, WFL, Loc, FoT, B, S} +struct TwoDZipperBuffer{FL, WFL, Loc, TY, B, S} send :: B recv :: B sign :: S end -struct ZipperCornerBuffer{FL, WFL, Loc, FoT, B, S} +struct ZipperCornerBuffer{FL, WFL, Loc, TY, B, S} send :: B recv :: B sign :: S end # Value-argument constructors: all types inferred -TwoDZipperBuffer(loc, fot, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {FL, WFL} = - TwoDZipperBuffer{FL, WFL, typeof(loc), typeof(fot), typeof(send), typeof(sign)}(send, recv, sign) -ZipperCornerBuffer(loc, fot, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {FL, WFL} = - ZipperCornerBuffer{FL, WFL, typeof(loc), typeof(fot), typeof(send), typeof(sign)}(send, recv, sign) -ZipperCornerBuffer(::Type{Loc}, ::Type{FoT}, send, recv, sign, ::Val{FL}, ::Val{WFL}) where {Loc, FoT, FL, WFL} = - ZipperCornerBuffer{FL, WFL, Loc, FoT, typeof(send), typeof(sign)}(send, recv, sign) +TwoDZipperBuffer(loc::Loc, fold_topology::TY, send::B, recv::B, sign::S, ::Val{FL}, ::Val{WFL}) where {FL, WFL, Loc, TY, B, S} = + TwoDZipperBuffer{FL, WFL, Loc, TY, B, S}(send, recv, sign) +ZipperCornerBuffer(loc::Loc, fold_topology::TY, send::B, recv::B, sign::S, ::Val{FL}, ::Val{WFL}) where {FL, WFL, Loc, TY, B, S} = + ZipperCornerBuffer{FL, WFL, Loc, TY, B, S}(send, recv, sign) +ZipperCornerBuffer(::Type{Loc}, ::Type{TY}, send::B, recv::B, sign::S, ::Val{FL}, ::Val{WFL}) where {Loc, TY, FL, WFL, B, S} = + ZipperCornerBuffer{FL, WFL, Loc, TY, B, S}(send, recv, sign) Adapt.adapt_structure(to, buff::TwoDZipperBuffer) = nothing Adapt.adapt_structure(to, buff::ZipperCornerBuffer) = nothing @@ -106,8 +106,8 @@ end Adapt.adapt_structure(to, buff::TripolarXBuffer) = nothing -TripolarXBuffer(send, recv, ::Val{FL}, ::Val{WFL}) where {FL, WFL} = - TripolarXBuffer{typeof(send), FL, WFL}(send, recv) +TripolarXBuffer(send::B, recv::B, ::Val{FL}, ::Val{WFL}) where {B, FL, WFL} = + TripolarXBuffer{B, FL, WFL}(send, recv) # X-buffers WFL: complement of the ADJACENT corner. west_writes_fold_line(arch) = !northwest_writes_fold_line(arch) @@ -141,10 +141,10 @@ _recv_from_east_buffer!(c, buff::TripolarXBuffer{<:Any, true, false}, Hx, Hy, Nx # Fallback for non-zipper north (south ranks): standard x-communication. # Separate west/east constructors since they complement different corners. function west_tripolar_buffer(arch, grid, data, Hx, bc, loc, - north::TwoDZipperBuffer{<:Any, <:Any, Loc, FoT}) where {Loc, FoT} + north::TwoDZipperBuffer{<:Any, <:Any, Loc, TY}) where {Loc, TY} loc_y = Loc.parameters[2]() - fl = has_fold_line(FoT, loc_y) - Ny_buf = length(loc_y, FoT(), size(grid, 2)) + fl = has_fold_line(TY, loc_y) + Ny_buf = length(loc_y, TY(), size(grid, 2)) wfl = fl && west_writes_fold_line(arch) _, _, Tz = size(parent(data)) FT = eltype(data) @@ -154,10 +154,10 @@ function west_tripolar_buffer(arch, grid, data, Hx, bc, loc, end function east_tripolar_buffer(arch, grid, data, Hx, bc, loc, - north::TwoDZipperBuffer{<:Any, <:Any, Loc, FoT}) where {Loc, FoT} + north::TwoDZipperBuffer{<:Any, <:Any, Loc, TY}) where {Loc, TY} loc_y = Loc.parameters[2]() - fl = has_fold_line(FoT, loc_y) - Ny_buf = length(loc_y, FoT(), size(grid, 2)) + fl = has_fold_line(TY, loc_y) + Ny_buf = length(loc_y, TY(), size(grid, 2)) wfl = fl && east_writes_fold_line(arch) _, _, Tz = size(parent(data)) FT = eltype(data) @@ -222,29 +222,29 @@ northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge) = corner_commu # based on whether the corner's global x-position is past the pivot. function northwest_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, - yedge::TwoDZipperBuffer{<:Any, <:Any, Loc, FoT}) where {Loc, FoT} + yedge::TwoDZipperBuffer{<:Any, <:Any, Loc, TY}) where {Loc, TY} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign loc_x_inst = instantiate(Loc.parameters[1]) - fl = has_fold_line(FoT, Loc.parameters[2]()) + fl = has_fold_line(TY, Loc.parameters[2]()) Hy′ = fl ? Hy + 1 : Hy wfl = fl && northwest_writes_fold_line(arch) Hx′ = nw_corner_nx(loc_x_inst, Hx) send = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) - return ZipperCornerBuffer(Loc, FoT, send, recv, sgn, Val(fl), Val(wfl)) + return ZipperCornerBuffer(Loc, TY, send, recv, sgn, Val(fl), Val(wfl)) end function northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, - yedge::TwoDZipperBuffer{<:Any, <:Any, Loc, FoT}) where {Loc, FoT} + yedge::TwoDZipperBuffer{<:Any, <:Any, Loc, TY}) where {Loc, TY} Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign loc_x_inst = instantiate(Loc.parameters[1]) - fl = has_fold_line(FoT, Loc.parameters[2]()) + fl = has_fold_line(TY, Loc.parameters[2]()) Hy′ = fl ? Hy + 1 : Hy wfl = fl && northeast_writes_fold_line(arch) Hx′ = ne_corner_nx(loc_x_inst, Hx) send = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) - return ZipperCornerBuffer(Loc, FoT, send, recv, sgn, Val(fl), Val(wfl)) + return ZipperCornerBuffer(Loc, TY, send, recv, sgn, Val(fl), Val(wfl)) end ##### @@ -263,11 +263,11 @@ const FF = Tuple{<:Face, <:Face, <:Any} # Type-parameter accessors @inline loc_x(::TwoDZipperBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[1]) @inline loc_y(::TwoDZipperBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[2]) -@inline fold_topo(::TwoDZipperBuffer{<:Any, <:Any, <:Any, FoT}) where FoT = FoT() +@inline fold_topo(::TwoDZipperBuffer{<:Any, <:Any, <:Any, TY}) where TY = TY() @inline loc_x(::ZipperCornerBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[1]) @inline loc_y(::ZipperCornerBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[2]) -@inline fold_topo(::ZipperCornerBuffer{<:Any, <:Any, <:Any, FoT}) where FoT = FoT() +@inline fold_topo(::ZipperCornerBuffer{<:Any, <:Any, <:Any, TY}) where TY = TY() # Send y-ranges (source rows, reversed for fold) @inline send_fold_y(::UPivotTopology, ::Center, Hy, Ny) = Ny + Hy From d4d24326222f73a04296f8c9304129b772ee0a07 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Wed, 8 Apr 2026 17:44:48 +1000 Subject: [PATCH 26/28] Remove misplaced comment --- src/OrthogonalSphericalShellGrids/distributed_zipper.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 8a72b43fbab..10b8ef09e55 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -138,7 +138,6 @@ _recv_from_east_buffer!(c, buff::TripolarXBuffer{<:Any, true, false}, Hx, Hy, Nx view(c, 1+Nx+Hx:Nx+2Hx, 1+Hy:size(buff.recv,2)-1+Hy, :) .= view(buff.recv, :, 1:size(buff.recv,2)-1, :) # Fold-aware x-buffer constructors. -# Fallback for non-zipper north (south ranks): standard x-communication. # Separate west/east constructors since they complement different corners. function west_tripolar_buffer(arch, grid, data, Hx, bc, loc, north::TwoDZipperBuffer{<:Any, <:Any, Loc, TY}) where {Loc, TY} From 70a4d2b79f9c1f0abf0b2b9c76ff546e01a76d20 Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Thu, 9 Apr 2026 00:33:20 +1000 Subject: [PATCH 27/28] declutter + clearer helper function names + simplify dispatch --- .../distributed_zipper.jl | 261 +++++++++--------- 1 file changed, 138 insertions(+), 123 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 10b8ef09e55..3d7455291ce 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -15,9 +15,6 @@ import Oceananigans.DistributedComputations: _recv_from_west_buffer!, _recv_from_east_buffer! import Oceananigans.Fields: communication_buffers -@inline instantiate(T::DataType) = T() -@inline instantiate(T) = T - const DistributedZipper = BoundaryCondition{<:DistributedCommunication, <:ZipperHaloCommunicationRanks} ##### @@ -71,25 +68,25 @@ end ##### Zipper communication buffers with fold-aware packing ##### -struct TwoDZipperBuffer{FL, WFL, Loc, TY, B, S} +struct TwoDZipperBuffer{FL, WFL, TY, Loc, B, S} + loc :: Loc send :: B recv :: B sign :: S end -struct ZipperCornerBuffer{FL, WFL, Loc, TY, B, S} +struct ZipperCornerBuffer{FL, WFL, TY, Loc, B, S} + loc :: Loc send :: B recv :: B sign :: S end -# Value-argument constructors: all types inferred -TwoDZipperBuffer(loc::Loc, fold_topology::TY, send::B, recv::B, sign::S, ::Val{FL}, ::Val{WFL}) where {FL, WFL, Loc, TY, B, S} = - TwoDZipperBuffer{FL, WFL, Loc, TY, B, S}(send, recv, sign) -ZipperCornerBuffer(loc::Loc, fold_topology::TY, send::B, recv::B, sign::S, ::Val{FL}, ::Val{WFL}) where {FL, WFL, Loc, TY, B, S} = - ZipperCornerBuffer{FL, WFL, Loc, TY, B, S}(send, recv, sign) -ZipperCornerBuffer(::Type{Loc}, ::Type{TY}, send::B, recv::B, sign::S, ::Val{FL}, ::Val{WFL}) where {Loc, TY, FL, WFL, B, S} = - ZipperCornerBuffer{FL, WFL, Loc, TY, B, S}(send, recv, sign) +# Partial type-parameter constructors: FL, WFL, TY specified; Loc, B, S inferred +TwoDZipperBuffer{FL, WFL, TY}(loc::Loc, send::B, recv::B, sign::S) where {FL, WFL, TY, Loc, B, S} = + TwoDZipperBuffer{FL, WFL, TY, Loc, B, S}(loc, send, recv, sign) +ZipperCornerBuffer{FL, WFL, TY}(loc::Loc, send::B, recv::B, sign::S) where {FL, WFL, TY, Loc, B, S} = + ZipperCornerBuffer{FL, WFL, TY, Loc, B, S}(loc, send, recv, sign) Adapt.adapt_structure(to, buff::TwoDZipperBuffer) = nothing Adapt.adapt_structure(to, buff::ZipperCornerBuffer) = nothing @@ -99,15 +96,15 @@ Adapt.adapt_structure(to, buff::ZipperCornerBuffer) = nothing # - FL: buffer has fold-line row (same as corners, for MPI size matching) # - WFL: recv writes the fold-line row (complement of the adjacent corner) # West WFL = !NW corner WFL, East WFL = !NE corner WFL. -struct TripolarXBuffer{B, FL, WFL} +struct TripolarXBuffer{FL, WFL, B} send :: B recv :: B end Adapt.adapt_structure(to, buff::TripolarXBuffer) = nothing -TripolarXBuffer(send::B, recv::B, ::Val{FL}, ::Val{WFL}) where {B, FL, WFL} = - TripolarXBuffer{B, FL, WFL}(send, recv) +TripolarXBuffer{FL, WFL}(send::B, recv::B) where {FL, WFL, B} = + TripolarXBuffer{FL, WFL, B}(send, recv) # X-buffers WFL: complement of the ADJACENT corner. west_writes_fold_line(arch) = !northwest_writes_fold_line(arch) @@ -120,49 +117,51 @@ _fill_east_send_buffer!(c, buff::TripolarXBuffer, Hx, Hy, Nx, Ny) = buff.send .= view(c, 1+Nx:Nx+Hx, 1+Hy:size(buff.send,2)+Hy, :) # TripolarXBuffer recv FL=false: no fold line, write all Ny_buf rows -_recv_from_west_buffer!(c, buff::TripolarXBuffer{<:Any, false}, Hx, Hy, Nx, Ny) = +_recv_from_west_buffer!(c, buff::TripolarXBuffer{false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Hy:size(buff.recv,2)+Hy, :) .= buff.recv -_recv_from_east_buffer!(c, buff::TripolarXBuffer{<:Any, false}, Hx, Hy, Nx, Ny) = +_recv_from_east_buffer!(c, buff::TripolarXBuffer{false}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Hy:size(buff.recv,2)+Hy, :) .= buff.recv # TripolarXBuffer recv FL=true, WFL=true: write all Ny_buf rows (fold line included) -_recv_from_west_buffer!(c, buff::TripolarXBuffer{<:Any, true, true}, Hx, Hy, Nx, Ny) = +_recv_from_west_buffer!(c, buff::TripolarXBuffer{true, true}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Hy:size(buff.recv,2)+Hy, :) .= buff.recv -_recv_from_east_buffer!(c, buff::TripolarXBuffer{<:Any, true, true}, Hx, Hy, Nx, Ny) = +_recv_from_east_buffer!(c, buff::TripolarXBuffer{true, true}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Hy:size(buff.recv,2)+Hy, :) .= buff.recv # TripolarXBuffer recv FL=true, WFL=false: skip last row (the fold line) -_recv_from_west_buffer!(c, buff::TripolarXBuffer{<:Any, true, false}, Hx, Hy, Nx, Ny) = +_recv_from_west_buffer!(c, buff::TripolarXBuffer{true, false}, Hx, Hy, Nx, Ny) = view(c, 1:Hx, 1+Hy:size(buff.recv,2)-1+Hy, :) .= view(buff.recv, :, 1:size(buff.recv,2)-1, :) -_recv_from_east_buffer!(c, buff::TripolarXBuffer{<:Any, true, false}, Hx, Hy, Nx, Ny) = +_recv_from_east_buffer!(c, buff::TripolarXBuffer{true, false}, Hx, Hy, Nx, Ny) = view(c, 1+Nx+Hx:Nx+2Hx, 1+Hy:size(buff.recv,2)-1+Hy, :) .= view(buff.recv, :, 1:size(buff.recv,2)-1, :) # Fold-aware x-buffer constructors. # Separate west/east constructors since they complement different corners. function west_tripolar_buffer(arch, grid, data, Hx, bc, loc, - north::TwoDZipperBuffer{<:Any, <:Any, Loc, TY}) where {Loc, TY} - loc_y = Loc.parameters[2]() - fl = has_fold_line(TY, loc_y) - Ny_buf = length(loc_y, TY(), size(grid, 2)) + north::TwoDZipperBuffer{<:Any, <:Any, TY}) where TY + ly = north.loc[2] + topo = fold_topo(north) + fl = has_fold_line(TY, ly) + Ny_buf = length(ly, topo, size(grid, 2)) wfl = fl && west_writes_fold_line(arch) _, _, Tz = size(parent(data)) FT = eltype(data) send = on_architecture(arch, zeros(FT, Hx, Ny_buf, Tz)) recv = on_architecture(arch, zeros(FT, Hx, Ny_buf, Tz)) - return TripolarXBuffer(send, recv, Val(fl), Val(wfl)) + return TripolarXBuffer{fl, wfl}(send, recv) end function east_tripolar_buffer(arch, grid, data, Hx, bc, loc, - north::TwoDZipperBuffer{<:Any, <:Any, Loc, TY}) where {Loc, TY} - loc_y = Loc.parameters[2]() - fl = has_fold_line(TY, loc_y) - Ny_buf = length(loc_y, TY(), size(grid, 2)) + north::TwoDZipperBuffer{<:Any, <:Any, TY}) where TY + ly = north.loc[2] + topo = fold_topo(north) + fl = has_fold_line(TY, ly) + Ny_buf = length(ly, topo, size(grid, 2)) wfl = fl && east_writes_fold_line(arch) _, _, Tz = size(parent(data)) FT = eltype(data) send = on_architecture(arch, zeros(FT, Hx, Ny_buf, Tz)) recv = on_architecture(arch, zeros(FT, Hx, Ny_buf, Tz)) - return TripolarXBuffer(send, recv, Val(fl), Val(wfl)) + return TripolarXBuffer{fl, wfl}(send, recv) end # Fallback when north is not a zipper (south ranks) @@ -208,7 +207,7 @@ function y_tripolar_buffer(arch, grid::AbstractGrid{<:Any, <:Any, Topo}, wfl = fl && north_writes_fold_line(arch) send = on_architecture(arch, zeros(FT, Nx, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Nx, Hy′, Tz)) - return TwoDZipperBuffer(loc, Topo(), send, recv, sgn, Val(fl), Val(wfl)) + return TwoDZipperBuffer{fl, wfl, Topo}(loc, send, recv, sgn) end # Fallbacks: non-zipper corners @@ -221,29 +220,29 @@ northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge) = corner_commu # based on whether the corner's global x-position is past the pivot. function northwest_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, - yedge::TwoDZipperBuffer{<:Any, <:Any, Loc, TY}) where {Loc, TY} + yedge::TwoDZipperBuffer{<:Any, <:Any, TY}) where TY Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - loc_x_inst = instantiate(Loc.parameters[1]) - fl = has_fold_line(TY, Loc.parameters[2]()) + lx, ly = yedge.loc[1], yedge.loc[2] + fl = has_fold_line(TY, ly) Hy′ = fl ? Hy + 1 : Hy wfl = fl && northwest_writes_fold_line(arch) - Hx′ = nw_corner_nx(loc_x_inst, Hx) + Hx′ = northwest_x_size(lx, Hx) send = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) - return ZipperCornerBuffer(Loc, TY, send, recv, sgn, Val(fl), Val(wfl)) + return ZipperCornerBuffer{fl, wfl, TY}(yedge.loc, send, recv, sgn) end function northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, - yedge::TwoDZipperBuffer{<:Any, <:Any, Loc, TY}) where {Loc, TY} + yedge::TwoDZipperBuffer{<:Any, <:Any, TY}) where TY Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - loc_x_inst = instantiate(Loc.parameters[1]) - fl = has_fold_line(TY, Loc.parameters[2]()) + lx, ly = yedge.loc[1], yedge.loc[2] + fl = has_fold_line(TY, ly) Hy′ = fl ? Hy + 1 : Hy wfl = fl && northeast_writes_fold_line(arch) - Hx′ = ne_corner_nx(loc_x_inst, Hx) + Hx′ = northeast_x_size(lx, Hx) send = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) - return ZipperCornerBuffer(Loc, TY, send, recv, sgn, Val(fl), Val(wfl)) + return ZipperCornerBuffer{fl, wfl, TY}(yedge.loc, send, recv, sgn) end ##### @@ -259,65 +258,62 @@ const FF = Tuple{<:Face, <:Face, <:Any} ##### Helper functions: compute ranges from location and topology ##### -# Type-parameter accessors -@inline loc_x(::TwoDZipperBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[1]) -@inline loc_y(::TwoDZipperBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[2]) -@inline fold_topo(::TwoDZipperBuffer{<:Any, <:Any, <:Any, TY}) where TY = TY() - -@inline loc_x(::ZipperCornerBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[1]) -@inline loc_y(::ZipperCornerBuffer{<:Any, <:Any, Loc}) where Loc = instantiate(Loc.parameters[2]) -@inline fold_topo(::ZipperCornerBuffer{<:Any, <:Any, <:Any, TY}) where TY = TY() +# Field accessors +@inline fold_topo(::TwoDZipperBuffer{<:Any, <:Any, TY}) where TY = TY() +@inline fold_topo(::ZipperCornerBuffer{<:Any, <:Any, TY}) where TY = TY() -# Send y-ranges (source rows, reversed for fold) -@inline send_fold_y(::UPivotTopology, ::Center, Hy, Ny) = Ny + Hy -@inline send_fold_y(::FPivotTopology, ::Face, Hy, Ny) = Ny + 1 + Hy +# North buffer: fold-line y-index (source/destination row at the fold) +@inline north_foldline_send_y_index(::UPivotTopology, ::Center, Hy, Ny) = Ny + Hy +@inline north_foldline_send_y_index(::FPivotTopology, ::Face, Hy, Ny) = Ny + 1 + Hy -@inline send_halo_y(::UPivotTopology, ::Center, Hy, Ny) = Ny+Hy-1:-1:Ny -@inline send_halo_y(topo, loc_y, Hy, Ny) = Ny+Hy:-1:Ny+1 +@inline north_foldline_recv_y_index(::UPivotTopology, ::Center, Hy, Ny) = Ny + Hy +@inline north_foldline_recv_y_index(::FPivotTopology, ::Face, Hy, Ny) = Ny + 1 + Hy -# Recv y-ranges (destination rows in parent) -@inline recv_fold_y(::UPivotTopology, ::Center, Hy, Ny) = Ny + Hy -@inline recv_fold_y(::FPivotTopology, ::Face, Hy, Ny) = Ny + 1 + Hy +# North buffer: halo y-range (reversed for send, forward for recv) +@inline north_halo_send_y_range(::UPivotTopology, ::Center, Hy, Ny) = Ny+Hy-1:-1:Ny +@inline north_halo_send_y_range(topo, loc_y, Hy, Ny) = Ny+Hy:-1:Ny+1 -@inline recv_halo_y(::FPivotTopology, ::Face, Hy, Ny) = 2+Ny+Hy:1+Ny+2Hy -@inline recv_halo_y(topo, loc_y, Hy, Ny) = 1+Ny+Hy:Ny+2Hy +@inline north_halo_recv_y_range(::FPivotTopology, ::Face, Hy, Ny) = 2+Ny+Hy:1+Ny+2Hy +@inline north_halo_recv_y_range(topo, loc_y, Hy, Ny) = 1+Ny+Hy:Ny+2Hy -# Recv x-range for north buffer -@inline recv_x_range(::Center, Hx, Nx) = 1+Hx:Nx+Hx -@inline recv_x_range(::Face, Hx, Nx) = 2+Hx:Nx+Hx+1 +# North buffer: recv x-range +@inline north_recv_x_range(::Center, Hx, Nx) = 1+Hx:Nx+Hx +@inline north_recv_x_range(::Face, Hx, Nx) = 2+Hx:Nx+Hx+1 -# Corner send x-ranges (reversed for fold) -@inline nw_send_x(::Center, Hx, Nx) = 2Hx:-1:1+Hx -@inline nw_send_x(::Face, Hx, Nx) = 2Hx+1:-1:1+Hx -@inline ne_send_x(::Center, Hx, Nx) = Nx+Hx:-1:1+Nx -@inline ne_send_x(::Face, Hx, Nx) = Nx+Hx:-1:Nx+2 +# Corner buffers: send x-ranges (reversed for fold) +@inline northwest_send_x_range(::Center, Hx, Nx) = 2Hx:-1:1+Hx +@inline northwest_send_x_range(::Face, Hx, Nx) = 2Hx+1:-1:1+Hx +@inline northeast_send_x_range(::Center, Hx, Nx) = Nx+Hx:-1:1+Nx +@inline northeast_send_x_range(::Face, Hx, Nx) = Nx+Hx:-1:Nx+2 -# Corner recv x-ranges -@inline nw_recv_x(::Center, Hx, Nx) = 1:Hx -@inline nw_recv_x(::Face, Hx, Nx) = 1:Hx+1 -@inline ne_recv_x(::Center, Hx, Nx) = 1+Nx+Hx:Nx+2Hx -@inline ne_recv_x(::Face, Hx, Nx) = 2+Nx+Hx:Nx+2Hx +# Corner buffers: recv x-ranges +@inline northwest_recv_x_range(::Center, Hx, Nx) = 1:Hx +@inline northwest_recv_x_range(::Face, Hx, Nx) = 1:Hx+1 +@inline northeast_recv_x_range(::Center, Hx, Nx) = 1+Nx+Hx:Nx+2Hx +@inline northeast_recv_x_range(::Face, Hx, Nx) = 2+Nx+Hx:Nx+2Hx -# Corner buffer x-size (for constructors) -@inline nw_corner_nx(::Center, Hx) = Hx -@inline nw_corner_nx(::Face, Hx) = Hx + 1 -@inline ne_corner_nx(::Center, Hx) = Hx -@inline ne_corner_nx(::Face, Hx) = Hx - 1 +# Corner buffers: x-size (for constructors) +@inline northwest_x_size(::Center, Hx) = Hx +@inline northwest_x_size(::Face, Hx) = Hx + 1 +@inline northeast_x_size(::Center, Hx) = Hx +@inline northeast_x_size(::Face, Hx) = Hx - 1 ##### ##### TwoDZipperBuffer: north send (FL=true has fold line, FL=false does not) ##### function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{true}, Hx, Hy, Nx, Ny) - topo, ly = fold_topo(b), loc_y(b) - fy = send_fold_y(topo, ly, Hy, Ny) - view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, fy:fy, :) - view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, send_halo_y(topo, ly, Hy, Ny), :) + topo, ly = fold_topo(b), b.loc[2] + j = north_foldline_send_y_index(topo, ly, Hy, Ny) + jrange = north_halo_send_y_range(topo, ly, Hy, Ny) + view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, j:j, :) + view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, jrange, :) end function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{false}, Hx, Hy, Nx, Ny) - topo, ly = fold_topo(b), loc_y(b) - b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, send_halo_y(topo, ly, Hy, Ny), :) + topo, ly = fold_topo(b), b.loc[2] + jrange = north_halo_send_y_range(topo, ly, Hy, Ny) + b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, jrange, :) end ##### @@ -325,21 +321,24 @@ end ##### function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{true, true}, Hx, Hy, Nx, Ny) - xr = recv_x_range(loc_x(buff), Hx, Nx) - topo, ly = fold_topo(buff), loc_y(buff) - fy = recv_fold_y(topo, ly, Hy, Ny) - view(c, xr, fy:fy , :) .= view(buff.recv, :, 1:1 , :) - view(c, xr, recv_halo_y(topo, ly, Hy, Ny), :) .= view(buff.recv, :, 2:Hy+1, :) + topo, ly = fold_topo(buff), buff.loc[2] + j = north_foldline_recv_y_index(topo, ly, Hy, Ny) + xr = north_recv_x_range(buff.loc[1], Hx, Nx) + yr = north_halo_recv_y_range(topo, ly, Hy, Ny) + view(c, xr, j:j , :) .= view(buff.recv, :, 1:1 , :) + view(c, xr, yr, :) .= view(buff.recv, :, 2:Hy+1, :) end function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{true, false}, Hx, Hy, Nx, Ny) - xr = recv_x_range(loc_x(buff), Hx, Nx) - view(c, xr, recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= view(buff.recv, :, 2:Hy+1, :) + xr = north_recv_x_range(buff.loc[1], Hx, Nx) + yr = north_halo_recv_y_range(fold_topo(buff), buff.loc[2], Hy, Ny) + view(c, xr, yr, :) .= view(buff.recv, :, 2:Hy+1, :) end function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{false}, Hx, Hy, Nx, Ny) - xr = recv_x_range(loc_x(buff), Hx, Nx) - view(c, xr, recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= buff.recv + xr = north_recv_x_range(buff.loc[1], Hx, Nx) + yr = north_halo_recv_y_range(fold_topo(buff), buff.loc[2], Hy, Ny) + view(c, xr, yr, :) .= buff.recv end ##### @@ -348,28 +347,32 @@ end function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{true}, Hx, Hy, Nx, Ny) topo = fold_topo(b) - ly = loc_y(b) - fy = send_fold_y(topo, ly, Hy, Ny) - b.send .= b.sign .* view(c, nw_send_x(loc_x(b), Hx, Nx), fy:-1:fy-Hy, :) + ly = b.loc[2] + j = north_foldline_send_y_index(topo, ly, Hy, Ny) + xr = northwest_send_x_range(b.loc[1], Hx, Nx) + b.send .= b.sign .* view(c, xr, j:-1:j-Hy, :) end function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{false}, Hx, Hy, Nx, Ny) topo = fold_topo(b) - fy = send_halo_y(topo, loc_y(b), Hy, Ny) - b.send .= b.sign .* view(c, nw_send_x(loc_x(b), Hx, Nx), fy, :) + yr = north_halo_send_y_range(topo, b.loc[2], Hy, Ny) + xr = northwest_send_x_range(b.loc[1], Hx, Nx) + b.send .= b.sign .* view(c, xr, yr, :) end function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{true}, Hx, Hy, Nx, Ny) topo = fold_topo(b) - ly = loc_y(b) - fy = send_fold_y(topo, ly, Hy, Ny) - b.send .= b.sign .* view(c, ne_send_x(loc_x(b), Hx, Nx), fy:-1:fy-Hy, :) + ly = b.loc[2] + j = north_foldline_send_y_index(topo, ly, Hy, Ny) + xr = northeast_send_x_range(b.loc[1], Hx, Nx) + b.send .= b.sign .* view(c, xr, j:-1:j-Hy, :) end function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{false}, Hx, Hy, Nx, Ny) topo = fold_topo(b) - fy = send_halo_y(topo, loc_y(b), Hy, Ny) - b.send .= b.sign .* view(c, ne_send_x(loc_x(b), Hx, Nx), fy, :) + yr = north_halo_send_y_range(topo, b.loc[2], Hy, Ny) + xr = northeast_send_x_range(b.loc[1], Hx, Nx) + b.send .= b.sign .* view(c, xr, yr, :) end ##### @@ -377,31 +380,43 @@ end ##### function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{true, true}, Hx, Hy, Nx, Ny) - xr = nw_recv_x(loc_x(buff), Hx, Nx) - topo, ly = fold_topo(buff), loc_y(buff) - fy = recv_fold_y(topo, ly, Hy, Ny) - view(c, xr, fy:fy, :) .= view(buff.recv, :, 1:1, :) - view(c, xr, recv_halo_y(topo, ly, Hy, Ny), :) .= view(buff.recv, :, 2:size(buff.recv,2), :) + topo, ly = fold_topo(buff), buff.loc[2] + j = north_foldline_recv_y_index(topo, ly, Hy, Ny) + xr = northwest_recv_x_range(buff.loc[1], Hx, Nx) + yr = north_halo_recv_y_range(topo, ly, Hy, Ny) + view(c, xr, j:j, :) .= view(buff.recv, :, 1:1, :) + view(c, xr, yr, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) end -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{true, false}, Hx, Hy, Nx, Ny) = - view(c, nw_recv_x(loc_x(buff), Hx, Nx), recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= - view(buff.recv, :, 2:size(buff.recv,2), :) +function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{true, false}, Hx, Hy, Nx, Ny) + xr = northwest_recv_x_range(buff.loc[1], Hx, Nx) + yr = north_halo_recv_y_range(fold_topo(buff), buff.loc[2], Hy, Ny) + view(c, xr, yr, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +end -_recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{false}, Hx, Hy, Nx, Ny) = - view(c, nw_recv_x(loc_x(buff), Hx, Nx), recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= buff.recv +function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{false}, Hx, Hy, Nx, Ny) + xr = northwest_recv_x_range(buff.loc[1], Hx, Nx) + yr = north_halo_recv_y_range(fold_topo(buff), buff.loc[2], Hy, Ny) + view(c, xr, yr, :) .= buff.recv +end function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{true, true}, Hx, Hy, Nx, Ny) - xr = ne_recv_x(loc_x(buff), Hx, Nx) - topo, ly = fold_topo(buff), loc_y(buff) - fy = recv_fold_y(topo, ly, Hy, Ny) - view(c, xr, fy:fy, :) .= view(buff.recv, :, 1:1, :) - view(c, xr, recv_halo_y(topo, ly, Hy, Ny), :) .= view(buff.recv, :, 2:size(buff.recv,2), :) + topo, ly = fold_topo(buff), buff.loc[2] + xr = northeast_recv_x_range(buff.loc[1], Hx, Nx) + j = north_foldline_recv_y_index(topo, ly, Hy, Ny) + yr = north_halo_recv_y_range(topo, ly, Hy, Ny) + view(c, xr, j:j, :) .= view(buff.recv, :, 1:1, :) + view(c, xr, yr , :) .= view(buff.recv, :, 2:size(buff.recv,2), :) end -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{true, false}, Hx, Hy, Nx, Ny) = - view(c, ne_recv_x(loc_x(buff), Hx, Nx), recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= - view(buff.recv, :, 2:size(buff.recv,2), :) +function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{true, false}, Hx, Hy, Nx, Ny) + xr = northeast_recv_x_range(buff.loc[1], Hx, Nx) + yr = north_halo_recv_y_range(fold_topo(buff), buff.loc[2], Hy, Ny) + view(c, xr, yr, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) +end -_recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{false}, Hx, Hy, Nx, Ny) = - view(c, ne_recv_x(loc_x(buff), Hx, Nx), recv_halo_y(fold_topo(buff), loc_y(buff), Hy, Ny), :) .= buff.recv +function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{false}, Hx, Hy, Nx, Ny) + xr = northeast_recv_x_range(buff.loc[1], Hx, Nx) + yr = north_halo_recv_y_range(fold_topo(buff), buff.loc[2], Hy, Ny) + view(c, xr, yr, :) .= buff.recv +end From 658e18f293eb539f1410dca82b11c1f5aaaaa59f Mon Sep 17 00:00:00 2001 From: Benoit Pasquier Date: Thu, 9 Apr 2026 01:36:41 +1000 Subject: [PATCH 28/28] =?UTF-8?q?`lx`=20->=20`=E2=84=93x`=20and=20`Topo`?= =?UTF-8?q?=20->=20`TY`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../distributed_zipper.jl | 68 +++++++++---------- 1 file changed, 34 insertions(+), 34 deletions(-) diff --git a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl index 3d7455291ce..f6fc323a437 100644 --- a/src/OrthogonalSphericalShellGrids/distributed_zipper.jl +++ b/src/OrthogonalSphericalShellGrids/distributed_zipper.jl @@ -138,10 +138,10 @@ _recv_from_east_buffer!(c, buff::TripolarXBuffer{true, false}, Hx, Hy, Nx, Ny) = # Separate west/east constructors since they complement different corners. function west_tripolar_buffer(arch, grid, data, Hx, bc, loc, north::TwoDZipperBuffer{<:Any, <:Any, TY}) where TY - ly = north.loc[2] + ℓy = north.loc[2] topo = fold_topo(north) - fl = has_fold_line(TY, ly) - Ny_buf = length(ly, topo, size(grid, 2)) + fl = has_fold_line(TY, ℓy) + Ny_buf = length(ℓy, topo, size(grid, 2)) wfl = fl && west_writes_fold_line(arch) _, _, Tz = size(parent(data)) FT = eltype(data) @@ -152,10 +152,10 @@ end function east_tripolar_buffer(arch, grid, data, Hx, bc, loc, north::TwoDZipperBuffer{<:Any, <:Any, TY}) where TY - ly = north.loc[2] + ℓy = north.loc[2] topo = fold_topo(north) - fl = has_fold_line(TY, ly) - Ny_buf = length(ly, topo, size(grid, 2)) + fl = has_fold_line(TY, ℓy) + Ny_buf = length(ℓy, topo, size(grid, 2)) wfl = fl && east_writes_fold_line(arch) _, _, Tz = size(parent(data)) FT = eltype(data) @@ -196,18 +196,18 @@ end y_tripolar_buffer(arch, grid, data, Hy, bc, loc) = y_communication_buffer(arch, grid, data, Hy, bc) # 2D fold (MxN) → TwoDZipperBuffer (interior-width, Hy′ = Hy or Hy+1 rows for fold line) -function y_tripolar_buffer(arch, grid::AbstractGrid{<:Any, <:Any, Topo}, - data, Hy, bc::DistributedZipper, loc::Loc) where {Topo <: TwoDFoldTopology, Loc} +function y_tripolar_buffer(arch, grid::AbstractGrid{<:Any, <:Any, TY}, + data, Hy, bc::DistributedZipper, loc::Loc) where {TY <: TwoDFoldTopology, Loc} Nx = size(grid, 1) _, _, Tz = size(parent(data)) FT = eltype(data) sgn = bc.condition.sign - fl = has_fold_line(Topo, loc[2]) + fl = has_fold_line(TY, loc[2]) Hy′ = fl ? Hy + 1 : Hy wfl = fl && north_writes_fold_line(arch) send = on_architecture(arch, zeros(FT, Nx, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Nx, Hy′, Tz)) - return TwoDZipperBuffer{fl, wfl, Topo}(loc, send, recv, sgn) + return TwoDZipperBuffer{fl, wfl, TY}(loc, send, recv, sgn) end # Fallbacks: non-zipper corners @@ -222,11 +222,11 @@ northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge) = corner_commu function northwest_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge::TwoDZipperBuffer{<:Any, <:Any, TY}) where TY Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - lx, ly = yedge.loc[1], yedge.loc[2] - fl = has_fold_line(TY, ly) + ℓx, ℓy = yedge.loc[1], yedge.loc[2] + fl = has_fold_line(TY, ℓy) Hy′ = fl ? Hy + 1 : Hy wfl = fl && northwest_writes_fold_line(arch) - Hx′ = northwest_x_size(lx, Hx) + Hx′ = northwest_x_size(ℓx, Hx) send = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) return ZipperCornerBuffer{fl, wfl, TY}(yedge.loc, send, recv, sgn) @@ -235,11 +235,11 @@ end function northeast_tripolar_buffer(arch, grid, data, Hx, Hy, xedge, yedge::TwoDZipperBuffer{<:Any, <:Any, TY}) where TY Tz = size(parent(data), 3); FT = eltype(data); sgn = yedge.sign - lx, ly = yedge.loc[1], yedge.loc[2] - fl = has_fold_line(TY, ly) + ℓx, ℓy = yedge.loc[1], yedge.loc[2] + fl = has_fold_line(TY, ℓy) Hy′ = fl ? Hy + 1 : Hy wfl = fl && northeast_writes_fold_line(arch) - Hx′ = northeast_x_size(lx, Hx) + Hx′ = northeast_x_size(ℓx, Hx) send = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) recv = on_architecture(arch, zeros(FT, Hx′, Hy′, Tz)) return ZipperCornerBuffer{fl, wfl, TY}(yedge.loc, send, recv, sgn) @@ -303,16 +303,16 @@ const FF = Tuple{<:Face, <:Face, <:Any} ##### function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{true}, Hx, Hy, Nx, Ny) - topo, ly = fold_topo(b), b.loc[2] - j = north_foldline_send_y_index(topo, ly, Hy, Ny) - jrange = north_halo_send_y_range(topo, ly, Hy, Ny) + topo, ℓy = fold_topo(b), b.loc[2] + j = north_foldline_send_y_index(topo, ℓy, Hy, Ny) + jrange = north_halo_send_y_range(topo, ℓy, Hy, Ny) view(b.send, :, 1:1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, j:j, :) view(b.send, :, 2:Hy+1, :) .= b.sign .* view(c, Nx+Hx:-1:1+Hx, jrange, :) end function _fill_north_send_buffer!(c, b::TwoDZipperBuffer{false}, Hx, Hy, Nx, Ny) - topo, ly = fold_topo(b), b.loc[2] - jrange = north_halo_send_y_range(topo, ly, Hy, Ny) + topo, ℓy = fold_topo(b), b.loc[2] + jrange = north_halo_send_y_range(topo, ℓy, Hy, Ny) b.send .= b.sign .* view(c, Nx+Hx:-1:1+Hx, jrange, :) end @@ -321,10 +321,10 @@ end ##### function _recv_from_north_buffer!(c, buff::TwoDZipperBuffer{true, true}, Hx, Hy, Nx, Ny) - topo, ly = fold_topo(buff), buff.loc[2] - j = north_foldline_recv_y_index(topo, ly, Hy, Ny) + topo, ℓy = fold_topo(buff), buff.loc[2] + j = north_foldline_recv_y_index(topo, ℓy, Hy, Ny) xr = north_recv_x_range(buff.loc[1], Hx, Nx) - yr = north_halo_recv_y_range(topo, ly, Hy, Ny) + yr = north_halo_recv_y_range(topo, ℓy, Hy, Ny) view(c, xr, j:j , :) .= view(buff.recv, :, 1:1 , :) view(c, xr, yr, :) .= view(buff.recv, :, 2:Hy+1, :) end @@ -347,8 +347,8 @@ end function _fill_northwest_send_buffer!(c, b::ZipperCornerBuffer{true}, Hx, Hy, Nx, Ny) topo = fold_topo(b) - ly = b.loc[2] - j = north_foldline_send_y_index(topo, ly, Hy, Ny) + ℓy = b.loc[2] + j = north_foldline_send_y_index(topo, ℓy, Hy, Ny) xr = northwest_send_x_range(b.loc[1], Hx, Nx) b.send .= b.sign .* view(c, xr, j:-1:j-Hy, :) end @@ -362,8 +362,8 @@ end function _fill_northeast_send_buffer!(c, b::ZipperCornerBuffer{true}, Hx, Hy, Nx, Ny) topo = fold_topo(b) - ly = b.loc[2] - j = north_foldline_send_y_index(topo, ly, Hy, Ny) + ℓy = b.loc[2] + j = north_foldline_send_y_index(topo, ℓy, Hy, Ny) xr = northeast_send_x_range(b.loc[1], Hx, Nx) b.send .= b.sign .* view(c, xr, j:-1:j-Hy, :) end @@ -380,10 +380,10 @@ end ##### function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{true, true}, Hx, Hy, Nx, Ny) - topo, ly = fold_topo(buff), buff.loc[2] - j = north_foldline_recv_y_index(topo, ly, Hy, Ny) + topo, ℓy = fold_topo(buff), buff.loc[2] + j = north_foldline_recv_y_index(topo, ℓy, Hy, Ny) xr = northwest_recv_x_range(buff.loc[1], Hx, Nx) - yr = north_halo_recv_y_range(topo, ly, Hy, Ny) + yr = north_halo_recv_y_range(topo, ℓy, Hy, Ny) view(c, xr, j:j, :) .= view(buff.recv, :, 1:1, :) view(c, xr, yr, :) .= view(buff.recv, :, 2:size(buff.recv,2), :) end @@ -401,10 +401,10 @@ function _recv_from_northwest_buffer!(c, buff::ZipperCornerBuffer{false}, Hx, Hy end function _recv_from_northeast_buffer!(c, buff::ZipperCornerBuffer{true, true}, Hx, Hy, Nx, Ny) - topo, ly = fold_topo(buff), buff.loc[2] + topo, ℓy = fold_topo(buff), buff.loc[2] xr = northeast_recv_x_range(buff.loc[1], Hx, Nx) - j = north_foldline_recv_y_index(topo, ly, Hy, Ny) - yr = north_halo_recv_y_range(topo, ly, Hy, Ny) + j = north_foldline_recv_y_index(topo, ℓy, Hy, Ny) + yr = north_halo_recv_y_range(topo, ℓy, Hy, Ny) view(c, xr, j:j, :) .= view(buff.recv, :, 1:1, :) view(c, xr, yr , :) .= view(buff.recv, :, 2:size(buff.recv,2), :) end