support for active_cells_map in kernels

src/DistributedComputations/distributed_immersed_boundaries.jl

@@ -121,23 +121,26 @@ function map_interior_active_cells(ibg::ImmersedBoundaryGrid{<:Any, <:Any, <:Any
     Nx, Ny, Nz = size(ibg)
     Hx, Hy, _  = halo_size(ibg)
 
-    x_boundary = (Hx, Ny, Nz)
-    y_boundary = (Nx, Hy, Nz)
-
-    left_offsets    = (0,  0,  0)
-    right_x_offsets = (Nx-Hx, 0,     0)
-    right_y_offsets = (0,     Ny-Hy, 0)
+    west_boundary  = (1:Hx,       1:Ny, 1:Nz)
+    east_boundary  = (Nx-Hx+1:Nx, 1:Ny, 1:Nz)
+    south_boundary = (1:Nx, 1:Hy,       1:Nz)
+    north_boundary = (1:Nx, Ny-Hy+1:Ny, 1:Nz)
 
     include_west  = !isa(ibg, XFlatGrid) && (Rx != 1) && !(Tx == RightConnected)
     include_east  = !isa(ibg, XFlatGrid) && (Rx != 1) && !(Tx == LeftConnected)
     include_south = !isa(ibg, YFlatGrid) && (Ry != 1) && !(Ty == RightConnected)
     include_north = !isa(ibg, YFlatGrid) && (Ry != 1) && !(Ty == LeftConnected)
 
-    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
+    west_halo_dependent_cells  = interior_active_indices(ibg; parameters = KernelParameters(west_boundary...))
+    east_halo_dependent_cells  = interior_active_indices(ibg; parameters = KernelParameters(east_boundary...))
+    south_halo_dependent_cells = interior_active_indices(ibg; parameters = KernelParameters(south_boundary...))
+    north_halo_dependent_cells = interior_active_indices(ibg; parameters = KernelParameters(north_boundary...))
 
+    west_halo_dependent_cells  = ifelse(include_west,  west_halo_dependent_cells,  nothing)
+    east_halo_dependent_cells  = ifelse(include_east,  east_halo_dependent_cells,  nothing)
+    south_halo_dependent_cells = ifelse(include_south, south_halo_dependent_cells, nothing)
+    north_halo_dependent_cells = ifelse(include_north, north_halo_dependent_cells, 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)
 

src/ImmersedBoundaries/immersed_grid_metrics.jl

@@ -1,7 +1,6 @@
 using Oceananigans.AbstractOperations: GridMetricOperation
 
-import Oceananigans.Grids: coordinates
-import Oceananigans.Operators: Δrᵃᵃᶠ, Δrᵃᵃᶜ
+import Oceananigans.Operators: Δrᵃᵃᶠ, Δrᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ
 import Oceananigans.Operators: Δxᶠᵃᵃ, Δxᶜᵃᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ
 import Oceananigans.Operators: Δyᵃᶠᵃ, Δyᵃᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶠᵃ, Δyᶜᶜᵃ
 import Oceananigans.Operators: Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl

@@ -1,30 +1,19 @@
 # Kernels to compute the vertical integral of the velocities
-@kernel function _barotropic_mode_kernel!(U, V, grid, ::Nothing, u, v)
+@kernel function _barotropic_mode_kernel!(U, V, grid, u, v)
     i, j  = @index(Global, NTuple)
-    barotropic_mode_kernel!(U, V, i, j, grid, u, v)
-end
-
-@kernel function _barotropic_mode_kernel!(U, V, grid, active_cells_map, u, v)
-    idx = @index(Global, Linear)
-    i, j = active_linear_index_to_tuple(idx, active_cells_map)
-    barotropic_mode_kernel!(U, V, i, j, grid, u, v)
-end
 
-@inline function barotropic_mode_kernel!(U, V, i, j, grid, u, v)
     @inbounds U[i, j, 1] = Δzᶠᶜᶜ(i, j, 1, grid) * u[i, j, 1]
     @inbounds V[i, j, 1] = Δzᶜᶠᶜ(i, j, 1, grid) * v[i, j, 1]
 
     for k in 2:grid.Nz
         @inbounds U[i, j, 1] += Δzᶠᶜᶜ(i, j, k, grid) * u[i, j, k]
         @inbounds V[i, j, 1] += Δzᶜᶠᶜ(i, j, k, grid) * v[i, j, k]
     end
-
-    return nothing
 end
 
 @inline function compute_barotropic_mode!(U, V, grid, u, v)
     active_cells_map = retrieve_surface_active_cells_map(grid)
-    launch!(architecture(grid), grid, :xy, _barotropic_mode_kernel!, U, V, grid, active_cells_map, u, v; active_cells_map)
+    launch!(architecture(grid), grid, :xy, _barotropic_mode_kernel!, U, V, grid, u, v; active_cells_map)
     return nothing
 end
 

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl

@@ -3,18 +3,9 @@
 #####
 
 # Calculate RHS for the barotropic time step.
-@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, ::Nothing, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
     i, j  = @index(Global, NTuple)
-    ab2_integrate_tendencies!(Gᵁ, Gⱽ, i, j, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
-end
-
-@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
-    idx = @index(Global, Linear)
-    i, j = active_linear_index_to_tuple(idx, active_cells_map)
-    ab2_integrate_tendencies!(Gᵁ, Gⱽ, i, j, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
-end
 
-@inline function ab2_integrate_tendencies!(Gᵁ, Gⱽ, i, j, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
     locU = (Face(), Center(), Center())
     locV = (Center(), Face(), Center())
 
@@ -49,7 +40,7 @@ end
     Gv⁻ = timestepper.G⁻.v
 
     launch!(architecture(grid), grid, :xy, _compute_integrated_ab2_tendencies!, GUⁿ, GVⁿ, grid,
-            active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, timestepper.χ; active_cells_map)
+            Gu⁻, Gv⁻, Guⁿ, Gvⁿ, timestepper.χ; active_cells_map)
 
     return nothing
 end
@@ -71,13 +62,7 @@ end
     return Gⁿ⁺¹
 end
 
-@kernel function _compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, grid, active_cells_map, Guⁿ, Gvⁿ, stage)
-    idx = @index(Global, Linear)
-    i, j = active_linear_index_to_tuple(idx, active_cells_map)
-    compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, stage)
-end
-
-@kernel function _compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, grid, ::Nothing, Guⁿ, Gvⁿ, stage)
+@kernel function _compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, grid, Guⁿ, Gvⁿ, stage)
     i, j = @index(Global, NTuple)
     compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, stage)
 end
