diff --git a/src/DistributedComputations/distributed_immersed_boundaries.jl b/src/DistributedComputations/distributed_immersed_boundaries.jl
index 3395d01fd2..8948e568b6 100644
--- a/src/DistributedComputations/distributed_immersed_boundaries.jl
+++ b/src/DistributedComputations/distributed_immersed_boundaries.jl
@@ -140,7 +140,7 @@ function map_interior_active_cells(ibg::ImmersedBoundaryGrid{<:Any, <:Any, <:Any
     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)
 
diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl
index 5b31e5f98c..3de75cc256 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl
+++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl
@@ -1,16 +1,6 @@
 # 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, η)
     k_top  = size(grid, 3) + 1
 
     hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid)
@@ -32,13 +22,11 @@ end
         @inbounds U̅[i, j, 1] += Δrᶠᶜᶜ(i, j, k, grid) * u[i, j, k] * σᶠᶜ
         @inbounds V̅[i, j, 1] += Δrᶜᶠᶜ(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
 
diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl
index b87d1a09be..9cc29a5b04 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl
+++ b/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
@@ -116,7 +101,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 
diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl
index 481b962df0..f09b7aed62 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl
+++ b/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
@@ -148,13 +147,11 @@ 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;
-            active_cells_map)
+            u_kernel_args; active_cells_map)
 
     launch!(arch, grid, kernel_parameters,
             compute_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid, 
-            active_cells_map, v_kernel_args;
-            active_cells_map)
+            v_kernel_args; active_cells_map)
 
     return nothing
 end
@@ -183,41 +180,23 @@ end
 #####
 
 """ Calculate the right-hand-side of the u-velocity equation. """
-@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, ::Nothing, args)
+@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, args)
     i, j, k = @index(Global, NTuple)
     @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...)
 end
 
-@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, active_cells_map, args)
-    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...)
-end
-
 """ Calculate the right-hand-side of the v-velocity equation. """
-@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, ::Nothing, args)
+@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, args)
     i, j, k = @index(Global, NTuple)
     @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...)
 end
 
-@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, active_cells_map, args)
-    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...)
-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
\ No newline at end of file
diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl
index 1f7fa9a40f..54b24fd565 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl
+++ b/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!
 
diff --git a/src/Models/HydrostaticFreeSurfaceModels/store_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/store_hydrostatic_free_surface_tendencies.jl
index 6c6fe5395c..41c16cd626 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/store_hydrostatic_free_surface_tendencies.jl
+++ b/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!
 
diff --git a/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl b/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl
index cb0947b4be..4e50f9c015 100644
--- a/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl
+++ b/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
 #####
diff --git a/src/Utils/kernel_launching.jl b/src/Utils/kernel_launching.jl
index b1d539c21c..27a8796028 100644
--- a/src/Utils/kernel_launching.jl
+++ b/src/Utils/kernel_launching.jl
@@ -6,10 +6,13 @@ using Oceananigans: location
 using Oceananigans.Architectures
 using Oceananigans.Grids
 using Oceananigans.Grids: AbstractGrid
+using Adapt
 using Base: @pure
+using KernelAbstractions: Kernel
 
 import Oceananigans
 import KernelAbstractions: get, expand
+import Base
 
 struct KernelParameters{S, O} end
 
@@ -83,6 +86,12 @@ KernelParameters(args::Tuple) = KernelParameters(args...)
 contiguousrange(range::NTuple{N, Int}, offset::NTuple{N, Int}) where N = Tuple(1+o:r+o for (r, o) in zip(range, offset))
 flatten_reduced_dimensions(worksize, dims) = Tuple(d ∈ dims ? 1 : worksize[d] for d = 1:3)
 
+# Internal utility to launch a function mapped on an index_map
+struct MappedFunction{F, M} <: Function
+    func::F
+    index_map::M
+end
+
 # Support for 1D
 heuristic_workgroup(Wx) = min(Wx, 256)
 
@@ -230,6 +239,7 @@ the architecture `arch`.
                                   location = nothing,
                                   active_cells_map = nothing)
 
