nabu 2024.1.0rc3__py3-none-any.whl → 2024.1.1__py3-none-any.whl

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/processing/muladd.py

@@ -0,0 +1,29 @@
+import numpy as np
+from .processing_base import ProcessingBase
+
+
+class MulAdd:
+    processing_cls = ProcessingBase
+
+    def __init__(self, **backend_options):
+        self.processing = self.processing_cls(**(backend_options or {}))
+        self._init_finalize()
+
+    def _init_finalize(self):
+        pass
+
+    def mul_add(self, dst, other, fac_dst, fac_other, dst_region=None, other_region=None):
+        if dst_region is None:
+            dst_slice_y = dst_slice_x = slice(None, None)
+        else:
+            dst_slice_y, dst_slice_x = dst_region
+        if other_region is None:
+            other_slice_y = other_slice_x = slice(None, None)
+        else:
+            other_slice_y, other_slice_x = other_region
+
+        dst[dst_slice_y, dst_slice_x] = (
+            fac_dst * dst[dst_slice_y, dst_slice_x] + fac_other * other[other_slice_y, other_slice_x]
+        )
+
+    __call__ = mul_add

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/processing/muladd_cuda.py

@@ -0,0 +1,68 @@
+import numpy as np
+from nabu.utils import get_cuda_srcfile, updiv
+from .muladd import MulAdd
+from ..cuda.utils import __has_pycuda__
+from ..cuda.processing import CudaProcessing
+
+if __has_pycuda__:
+    import pycuda.gpuarray as garray
+
+
+class CudaMulAdd(MulAdd):
+    processing_cls = CudaProcessing
+
+    def _init_finalize(self):
+        self._init_kernel()
+
+    def _init_kernel(self):
+        self.muladd_kernel = self.processing.kernel(
+            "mul_add",
+            filename=get_cuda_srcfile("ElementOp.cu"),  # signature="PPiiffiiii"
+        )
+
+    def mul_add(self, dst, other, fac_dst, fac_other, dst_region=None, other_region=None):
+        """
+        'region' should be a tuple (slice(y_start, y_end), slice(x_start, x_end))
+        """
+
+        if dst_region is None:
+            dst_coords = (0, dst.shape[1], 0, dst.shape[0])
+        else:
+            dst_coords = (dst_region[1].start, dst_region[1].stop, dst_region[0].start, dst_region[0].stop)
+        if other_region is None:
+            other_coords = (0, other.shape[1], 0, other.shape[0])
+        else:
+            other_coords = (other_region[1].start, other_region[1].stop, other_region[0].start, other_region[0].stop)
+
+        delta_x = np.diff(dst_coords[:2])
+        delta_y = np.diff(dst_coords[2:])
+        if delta_x != np.diff(other_coords[:2]) or delta_y != np.diff(other_coords[2:]):
+            raise ValueError("Invalid dst_region and other_region provided. Regions must have the same size")
+        if delta_x == 0 or delta_y == 0:
+            raise ValueError("delta_x or delta_y is 0")
+
+        # can't use "int4" in pycuda ? int2 seems fine. Go figure
+        dst_x_range = np.array(dst_coords[:2], dtype=garray.vec.int2)
+        dst_y_range = np.array(dst_coords[2:], dtype=garray.vec.int2)
+        other_x_range = np.array(other_coords[:2], dtype=garray.vec.int2)
+        other_y_range = np.array(other_coords[2:], dtype=garray.vec.int2)
+
+        block = (32, 32, 1)
+        grid = [updiv(length, b) for (length, b) in zip((delta_x[0], delta_y[0]), block)]
+
+        self.muladd_kernel(
+            dst,
+            other,
+            np.int32(dst.shape[1]),
+            np.int32(other.shape[1]),
+            np.float32(fac_dst),
+            np.float32(fac_other),
+            dst_x_range,
+            dst_y_range,
+            other_x_range,
+            other_y_range,
+            grid=grid,
+            block=block,
+        )
+
+    __call__ = mul_add

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/processing/padding_base.py

