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]

you can also add a constructor for tuples

[glwagner]

for KernelParameters

[glwagner]

It's much better! Possibly, @vchuravy may want to look at the implementation

[simone-silvestri]

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?

[simone-silvestri]

Or, maybe I don't completely understand the @device_override macro. Would specializing the functions with

CUDA.@device_override function
    cuda specific body...
end

Metal.@device_override function
    metal specific body...
end

work to specialize the function for different backends?