Exploring GPU acceleration

[Dobbermann2]

For my physics bachelor's thesis, I will be exploring the use of GPUs to accelerate simulations in Parcels. This issue will serve as a logbook where I will collect interesting resources, findings and potential issues.

Feel free to comment with suggestions, questions or general discussion!

[Dobbermann2]

GPU Accelerated Array Libraries

Library	Supported GPUs	Notes
CuPy	NVIDIA, AMD (exp.)	Closest drop-in for NumPy
PyTorch	NVIDIA	Optimized for machine learning
JAX	NVIDIA, AMD, TPU	Only CPU support for Windows

[Dobbermann2]

Array Backend Interoperability

Nowadays there are many different Python array libraries such as NumPy, CuPy and JAX. The goal is to make Parcels seamlessly support multiple array libraries without requiring the additional development hours needed to explicitly implement support for each library separately

The suggested way to approach this is by using the Array API Standard Ecosystem. The Array API Standard defines a common interface for array libraries. This allows Parcels to write code that supports any Array API Standard compatible array implementation, giving the user the freedom to switch between array libraries.

A great talk about the Array API Standard by Lucas Colley, A Hitchhiker's Guide to the Array API Standard Ecosystem