GPU¶

Ground-state calculations on a GPU is an new feature with some limitations:

only PW-mode
it has only been implemented in the new GPAW code
only parallelization over k-points

See gpaw/test/gpu/test_pw.py for an example.

Tip

>>> import numpy as np
>>> from gpaw.gpu import cupy as cp
>>> a_cpu = np.zeros(...)
>>> a_gpu = cp.asarray(a_cpu)  # from CPU to GPU
>>> b_cpu = a_gpu.get()  # from GPU to CPU

The gpaw.gpu module¶

gpaw.gpu.cupy¶: cupy module (or gpaw.gpu.cpupy if cupy is not available)

gpaw.gpu.cupyx¶: cupyx module (or gpaw.gpu.cpupyx if cupyx is not available)

gpaw.gpu.cupy_is_fake = True¶: True if cupy has been replaced by gpaw.gpu.cpupy

gpaw.gpu.is_hip = False¶: True if we are using HIP

gpaw.gpu.as_np(array)[source]¶

Transfer array to CPU (if not already there).

Parameters:: array (numpy.ndarray | gpaw.gpu.cpupy.ndarray) – Numpy or CuPy array.

gpaw.gpu.as_xp(array, xp)[source]¶

Transfer array to CPU or GPU (if not already there).

Parameters:

array – Numpy or CuPy array.
xp – numpy or cupy.

gpaw.gpu.cupy_eigh(a, UPLO)[source]¶

Wrapper for eigh().

HIP-GPU version is too slow for now so we do it on the CPU.

Fake cupy library¶

The implementation uses cupy. In the code, we don’t do import cupy as cp. Instead we use from gpaw.gpu import cupy as cp. This allows us to use a fake cupy implementation so that we can run GPAW’s cupy code without having a physical GPU. To enable the fake cupy module, do:

GPAW_CPUPY=1 python ...

This allows users without a GPU to find out if their code interferes with the GPU implementation, simply by running the tests.

CuPy enabled container objects¶

The following objects:

can have their data (.data attribute) stored in a cupy.ndarray array instead of, as normal, a numpy.ndarray array. In additions, these objects now have an xp attribute that can be numpy or cupy.

Also, the AtomCenteredFunctions object can do its operations on the GPU.

GPU-aware MPI¶

Use a GPU-aware MPI implementation and set the GPAW_GPU when compiling GPAW’s C-extension.

GPAW_GPU¶: Add support for passing cupy.ndarray objects to MPI