Remove Immersed map from Adapting the ImmersedBoundaryGrid (#3690)

* add new branch

* add some comments

* add comments

* adapt non hydrostatic model

* add inbounds where needed

* adapt barotropic kernel

* last bugfix?

* add some explanation and name changes

* adding some explanation

* typo

* better naming

* better comment

* add comments

* add comments

* add another explanation

* change names

* some more explanation

* add some more explanation

* better naming

* more description

* change some naming

* typo

* bugfix

* Update src/ImmersedBoundaries/active_cells_map.jl

* Update src/ImmersedBoundaries/active_cells_map.jl

* make sure synchronized architectures do not split the active map

* adding some immersed map tests

* no active cells map for grid fitted boundary?

* active cells map for gridfitted boundary

* revert to previous implementation

src/Grids/Grids.jl

@@ -155,8 +155,10 @@ const YZFlatGrid = AbstractGrid{<:Any, <:Any, Flat, Flat}
 const XYZFlatGrid = AbstractGrid{<:Any, Flat, Flat, Flat}
 
 isrectilinear(grid) = false
-@inline active_surface_map(::AbstractGrid) = nothing
-@inline active_interior_map(::AbstractGrid) = nothing
+
+# Fallback
+@inline  retrieve_surface_active_cells_map(::AbstractGrid) = nothing
+@inline retrieve_interior_active_cells_map(::AbstractGrid, any_map_type) = nothing
 
 include("grid_utils.jl")
 include("nodes_and_spacings.jl")

src/ImmersedBoundaries/ImmersedBoundaries.jl

@@ -136,7 +136,7 @@ const IBG = ImmersedBoundaryGrid
 @inline get_ibg_property(ibg::IBG, ::Val{:immersed_boundary})      = getfield(ibg, :immersed_boundary)
 @inline get_ibg_property(ibg::IBG, ::Val{:underlying_grid})        = getfield(ibg, :underlying_grid)
 @inline get_ibg_property(ibg::IBG, ::Val{:interior_active_cells})  = getfield(ibg, :interior_active_cells)
-@inline get_ibg_property(ibg::IBG, ::Val{:active_z_columns})   = getfield(ibg, :active_z_columns)
+@inline get_ibg_property(ibg::IBG, ::Val{:active_z_columns})       = getfield(ibg, :active_z_columns)
 
 @inline architecture(ibg::IBG) = architecture(ibg.underlying_grid)
 
@@ -147,8 +147,8 @@ const IBG = ImmersedBoundaryGrid
 Adapt.adapt_structure(to, ibg::IBG{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} =
     ImmersedBoundaryGrid{TX, TY, TZ}(adapt(to, ibg.underlying_grid), 
                                      adapt(to, ibg.immersed_boundary), 
-                                     adapt(to, ibg.interior_active_cells), 
-                                     adapt(to, ibg.active_z_columns))
+                                     nothing, 
+                                     nothing)
 
 with_halo(halo, ibg::ImmersedBoundaryGrid) =
     ImmersedBoundaryGrid(with_halo(halo, ibg.underlying_grid), ibg.immersed_boundary)

src/ImmersedBoundaries/active_cells_map.jl

@@ -4,45 +4,48 @@ using Oceananigans.Grids: AbstractGrid
 
 using KernelAbstractions: @kernel, @index
 
-import Oceananigans.Grids: active_surface_map, active_interior_map
+import Oceananigans.Grids: retrieve_surface_active_cells_map, retrieve_interior_active_cells_map
 import Oceananigans.Utils: active_cells_work_layout
 
 using Oceananigans.Solvers: solve_batched_tridiagonal_system_z!, ZDirection
-using Oceananigans.DistributedComputations: DistributedGrid
+using Oceananigans.DistributedComputations: DistributedGrid, SynchronizedDistributed
 
 import Oceananigans.Solvers: solve_batched_tridiagonal_system_kernel!
 
-const DistributedActiveCellsIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:DistributedGrid, <:Any, <:NamedTuple} # Cannot be used to dispatch in kernels!!!
-const ArrayActiveCellsIBG       = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractArray}
-const NamedTupleActiveCellsIBG  = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:NamedTuple}
+# REMEMBER: since the active map is stripped out of the grid when `Adapt`ing to the GPU, 
+# The following types cannot be used to dispatch in kernels!!!
+
+# An IBG with a single interior active cells map that includes the whole :xyz domain
+const WholeActiveCellsMapIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractArray}
+
+# An IBG with an interior active cells map subdivided in 5 different sub-maps.
+# Only used (for the moment) in the case of distributed architectures where the boundary adjacent region 
+# has to be computed separately, these maps hold the active region in the "halo-independent" part of the domain
+# (; halo_independent_cells), and the "halo-dependent" regions in the west, east, north, and south, respectively
+const SplitActiveCellsMapIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:NamedTuple}
+
+# A distributed grid with split interior map
+const DistributedActiveCellsIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:DistributedGrid, <:Any, <:NamedTuple} 
 
 """
 A constant representing an immersed boundary grid, where interior active cells are mapped to linear indices in grid.interior_active_cells
 """