@@ -118,7 +103,7 @@ end
 
     active_cells_map = retrieve_surface_active_cells_map(grid)    
     launch!(architecture(grid), grid, :xy, _compute_integrated_rk3_tendencies!, 
-            GUⁿ, GVⁿ, GU⁻, GV⁻, grid, active_cells_map, Guⁿ, Gvⁿ, stage; active_cells_map)
+            GUⁿ, GVⁿ, GU⁻, GV⁻, grid, Guⁿ, Gvⁿ, stage; active_cells_map)
 
     return nothing
 end 

src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_buffers.jl

@@ -35,17 +35,16 @@ end
 function compute_buffer_tendency_contributions!(grid::DistributedActiveCellsIBG, arch, model)
     maps = grid.interior_active_cells
 
-    for (name, map) in zip(keys(maps), maps)
+    for name in (:west_halo_dependent_cells, 
+                 :east_halo_dependent_cells, 
+                 :south_halo_dependent_cells, 
+                 :north_halo_dependent_cells)
 
-        # If there exists a buffer map, then we compute the buffer contributions. If not, the 
-        # buffer contributions have already been calculated. We exclude the interior because it has
-        # already been calculated
-        compute_buffer = (name != :interior) && !isnothing(map) 
-
-        if compute_buffer
-            active_cells_map = retrieve_interior_active_cells_map(grid, Val(name))
-            compute_hydrostatic_free_surface_tendency_contributions!(model, :xyz; active_cells_map)
-        end
+        active_cells_map = @inbounds maps[name]
+
+        # If the map == nothing, we don't need to compute the buffer because 
+        # the buffer is not adjacent to a processor boundary
+        !isnothing(map) && compute_hydrostatic_free_surface_tendency_contributions!(model, :xyz; active_cells_map)
     end
 
     return nothing

