Repository: CliMA/Oceananigans.jl
Pull request #3920: support for active_cells_map in kernels
Author: login: simone-silvestri
Labels: name: cleanup 🧹
description: Paying off technical debt
name: 🚨 DO NOT MERGE 🚨
description: IN BIG BOLD RED CAPS WITH FLASHING SIRENS
name: benchmark performance
description: runs preconfigured benchamarks and spits out timing
Created: 2024-11-12T15:44:28Z | Merged: 2025-04-29T15:02:09Z

I am not sure how hacky this is or if it will work throughout the board, but this is supposed to achieve what was discussed in #3918:

for example, the current script:
```
using Oceananigans
grid = RectilinearGrid(size = (3, 3, 3), extent = (1, 1, 1))

using Oceananigans.Utils
using KernelAbstractions: @kernel, @index

@kernel function _test_indices(a)
    i, j, k = @index(Global, NTuple)
    @show a[i, j, k]
end

a = [(i, j, k) for i in 1:3, j in 1:3, k in 1:3]
```

returns



```
julia> launch!(CPU(), grid, :xyz, _test_indices, a)
a[i, j, k] = (1, 1, 1)
a[i, j, k] = (2, 1, 1)
a[i, j, k] = (3, 1, 1)
a[i, j, k] = (1, 2, 1)
a[i, j, k] = (2, 2, 1)
a[i, j, k] = (3, 2, 1)
a[i, j, k] = (1, 3, 1)
a[i, j, k] = (2, 3, 1)
a[i, j, k] = (3, 3, 1)
a[i, j, k] = (1, 1, 2)
a[i, j, k] = (2, 1, 2)
a[i, j, k] = (3, 1, 2)
a[i, j, k] = (1, 2, 2)
a[i, j, k] = (2, 2, 2)
a[i, j, k] = (3, 2, 2)
a[i, j, k] = (1, 3, 2)
a[i, j, k] = (2, 3, 2)
a[i, j, k] = (3, 3, 2)
a[i, j, k] = (1, 1, 3)
a[i, j, k] = (2, 1, 3)
a[i, j, k] = (3, 1, 3)
a[i, j, k] = (1, 2, 3)
a[i, j, k] = (2, 2, 3)
a[i, j, k] = (3, 2, 3)
a[i, j, k] = (1, 3, 3)
a[i, j, k] = (2, 3, 3)
a[i, j, k] = (3, 3, 3)

julia> launch!(CPU(), grid, :xyz, _test_indices, a; active_cells_map = [(1, 2, 3), (3, 2, 1)])
a[i, j, k] = (1, 2, 3)
a[i, j, k] = (3, 2, 1)
```


Diff:
diff --git a/ext/OceananigansAMDGPUExt.jl b/ext/OceananigansAMDGPUExt.jl
index e034ee7b4f0..e89791da44d 100644
--- a/ext/OceananigansAMDGPUExt.jl
+++ b/ext/OceananigansAMDGPUExt.jl
@@ -2,6 +2,9 @@ module OceananigansAMDGPUExt
 
 using AMDGPU
 using Oceananigans
+using Oceananigans.Utils: linear_expand, __linear_ndrange, MappedCompilerMetadata
+using KernelAbstractions: __dynamic_checkbounds, __iterspace
+import KernelAbstractions: __validindex
 
 import Oceananigans.Architectures:
     architecture,
@@ -9,7 +12,7 @@ import Oceananigans.Architectures:
     on_architecture
 
 const ROCGPU = GPU{<:AMDGPU.ROCBackend}
-ROCGPU() = GPU(Metal.AMDGPU.ROCBackend())
+ROCGPU() = GPU(AMDGPU.ROCBackend())
 
 architecture(::ROCArray) = ROCGPU()
 Base.summary(::ROCGPU) = "ROCGPU"
@@ -27,4 +30,13 @@ on_architecture(::ROCGPU, a::StepRangeLen) = a
 @inline convert_to_device(::ROCGPU, args) = AMDGPU.rocconvert(args)
 @inline convert_to_device(::ROCGPU, args::Tuple) = map(AMDGPU.rocconvert, args)
 
+AMDGPU.Device.@device_override @inline function __validindex(ctx::MappedCompilerMetadata)
+    if __dynamic_checkbounds(ctx)
+        I = @inbounds linear_expand(__iterspace(ctx), AMDGPU.Device.blockIdx().x, AMDGPU.Device.threadIdx().x)
+        return I in __linear_ndrange(ctx)
+    else
+        return true
+    end
+end
+
 end # module
