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

[simone-silvestri]

This PR might be able to be merged, I am just not sure why the OceananigansAMDGPUExt extension is not precompiling

[simone-silvestri]

I guess we can merge this because it seems to work. It complicates kernel_launching.jl, but it simplifies the rest of the code

[glwagner]

how about benchmarks

[simone-silvestri]

Should not play a role in performance, we can check the immersed benchmarks, which are the only ones affected.

[simone-silvestri]

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