+
     if !isnothing(active_cells_map) # everything else is irrelevant
         workgroup = min(length(active_cells_map), 256)
         worksize = length(active_cells_map)
@@ -241,9 +251,17 @@ the architecture `arch`.
 
     dev = Architectures.device(arch)
     loop = kernel!(dev, workgroup, worksize)
+
+    # Map out the function to use active_cells_map as an index map
+    if !isnothing(active_cells_map)
+        func = MappedFunction(loop.f, active_cells_map)
+        loop = Kernel{get_kernel_parameters(loop)..., typeof(func)}(dev, func)
+    end
+
     return loop, worksize
 end
 
+@inline get_kernel_parameters(::Kernel{Dev, B, W}) where {Dev, B, W} = Dev, B, W
        
 """
     launch!(arch, grid, workspec, kernel!, kernel_args...; kw...)
@@ -315,11 +333,20 @@ end
 #####
 
 # TODO: when offsets are implemented in KA so that we can call `kernel(dev, group, size, offsets)`, remove all of this
-
-using KernelAbstractions: Kernel
+using CUDA: @device_override, blockIdx, threadIdx
 using KernelAbstractions.NDIteration: _Size, StaticSize
 using KernelAbstractions.NDIteration: NDRange
 
+using KernelAbstractions.NDIteration
+using KernelAbstractions: ndrange, workgroupsize
+
+using KernelAbstractions: __iterspace, __groupindex, __dynamic_checkbounds
+using KernelAbstractions: CompilerMetadata
+
+import KernelAbstractions: partition
+import KernelAbstractions: __ndrange, __groupsize
+import KernelAbstractions: __validindex
+
 struct OffsetStaticSize{S} <: _Size
     function OffsetStaticSize{S}() where S
         new{S::Tuple{Vararg}}()
@@ -369,13 +396,6 @@ const OffsetNDRange{N} = NDRange{N, <:StaticSize, <:StaticSize, <:Any, <:KernelO
     return CartesianIndex(nI)
 end
 
-using KernelAbstractions.NDIteration
-using KernelAbstractions: ndrange, workgroupsize
-import KernelAbstractions: partition
-
-using KernelAbstractions: CompilerMetadata
-import KernelAbstractions: __ndrange, __groupsize
-
 @inline __ndrange(::CompilerMetadata{NDRange}) where {NDRange<:OffsetStaticSize}  = CartesianIndices(get(NDRange))
 @inline __groupsize(cm::CompilerMetadata{NDRange}) where {NDRange<:OffsetStaticSize} = size(__ndrange(cm))
 
@@ -413,3 +433,109 @@ function partition(kernel::OffsetKernel, inrange, ingroupsize)
     return iterspace, dynamic
 end
 
+#####
+##### Utilities for Mapped kernels
+#####
+
+struct IndexMap end
+
+const MappedNDRange{N} = NDRange{N, <:StaticSize, <:StaticSize, <:IndexMap, <:AbstractArray} where N
+
+# TODO: maybe don't do this
+# NDRange has been modified to include an index_map in place of workitems.
+# Remember, dynamic kernels are not possible in combination with this extension!!
+# Also, mapped kernels work only with a 1D kernel and a 1D map, it is not possible to launch a ND kernel.
+@inline function expand(ndrange::MappedNDRange, groupidx::CartesianIndex{N}, idx::CartesianIndex{N}) where N
+    nI = ntuple(Val(N)) do I
+        Base.@_inline_meta
+        offsets = workitems(ndrange)
+        stride = size(offsets, I)
+        gidx = groupidx.I[I]
+        ndrange.workitems[(gidx - 1) * stride + idx.I[I]]
+    end
+    return CartesianIndex(nI...)
+end
+
+const MappedKernel{D} = Kernel{D, <:Any, <:Any, <:MappedFunction} where D
+
+# Override the getproperty to make sure we launch the correct function in the kernel
+@inline Base.getproperty(k::MappedKernel, prop::Symbol) = get_mapped_kernel_property(k, Val(prop))
+
+@inline get_mapped_kernel_property(k, ::Val{prop}) where prop = getfield(k, prop)
+@inline get_mapped_kernel_property(k, ::Val{:index_map}) = getfield(getfield(k, :f), :index_map)
+@inline get_mapped_kernel_property(k, ::Val{:f})         = getfield(getfield(k, :f), :func)
+
+Adapt.adapt_structure(to, cm::CompilerMetadata{N, C}) where {N, C} = 
+    CompilerMetadata{N, C}(Adapt.adapt(to, cm.groupindex), 
+                           Adapt.adapt(to, cm.ndrange),
+                           Adapt.adapt(to, cm.iterspace))
+
+Adapt.adapt_structure(to, ndrange::NDRange{N, B, W}) where {N, B, W} = 
+    NDRange{N, B, W}(Adapt.adapt(to, ndrange.blocks), Adapt.adapt(to, ndrange.workitems))
+
+# Extending the partition function to include the index_map in NDRange: note that in this case the 
+# index_map takes the place of the DynamicWorkitems which we assume is not needed in static kernels
+function partition(kernel::MappedKernel, inrange, ingroupsize)
+    static_workgroupsize = workgroupsize(kernel)
+    
+    # Calculate the static NDRange and WorkgroupSize
+    index_map = kernel.index_map
+    range = length(index_map)
+    groupsize = get(static_workgroupsize)
+
+    blocks, groupsize, dynamic = NDIteration.partition(range, groupsize)
+
+    static_blocks = StaticSize{blocks}
+    static_workgroupsize = StaticSize{groupsize} # we might have padded workgroupsize
+    
+    iterspace = NDRange{length(range), static_blocks, static_workgroupsize}(IndexMap(), index_map)
+
+    return iterspace, dynamic
+end
+
+#####
+##### Extend the valid index function to check whether the index is valid in the index map
+#####
+
+const MappedCompilerMetadata = CompilerMetadata{<:StaticSize, <:Any, <:Any, <:Any, <:MappedNDRange}
+
+@inline __linear_ndrange(ctx::MappedCompilerMetadata) = length(__iterspace(ctx).workitems)
+
+# Mapped kernels are always 1D
+Base.@propagate_inbounds function linear_expand(ndrange::MappedNDRange, gidx::Integer, idx::Integer)
+    offsets = workitems(ndrange)
+    stride = size(offsets, 1)
+    return (gidx - 1) * stride + idx
+end
+
+# Mapped kernels are always 1D
+Base.@propagate_inbounds function linear_expand(ndrange::MappedNDRange, groupidx::CartesianIndex{1}, idx::CartesianIndex{1})
+    offsets = workitems(ndrange)
+    stride = size(offsets, 1)
+    gidx = groupidx.I[1]
+    return (gidx - 1) * stride + idx.I[1]
+end
+
+# To check whether the index is valid in the index map, we need to 
+# check whether the linear index is smaller than the size of the index map
+
+# CPU version, the index is passed explicitly
+@inline function __validindex(ctx::MappedCompilerMetadata, idx::CartesianIndex)
+    # Turns this into a noop for code where we can turn of checkbounds of
+    if __dynamic_checkbounds(ctx)
+        index = @inbounds linear_expand(__iterspace(ctx), __groupindex(ctx), idx)
+        return index ≤ __linear_ndrange(ctx)
+    else
+        return true
+    end
+end
+
+# GPU version, the indices are passed implicitly
+@device_override @inline function __validindex(ctx::MappedCompilerMetadata)
+    if __dynamic_checkbounds(ctx)
+        index = @inbounds linear_expand(__iterspace(ctx), blockIdx().x, threadIdx().x)
+        return index ≤ __linear_ndrange(ctx)
+    else
+        return true
+    end
+end
diff --git a/test/test_active_cells_map.jl b/test/test_active_cells_map.jl
index d1533b9451..b00772515d 100644
--- a/test/test_active_cells_map.jl
+++ b/test/test_active_cells_map.jl
@@ -4,7 +4,6 @@ using Oceananigans.Operators: hack_cosd
 using Oceananigans.ImmersedBoundaries: retrieve_surface_active_cells_map, 
                                        retrieve_interior_active_cells_map,
                                        immersed_cell
-
 function Δ_min(grid) 
     Δx_min = minimum_xspacing(grid, Center(), Center(), Center())
     Δy_min = minimum_yspacing(grid, Center(), Center(), Center())
diff --git a/test/test_distributed_hydrostatic_model.jl b/test/test_distributed_hydrostatic_model.jl
index 68ce34424a..c1a85f0135 100644
--- a/test/test_distributed_hydrostatic_model.jl
+++ b/test/test_distributed_hydrostatic_model.jl
@@ -81,7 +81,7 @@ function rotation_with_shear_test(grid, closure=nothing)
 
     return model
 end
-
+  
 Nx = 32
 Ny = 32 
 
diff --git a/test/utils_for_runtests.jl b/test/utils_for_runtests.jl
index 36fb02ce6a..04d5e67c8a 100644
--- a/test/utils_for_runtests.jl
+++ b/test/utils_for_runtests.jl
@@ -1,5 +1,6 @@
 using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, RungeKutta3TimeStepper, update_state!
-using Oceananigans.DistributedComputations: Distributed, Partition, child_architecture, Fractional, Equal
+using Oceananigans.DistributedComputations: Distributed, Partition, child_architecture, Fractional, Equal, all_reduce
+using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization
 
 import Oceananigans.Fields: interior
 
@@ -201,6 +202,52 @@ function run_script(replace_strings, script_name, script_filepath, module_suffix
     return true
 end
 
+function Δ_min(grid) 
+    Δx_min = minimum_xspacing(grid, Center(), Center(), Center())
+    Δy_min = minimum_yspacing(grid, Center(), Center(), Center())
+    return min(Δx_min, Δy_min)
+end
+
+@inline Gaussian(x, y, L) = exp(-(x^2 + y^2) / L^2)
+
+function solid_body_rotation_test(grid)
+
+    free_surface = SplitExplicitFreeSurface(grid; substeps = 10, gravitational_acceleration = 1)
+    coriolis     = HydrostaticSphericalCoriolis(rotation_rate = 1)
+    closure      = VerticalScalarDiffusivity(VerticallyImplicitTimeDiscretization(), ν = 0.01, κ = 0.01)
+
+    model = HydrostaticFreeSurfaceModel(; grid,
+                                        momentum_advection = VectorInvariant(),
+                                        free_surface = free_surface,
+                                        coriolis = coriolis,
+                                        tracers = :c,
+                                        tracer_advection = WENO(),
+                                        buoyancy = nothing,
+                                        closure)
+
+    g = model.free_surface.gravitational_acceleration
+    R = grid.radius
+    Ω = model.coriolis.rotation_rate
+
+    uᵢ(λ, φ, z) = 0.1 * cosd(φ) * sind(λ)
+    ηᵢ(λ, φ, z) = (R * Ω * 0.1 + 0.1^2 / 2) * sind(φ)^2 / g * sind(λ)
+
+    # Gaussian leads to values with O(1e-60),
+    # too small for repetible testing. We cap it at 0.1
+    cᵢ(λ, φ, z) = max(Gaussian(λ, φ - 5, 10), 0.1)
+    vᵢ(λ, φ, z) = 0.1
+
+    set!(model, u=uᵢ, η=ηᵢ, c=cᵢ)
+
+    @show Δt_local = 0.1 * Δ_min(grid) / sqrt(g * grid.Lz) 
+    @show Δt = all_reduce(min, Δt_local, architecture(grid))
+
+    simulation = Simulation(model; Δt, stop_iteration = 10)
+    run!(simulation)
+
+    return merge(model.velocities, model.tracers, (; η = model.free_surface.η))
+end
+
 #####
 ##### Boundary condition utils
 #####