-const ActiveCellsIBG = Union{DistributedActiveCellsIBG, ArrayActiveCellsIBG, NamedTupleActiveCellsIBG}
+const ActiveCellsIBG = Union{DistributedActiveCellsIBG, WholeActiveCellsMapIBG, SplitActiveCellsMapIBG}
 
 """
 A constant representing an immersed boundary grid, where active columns in the Z-direction are mapped to linear indices in grid.active_z_columns
 """
 const ActiveZColumnsIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractArray}
 
-struct InteriorMap end
-struct ZColumnMap end
-
-struct WestMap  end
-struct EastMap  end
-struct SouthMap end
-struct NorthMap end
+@inline retrieve_surface_active_cells_map(grid::ActiveZColumnsIBG) = grid.active_z_columns
 
-@inline active_surface_map(::ActiveZColumnsIBG) = ZColumnMap()
-
-@inline active_interior_map(::Val{:west})  = WestMap()
-@inline active_interior_map(::Val{:east})  = EastMap()
-@inline active_interior_map(::Val{:south}) = SouthMap()
-@inline active_interior_map(::Val{:north}) = NorthMap()
-
-@inline active_interior_map(::ActiveCellsIBG)            = InteriorMap()
-@inline active_interior_map(::DistributedActiveCellsIBG) = InteriorMap()
+@inline retrieve_interior_active_cells_map(grid::WholeActiveCellsMapIBG, ::Val{:interior}) = grid.interior_active_cells
+@inline retrieve_interior_active_cells_map(grid::SplitActiveCellsMapIBG, ::Val{:interior}) = grid.interior_active_cells.halo_independent_cells
+@inline retrieve_interior_active_cells_map(grid::SplitActiveCellsMapIBG, ::Val{:west})     = grid.interior_active_cells.west_halo_dependent_cells
+@inline retrieve_interior_active_cells_map(grid::SplitActiveCellsMapIBG, ::Val{:east})     = grid.interior_active_cells.east_halo_dependent_cells
+@inline retrieve_interior_active_cells_map(grid::SplitActiveCellsMapIBG, ::Val{:south})    = grid.interior_active_cells.south_halo_dependent_cells
+@inline retrieve_interior_active_cells_map(grid::SplitActiveCellsMapIBG, ::Val{:north})    = grid.interior_active_cells.north_halo_dependent_cells
+@inline retrieve_interior_active_cells_map(grid::ActiveZColumnsIBG,      ::Val{:surface})  = grid.active_z_columns
 
 """
     active_cells_work_layout(group, size, map_type, grid)
@@ -52,19 +55,12 @@ Compute the work layout for active cells based on the given map type and grid.
 # Arguments
 - `group`: The previous workgroup.
 - `size`: The previous worksize.
-- `map_type`: The type of map (e.g., `InteriorMap`, `WestMap`, `EastMap`, `SouthMap`, `NorthMap`).
-- `grid`: The grid containing the active cells.
+- `active_cells_map`: The map containing the index of the active cells
 
 # Returns
 - A tuple `(workgroup, worksize)` representing the work layout for active cells.
 """
