GPU

Ground-state calculations on a GPU is an experimental feature at the moment with some limitations:

  • only PW-mode

  • it has only been implemented in the new GPAW

  • 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