@@ -0,0 +1,71 @@
+import numpy as np
+from ..utils import check_supported
+
+
+class PaddingBase:
+    """
+    A class for performing padding based on coordinate transform.
+    The Cuda and OpenCL backends will subclass this class.
+    """
+
+    supported_modes = ["constant", "edge", "reflect", "symmetric", "wrap"]
+
+    def __init__(self, shape, pad_width, mode="constant", **kwargs):
+        """
+        Initialize a Padding object.
+
+        Parameters
+        ----------
+        shape: tuple
+            Image shape
+        pad_width: tuple
+            Padding width for each axis. Please see the documentation of numpy.pad().
+        mode: str
+            Padding mode
+
+        Other parameters
+        ----------------
+        constant_values: tuple
+            Tuple containing the values to fill when mode="constant" (as in numpy.pad)
+        """
+        if len(shape) != 2:
+            raise ValueError("This class only works on images")
+        self.shape = shape
+        self._set_mode(mode, **kwargs)
+        self._get_padding_arrays(pad_width)
+
+    def _set_mode(self, mode, **kwargs):
+        # COMPAT.
+        if mode == "edges":
+            mode = "edge"
+        #
+        check_supported(mode, self.supported_modes, "padding mode")
+        self.mode = mode
+        self._kwargs = kwargs
+
+    def _get_padding_arrays(self, pad_width):
+        self.pad_width = pad_width
+        if isinstance(pad_width, tuple) and isinstance(pad_width[0], np.ndarray):
+            # user-defined coordinate transform
+            err_msg = "pad_width must be either a scalar, a tuple in the form ((a, b), (c, d)), or a tuple of two one-dimensional numpy arrays (eg. use numpy.indices(..., sparse=True))"
+            if len(pad_width) != 2:
+                raise ValueError(err_msg)
+            if any([np.squeeze(pw).ndim > 1 for pw in pad_width]):
+                raise ValueError(err_msg)
+            if self.mode == "constant":
+                raise ValueError("Custom coordinate transform does not work with mode='constant'")
+            self.mode = "custom"
+            self.coords_rows, self.coords_cols = pad_width
+        else:
+            if self.mode == "constant":
+                # no need for coordinate transform here
+                constant_values = self._kwargs.get("constant_values", 0)
+                self.padded_array_constant = np.pad(
+                    np.zeros(self.shape, dtype="f"), self.pad_width, mode="constant", constant_values=constant_values
+                )
+                self.padded_shape = self.padded_array_constant.shape
+                return
+            R, C = np.indices(self.shape, dtype=np.int32, sparse=True)
+            self.coords_rows = np.pad(R.ravel(), self.pad_width[0], mode=self.mode)
+            self.coords_cols = np.pad(C.ravel(), self.pad_width[1], mode=self.mode)
+        self.padded_shape = (self.coords_rows.size, self.coords_cols.size)

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/processing/padding_cuda.py