-@inline active_cells_work_layout(group, size, ::InteriorMap, grid::ArrayActiveCellsIBG)      = min(length(grid.interior_active_cells), 256),          length(grid.interior_active_cells)
-@inline active_cells_work_layout(group, size, ::InteriorMap, grid::NamedTupleActiveCellsIBG) = min(length(grid.interior_active_cells.interior), 256), length(grid.interior_active_cells.interior)
-@inline active_cells_work_layout(group, size, ::WestMap,     grid::NamedTupleActiveCellsIBG) = min(length(grid.interior_active_cells.west),     256), length(grid.interior_active_cells.west)
-@inline active_cells_work_layout(group, size, ::EastMap,     grid::NamedTupleActiveCellsIBG) = min(length(grid.interior_active_cells.east),     256), length(grid.interior_active_cells.east)
-@inline active_cells_work_layout(group, size, ::SouthMap,    grid::NamedTupleActiveCellsIBG) = min(length(grid.interior_active_cells.south),    256), length(grid.interior_active_cells.south)
-@inline active_cells_work_layout(group, size, ::NorthMap,    grid::NamedTupleActiveCellsIBG) = min(length(grid.interior_active_cells.north),    256), length(grid.interior_active_cells.north)
-@inline active_cells_work_layout(group, size, ::ZColumnMap,  grid::ActiveZColumnsIBG)        = min(length(grid.active_z_columns),  256),              length(grid.active_z_columns)
+@inline active_cells_work_layout(group, size, active_cells_map) = min(length(active_cells_map), 256), length(active_cells_map)
 
 """
     active_linear_index_to_tuple(idx, map, grid)
@@ -73,19 +69,12 @@ Converts a linear index to a tuple of indices based on the given map and grid.
 
 # Arguments
 - `idx`: The linear index to convert.
-- `map`: The map indicating the type of index conversion to perform.
-- `grid`: The grid containing the active cells.
+- `active_cells_map`: The map containing the N-dimensional index of the active cells
 
 # Returns
 A tuple of indices corresponding to the linear index.
 """
-@inline active_linear_index_to_tuple(idx, ::InteriorMap, grid::ArrayActiveCellsIBG)      = Base.map(Int, grid.interior_active_cells[idx])
-@inline active_linear_index_to_tuple(idx, ::InteriorMap, grid::NamedTupleActiveCellsIBG) = Base.map(Int, grid.interior_active_cells.interior[idx])
-@inline active_linear_index_to_tuple(idx, ::WestMap,     grid::NamedTupleActiveCellsIBG) = Base.map(Int, grid.interior_active_cells.west[idx])
-@inline active_linear_index_to_tuple(idx, ::EastMap,     grid::NamedTupleActiveCellsIBG) = Base.map(Int, grid.interior_active_cells.east[idx])
-@inline active_linear_index_to_tuple(idx, ::SouthMap,    grid::NamedTupleActiveCellsIBG) = Base.map(Int, grid.interior_active_cells.south[idx])
-@inline active_linear_index_to_tuple(idx, ::NorthMap,    grid::NamedTupleActiveCellsIBG) = Base.map(Int, grid.interior_active_cells.north[idx])
-@inline active_linear_index_to_tuple(idx, ::ZColumnMap,  grid::ActiveZColumnsIBG)         = Base.map(Int, grid.active_z_columns[idx])
+@inline active_linear_index_to_tuple(idx, active_cells_map) = @inbounds Base.map(Int, active_cells_map[idx])
 
 function ImmersedBoundaryGrid(grid, ib; active_cells_map::Bool = true) 
 
@@ -126,22 +115,33 @@ function compute_interior_active_cells(ibg; parameters = :xyz)
 end
 
 function compute_active_z_columns(ibg)
