search

CliMA / Oceananigans.jl

🌊 Julia software for fast, friendly, flexible, ocean-flavored fluid dynamics on CPUs and GPUs
https://clima.github.io/OceananigansDocumentation/stable
MIT License
999 stars 196 forks source link

Improve active cells map kernels #3918

Open simone-silvestri opened 2 weeks ago

simone-silvestri commented 2 weeks ago

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: https://github.com/CliMA/Oceananigans.jl/blob/9111a8f3aa3742c57422079d78692f6114e33ea1/src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl#L145-L155

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 commented 2 weeks ago

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

https://github.com/CliMA/Oceananigans.jl/blob/9111a8f3aa3742c57422079d78692f6114e33ea1/src/Utils/kernel_launching.jl#L357

glwagner commented 2 weeks ago

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 commented 2 weeks ago

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 commented 2 weeks ago

Why is that, can you elaborate?

simone-silvestri commented 2 weeks ago

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 commented 2 weeks ago

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

That clarifies, thanks.