Improve active cells map kernels

[simone-silvestri]

Currently, if we build an immersed grid with active_cells_map = true, specific kernels compute only on the map provided by the grid. This is very beneficial, but the implementation is a bit clunky, requiring to pass the active_cells_map and define two different kernels, for example:

Oceananigans.jl/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl

Lines 145 to 155 in 9111a8f

	""" Calculate the right-hand-side of the u-velocity equation. """
	@kernel function compute_Gu!(Gu, grid, ::Nothing, 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

Since the implementation is a bit clunky, we restrict this behavior to the largest performance-critical kernels that are seldom called during time stepping. Otherwise, the code base would explode.

However, launching kernels only on the active cells is a low-hanging fruit for really speeding up the code, and most kernels in Oceananigans could benefit from it. This calls for a more implicit implementation that does not require defining two different kernels. There might be some way to pass a vector of indices to the kernel configuration and launch a one-dimensional kernel that iterates over that map.

That way, we could automatically compute everything on active cells if active_cells_map = true without bloating the code base. We could even avoid masking fields that live on a grid with an active map.

I think @vchuravy mentioned launching kernels on a map some time ago.

This is probably more of an issue for KernelAbstraction than for Oceananigans, but I just want to see what people think.

[glwagner]

Can you implement such a feature like you implemented offsets, eg defining an IndexMap (rather than KernelOffsets) and MappedNDRange instead of OffsetNDRange?

Oceananigans.jl/src/Utils/kernel_launching.jl

Line 357 in 9111a8f

const OffsetNDRange{N} = NDRange{N, <:StaticSize, <:StaticSize, <:Any, <:KernelOffsets} where N

[glwagner]

We could even avoid masking fields that live on a grid with an active map.

I'm not even sure we need masking as is, but agree, it would be nice to consider masking as a convenience for output rather than something we need to do while the simulation runs.

[simone-silvestri]

Can you implement such a feature like you implemented offsets, eg defining an IndexMap (rather than KernelOffsets) and MappedNDRange instead of OffsetNDRange?

We could try, but the map is an AbstractArray, so I suspect it has to be stored in the Kernel and passed as an additional argument to the kernel.

[glwagner]

Why is that, can you elaborate?

[simone-silvestri]

The offsets is passed as a type parameter in the NDRange argument of the kernel here
https://github.com/JuliaGPU/KernelAbstractions.jl/blob/265e5b82a047aca1b7bcf50966d14f540b6ba57e/src/KernelAbstractions.jl#L588

The problem is that an active map is an AbstractArray so we cannot pass it through type parameters, but it should be stored in the Kernel type I think

[glwagner]

https://github.com/JuliaGPU/KernelAbstractions.jl/blob/265e5b82a047aca1b7bcf50966d14f540b6ba57e/src/KernelAbstractions.jl#L588

That clarifies, thanks.