-    one_field = ConditionalOperation{Center, Center, Center}(OneField(Int), identity, ibg, NotImmersed(truefunc), 0)
-    column    = sum(one_field, dims = 3)
-    is_immersed_column = KernelFunctionOperation{Center, Center, Nothing}(active_column, ibg, column)
-    active_cells_field = Field{Center, Center, Nothing}(ibg, Bool)
-    set!(active_cells_field, is_immersed_column)
-    return active_cells_field
+    one_field = OneField(Int)
+    condition = NotImmersed(truefunc)
+    mask = 0
+
+    # Compute all the active cells in a z-column using a ConditionalOperation
+    conditional_active_cells = ConditionalOperation{Center, Center, Center}(one_field, identity, ibg, condition, mask)
+    active_cells_in_column   = sum(conditional_active_cells, dims = 3)
+
+    # Check whether the column ``i, j`` is immersed, which would correspond to `active_cells_in_column[i, j, 1] == 0`
+    is_immersed_column = KernelFunctionOperation{Center, Center, Nothing}(active_column, ibg, active_cells_in_column)
+    active_z_columns = Field{Center, Center, Nothing}(ibg, Bool)
+    set!(active_z_columns, is_immersed_column)
+
+    return active_z_columns
 end
 
+# Maximum integer represented by the 
+# `UInt8`, `UInt16` and `UInt32` types
 const MAXUInt8  = 2^8  - 1
 const MAXUInt16 = 2^16 - 1
 const MAXUInt32 = 2^32 - 1
 
 """
     interior_active_indices(ibg; parameters = :xyz)
 
-Compute the indices of the active interior cells in the given immersed boundary grid.
+Compute the indices of the active interior cells in the given immersed boundary grid within the indices
+specified by the `parameters` keyword argument
 
 # Arguments
 - `ibg`: The immersed boundary grid.
@@ -158,7 +158,9 @@ function interior_active_indices(ibg; parameters = :xyz)
 
     IndicesType = Tuple{IntType, IntType, IntType}
 
-    # Cannot findall on the entire field because we incur on OOM errors
+    # Cannot findall on the entire field because we could incur on OOM errors
+    # For this reason, we split the computation in vertical levels and `findall` the active indices in 
+    # subsequent xy planes, then stitch them back together
     active_indices = IndicesType[]
     active_indices = findall_active_indices!(active_indices, active_cells_field, ibg, IndicesType)
     active_indices = on_architecture(architecture(ibg), active_indices)
@@ -187,15 +189,32 @@ function convert_interior_indices(interior_indices, k, IndicesType)
     return interior_indices
 end
 
-@inline add_3rd_index(t::Tuple, k) = (t[1], t[2], k) 
+@inline add_3rd_index(ij::Tuple, k) = (ij[1], ij[2], k) 
 
+# In case of a serial grid, the interior computations are performed over the whole three-dimensional
+# domain. Therefore, the `interior_active_cells` field contains the indices of all the active cells in 
+# the range 1:Nx, 1:Ny and 1:Nz (i.e., we construct the map with parameters :xyz)
 map_interior_active_cells(ibg) = interior_active_indices(ibg; parameters = :xyz)
 
-# In case of a `DistributedGrid` we want to have different maps depending on the 
-# partitioning of the domain
+# In case of a `DistributedGrid` we want to have different maps depending on the partitioning of the domain:
+#
+# If we partition the domain in the x-direction, we typically want to have the option to split three-dimensional 
+# kernels in a `halo-independent` part in the range Hx+1:Nx-Hx, 1:Ny, 1:Nz and two `halo-dependent` computations:
+# a west one spanning 1:Hx, 1:Ny, 1:Nz and an east one spanning Nx-Hx+1:Nx, 1:Ny, 1:Nz. 
+# For this reason we need three different maps, one containing the `halo_independent` active region, a `west` map and an `east` map. 
+# For the same reason we need to construct `south` and `north` maps if we partition the domain in the y-direction.
+# Therefore, the `interior_active_cells` in this case is a `NamedTuple` containing 5 elements.
+# Note that boundary-adjacent maps corresponding to non-partitioned directions are set to `nothing`
 function map_interior_active_cells(ibg::ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:DistributedGrid})
 
     arch = architecture(ibg)
+
+    # If we using a synchronized architecture, nothing
+    # changes with serial execution.
+    if arch isa SynchronizedDistributed
+        return interior_active_indices(ibg; parameters = :xyz)
+    end
+
     Rx, Ry, _  = arch.ranks
     Tx, Ty, _  = topology(ibg)
     Nx, Ny, Nz = size(ibg)
@@ -213,20 +232,24 @@ function map_interior_active_cells(ibg::ImmersedBoundaryGrid{<:Any, <:Any, <:Any
     include_south = !isa(ibg, YFlatGrid) && (Ry != 1) && !(Ty == RightConnected)
     include_north = !isa(ibg, YFlatGrid) && (Ry != 1) && !(Ty == LeftConnected)
 
-    west  = include_west  ? interior_active_indices(ibg; parameters = KernelParameters(x_boundary, left_offsets))    : nothing
-    east  = include_east  ? interior_active_indices(ibg; parameters = KernelParameters(x_boundary, right_x_offsets)) : nothing
-    south = include_south ? interior_active_indices(ibg; parameters = KernelParameters(y_boundary, left_offsets))    : nothing
-    north = include_north ? interior_active_indices(ibg; parameters = KernelParameters(y_boundary, right_y_offsets)) : nothing
+    west_halo_dependent_cells  = include_west  ? interior_active_indices(ibg; parameters = KernelParameters(x_boundary, left_offsets))    : nothing
+    east_halo_dependent_cells  = include_east  ? interior_active_indices(ibg; parameters = KernelParameters(x_boundary, right_x_offsets)) : nothing
+    south_halo_dependent_cells = include_south ? interior_active_indices(ibg; parameters = KernelParameters(y_boundary, left_offsets))    : nothing
+    north_halo_dependent_cells = include_north ? interior_active_indices(ibg; parameters = KernelParameters(y_boundary, right_y_offsets)) : nothing
 
     nx = Rx == 1 ? Nx : (Tx == RightConnected || Tx == LeftConnected ? Nx - Hx : Nx - 2Hx)
     ny = Ry == 1 ? Ny : (Ty == RightConnected || Ty == LeftConnected ? Ny - Hy : Ny - 2Hy)
 
     ox = Rx == 1 || Tx == RightConnected ? 0 : Hx
     oy = Ry == 1 || Ty == RightConnected ? 0 : Hy
 
-    interior = interior_active_indices(ibg; parameters = KernelParameters((nx, ny, Nz), (ox, oy, 0)))
+    halo_independent_cells = interior_active_indices(ibg; parameters = KernelParameters((nx, ny, Nz), (ox, oy, 0)))
 
-    return (; interior, west, east, south, north)
+    return (; halo_independent_cells, 
+              west_halo_dependent_cells, 
+              east_halo_dependent_cells, 
+              south_halo_dependent_cells, 
+              north_halo_dependent_cells)
 end
 
 # If we eventually want to perform also barotropic step, `w` computation and `p` 

src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_boundary_tendencies.jl

@@ -2,7 +2,7 @@ import Oceananigans.Models: compute_boundary_tendencies!
 import Oceananigans.Models: compute_boundary_tendencies!
 
 using Oceananigans.Grids: halo_size
-using Oceananigans.ImmersedBoundaries: active_interior_map, DistributedActiveCellsIBG
+using Oceananigans.ImmersedBoundaries: retrieve_interior_active_cells_map, DistributedActiveCellsIBG
 using Oceananigans.Models.NonhydrostaticModels: boundary_tendency_kernel_parameters,
                                                 boundary_p_kernel_parameters, 
                                                 boundary_κ_kernel_parameters,
@@ -38,13 +38,15 @@ function compute_boundary_tendency_contributions!(grid::DistributedActiveCellsIB
     maps = grid.interior_active_cells
 
     for (name, map) in zip(keys(maps), maps)
-        compute_boundary = (name != :interior) && !isnothing(map) 
 
-        # If there exists a boundary map, then we compute the boundary contributions
+        # If there exists a boundary map, then we compute the boundary contributions. If not, the 
+        # boundary contributions have already been calculated. We exclude the interior because it has
+        # already been calculated
+        compute_boundary = (name != :interior) && !isnothing(map) 
+
         if compute_boundary
-            active_boundary_map = active_interior_map(Val(name))
-            compute_hydrostatic_free_surface_tendency_contributions!(model, tuple(:xyz); 
-                                                                     active_cells_map = active_boundary_map)
+            active_cells_map = retrieve_interior_active_cells_map(grid, Val(name))
+            compute_hydrostatic_free_surface_tendency_contributions!(model, tuple(:xyz); active_cells_map)
         end
     end