diff --git a/ext/OceananigansMetalExt.jl b/ext/OceananigansMetalExt.jl
index 7c06eb55a84..1c6a2634ffc 100644
--- a/ext/OceananigansMetalExt.jl
+++ b/ext/OceananigansMetalExt.jl
@@ -3,6 +3,11 @@ module OceananigansMetalExt
 using Metal
 using Oceananigans
 
+using Metal: thread_position_in_threadgroup_1d, threadgroup_position_in_grid_1d
+using Oceananigans.Utils: linear_expand, __linear_ndrange, MappedCompilerMetadata
+using KernelAbstractions: __dynamic_checkbounds, __iterspace
+import KernelAbstractions: __validindex
+
 import Oceananigans.Architectures:
     architecture,
     convert_to_device,
@@ -35,5 +40,15 @@ end
 @inline convert_to_device(::MetalGPU, args) = Metal.mtlconvert(args)
 @inline convert_to_device(::MetalGPU, args::Tuple) = map(Metal.mtlconvert, args)
 
+Metal.@device_override @inline function __validindex(ctx::MappedCompilerMetadata)
+    if __dynamic_checkbounds(ctx)
+        I = @inbounds linear_expand(__iterspace(ctx), threadgroup_position_in_grid_1d(),
+                                thread_position_in_threadgroup_1d())
+        return I in __linear_ndrange(ctx)
+    else
+        return true
+    end
+end
+
 end # module
 
diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl
index a578afec9fe..2e2f3f2b7d2 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 _compute_barotropic_mode!(U̅, V̅, grid, ::Nothing, u, v, η)
+@kernel function _compute_barotropic_mode!(U̅, V̅, grid, u, v, η)
     i, j  = @index(Global, NTuple)
-    integrate_barotropic_mode!(U̅, V̅, i, j, grid, u, v, η)
-end
-
-@kernel function _compute_barotropic_mode!(U̅, V̅, grid, active_cells_map, u, v, η)
-    idx = @index(Global, Linear)
-    i, j = linear_index_to_tuple(idx, active_cells_map)
-    integrate_barotropic_mode!(U̅, V̅, i, j, grid, u, v, η)
-end
-
-@inline function integrate_barotropic_mode!(U̅, V̅, i, j, grid, u, v, η)
     k_top  = size(grid, 3) + 1
 
     hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid)
@@ -32,17 +22,15 @@ 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
 
 # Note: this function is also used during initialization
 function compute_barotropic_mode!(U̅, V̅, grid, u, v, η)
-    active_columns = get_active_column_map(grid) # may be nothing
+    active_cells_map = get_active_column_map(grid) # may be nothing
 
     launch!(architecture(grid), grid, :xy,
             _compute_barotropic_mode!,
-            U̅, V̅, grid, active_columns, u, v, η; active_cells_map=active_columns)
+            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 fe0392c2b2d..006ad6804a5 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 = 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 = 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
@@ -114,9 +99,9 @@ end
     GU⁻ = timestepper.G⁻.U
     GV⁻ = timestepper.G⁻.V
 
-    active_cells_map = get_active_column_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)
+    active_cells_map = get_active_column_map(grid)    
+    launch!(architecture(grid), grid, :xy, _compute_integrated_rk3_tendencies!, 
+            GUⁿ, GVⁿ, GU⁻, GV⁻, grid, Guⁿ, Gvⁿ, stage; active_cells_map)
 
     return nothing
 end
diff --git a/src/Models/HydrostaticFreeSurfaceModels/cache_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/cache_hydrostatic_free_surface_tendencies.jl
index a11f3f87ebd..8ada686188a 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/cache_hydrostatic_free_surface_tendencies.jl
+++ b/src/Models/HydrostaticFreeSurfaceModels/cache_hydrostatic_free_surface_tendencies.jl
@@ -4,7 +4,6 @@ using Oceananigans.TimeSteppers: _cache_field_tendencies!
 
 using Oceananigans: prognostic_fields
 using Oceananigans.Grids: AbstractGrid
-using Oceananigans.ImmersedBoundaries: get_active_cells_map
 
 using Oceananigans.Utils: launch!
 
diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl
index 36cc7aa7f28..64bbc530be6 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
@@ -147,14 +146,12 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para
     v_kernel_args = tuple(start_momentum_kernel_args..., v_immersed_bc, end_momentum_kernel_args..., v_forcing)
 
     launch!(arch, grid, kernel_parameters,
-            compute_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid,
-            active_cells_map, u_kernel_args;
-            active_cells_map)
+            compute_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid, 
+            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)
+            compute_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid, 
+            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 = 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 = 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 = linear_index_to_tuple(idx, active_cells_map)
-    @inbounds Gc[i, j, k] = hydrostatic_free_surface_tracer_tendency(i, j, k, grid, args...)
-end
diff --git a/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl b/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl
index 8091e85b3fa..a807be4d1ff 100644
--- a/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl
+++ b/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl
@@ -102,16 +102,16 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti
                           hydrostatic_pressure, clock, forcings.w)
 
     exclude_periphery = true
-    launch!(arch, grid, kernel_parameters, compute_Gu!,
-            tendencies.u, grid, active_cells_map, u_kernel_args;
+    launch!(arch, grid, kernel_parameters, compute_Gu!, 
+            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;
+    launch!(arch, grid, kernel_parameters, compute_Gv!, 
+            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;
+    launch!(arch, grid, kernel_parameters, compute_Gw!, 
+            tendencies.w, grid, w_kernel_args;
             active_cells_map, exclude_periphery)
 
     start_tracer_kernel_args = (advection, closure)
@@ -130,8 +130,8 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti
                      end_tracer_kernel_args...,
                      clock, forcing)
 
-        launch!(arch, grid, kernel_parameters, compute_Gc!,
-                c_tendency, grid, active_cells_map, args;
+        launch!(arch, grid, kernel_parameters, compute_Gc!, 
+                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 = 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 = 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 = 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 = 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 d07f1e3d2ba..1dee5457d00 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,22 @@ 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)
 
+"""
+    MappedFunction(func, index_map)
+
+A `MappedFunction` is a wrapper around a function `func` of a kernel that is mapped over an `index_map`.
+The `index_map` is a one-dimensional `AbstractArray` where the elements are tuple of indices `(i, j, k, ....)`.
+
+A kernel launched over a `MappedFunction` **needs** to be launched with a one-dimensional **static** workgroup and worksize.
+If using `launch!` with a non-nothing `active_cells_map` keyword argument, the kernel function will be automatically wrapped 
+in a `MappedFunction` with `index_map = active_cells_map` and the resulting kernel will be launched with a 
+one-dimensional workgroup and worksize equal  to the length of the `active_cells_map`.
+"""
+struct MappedFunction{F, M} <: Function
+    func :: F
+    index_map :: M
+end
+
 # Support for 1D
 heuristic_workgroup(Wx) = min(Wx, 256)
 
@@ -230,6 +249,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)
@@ -239,12 +259,23 @@ the architecture `arch`.
         workgroup, worksize = work_layout(grid, workspec, reduced_dimensions)
     end
 
-    dev = Architectures.device(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)
+        loop = mapped_kernel(loop, dev, active_cells_map)
+    end
+
     return loop, worksize
 end
 
-
+@inline function mapped_kernel(kernel::Kernel{Dev, B, W}, dev, map) where {Dev, B, W}
+    f  = kernel.f
+    mf = MappedFunction(f, map)
+    return Kernel{Dev, B, W, typeof(mf)}(dev, mf)
+end
+       
 """
     launch!(arch, grid, workspec, kernel!, kernel_args...; kw...)
 
@@ -315,11 +346,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}}()
@@ -362,13 +402,6 @@ const OffsetNDRange{N, S} = NDRange{N, <:StaticSize, <:StaticSize, <:Any, <:Offs
     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))
 
@@ -406,3 +439,109 @@ function partition(kernel::OffsetKernel, inrange, ingroupsize)
     return iterspace, dynamic
 end
 
+#####
+##### Utilities for Mapped kernels
+#####
+
+struct IndexMap end
+
+const MappedNDRange{N, B, W} = NDRange{N, B, W, <:IndexMap, <:AbstractArray} where {N, B<:StaticSize, W<:StaticSize}
+
+# 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, ndrange::MappedNDRange{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{N, C} = CompilerMetadata{N, C, <:Any, <:Any, <:MappedNDRange} where {N<:StaticSize, C}
+
+Adapt.adapt_structure(to, cm::MappedCompilerMetadata{N, C}) where {N, C} = 
+    CompilerMetadata{N, C}(Adapt.adapt(to, cm.groupindex), 
+                           Adapt.adapt(to, cm.ndrange),
+                           Adapt.adapt(to, cm.iterspace))
+
+@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
+CUDA.@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


Reviews:
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2024-11-12T18:58:38Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2024-11-12T19:00:34Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2024-11-12T19:00:42Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2024-11-12T19:01:01Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2024-11-12T19:02:41Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2024-11-12T20:48:05Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2024-11-13T22:43:45Z
user: login: simone-silvestri
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2024-11-28T09:22:12Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2024-12-16T16:12:11Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2024-12-16T16:12:18Z
user: login: glwagner
body: It's much better! Possibly, @vchuravy may want to look at the implementation
state: APPROVED
author_association: MEMBER
submitted_at: 2024-12-16T16:13:29Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-26T13:31:36Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-26T13:33:04Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-26T13:33:44Z
user: login: glwagner
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-26T13:35:11Z
user: login: simone-silvestri
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-26T14:15:12Z
user: login: simone-silvestri
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-26T14:16:57Z

Review comments:
diff_hunk: @@ -413,3 +445,58 @@ function partition(kernel::OffsetKernel, inrange, ingroupsize)
     return iterspace, dynamic
 end
 
+#####
+##### Utilities for Mapped kernels
+#####
+
+struct IndexMap{M}
+    imap :: M
path: src/Utils/kernel_launching.jl
user: login: glwagner
body: ```suggestion
    index_map :: M
```
created_at: 2024-11-12T18:58:37Z
author_association: MEMBER
diff_hunk: @@ -80,6 +82,28 @@ end
 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
+    f::F
+    imap::M
path: src/Utils/kernel_launching.jl
user: login: glwagner
body: index_map?
created_at: 2024-11-12T19:00:34Z
author_association: MEMBER
diff_hunk: @@ -80,6 +82,28 @@ end
 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
+    f::F
path: src/Utils/kernel_launching.jl
user: login: glwagner
body: ```suggestion
    func::F
```
created_at: 2024-11-12T19:00:42Z
author_association: MEMBER
diff_hunk: @@ -80,6 +82,28 @@ end
 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
+    f::F
+    imap::M
+end
+
+Adapt.adapt_structure(to, m::MappedFunction) = 
+    MappedFunction(Adapt.adapt(to, m.f), Adapt.adapt(to, m.imap))
+
+@inline function (m::MappedFunction)(_ctx_)  
+    m.f(_ctx_)
+    return nothing
+end
path: src/Utils/kernel_launching.jl
user: login: glwagner
body: isn't this case covered by the case with `args...`?
created_at: 2024-11-12T19:01:01Z
author_association: MEMBER
diff_hunk: @@ -238,9 +262,20 @@ 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)
+        param = get_kernel_parameters(loop)
+        M     = typeof(func)
+        loop  = Kernel{param..., M}(dev, func)
+    end
path: src/Utils/kernel_launching.jl
user: login: glwagner
body: doesn't it make more sense to put this in the first if-statement in this function?
created_at: 2024-11-12T19:02:41Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,18 @@
+using Revise
+using Oceananigans
+using Oceananigans.Utils
+
+using KernelAbstractions: @index, @kernel
+
+arch = CPU()
+
+@kernel function _test_indices(a)
+    i, j, k = @index(Global, NTuple)
+    a[i, j, k] = i + j + k
+end
+
+grid  = RectilinearGrid(arch, size = (3, 3, 3), extent = (1, 1, 1))
+array = zeros(arch, 3, 3, 3)
+imap  = on_architecture(arch, [(i, j, k) for i in 1:3, j in 1:3, k in 1:2])
path: test.jl
user: login: glwagner
body: ```suggestion
imap  = on_architecture(arch, [(i, j, k) for i in 1:3, j in 1:3, k in 1:3])
```
created_at: 2024-11-12T20:48:05Z
author_association: MEMBER
diff_hunk: @@ -440,54 +430,85 @@ end
 ##### Utilities for Mapped kernels
 #####
 
-struct IndexMap{M}
-    index_map :: M
-end
-
-Adapt.adapt_structure(to, m::IndexMap) = IndexMap(Adapt.adapt(to, m.index_map))
+struct IndexMap end
 
-const MappedNDRange{N} = NDRange{N, <:StaticSize, <:StaticSize, <:Any, <:IndexMap} where N
+const MappedNDRange{N} = NDRange{N, <:StaticSize, <:StaticSize, <:IndexMap, <:AbstractArray} where N
 
 # NDRange has been modified to include an index_map in place of workitems.
 # Remember, dynamic offset kernels are not possible 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.
 # TODO: maybe don't do this
-@inline function expand(ndrange::MappedNDRange, groupidx::CartesianIndex, idx::CartesianIndex) 
-    Base.@_inline_meta
-    offsets = workitems(ndrange)
-    stride = size(offsets, 1)
-    gidx = groupidx.I[1]
-    tI = (gidx - 1) * stride + idx.I[1]
-    nI = ndrange.workitems.index_map[tI]
-    return CartesianIndex(nI)
+@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 get the correct properties
-@inline getproperty(k::MappedKernel, prop::Symbol) = get_mapped_property(k, Val(prop))
+# 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)
 
-@inline get_mapped_property(k, ::Val{:index_map}) = k.f.index_map
-@inline get_mapped_property(k, ::Val{:func}) = 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 offsets in NDRange: note that in this case the 
 # offsets take 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 = getproperty(kernel, :index_map)
+    index_map = kernel.index_map
     range = length(index_map)
-    arch  = Oceananigans.Architectures.architecture(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
     
-    index_map = Oceananigans.Architectures.convert_args(arch, index_map)
-    iterspace = NDRange{length(range), static_blocks, static_workgroupsize}(blocks, IndexMap(index_map))
+    iterspace = NDRange{length(range), static_blocks, static_workgroupsize}(IndexMap(), index_map)
 
     return iterspace, dynamic
-end
\ No newline at end of file
+end
+
+using KernelAbstractions: CompilerMetadata, NoDynamicCheck
+using CUDA: CUDABackend
+
+import KernelAbstractions: mkcontext, __dynamic_checkbounds, __validindex
+
+# Very dangerous override of mkcontext which will not work if we are not 
+# carefull with making sure that indices are correct when launching a `MappedKernel`
path: src/Utils/kernel_launching.jl
user: login: glwagner
body: ```suggestion
# careful with making sure that indices are correct when launching a `MappedKernel`
```
created_at: 2024-11-13T22:43:44Z
author_association: MEMBER
diff_hunk: @@ -238,9 +262,20 @@ 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)
+        param = get_kernel_parameters(loop)
+        M     = typeof(func)
+        loop  = Kernel{param..., M}(dev, func)
+    end
path: src/Utils/kernel_launching.jl
user: login: simone-silvestri
body: we need to build the kernel first to be able to retrieve the function
created_at: 2024-11-28T09:22:12Z
author_association: COLLABORATOR
diff_hunk: @@ -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...))
path: src/DistributedComputations/distributed_immersed_boundaries.jl
user: login: glwagner
body: you can also add a constructor for tuples
created_at: 2024-12-16T16:12:11Z
author_association: MEMBER
diff_hunk: @@ -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...))
path: src/DistributedComputations/distributed_immersed_boundaries.jl
user: login: glwagner
body: for `KernelParameters`
created_at: 2024-12-16T16:12:17Z
author_association: MEMBER
diff_hunk: @@ -2,12 +2,15 @@ module OceananigansAMDGPUExt
 
 using AMDGPU
 using Oceananigans
+using KernelAbstractions: __dynamic_checkbounds, __iterspace
 
 import Oceananigans.Architectures:
     architecture,
     convert_to_device,
     on_architecture
 
+using Ocenanigans.Utils: linear_expand, __linear_ndrange, MappedCompilerMetadata
path: ext/OceananigansAMDGPUExt.jl
user: login: glwagner
body: ```suggestion
using Oceananigans.Utils: linear_expand, __linear_ndrange, MappedCompilerMetadata
```
created_at: 2025-04-26T13:31:36Z
author_association: MEMBER
diff_hunk: @@ -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
path: src/Utils/kernel_launching.jl
user: login: glwagner
body: ```suggestion
    func :: F
    index_map :: M
end
```
created_at: 2025-04-26T13:33:03Z
author_association: MEMBER
diff_hunk: @@ -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
path: src/Utils/kernel_launching.jl
user: login: glwagner
body: Might make sense to write a longer comment here. `index_map` is the key feature that is utilized by active_cells_map, right?
created_at: 2025-04-26T13:33:44Z
author_association: MEMBER
diff_hunk: @@ -241,10 +251,18 @@ 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
path: src/Utils/kernel_launching.jl
user: login: glwagner
body: This is poorly named because this file includes a type called `KernelParameters`.

Can you define a constructor for `Kernel` that takes in `loop`, `dev`, `func`?
created_at: 2025-04-26T13:35:11Z
author_association: MEMBER
diff_hunk: @@ -2,12 +2,15 @@ module OceananigansAMDGPUExt
 
 using AMDGPU
 using Oceananigans
+using KernelAbstractions: __dynamic_checkbounds, __iterspace
 
 import Oceananigans.Architectures:
     architecture,
     convert_to_device,
     on_architecture
 
+using Ocenanigans.Utils: linear_expand, __linear_ndrange, MappedCompilerMetadata
path: ext/OceananigansAMDGPUExt.jl
user: login: simone-silvestri
body: that's why 😅 
created_at: 2025-04-26T14:15:12Z
author_association: COLLABORATOR
diff_hunk: @@ -241,10 +251,18 @@ 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
path: src/Utils/kernel_launching.jl
user: login: simone-silvestri
body: We can do it, It would have to take in `{Dev, B, W}` however, because these are parametric types that do not correspond to a field of `Kernel`.
created_at: 2025-04-26T14:16:57Z
author_association: COLLABORATOR

PR comments:
user: login: simone-silvestri
created_at: 2024-11-12T17:31:22Z
body: Probably this is very hacky, a better way to go about it would be to add a map to the `Kernel`
author_association: COLLABORATOR
user: login: simone-silvestri
created_at: 2024-12-11T12:29:54Z
body: This solution seems to work.

PROS: makes the codebase much simpler and allows to use of `active_cells_map` in many more kernels with very little change
CONS: The solution is a bit convoluted in the `kernel_launching.jl` file

I am in favor of merging, it will make future low-hanging fruit optimization (using active cells map in `compute_diffusivities`, the catke TKE time stepping, and the split explicit free surface as well as kernels in other packages like flux computation in ClimaOcean) much more straightforward.

but I am happy to see what people think

author_association: COLLABORATOR
user: login: simone-silvestri
created_at: 2025-01-27T09:12:28Z
body: I don't think this is correct, or at least, it works but it is very ad-hoc because it is adding a CUDA-specific method (with the cuda-specific `@device_override` macro which would be, for example, different for `Metal` and `AMDGPU`). This might hamper our possibility of making the code compatible with other backends so I would prefer not to merge it.

@vchuravy, what would be the best way to add something like this to KA? Would adding a new field to the `Kernel` struct (for example `index_map`) work? 
author_association: COLLABORATOR
user: login: simone-silvestri
created_at: 2025-01-27T09:18:28Z
body: Or, maybe I don't completely understand the `@device_override` macro. Would specializing the functions with 
```julia
CUDA.@device_override function
    cuda specific body...
end

Metal.@device_override function
    metal specific body...
end
```
work to specialize the function for different backends?
author_association: COLLABORATOR
user: login: glwagner
created_at: 2025-03-31T20:20:59Z
body: lots of conflicts here unfortunately. Also it might be nice to get @vchuravy's thoughts here. Since a better implementation here is all the more useful someday when KA is ready
author_association: MEMBER
user: login: simone-silvestri
created_at: 2025-04-26T11:11:00Z
body: This PR might be able to be merged, I am just not sure why the `OceananigansAMDGPUExt` extension is not precompiling
author_association: COLLABORATOR
user: login: simone-silvestri
created_at: 2025-04-29T05:26:11Z
body: I guess we can merge this because it seems to work. It complicates `kernel_launching.jl`, but it simplifies the rest of the code
author_association: COLLABORATOR
user: login: glwagner
created_at: 2025-04-29T05:48:39Z
body: how about benchmarks
author_association: MEMBER
user: login: simone-silvestri
created_at: 2025-04-29T07:04:48Z
body: Should not play a role in performance, we can check the immersed benchmarks, which are the only ones affected.
author_association: COLLABORATOR
user: login: simone-silvestri
created_at: 2025-04-29T08:29:35Z
body: In the `diff_immersed_output.txt` there seems to be something wrong in the data scraped by the main build, it looks like it is taking the previous main before the performance improvements of using the reciprocals, maybe updating the branch might fix this issue. However, it seems that this PR does not affect performance looking at the results on this branch and (the correct ones) on main. I ll update the branch to confirm
author_association: COLLABORATOR