@@ -0,0 +1,75 @@
+import numpy as np
+from ..utils import get_cuda_srcfile, updiv
+from ..cuda.processing import CudaProcessing
+from ..cuda.utils import __has_pycuda__
+from .padding_base import PaddingBase
+
+
+class CudaPadding(PaddingBase):
+    """
+    A class for performing padding on GPU using Cuda
+    """
+
+    backend = "cuda"
+
+    # TODO docstring from base class
+    def __init__(self, shape, pad_width, mode="constant", cuda_options=None, **kwargs):
+        super().__init__(shape, pad_width, mode=mode, **kwargs)
+        self.cuda_processing = self.processing = CudaProcessing(**(cuda_options or {}))
+        self._init_cuda_coordinate_transform()
+
+    def _init_cuda_coordinate_transform(self):
+        if self.mode == "constant":
+            self.d_padded_array_constant = self.processing.to_device(
+                "d_padded_array_constant", self.padded_array_constant
+            )
+            return
+        self._coords_transform_kernel = self.processing.kernel(
+            "coordinate_transform",
+            filename=get_cuda_srcfile("padding.cu"),
+            signature="PPPPiii",
+        )
+        self._coords_transform_block = (32, 32, 1)
+        self._coords_transform_grid = [
+            updiv(a, b) for a, b in zip(self.padded_shape[::-1], self._coords_transform_block)
+        ]
+        self.d_coords_rows = self.processing.to_device("d_coords_rows", self.coords_rows)
+        self.d_coords_cols = self.processing.to_device("d_coords_cols", self.coords_cols)
+
+    def _pad_constant(self, image, output):
+        pad_y, pad_x = self.pad_width
+        self.d_padded_array_constant[pad_y[0] : pad_y[0] + self.shape[0], pad_x[0] : pad_x[0] + self.shape[1]] = image[
+            :
+        ]
+        output[:] = self.d_padded_array_constant[:]
+        return output
+
+    def pad(self, image, output=None):
+        """
+        Pad an array.
+
+        Parameters
+        ----------
+        image: pycuda.gpuarray.GPUArray
+            Image to pad
+        output: pycuda.gpuarray.GPUArray, optional
+            Output image. If provided, must be in the expected shape.
+        """
+        if output is None:
+            output = self.processing.allocate_array("d_output", self.padded_shape)
+        if self.mode == "constant":
+            return self._pad_constant(image, output)
+        self._coords_transform_kernel(
+            image,
+            output,
+            self.d_coords_cols,
+            self.d_coords_rows,
+            np.int32(self.shape[1]),
+            np.int32(self.padded_shape[1]),
+            np.int32(self.padded_shape[0]),
+            grid=self._coords_transform_grid,
+            block=self._coords_transform_block,
+        )
+        return output
+
+    __call__ = pad

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/processing/padding_opencl.py

@@ -0,0 +1,77 @@
+import numpy as np
+from ..utils import get_opencl_srcfile
+from ..opencl.processing import OpenCLProcessing
+from .padding_base import PaddingBase
+from ..opencl.utils import __has_pyopencl__
+
+if __has_pyopencl__:
+    from ..opencl.memcpy import OpenCLMemcpy2D
+
+
+class OpenCLPadding(PaddingBase):
+    """
+    A class for performing padding on GPU using OpenCL
+    """
+
+    backend = "opencl"
+
+    # TODO docstring from base class
+    def __init__(self, shape, pad_width, mode="constant", opencl_options=None, **kwargs):
+        super().__init__(shape, pad_width, mode=mode, **kwargs)
+        self.opencl_processing = self.processing = OpenCLProcessing(**(opencl_options or {}))
+        self.queue = self.opencl_processing.queue
+        self._init_opencl_coordinate_transform()
+
+    def _init_opencl_coordinate_transform(self):
+        if self.mode == "constant":
+            self.d_padded_array_constant = self.processing.to_device(
+                "d_padded_array_constant", self.padded_array_constant
+            )
+            self.memcpy2D = OpenCLMemcpy2D(ctx=self.processing.ctx, queue=self.queue)
+            return
+        self._coords_transform_kernel = self.processing.kernel(
+            "coordinate_transform",
+            filename=get_opencl_srcfile("padding.cl"),
+        )
+        self._coords_transform_global_size = self.padded_shape[::-1]
+        self.d_coords_rows = self.processing.to_device("d_coords_rows", self.coords_rows)
+        self.d_coords_cols = self.processing.to_device("d_coords_cols", self.coords_cols)
+
+    def _pad_constant(self, image, output):
+        pad_y, pad_x = self.pad_width
+        # the following line is not implemented in pyopencl
+        # self.d_padded_array_constant[pad_y[0] : pad_y[0] + self.shape[0], pad_x[0] : pad_x[0] + self.shape[1]] = image[:]
+        # cl.enqueue_copy is too cumbersome to use for Buffer <-> Buffer.
+        # Use a dedicated kernel instead.
+        # This is not optimal (two copies) - TODO write a constant padding kernel
+        self.memcpy2D(self.d_padded_array_constant, image, image.shape[::-1], dst_offset_xy=(pad_x[0], pad_y[0]))
+        output[:] = self.d_padded_array_constant[:]
+        return output
+
+    def pad(self, image, output=None):
+        """
+        Pad an array.
+
+        Parameters
+        ----------
+        image: pyopencl array
+            Image to pad
+        output: pyopencl array
+            Output image. If provided, must be in the expected shape.
+        """
+        if output is None:
+            output = self.processing.allocate_array("d_output", self.padded_shape)
+        if self.mode == "constant":
+            return self._pad_constant(image, output)
+        self._coords_transform_kernel(
+            self.queue,
+            image,
+            output,
+            self.d_coords_cols,
+            self.d_coords_rows,
+            np.int32(self.shape[1]),
+            np.int32(self.padded_shape[1]),
+            np.int32(self.padded_shape[0]),
+            global_size=self._coords_transform_global_size,
+        )
+        return output

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/processing/processing_base.py