src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl

@@ -98,7 +98,6 @@ function compute_hydrostatic_free_surface_tendency_contributions!(model, kernel_
                 compute_hydrostatic_free_surface_Gc!,
                 c_tendency,
                 grid,
-                active_cells_map,
                 args;
                 active_cells_map)
     end
@@ -144,12 +143,12 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para
 
     launch!(arch, grid, kernel_parameters,
             compute_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid, 
-            active_cells_map, u_kernel_args, model.forcing.u;
+            u_kernel_args, model.forcing.u;
             active_cells_map)
 
     launch!(arch, grid, kernel_parameters,
             compute_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid, 
-            active_cells_map, v_kernel_args, model.forcing.v;
+            v_kernel_args, model.forcing.v;
             active_cells_map)
 
     return nothing
@@ -179,41 +178,23 @@ end
 #####
 
 """ Calculate the right-hand-side of the u-velocity equation. """
-@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, ::Nothing, args, forcing)
+@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, args, forcing)
     i, j, k = @index(Global, NTuple)
     @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args..., forcing)
 end
 
-@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, active_cells_map, args, forcing)
-    idx = @index(Global, Linear)
-    i, j, k = active_linear_index_to_tuple(idx, active_cells_map)
-    @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args..., forcing)
-end
-
 """ Calculate the right-hand-side of the v-velocity equation. """
-@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, ::Nothing, args, forcing)
+@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, args, forcing)
     i, j, k = @index(Global, NTuple)
     @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args..., forcing)
 end
 
-@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, active_cells_map, args, forcing)
-    idx = @index(Global, Linear)
-    i, j, k = active_linear_index_to_tuple(idx, active_cells_map)
-    @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args..., forcing)
-end
-
 #####
 ##### Tendency calculators for tracers
 #####
 
 """ Calculate the right-hand-side of the tracer advection-diffusion equation. """
-@kernel function compute_hydrostatic_free_surface_Gc!(Gc, grid, ::Nothing, args)
+@kernel function compute_hydrostatic_free_surface_Gc!(Gc, grid, args)
     i, j, k = @index(Global, NTuple)
     @inbounds Gc[i, j, k] = hydrostatic_free_surface_tracer_tendency(i, j, k, grid, args...)
-end
-
-@kernel function compute_hydrostatic_free_surface_Gc!(Gc, grid, active_cells_map, args)
-    idx = @index(Global, Linear)
-    i, j, k = active_linear_index_to_tuple(idx, active_cells_map)
-    @inbounds Gc[i, j, k] = hydrostatic_free_surface_tracer_tendency(i, j, k, grid, args...)
-end
+end

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl

@@ -1,7 +1,6 @@
 using Oceananigans.Fields: location
 using Oceananigans.TimeSteppers: ab2_step_field!
 using Oceananigans.TurbulenceClosures: implicit_step!
-using Oceananigans.ImmersedBoundaries: retrieve_interior_active_cells_map, retrieve_surface_active_cells_map
 
 import Oceananigans.TimeSteppers: ab2_step!
 

src/Models/HydrostaticFreeSurfaceModels/store_hydrostatic_free_surface_tendencies.jl

@@ -4,7 +4,6 @@ using Oceananigans.TimeSteppers: store_field_tendencies!
 
 using Oceananigans: prognostic_fields
 using Oceananigans.Grids: AbstractGrid
-using Oceananigans.ImmersedBoundaries: retrieve_interior_active_cells_map
 
 using Oceananigans.Utils: launch!
 

src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl

