Investigations in using GPU for pyresample

[mraspaud]

Triggered by the recent availability of GPUs in computing resources I had access too, I started investigation how feasible it would be to use the GPU to speed up resampling of satellite imagery.

In order for others to see how far we are on this, I thought I would open this issue to have some visibility of the investigations and work that has been done here. Feel free to complement with further investigations in the comments.

Cuproj, transforming coordinates

One requirement to be able to resample is to have the possibility to convert coordinates, as we do with pyproj at the moment. The rapidsai project has a cuproj library, that provide equivalent interface. However they only provide support for epsg:4326 and utm-based projection.

Testing this shows about a factor 100 speed up with the gpu.

print("Numpy/CPU")
import numpy as np
from pyproj.transformer import Transformer

lons, lats = np.meshgrid(np.linspace(-180, 180, 20000), np.linspace(-90, 90, 10000))

print("creating transform")
tr = Transformer.from_crs(4326, 32630)
print("transforming")
x, y = tr.transform(lons, lats)

and

print("Cupy/GPU")
import cupy as cp
from cuproj.transformer import Transformer

lons, lats = cp.meshgrid(cp.linspace(-180, 180, 20000), cp.linspace(-90, 90, 10000))

print("creating transform")
tr = Transformer.from_crs("EPSG:4326", "EPSG:32630")
print("transforming")
x, y = tr.transform(lons.reshape(-1), lats.reshape(-1))
x = x.reshape((20000, 10000))
y = y.reshape((20000, 10000))

outputs repectively (using time)

Numpy/CPU
creating transform
transforming

real    1m31.581s
user    1m31.588s
sys 0m0.644s

and

Cupy/GPU
creating transform
transforming

real    0m0.981s
user    0m1.503s
sys 0m0.142s

KDTree implementation

Cupy seems to have a GPU-optimized kdtree https://docs.cupy.dev/en/latest/reference/generated/cupyx.scipy.spatial.KDTree.html However at the time writing, this has not been released yet and would need manual building of cupy to try it out (which I don't have time for right now).

Gradient search

Cupy has the possibility to define custom kernels, where we could implement the gradient search. However, GPUs are good for doing things for each pixel in parallel, so we might need to implement a pixel-wise version of the algorithm. I haven't test this.

[djhoese]

On mobile right now so can'tfind the links easily, but I believe there are other GPU issues on this repository and/or pykdtree.

[mraspaud]

Indeed, here is a previous issue of yours on the topic of kdtrees https://github.com/pytroll/pyresample/issues/174