@@ -0,0 +1,123 @@
+import numpy as np
+from ..utils import BaseClassError
+
+"""
+Base class for OpenCLProcessing and CudaProcessing
+Should not be used directly
+"""
+
+
+class ProcessingBase:
+    array_class = None
+    dtype_to_ctype = BaseClassError
+
+    def __init__(self):
+        self._allocated = {}
+
+    def init_arrays_to_none(self, arrays_names):
+        """
+        Initialize arrays to None. After calling this method, the current instance will
+        have self.array_name = None, and self._old_array_name = None.
+
+        Parameters
+        ----------
+        arrays_names: list of str
+            List of arrays names.
+        """
+        for array_name in arrays_names:
+            setattr(self, array_name, None)
+            setattr(self, "_old_" + array_name, None)
+            self._allocated[array_name] = False
+
+    def recover_arrays_references(self, arrays_names):
+        """
+        Performs self._array_name = self._old_array_name,
+        for each array_name in arrays_names.
+
+        Parameters
+        ----------
+        arrays_names: list of str
+            List of array names
+        """
+        for array_name in arrays_names:
+            old_arr = getattr(self, "_old_" + array_name, None)
+            if old_arr is not None:
+                setattr(self, array_name, old_arr)
+
+    def _allocate_array_mem(self, shape, dtype):
+        raise ValueError("Base class")
+
+    def allocate_array(self, array_name, shape, dtype=np.float32):
+        """
+        Allocate a GPU array on the current context/stream/device,
+        and set 'self.array_name' to this array.
+
+        Parameters
+        ----------
+        array_name: str
+            Name of the array (for book-keeping)
+        shape: tuple of int
+            Array shape
+        dtype: numpy.dtype, optional
+            Data type. Default is float32.
+        """
+        if not self._allocated.get(array_name, False):
+            new_device_arr = self._allocate_array_mem(shape, dtype)
+            setattr(self, array_name, new_device_arr)
+            self._allocated[array_name] = True
+        return getattr(self, array_name)
+
+    def set_array(self, array_name, array_ref, dtype=np.float32):
+        """
+        Set the content of a device array.
+
+        Parameters
+        ----------
+        array_name: str
+            Array name. This method will look for self.array_name.
+        array_ref: array (numpy or GPU array)
+            Array containing the data to copy to 'array_name'.
+        dtype: numpy.dtype, optional
+            Data type. Default is float32.
+        """
+        if isinstance(array_ref, self.array_class):
+            current_arr = getattr(self, array_name, None)
+            setattr(self, "_old_" + array_name, current_arr)
+            setattr(self, array_name, array_ref)
+        elif isinstance(array_ref, np.ndarray):
+            self.allocate_array(array_name, array_ref.shape, dtype=dtype)
+            getattr(self, array_name).set(array_ref)
+        else:
+            raise ValueError("Expected numpy array or pycuda array")
+        return getattr(self, array_name)
+
+    def get_array(self, array_name):
+        return getattr(self, array_name, None)
+
+    # COMPAT.
+    _init_arrays_to_none = init_arrays_to_none
+    _recover_arrays_references = recover_arrays_references
+    _allocate_array = allocate_array
+    _set_array = set_array
+
+    def check_array(self, arr, expected_shape, expected_dtype="f", check_contiguous=True):
+        """
+        Check whether a given array is suitable for being processed (shape, dtype, contiguous)
+        """
+        if arr.shape != expected_shape:
+            raise ValueError("Expected shape %s but got %s" % (str(expected_shape), str(arr.shape)))
+        if arr.dtype != np.dtype(expected_dtype):
+            raise ValueError("Expected data type %s but got %s" % (str(expected_dtype), str(arr.dtype)))
+        if check_contiguous:
+            if isinstance(arr, np.ndarray) and not (arr.flags["C_CONTIGUOUS"]):
+                raise ValueError("Expected C-contiguous array")
+            if isinstance(arr, self.array_class) and not arr.flags.c_contiguous:
+                raise ValueError("Expected C-contiguous array")
+
+    def kernel(self, *args, **kwargs):
+        raise ValueError("Base class")
+
+    def to_device(self, array_name, array):
+        arr_ref = self.allocate_array(array_name, array.shape, dtype=array.dtype)
+        arr_ref.set(array)
+        return arr_ref

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/processing/roll_opencl.py

