Consider using KernelAbstractions.jl for GPU support

[SouthEndMusic]

Feature type

None

Improvement Description

See https://juliagpu.github.io/KernelAbstractions.jl/stable/kernels/. It provides a syntax to write kernels that can be compiled for either CPU or GPU.

Implementation Description

Benefits of using this:

• It allows for both CPU and GPU parallelizatiation without code duplication
• The transition can be done stepwise in the current code base instead of building a separate ‘Wflow for GPU’ code base from scratch
• You don’t have to align with the SciML data structure standard if you don’t want to
• I have some experience 🙂

Additional Context

No response

[BSchilperoort]

Hi Bart, we were also considering using KernelAbstractions for parallelizing parts of Wflow for the exact reasons you stated (as part of the ESiWACE project). Good to see that you've found this package too! We can try to keep you in the loop about our progress towards this.

[vers-w]

Just adding this as a reminder here (discussed as part of the ESiWACE project on 17 Feb. 2025):

Recommendation is to also check the CPU performance of KernelAbstractions during implementation as we have experienced a decrease in performance when using > 8 threads with Julia's native treading, hence the use of Polyester (> 8 threads). Additionally, we make use of LoopVectorization for local inertial routing. @SouthEndMusic provided the following with some info on the KernelAbstractions CPU backend. And there is a KernelAbstractions issue reporting a slowdown with CPU.