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]

doesn't it make more sense to put this in the first if-statement in this function?

[simone-silvestri]

we need to build the kernel first to be able to retrieve the function

[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])