@@ -0,0 +1,64 @@
+#
+# WIP !
+#
+import numpy as np
+from ..opencl.utils import __has_pyopencl__
+from ..utils import get_opencl_srcfile
+
+if __has_pyopencl__:
+    import pyopencl as cl
+    from ..opencl.processing import OpenCLProcessing
+    from ..opencl.kernel import OpenCLKernel
+    from pyopencl.tools import dtype_to_ctype as cl_dtype_to_ctype
+
+
+class OpenCLRoll:
+    def __init__(self, dtype, direction=1, offset=None, **processing_kwargs):
+        self.processing = OpenCLProcessing(queue=processing_kwargs.get("queue", None))
+        self.dtype = np.dtype(dtype)
+        compile_options = ["-DDTYPE=%s" % cl_dtype_to_ctype(self.dtype)]
+        self.offset = offset or 0
+        self.roll_kernel = OpenCLKernel(
+            "roll_forward_x",
+            None,
+            queue=self.processing.queue,
+            filename=get_opencl_srcfile("roll.cl"),
+            options=compile_options,
+        )
+        self.shmem = cl.LocalMemory(self.dtype.itemsize)
+        self.direction = direction
+        if self.direction < 0:
+            self.revert_kernel = OpenCLKernel(
+                "revert_array_x",
+                None,
+                queue=self.processing.queue,
+                filename=get_opencl_srcfile("roll.cl"),
+                options=compile_options,
+            )
+
+    def __call__(self, arr):
+        ny, nx = arr.shape
+        # Launch one big horizontal workgroup
+        wg_x = min((nx - self.offset) // 2, self.processing.queue.device.max_work_group_size)
+        local_size = (wg_x, 1, 1)
+        global_size = [wg_x, ny]
+        if self.direction < 0:
+            local_size2 = None
+            global_size2 = [nx - self.offset, ny]
+            self.revert_kernel(
+                arr, np.int32(nx), np.int32(ny), np.int32(self.offset), local_size=local_size2, global_size=global_size2
+            )
+        self.roll_kernel(
+            arr,
+            np.int32(nx),
+            np.int32(ny),
+            np.int32(self.offset),
+            self.shmem,
+            local_size=local_size,
+            global_size=global_size,
+        )
+        if self.direction < 0:
+            self.revert_kernel(
+                arr, np.int32(nx), np.int32(ny), np.int32(self.offset), local_size=local_size2, global_size=global_size2
+            )
+        return arr

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/processing/rotation.py

@@ -0,0 +1,63 @@
+try:
+    from skimage.transform import rotate
+
+    __have__skimage__ = True
+except ImportError:
+    __have__skimage__ = False
+
+
+class Rotation:
+    supported_modes = {
+        "constant": "constant",
+        "zeros": "constant",
+        "edge": "edge",
+        "edges": "edge",
+        "symmetric": "symmetric",
+        "sym": "symmetric",
+        "reflect": "reflect",
+        "wrap": "wrap",
+        "periodic": "wrap",
+    }
+
+    def __init__(self, shape, angle, center=None, mode="edge", reshape=False, **sk_kwargs):
+        """
+        Initiate a Rotation object.
+
+        Parameters
+        ----------
+        shape: tuple of int
+            Shape of the images to process
+        angle: float
+            Rotation angle in DEGREES
+        center: tuple of float, optional
+            Coordinates of the center of rotation, in the format (X, Y) (mind the non-python
+            convention !).
+            Default is ((Nx - 1)/2.0, (Ny - 1)/2.0)
+        mode: str, optional
+            Padding mode. Default is "edge".
+        reshape: bool, optional
+
+
+        Other Parameters
+        -----------------
+        All the other parameters are passed directly to scikit image 'rotate' function:
+        order, cval, clip, preserve_range.
+        """
+        self.shape = shape
+        self.angle = angle
+        self.center = center
+        self.mode = mode
+        self.reshape = reshape
+        self.sk_kwargs = sk_kwargs
+
+    def rotate(self, img, output=None):
+        if not __have__skimage__:
+            raise ValueError("scikit-image is needed for using rotate()")
+        res = rotate(img, self.angle, resize=self.reshape, center=self.center, mode=self.mode, **self.sk_kwargs)
+        if output is not None:
+            output[:] = res[:]
+            return output
+        else:
+            return res
+
+    __call__ = rotate

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/processing/rotation_cuda.py

@@ -0,0 +1,66 @@
+import numpy as np
+from .rotation import Rotation
+from ..utils import get_cuda_srcfile, updiv
+from ..cuda.utils import __has_pycuda__, copy_array
+from ..cuda.processing import CudaProcessing
+
+if __has_pycuda__:
+    from ..cuda.kernel import CudaKernel
+    import pycuda.driver as cuda
+
+
+class CudaRotation(Rotation):
+    def __init__(self, shape, angle, center=None, mode="edge", reshape=False, cuda_options=None, **sk_kwargs):
+        if center is None:
+            center = ((shape[1] - 1) / 2.0, (shape[0] - 1) / 2.0)
+        super().__init__(shape, angle, center=center, mode=mode, reshape=reshape, **sk_kwargs)
+        self._init_cuda_rotation(cuda_options)
+
+    def _init_cuda_rotation(self, cuda_options):
+        cuda_options = cuda_options or {}
+        self.cuda_processing = CudaProcessing(**cuda_options)
+        self._allocate_arrays()
+        self._init_rotation_kernel()
+
+    def _allocate_arrays(self):
+        self._d_image_cua = cuda.np_to_array(np.zeros(self.shape, "f"), "C")
+        self.cuda_processing.init_arrays_to_none(["d_output"])
+
+    def _init_rotation_kernel(self):
+        self.cuda_rotation_kernel = CudaKernel("rotate", get_cuda_srcfile("rotation.cu"))
+        self.texref_image = self.cuda_rotation_kernel.module.get_texref("tex_image")
+        self.texref_image.set_filter_mode(cuda.filter_mode.LINEAR)  # bilinear
+        self.texref_image.set_address_mode(0, cuda.address_mode.CLAMP)  # TODO tune
+        self.texref_image.set_address_mode(1, cuda.address_mode.CLAMP)  # TODO tune
+        self.cuda_rotation_kernel.prepare("Piiffff", [self.texref_image])
+        self.texref_image.set_array(self._d_image_cua)
+        self._cos_theta = np.cos(np.deg2rad(self.angle))
+        self._sin_theta = np.sin(np.deg2rad(self.angle))
+        self._Nx = np.int32(self.shape[1])
+        self._Ny = np.int32(self.shape[0])
+        self._center_x = np.float32(self.center[0])
+        self._center_y = np.float32(self.center[1])
+        self._block = (32, 32, 1)  # tune ?
+        self._grid = (updiv(self.shape[1], self._block[1]), updiv(self.shape[0], self._block[0]), 1)
+
+    def rotate(self, img, output=None, do_checks=True):
+        copy_array(self._d_image_cua, img, check=do_checks)
+        if output is not None:
+            d_out = output
+        else:
+            self.cuda_processing.allocate_array("d_output", self.shape, np.float32)
+            d_out = self.cuda_processing.d_output
+        self.cuda_rotation_kernel(
+            d_out,
+            self._Nx,
+            self._Ny,
+            self._cos_theta,
+            self._sin_theta,
+            self._center_x,
+            self._center_y,
+            grid=self._grid,
+            block=self._block,
+        )
+        return d_out
+
+    __call__ = rotate