@@ -103,15 +103,15 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti
 
     exclude_periphery = true
     launch!(arch, grid, kernel_parameters, compute_Gu!, 
-            tendencies.u, grid, active_cells_map, u_kernel_args;
+            tendencies.u, grid, u_kernel_args;
             active_cells_map, exclude_periphery)
 
     launch!(arch, grid, kernel_parameters, compute_Gv!, 
-            tendencies.v, grid, active_cells_map, v_kernel_args;
+            tendencies.v, grid, v_kernel_args;
             active_cells_map, exclude_periphery)
 
     launch!(arch, grid, kernel_parameters, compute_Gw!, 
-            tendencies.w, grid, active_cells_map, w_kernel_args;
+            tendencies.w, grid, w_kernel_args;
             active_cells_map, exclude_periphery)
 
     start_tracer_kernel_args = (advection, closure)
@@ -131,7 +131,7 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti
                      clock, forcing)
 
         launch!(arch, grid, kernel_parameters, compute_Gc!, 
-                c_tendency, grid, active_cells_map, args;
+                c_tendency, grid, args;
                 active_cells_map)
     end
 
@@ -143,57 +143,34 @@ end
 #####
 
 """ Calculate the right-hand-side of the u-velocity equation. """
-@kernel function compute_Gu!(Gu, grid, ::Nothing, args) 
+@kernel function compute_Gu!(Gu, grid, args) 
     i, j, k = @index(Global, NTuple)
     @inbounds Gu[i, j, k] = u_velocity_tendency(i, j, k, grid, args...)
 end
 
-@kernel function compute_Gu!(Gu, grid, interior_map, args) 
-    idx = @index(Global, Linear)
-    i, j, k = active_linear_index_to_tuple(idx, interior_map)
-    @inbounds Gu[i, j, k] = u_velocity_tendency(i, j, k, grid, args...)
-end
-
 """ Calculate the right-hand-side of the v-velocity equation. """
-@kernel function compute_Gv!(Gv, grid, ::Nothing, args) 
+@kernel function compute_Gv!(Gv, grid, args) 
     i, j, k = @index(Global, NTuple)
     @inbounds Gv[i, j, k] = v_velocity_tendency(i, j, k, grid, args...)
 end
 
-@kernel function compute_Gv!(Gv, grid, interior_map, args) 
-    idx = @index(Global, Linear)
-    i, j, k = active_linear_index_to_tuple(idx, interior_map)
-    @inbounds Gv[i, j, k] = v_velocity_tendency(i, j, k, grid, args...)
-end
-
 """ Calculate the right-hand-side of the w-velocity equation. """
-@kernel function compute_Gw!(Gw, grid, ::Nothing, args) 
+@kernel function compute_Gw!(Gw, grid, args) 
     i, j, k = @index(Global, NTuple)
     @inbounds Gw[i, j, k] = w_velocity_tendency(i, j, k, grid, args...)
 end
 
-@kernel function compute_Gw!(Gw, grid, interior_map, args)
-    idx = @index(Global, Linear)
-    i, j, k = active_linear_index_to_tuple(idx, interior_map)
-    @inbounds Gw[i, j, k] = w_velocity_tendency(i, j, k, grid, args...)
-end
 
 #####
 ##### Tracer(s)
 #####
 
 """ Calculate the right-hand-side of the tracer advection-diffusion equation. """
-@kernel function compute_Gc!(Gc, grid, ::Nothing, args)
+@kernel function compute_Gc!(Gc, grid, args)
     i, j, k = @index(Global, NTuple)
     @inbounds Gc[i, j, k] = tracer_tendency(i, j, k, grid, args...)
 end
 
-@kernel function compute_Gc!(Gc, grid, interior_map, args) 
-    idx = @index(Global, Linear)
-    i, j, k = active_linear_index_to_tuple(idx, interior_map)
-    @inbounds Gc[i, j, k] = tracer_tendency(i, j, k, grid, args...)
-end
-
 #####
 ##### Boundary contributions to tendencies due to user-prescribed fluxes
 #####