GPU Support

Overview

From v1.12.0, PyLops supports computations on GPUs powered by CuPy (cupy-cudaXX>=8.1.0) and cuSignal (cusignal>=0.16.0). They must be installed before PyLops is installed.

Note

Set environment variables CUPY_PYLOPS=0 and/or CUSIGNAL_PYLOPS=0 to force PyLops to ignore cupy and cusignal backends. This can be also used if a previous version of cupy or cusignal is installed in your system, otherwise you will get an error when importing PyLops.

Apart from a few exceptions, all operators and solvers in PyLops can seamlessly work with numpy arrays on CPU as well as with cupy arrays on GPU. Users do simply need to consistently create operators and provide data vectors to the solvers, e.g., when using pylops.MatrixMult the input matrix must be a cupy array if the data provided to a solver is also cupy array.

Warning

Some pylops.LinearOperator methods are currently on GPU:

Warning

Some solvers are currently not available on GPU:

  • pylops.optimization.sparsity.SPGL1

Example

Finally, let’s briefly look at an example. First we write a code snippet using numpy arrays which PyLops will run on your CPU:

ny, nx = 400, 400
G = np.random.normal(0, 1, (ny, nx)).astype(np.float32)
x = np.ones(nx, dtype=np.float32)

Gop = MatrixMult(G, dtype='float32')
y = Gop * x
xest = Gop / y

Now we write a code snippet using cupy arrays which PyLops will run on your GPU:

ny, nx = 400, 400
G = cp.random.normal(0, 1, (ny, nx)).astype(np.float32)
x = cp.ones(nx, dtype=np.float32)

Gop = MatrixMult(G, dtype='float32')
y = Gop * x
xest = Gop / y

The code is almost unchanged apart from the fact that we now use cupy arrays, PyLops will figure this out!

Note

The CuPy backend is in active development, with many examples not yet in the docs. You can find many other examples from the PyLops Notebooks repository.