support for active_cells_map in kernels

[simone-silvestri]

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)

[simone-silvestri]

Probably this is very hacky, a better way to go about it would be to add a map to the Kernel

[glwagner]

Suggested change

	imap :: M
	index_map :: M

[glwagner]

index_map?

[glwagner]

Suggested change

	imap = on_architecture(arch, [(i, j, k) for i in 1:3, j in 1:3, k in 1:2])
	imap = on_architecture(arch, [(i, j, k) for i in 1:3, j in 1:3, k in 1:3])

[simone-silvestri]

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