Source code for gpaw.gpu

from __future__ import annotations
import contextlib
from time import time
from typing import TYPE_CHECKING
from types import ModuleType

import numpy as np

cupy_is_fake = True
"""True if :mod:`cupy` has been replaced by ``gpaw.gpu.cpupy``"""

is_hip = False
"""True if we are using HIP"""

if TYPE_CHECKING:
    import gpaw.gpu.cpupy as cupy
    import gpaw.gpu.cpupyx as cupyx
else:
    try:
        import cupy
        import cupyx
        from cupy.cuda import runtime
        is_hip = runtime.is_hip
        cupy_is_fake = False
    except ImportError:
        import gpaw.gpu.cpupy as cupy
        import gpaw.gpu.cpupyx as cupyx

__all__ = ['cupy', 'cupyx', 'as_xp', 'as_np', 'synchronize']


def synchronize():
    if not cupy_is_fake:
        cupy.cuda.get_current_stream().synchronize()


def setup():
    if not cupy_is_fake:
        # select GPU device (round-robin based on MPI rank)
        # if not set, all MPI ranks will use the same default device
        from gpaw.mpi import rank
        device_id = rank % cupy.cuda.runtime.getDeviceCount()
        cupy.cuda.runtime.setDevice(device_id)
        # initialise C parameters and memory buffers
        import _gpaw
        _gpaw.gpaw_gpu_init()


[docs]def as_np(array: np.ndarray | cupy.ndarray) -> np.ndarray: """Transfer array to CPU (if not already there). Parameters ========== array: Numpy or CuPy array. """ if isinstance(array, np.ndarray): return array return cupy.asnumpy(array)
[docs]def as_xp(array, xp): """Transfer array to CPU or GPU (if not already there). Parameters ========== array: Numpy or CuPy array. xp: :mod:`numpy` or :mod:`cupy`. """ if xp is np: if isinstance(array, np.ndarray): return array return cupy.asnumpy(array) if isinstance(array, np.ndarray): return cupy.asarray(array) 1 / 0 return array
def einsum(subscripts, *operands, out): if isinstance(out, np.ndarray): np.einsum(subscripts, *operands, out=out) else: out[:] = cupy.einsum(subscripts, *operands)
[docs]def cupy_eigh(a: cupy.ndarray, UPLO: str) -> tuple[cupy.ndarray, cupy.ndarray]: """Wrapper for ``eigh()``. HIP-GPU version is too slow for now so we do it on the CPU. """ from scipy.linalg import eigh if not is_hip: return cupy.linalg.eigh(a, UPLO=UPLO) eigs, evals = eigh(cupy.asnumpy(a), lower=(UPLO == 'L'), check_finite=False) return cupy.asarray(eigs), cupy.asarray(evals)
class XP: """Class for adding xp attribute (numpy or cupy). Also implements pickling which will not work out of the box because a module can't be pickled. """ def __init__(self, xp: ModuleType): self.xp = xp def __getstate__(self): state = self.__dict__.copy() assert self.xp is np del state['xp'] return state def __setstate__(self, state): state['xp'] = np self.__dict__.update(state) @contextlib.contextmanager def T(): t1 = time() yield synchronize() t2 = time() print(f'{(t2 - t1) * 1e9:_.3f} ns')