context_cupy.py

# copyright ################################# #
# This file is part of the Xobjects Package.  #
# Copyright (c) CERN, 2021.                   #
# ########################################### #

import logging
from typing import Dict, List, Tuple

import numpy as np

from .context import (
    ModuleNotAvailable,
    SourceType,
    XBuffer,
    XContext,
    _concatenate_sources,
    available,
    classes_from_kernels,
    sort_classes,
    sources_from_classes,
    ModuleNotAvailableError,
)
from .linkedarray import BaseLinkedArray
from .specialize_source import specialize_source

log = logging.getLogger(__name__)

try:
    import cupy
    import cupyx.scipy
    import cupyx.scipy.interpolate
    import cupyx.scipy.signal
    import cupyx.scipy.special
    import cupyx.scipy.stats
    from cupyx.scipy import fftpack as cufftp

    _enabled = True
except ImportError:
    log.info("cupy is not installed, ContextCupy will not be available")
    cupy = ModuleNotAvailable(
        message=("cupy is not installed. " "ContextCupy is not available!")
    )
    cufftp = cupy
    _enabled = False

if _enabled:
    # order of base classes matters as it defines which __setitem__ is used
    class LinkedArrayCupy(BaseLinkedArray, cupy.ndarray):
        @classmethod
        def _build_view(cls, a):
            assert len(a.shape) == 1
            return cls(
                shape=a.shape,
                dtype=a.dtype,
                memptr=a.data,
                strides=a.strides,
                order="C",
            )

        def copy(self):
            res = cupy.zeros(shape=self.shape, dtype=self.dtype)
            res[:] = self[:]
            return res

        def _as_cupy(self):
            return cupy.ndarray(
                shape=self.shape,
                dtype=self.dtype,
                memptr=self.data,
                strides=self.strides,
            )

        def _basic_setitem(self, indx, val):
            if hasattr(val, "_as_cupy"):
                val = val._as_cupy()
            if hasattr(indx, "_as_cupy"):
                indx = indx._as_cupy()
            cupy.ndarray.__setitem__(self._as_cupy(), indx, val)

        def __getitem__(self, indx):
            if hasattr(indx, "_as_cupy"):
                indx = indx._as_cupy()
            return cupy.ndarray.__getitem__(self._as_cupy(), indx)

        #        # The following methods are generated by this code
        #        binops = [
        #            'add', 'radd', 'sub', 'rsub', 'mul', 'rmul', 'matmul', 'rmatmul',
        #            'truediv', 'rtruediv', 'floordiv', 'rfloordiv', 'mod', 'rmod',
        #            'pow', 'rpow', 'and_', 'rand', 'or_', 'ror', 'xor', 'rxor',
        #            'lt', 'rlt', 'le', 'rle', 'eq', 'req', 'ne', 'rne', 'ge',
        #            'rge', 'gt', 'rgt', 'rshift', 'rrshift', 'lshift']
        #
        #        for nn in binops:
        #            print(f"""
        #                 def __{nn}__(self, other):
        #                     if hasattr(other, '_as_cupy'):
        #                         other = other._as_cupy()
        #                     return cupy.ndarray.__{nn}__(self._as_cupy(), other)""")

        def __add__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__add__(self._as_cupy(), other)

        def __radd__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__radd__(self._as_cupy(), other)

        def __sub__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__sub__(self._as_cupy(), other)

        def __rsub__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rsub__(self._as_cupy(), other)

        def __mul__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__mul__(self._as_cupy(), other)

        def __rmul__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rmul__(self._as_cupy(), other)

        def __matmul__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__matmul__(self._as_cupy(), other)

        def __rmatmul__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rmatmul__(self._as_cupy(), other)

        def __truediv__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__truediv__(self._as_cupy(), other)

        def __rtruediv__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rtruediv__(self._as_cupy(), other)

        def __floordiv__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__floordiv__(self._as_cupy(), other)

        def __rfloordiv__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rfloordiv__(self._as_cupy(), other)

        def __mod__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__mod__(self._as_cupy(), other)

        def __rmod__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rmod__(self._as_cupy(), other)

        def __pow__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__pow__(self._as_cupy(), other)

        def __rpow__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rpow__(self._as_cupy(), other)

        def __and___(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__and___(self._as_cupy(), other)

        def __rand__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rand__(self._as_cupy(), other)

        def __or___(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__or___(self._as_cupy(), other)

        def __ror__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__ror__(self._as_cupy(), other)

        def __xor__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__xor__(self._as_cupy(), other)

        def __rxor__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rxor__(self._as_cupy(), other)

        def __lt__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__lt__(self._as_cupy(), other)

        def __rlt__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rlt__(self._as_cupy(), other)

        def __le__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__le__(self._as_cupy(), other)

        def __rle__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rle__(self._as_cupy(), other)

        def __eq__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__eq__(self._as_cupy(), other)

        def __req__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__req__(self._as_cupy(), other)

        def __ne__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__ne__(self._as_cupy(), other)

        def __rne__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rne__(self._as_cupy(), other)

        def __ge__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__ge__(self._as_cupy(), other)

        def __rge__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rge__(self._as_cupy(), other)

        def __gt__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__gt__(self._as_cupy(), other)

        def __rgt__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rgt__(self._as_cupy(), other)

        def __rshift__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rshift__(self._as_cupy(), other)

        def __rrshift__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__rrshift__(self._as_cupy(), other)

        def __lshift__(self, other):
            if hasattr(other, "_as_cupy"):
                other = other._as_cupy()
            return cupy.ndarray.__lshift__(self._as_cupy(), other)

        #        The following methods are generated by this code
        #        inplace_operators = ['add', 'floordiv', 'lshift', 'matmul', 'mod',
        #                             'mul', 'pow', 'rshift', 'sub', 'truediv', 'xor']
        #        for nn in inplace_operators:
        #            print(f"""
        #        def __i{nn}__(self, other):
        #            self._as_cupy()[:] = self.__{nn}__(other)
        #            return self""")

        def __iadd__(self, other):
            self._as_cupy()[:] = self.__add__(other)
            return self

        def __ifloordiv__(self, other):
            self._as_cupy()[:] = self.__floordiv__(other)
            return self

        def __ilshift__(self, other):
            self._as_cupy()[:] = self.__lshift__(other)
            return self

        def __imatmul__(self, other):
            self._as_cupy()[:] = self.__matmul__(other)
            return self

        def __imod__(self, other):
            self._as_cupy()[:] = self.__mod__(other)
            return self

        def __imul__(self, other):
            self._as_cupy()[:] = self.__mul__(other)
            return self

        def __ipow__(self, other):
            self._as_cupy()[:] = self.__pow__(other)
            return self

        def __irshift__(self, other):
            self._as_cupy()[:] = self.__rshift__(other)
            return self

        def __isub__(self, other):
            self._as_cupy()[:] = self.__sub__(other)
            return self

        def __itruediv__(self, other):
            self._as_cupy()[:] = self.__truediv__(other)
            return self

        def __ixor__(self, other):
            self._as_cupy()[:] = self.__xor__(other)
            return self

        #        # The following methods are generated by this code
        #        unops = ['neg', 'pos', 'invert']
        #        for nn in unops:
        #            print(f"""
        #                 def __{nn}__(self):
        #                     return cupy.ndarray.__{nn}__(self._as_cupy())""")
        def __neg__(self):
            return cupy.ndarray.__neg__(self._as_cupy())

        def __pos__(self):
            return cupy.ndarray.__pos__(self._as_cupy())

        def __invert__(self):
            return cupy.ndarray.__invert__(self._as_cupy())


cudaheader: List[SourceType] = ["""\
typedef signed int         int32_t;  //only_for_context cuda
typedef signed short       int16_t;  //only_for_context cuda
typedef signed char        int8_t;   //only_for_context cuda
typedef unsigned int       uint32_t; //only_for_context cuda
typedef unsigned short     uint16_t; //only_for_context cuda
typedef unsigned char      uint8_t;  //only_for_context cuda

#if defined(__CUDACC__) || defined(__HIPCC_RTC__)
typedef signed long long   int64_t;
typedef unsigned long long uint64_t;
#endif

#ifndef NULL
  #define NULL nullptr
#endif

"""]


def nplike_to_cupy(arr):
    return cupy.array(arr)


class ContextCupy(XContext):
    """
    Creates a Cupy Context object, that allows performing the computations
    on nVidia GPUs.

    Args:
        default_block_size (int):  CUDA thread size that is used by default
            for kernel execution in case a block size is not specified
            directly in the kernel object. The default value is 256.
        device (int): Identifier of the device to be used by the context.
    Returns:
        ContextCupy: context object.

    """

    @property
    def nplike_array_type(self):
        return cupy.ndarray

    @property
    def linked_array_type(self):
        return LinkedArrayCupy

    def __init__(
        self,
        default_block_size=256,
        default_shared_mem_size_bytes=0,
        device=None,
    ):
        if not _enabled:
            raise ModuleNotAvailableError(
                "cupy is not installed. " "ContextCupy is not available!"
                )
        
        if device is not None:
            cupy.cuda.Device(device).use()

        super().__init__()

        self.default_block_size = default_block_size
        self.default_shared_mem_size_bytes = default_shared_mem_size_bytes

    def _make_buffer(self, capacity):
        return BufferCupy(capacity=capacity, context=self)

    def build_kernels(
        self,
        sources,
        kernel_descriptions,
        specialize=True,
        apply_to_source=(),
        save_source_as=None,
        extra_compile_args=(),
        extra_cdef=None,
        extra_classes=(),
        extra_headers=(),
        compile=True,  # noqa
    ) -> Dict[Tuple[str, tuple], "KernelCupy"]:
        if not compile:
            raise NotImplementedError("compile=False available only on CPU.")
        classes = list(classes_from_kernels(kernel_descriptions))
        classes += list(extra_classes)
        classes = sort_classes(classes)

        # Update the kernel descriptions with the overriden classes
        cls_for_name = {cls.__name__: cls for cls in classes}
        for kernel_name, kernel in kernel_descriptions.items():
            for arg in kernel.args:
                arg.atype = cls_for_name.get(arg.atype.__name__, arg.atype)

        cls_sources = sources_from_classes(classes)

        headers = cudaheader + list(extra_headers)

        sources = headers + cls_sources + sources

        source, folders = _concatenate_sources(sources, apply_to_source)
        source = "\n".join(['extern "C"{', source, "}"])

        if specialize:
            # included files are searched in the same folders od the src_filed
            specialized_source = specialize_source(
                source, specialize_for="cuda", search_in_folders=folders
            )
        else:
            specialized_source = source

        if save_source_as is not None:
            with open(save_source_as, "w") as fid:
                fid.write(specialized_source)

        extra_include_paths = self.get_installed_c_source_paths()
        include_flags = [f"-I{path}" for path in extra_include_paths]
        extra_compile_args = (
            *extra_compile_args,
            *include_flags,
            "-DXO_CONTEXT_CUDA",
        )

        module = cupy.RawModule(
            code=specialized_source, options=extra_compile_args
        )

        out_kernels = {}
        for pyname, kernel in kernel_descriptions.items():
            if kernel.c_name is None:
                kernel.c_name = pyname
            out_kernels[pyname] = KernelCupy(
                function=module.get_function(kernel.c_name),
                description=kernel,
                block_size=self.default_block_size,
                context=self,
                shared_mem_size_bytes=self.default_shared_mem_size_bytes,
            )

            out_kernels[pyname].source = source
            out_kernels[pyname].specialized_source = specialized_source

        return out_kernels

    def __str__(self):
        return f"{type(self).__name__}:{cupy.cuda.get_device_id()}"

    def nparray_to_context_array(self, arr, copy=False):
        """
        Copies a numpy array to the device memory.

        Args:
            arr (numpy.ndarray): Array to be transferred
            copy (bool): This parameter is ignored for CUDA, as the data lives
                on a different device.

        Returns:
            cupy.ndarray:The same array copied to the device.
        """
        dev_arr = cupy.array(arr)
        return dev_arr

    def nparray_from_context_array(self, dev_arr, copy=False):
        """
        Copies an array to the device to a numpy array.

        Args:
            dev_arr (cupy.ndarray): Array to be transferred.
            copy (bool): This parameter is ignored for CUDA, as the data lives
                on a different device.

        Returns:
            numpy.ndarray: The same data copied to a numpy array.
        """
        return dev_arr.get()

    @property
    def nplike_lib(self):
        """
        Module containing all the numpy features supported by cupy.
        """
        return cupy

    @property
    def splike_lib(self):
        """
        Module containing all the scipy features supported by cupy.
        """
        return cupyx.scipy

    def synchronize(self):
        """
        Ensures that all computations submitted to the context are completed.
        Equivalent to ``cupy.cuda.stream.get_current_stream().synchronize()``
        """
        cupy.cuda.stream.get_current_stream().synchronize()

    def zeros(self, *args, **kwargs):
        """
        Allocates an array of zeros on the device. The function has the same
        interface of numpy.zeros"""
        return self.nplike_lib.zeros(*args, **kwargs)

    def plan_FFT(
        self,
        data,
        axes,
    ):
        """
        Generates an FFT plan object to be executed on the context.

        Args:
            data (cupy.ndarray): Array having type and shape for which the FFT
                needs to be planned.
            axes (sequence of ints): Axes along which the FFT needs to be
                performed.
        Returns:
            FFTCupy: FFT plan for the required array shape, type and axes.

        Example:

        .. code-block:: python

            plan = context.plan_FFT(data, axes=(0,1))

            data2 = 2*data

            # Forward tranform (in place)
            plan.transform(data2)

            # Inverse tranform (in place)
            plan.itransform(data2)
        """
        return FFTCupy(self, data, axes)

    @property
    def kernels(self):
        """
        Dictionary containing all the kernels that have been imported to the context.
        The syntax ``context.kernels.mykernel`` can also be used.
        """

        return self._kernels


class BufferCupy(XBuffer):
    def _make_context(self):
        return ContextCupy()

    def _new_buffer(self, capacity):
        return cupy.zeros(shape=(capacity,), dtype=cupy.uint8)

    def update_from_native(self, offset, source, source_offset, nbytes):
        """Copy data from native buffer into self.buffer starting from offset"""
        self.buffer[offset : offset + nbytes] = source[
            source_offset : source_offset + nbytes
        ]

    def to_native(self, offset, nbytes):
        """Return a new cupy buffer with data from offset"""
        return self.buffer[offset : offset + nbytes].copy()

    def copy_to_native(self, dest, dest_offset, source_offset, nbytes):
        """copy data from self to source from offset and nbytes"""
        dest[dest_offset : dest_offset + nbytes] = self.buffer[
            source_offset : source_offset + nbytes
        ]

    def update_from_buffer(self, offset, source):
        """Copy data from python buffer such as bytearray, bytes, memoryview, numpy array.data"""
        nbytes = len(source)
        self.buffer[offset : offset + nbytes] = cupy.array(
            np.frombuffer(source, dtype=np.uint8)
        )

    def to_nplike(self, offset, dtype, shape):
        """view in nplike"""
        nbytes = np.prod(shape) * np.dtype(dtype).itemsize
        return (
            self.buffer[offset : offset + nbytes]
            .view(dtype=dtype)
            .reshape(*shape)
        )

    def to_nparray(self, offset, dtype, shape):
        return self.to_nplike(offset, dtype, shape).get()

    def update_from_nplike(self, offset, dest_dtype, value):
        if dest_dtype != value.dtype:
            value = value.astype(dtype=dest_dtype)  # make a copy
        src = value.view("int8")
        self.buffer[offset : offset + src.nbytes] = value.flatten().view(
            "int8"
        )

    def to_bytearray(self, offset, nbytes):
        """copy in byte array: used in update_from_xbuffer"""
        return self.buffer[offset : offset + nbytes].get().tobytes()

    def to_pointer_arg(self, offset, nbytes):
        """return data that can be used as argument in kernel"""
        return self.buffer[offset : offset + nbytes]


class KernelCupy(object):
    def __init__(
        self, function, description, block_size, context, shared_mem_size_bytes
    ):
        self.function = function
        self.description = description
        self.block_size = block_size
        self.shared_mem_size_bytes = shared_mem_size_bytes
        self.context = context

    def to_function_arg(self, arg, value):
        if arg.pointer:
            if hasattr(arg.atype, "_dtype"):  # it is numerical scalar
                if hasattr(value, "dtype"):  # nparray
                    assert isinstance(value, cupy.ndarray)
                    return value.data
                elif hasattr(value, "_shape"):  # xobject array
                    raise NotImplementedError
            else:
                raise ValueError(
                    f"Invalid value {value} for argument {arg.name} "
                    f"of kernel {self.description.pyname}"
                )
        else:
            if hasattr(arg.atype, "_dtype"):  # it is numerical scalar
                return arg.atype(value)  # try to return a numpy scalar
            elif hasattr(arg.atype, "_size"):  # it is a compound xobject
                assert (
                    value._buffer.context is self.context
                ), f"Incompatible context for argument `{arg.name}`"
                return value._buffer.buffer[value._offset :]
            else:
                raise ValueError(
                    f"Invalid value {value} for argument {arg.name} of kernel {self.description.pyname}"
                )

    @property
    def num_args(self):
        return len(self.description.args)

    def __call__(self, **kwargs):
        assert len(kwargs.keys()) == self.num_args
        arg_list = []
        for arg in self.description.args:
            vv = kwargs[arg.name]
            arg_list.append(self.to_function_arg(arg, vv))

        if isinstance(self.description.n_threads, str):
            n_threads = kwargs[self.description.n_threads]
        else:
            n_threads = self.description.n_threads

        if "shared_mem_size_bytes" in kwargs.keys():
            shared_mem_size_bytes = kwargs["shared_mem_size_bytes"]
        else:
            shared_mem_size_bytes = self.shared_mem_size_bytes

        grid_size = int(np.ceil(n_threads / self.block_size))
        self.function(
            (grid_size,),
            (self.block_size,),
            arg_list,
            shared_mem=shared_mem_size_bytes,
        )


class FFTCupy(object):
    def __init__(self, context, data, axes):
        self.context = context
        self.axes = axes

        assert len(data.shape) > max(axes)

        from cupyx.scipy import fftpack as cufftp

        if data.flags.f_contiguous:
            self._ax = [data.ndim - 1 - aa for aa in axes]
            _dat = data.T
            self.f_contiguous = True
        else:
            self._ax = axes
            _dat = data
            self.f_contiguous = False

        self._fftplan = cufftp.get_fft_plan(
            _dat, axes=self._ax, value_type="C2C"
        )

    def transform(self, data):
        if self.f_contiguous:
            _dat = data.T
        else:
            _dat = data
        _dat[:] = cufftp.fftn(_dat, axes=self._ax, plan=self._fftplan)[:]
        """The transform is done inplace"""

    def itransform(self, data):
        """The transform is done inplace"""
        if self.f_contiguous:
            _dat = data.T
        else:
            _dat = data
        _dat[:] = cufftp.ifftn(_dat, axes=self._ax, plan=self._fftplan)[:]


if _enabled:
    available.append(ContextCupy)