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/nabu/misc/utils.py

@@ -0,0 +1,117 @@
+import numpy as np
+
+
+def rescale_data(data, new_min, new_max, data_min=None, data_max=None):
+    if data_min is None:
+        data_min = np.min(data)
+    if data_max is None:
+        data_max = np.max(data)
+    return (new_max - new_min) / (data_max - data_min) * (data - data_min) + new_min
+
+
+def get_dtype_range(dtype, normalize_floats=False):
+    if np.dtype(dtype).kind in ["u", "i"]:
+        dtype_range = (np.iinfo(dtype).min, np.iinfo(dtype).max)
+    else:
+        if normalize_floats:
+            dtype_range = (-1.0, 1.0)
+        else:
+            dtype_range = (np.finfo(dtype).min, np.finfo(dtype).max)
+    return dtype_range
+
+
+def psnr(img1, img2):
+    if img1.dtype != img2.dtype:
+        raise ValueError("both images should have the same data type")
+    dtype_range = get_dtype_range(img1.dtype, normalize_floats=True)
+    dtype_range = dtype_range[-1] - dtype_range[0]
+    if np.dtype(img1.dtype).kind in ["f", "c"]:
+        img1 = rescale_data(img1, -1.0, 1.0)
+        img2 = rescale_data(img2, -1.0, 1.0)
+    mse = np.mean((img1.astype(np.float64) - img2) ** 2)
+    return 10 * np.log10((dtype_range**2) / mse)
+
+
+#
+# silx.opencl.utils.ConvolutionInfos cannot be used as long as
+# silx.opencl instantiates the "ocl" singleton in __init__,
+# leaving opencl contexts all over the place in some cases.
+#
+# so for now: copypasta
+#
+
+
+class ConvolutionInfos(object):
+    allowed_axes = {
+        "1D": [None],
+        "separable_2D_1D_2D": [None, (0, 1), (1, 0)],
+        "batched_1D_2D": [(0,), (1,)],
+        "separable_3D_1D_3D": [None, (0, 1, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0), (1, 0, 2), (0, 2, 1)],
+        "batched_1D_3D": [(0,), (1,), (2,)],
+        "batched_separable_2D_1D_3D": [(0,), (1,), (2,)],  # unsupported (?)
+        "2D": [None],
+        "batched_2D_3D": [(0,), (1,), (2,)],
+        "separable_3D_2D_3D": [
+            (1, 0),
+            (0, 1),
+            (2, 0),
+            (0, 2),
+            (1, 2),
+            (2, 1),
+        ],
+        "3D": [None],
+    }
+    use_cases = {
+        (1, 1): {
+            "1D": {
+                "name": "1D convolution on 1D data",
+                "kernels": ["convol_1D_X"],
+            },
+        },
+        (2, 2): {
+            "2D": {
+                "name": "2D convolution on 2D data",
+                "kernels": ["convol_2D_XY"],
+            },
+        },
+        (3, 3): {
+            "3D": {
+                "name": "3D convolution on 3D data",
+                "kernels": ["convol_3D_XYZ"],
+            },
+        },
+        (2, 1): {
+            "separable_2D_1D_2D": {
+                "name": "Separable (2D->1D) convolution on 2D data",
+                "kernels": ["convol_1D_X", "convol_1D_Y"],
+            },
+            "batched_1D_2D": {
+                "name": "Batched 1D convolution on 2D data",
+                "kernels": ["convol_1D_X", "convol_1D_Y"],
+            },
+        },
+        (3, 1): {
+            "separable_3D_1D_3D": {
+                "name": "Separable (3D->1D) convolution on 3D data",
+                "kernels": ["convol_1D_X", "convol_1D_Y", "convol_1D_Z"],
+            },
+            "batched_1D_3D": {
+                "name": "Batched 1D convolution on 3D data",
+                "kernels": ["convol_1D_X", "convol_1D_Y", "convol_1D_Z"],
+            },
+            "batched_separable_2D_1D_3D": {
+                "name": "Batched separable (2D->1D) convolution on 3D data",
+                "kernels": ["convol_1D_X", "convol_1D_Y", "convol_1D_Z"],
+            },
+        },
+        (3, 2): {
+            "separable_3D_2D_3D": {
+                "name": "Separable (3D->2D) convolution on 3D data",
+                "kernels": ["convol_2D_XY", "convol_2D_XZ", "convol_2D_YZ"],
+            },
+            "batched_2D_3D": {
+                "name": "Batched 2D convolution on 3D data",
+                "kernels": ["convol_2D_XY", "convol_2D_XZ", "convol_2D_YZ"],
+            },
+        },
+    }

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/opencl/fft.py

@@ -0,0 +1,6 @@
+from ..processing.fft_opencl import *
+from ..utils import deprecation_warning
+
+deprecation_warning(
+    "nabu.opencl.fft has been moved to nabu.processing.fft_opencl", do_print=True, func_name="fft_opencl"
+)

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/opencl/kernel.py

@@ -0,0 +1,135 @@
+import pyopencl.array as parray
+from pyopencl import Program, CommandQueue, kernel_work_group_info
+from ..utils import deprecation_warning
+from ..processing.kernel_base import KernelBase
+
+
+class OpenCLKernel(KernelBase):
+    """
+    Helper class that wraps OpenCL kernel through pyopencl.
+
+    Parameters
+    -----------
+    kernel_name: str
+        Name of the OpenCL kernel.
+    ctx: pyopencl.Context
+        OpenCL context to use.
+    queue: pyopencl.CommandQueue
+        OpenCL queue to use. If provided, will use this queue's context instead of 'ctx'
+    filename: str, optional
+        Path to the file name containing kernels definitions
+    src: str, optional
+        Source code of kernels definitions
+    automation_params: dict, optional
+        Automation parameters, see below
+    build_kwargs: optional
+        Extra arguments to provide to pyopencl.Program.build(),
+    """
+
+    def __init__(
+        self,
+        kernel_name,
+        ctx,
+        queue=None,
+        filename=None,
+        src=None,
+        automation_params=None,
+        **build_kwargs,
+    ):
+        super().__init__(kernel_name, filename=filename, src=src, automation_params=automation_params)
+        if queue is not None:
+            self.ctx = queue.context
+            self.queue = queue
+        else:
+            self.ctx = ctx
+            self.queue = None
+        self.compile_kernel_source(kernel_name, build_kwargs)
+        self.get_kernel()
+
+    def compile_kernel_source(self, kernel_name, build_kwargs):
+        self.build_kwargs = build_kwargs
+        self.kernel_name = kernel_name
+        self.program = Program(self.ctx, self.src).build(**self.build_kwargs)
+
+    def get_kernel(self):
+        self.kernel = None
+        for kern in self.program.all_kernels():
+            if kern.function_name == self.kernel_name:
+                self.kernel = kern
+        if self.kernel is None:
+            raise ValueError(
+                "Could not find a kernel with function name '%s'. Available are: %s"
+                % (self.kernel_name, self.program.kernel_names)
+            )
+
+    # overwrite parent method
+    def guess_block_size(self, shape):
+        device = self.ctx.devices[0]
+        wg_max = device.max_work_group_size
+        wg_multiple = self.kernel.get_work_group_info(kernel_work_group_info.PREFERRED_WORK_GROUP_SIZE_MULTIPLE, device)
+        ndim = len(shape)
+        # Try to have workgroup relatively well-balanced in all dimensions,
+        # with more work items in x > y > z
+        if ndim == 1:
+            wg = (wg_max, 1, 1)
+        else:
+            w = (wg_max // wg_multiple, wg_multiple)
+            wg = w if w[0] > w[1] else w[::-1]
+            wg = wg + (1,)
+            if ndim == 3:
+                (wg[0] // 2, wg[1] // 4, 8)
+        return wg
+
+    def get_block_grid(self, *args, **kwargs):
+        local_size = None
+        global_size = block = None
+        # COMPAT.
+        block = kwargs.pop("block", None)
+        if block is not None:
+            deprecation_warning("Please use 'local_size' instead of 'block'")
+        grid = kwargs.pop("grid", None)
+        if grid is not None:
+            deprecation_warning("Please use 'global_size' instead of 'grid'")
+            global_size = tuple(g * b for g, b in zip(grid, block))
+        #
+
+        global_size = kwargs.pop("global_size", global_size)
+        local_size = kwargs.pop("local_size", block)
+
+        if global_size is None:
+            raise ValueError("Need to define global_size for kernel '%s'" % self.kernel_name)
+
+        if len(global_size) == 2 and local_size is not None and len(local_size) == 3:
+            local_size = local_size[:-1]  # TODO check that last dim is 1
+
+        self.last_block_size = local_size
+        self.last_grid_size = global_size
+        return local_size, global_size
+
+    def follow_device_arr(self, args):
+        args = list(args)
+        for i, arg in enumerate(args):
+            if isinstance(arg, parray.Array):
+                args[i] = arg.data
+        return tuple(args)
+
+    def call(self, *args, **kwargs):
+        if not isinstance(args[0], CommandQueue):
+            queue = self.queue
+            if queue is None:
+                raise ValueError(
+                    "First argument must be a pyopencl queue - otherwise provide OpenCLKernel(..., queue=queue)"
+                )
+        else:
+            queue = args[0]
+            args = args[1:]
+        global_size, local_size, args, kwargs = self._prepare_call(*args, **kwargs)
+
+        kwargs.pop("global_size", None)
+        kwargs.pop("local_size", None)
+        kwargs.pop("grid", None)
+        kwargs.pop("block", None)
+
+        return self.kernel(queue, global_size, local_size, *args, **kwargs)
+
+    __call__ = call

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/opencl/memcpy.py

@@ -0,0 +1,35 @@
+import numpy as np
+from ..utils import get_opencl_srcfile
+from .kernel import OpenCLKernel
+from .processing import OpenCLProcessing
+
+
+class OpenCLMemcpy2D(OpenCLProcessing):
+    """
+    A class for performing rectangular memory copies between pyopencl arrays.
+    It will only work for float32 arrays!
+    It was written as pyopencl.enqueue_copy is too cumbersome to use for buffers.
+    """
+
+    def __init__(self, ctx=None, device_type="GPU", queue=None, **kwargs):
+        super().__init__(ctx=ctx, device_type=device_type, queue=queue, **kwargs)
+        self.memcpy2D = OpenCLKernel("cpy2d", self.ctx, filename=get_opencl_srcfile("ElementOp.cl"))
+
+    def __call__(self, dst, src, transfer_shape_xy, dst_offset_xy=None, src_offset_xy=None, wait=True):
+        if dst_offset_xy is None:
+            dst_offset_xy = (0, 0)
+        if src_offset_xy is None:
+            src_offset_xy = (0, 0)
+        evt = self.memcpy2D(
+            self.queue,
+            dst,
+            src,
+            np.int32(dst.shape[-1]),
+            np.int32(src.shape[-1]),
+            np.int32(dst_offset_xy),
+            np.int32(src_offset_xy),
+            np.int32(transfer_shape_xy),
+            global_size=transfer_shape_xy,
+        )
+        if wait:
+            evt.wait()

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/opencl/padding.py

@@ -0,0 +1,6 @@
+from ..processing.padding_opencl import *
+from ..utils import deprecation_warning
+
+deprecation_warning(
+    "nabu.opencl.padding has been moved to nabu.processing.padding_opencl", do_print=True, func_name="padding_opencl"
+)

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/opencl/processing.py

@@ -0,0 +1,62 @@
+from ..processing.processing_base import ProcessingBase
+from ..utils import MissingComponentError
+from .utils import get_opencl_context, __has_pyopencl__
+
+if __has_pyopencl__:
+    from .kernel import OpenCLKernel
+    import pyopencl as cl
+    import pyopencl.array as parray
+    from pyopencl.tools import dtype_to_ctype
+
+    OpenCLArray = parray.Array
+else:
+    OpenCLArray = MissingComponentError("pyopencl")
+    dtype_to_ctype = MissingComponentError("pyopencl")
+
+
+class OpenCLProcessing(ProcessingBase):
+    array_class = OpenCLArray
+    dtype_to_ctype = dtype_to_ctype
+
+    def __init__(self, ctx=None, device_type="all", queue=None, profile=False, **kwargs):
+        """
+        Initialie a OpenCLProcessing instance.
+
+        Parameters
+        ----------
+        ctx: pycuda.driver.Context, optional
+            Existing context to use. If provided, do not create a new context.
+        cleanup_at_exit: bool, optional
+            Whether to clean-up the context at exit.
+            Ignored if ctx is not None.
+        """
+        super().__init__()
+        if queue is not None:
+            # re-use an existing queue. In this case the this instance is mostly for convenience
+            ctx = queue.context
+        if ctx is None:
+            self.ctx = get_opencl_context(device_type=device_type, **kwargs)
+        else:
+            self.ctx = ctx
+        if queue is None:
+            queue_init_kwargs = {}
+            if profile:
+                queue_init_kwargs = {"properties": cl.command_queue_properties.PROFILING_ENABLE}
+            queue = cl.CommandQueue(self.ctx, **queue_init_kwargs)
+        self.queue = queue
+
+    # TODO push_context, pop_context ?
+
+    def _allocate_array_mem(self, shape, dtype):
+        return parray.zeros(self.queue, shape, dtype)
+
+    def kernel(self, kernel_name, filename=None, src=None, automation_params=None, **build_kwargs):
+        return OpenCLKernel(
+            kernel_name,
+            None,
+            queue=self.queue,
+            filename=filename,
+            src=src,
+            automation_params=automation_params,
+            **build_kwargs,
+        )

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/opencl/utils.py

@@ -0,0 +1,292 @@
+import numpy as np
+from ..utils import check_supported
+
+try:
+    import pyopencl as cl
+    import pyopencl.array as parray
+
+    __has_pyopencl__ = True
+    __pyopencl_error_msg__ = None
+except ImportError as err:
+    __has_pyopencl__ = False
+    __pyopencl_error_msg__ = str(err)
+from ..resources.gpu import GPUDescription
+
+
+def get_opencl_devices(
+    device_type,
+    vendor=None,
+    name=None,
+    order_by="global_mem_size",
+    exclude_platforms=None,
+    exclude_vendors=None,
+    prefer_GPU=True,
+):
+    """
+    List available OpenCL devices.
+
+    Parameters
+    ----------
+    device_type: str
+        Type of device, can be "CPU", "GPU" or "all".
+    vendor: str, optional
+        Filter devices by vendor, eg. "NVIDIA"
+    name: Filter devices by names, eg. "GeForce RTX 3080"
+    order_by: str, optional
+        Order results in decreasing order of this value. Default is to sort by global memory size.
+    exclude_platforms: str, optional
+        Platforms to exclude, eg. "Portable Computing Language"
+    exclude_vendors: str, optional
+        Vendors to be excluded
+    prefer_GPU: bool, optional
+        Whether to put GPUs on top of the returned list, regardless of the "order_by" parameter.
+        This can be useful when sorting by global memory size, as CPUs often have a bigger memory size.
+
+    Returns
+    -------
+    devices: list of pyopencl.Device
+        List of OpenCL devices matching the criteria, and ordered by the 'order_by' parameter.
+        The list may be empty.
+
+    """
+    exclude_platforms = exclude_platforms or []
+    exclude_vendors = exclude_vendors or []
+    dev_type = {
+        "cpu": cl.device_type.CPU,
+        "gpu": cl.device_type.GPU,
+        "accelerator": cl.device_type.ACCELERATOR,
+        "all": cl.device_type.ALL,
+        "any": cl.device_type.ALL,
+    }
+    device_type = device_type.lower()
+    check_supported(device_type, dev_type.keys(), "device_type")
+
+    devices = []
+    for platform in cl.get_platforms():
+        if vendor is not None and vendor.lower() not in platform.vendor.lower():
+            continue
+        if any(excluded_platform.lower() in platform.name.lower() for excluded_platform in exclude_platforms):
+            continue
+        if any(excluded_vendor.lower() in platform.vendor.lower() for excluded_vendor in exclude_vendors):
+            continue
+        for device in platform.get_devices():
+            if device.type & dev_type[device_type] == 0:
+                continue
+            if name is not None and name.lower() not in device.name.lower():
+                continue
+            devices.append(device)
+    if order_by is not None:
+        devices = sorted(devices, key=lambda dev: getattr(dev, order_by), reverse=True)
+    if prefer_GPU:
+        # put GPUs devices on top of the list
+        devices = [dev for dev in devices if dev.type & dev_type["gpu"] > 0] + [
+            dev for dev in devices if dev.type & dev_type["gpu"] == 0
+        ]
+    return devices
+
+
+def usable_opencl_devices():
+    """
+    Test the available OpenCL platforms/devices.
+
+    Returns
+    --------
+    platforms: dict
+        Dictionary where the key is the platform name, and the value is a list
+        of `silx.opencl.common.Device` object.
+    """
+    platforms = {}
+    for platform in cl.get_platforms():
+        platforms[platform.name] = platform.get_devices()
+    return platforms
+
+
+def detect_opencl_gpus():
+    """
+    Get the available OpenCL-compatible GPUs.
+
+    Returns
+    --------
+    gpus: dict
+        Nested dictionary where the keys are OpenCL platform names,
+        values are dictionary of GPU IDs and `silx.opencl.common.Device` object.
+    error_msg: str
+        In the case where there is an error, the message is returned in this item.
+        Otherwise, it is a None object.
+    """
+    gpus = {}
+    error_msg = None
+    if not (__has_pyopencl__):
+        return {}, __pyopencl_error_msg__
+    try:
+        platforms = usable_opencl_devices()
+    except Exception as exc:
+        error_msg = str(exc)
+    if error_msg is not None:
+        return {}, error_msg
+    for platform_name, devices in platforms.items():
+        for d_id, device in enumerate(devices):
+            if device.type == cl.device_type.GPU:  # and bool(device.available):
+                if platform_name not in gpus:
+                    gpus[platform_name] = {}
+                gpus[platform_name][d_id] = device
+    return gpus, None
+
+
+def collect_opencl_gpus():
+    """
+    Return a dictionary of platforms and brief description of each OpenCL-compatible
+    GPU with a few fields
+    """
+    gpus, error_msg = detect_opencl_gpus()
+    if error_msg is not None:
+        return None
+    opencl_gpus = {}
+    for platform, gpus in gpus.items():
+        for gpu_id, gpu in gpus.items():
+            if platform not in opencl_gpus:
+                opencl_gpus[platform] = {}
+            opencl_gpus[platform][gpu_id] = GPUDescription(gpu, device_id=gpu_id).get_dict()
+            opencl_gpus[platform][gpu_id]["platform"] = platform
+    return opencl_gpus
+
+
+def collect_opencl_cpus():
+    """
+    Return a dictionary of platforms and brief description of each OpenCL-compatible
+    CPU with a few fields
+    """
+    opencl_cpus = {}
+    platforms = usable_opencl_devices()
+    for platform, devices in platforms.items():
+        if "cuda" in platform.lower():
+            continue
+        opencl_cpus[platform] = {}
+        for device_id, device in enumerate(devices):  # device_id might be inaccurate
+            if device.type != cl.device_type.CPU:
+                continue
+            opencl_cpus[platform][device_id] = GPUDescription(device).get_dict()
+            opencl_cpus[platform][device_id]["platform"] = platform
+    return opencl_cpus
+
+
+def get_opencl_context(device_type, **kwargs):
+    """
+    Create an OpenCL context. Please refer to 'get_opencl_devices' documentation
+    """
+    devices = get_opencl_devices(device_type, **kwargs)
+    if devices == []:
+        raise RuntimeError("No OpenCL device found for device_type='%s' and %s" % (device_type, str(kwargs)))
+    return cl.Context([devices[0]])
+
+
+def replace_array_memory(arr, new_shape):
+    """
+    Replace the underlying buffer data of a  `pyopencl.array.Array`.
+    This function is dangerous !
+    It should merely be used to clear memory, the array should not be used afterwise.
+    """
+    arr.data.release()
+    arr.base_data = cl.Buffer(arr.context, cl.mem_flags.READ_WRITE, np.prod(new_shape) * arr.dtype.itemsize)
+    arr.shape = new_shape
+    # strides seems to be updated by pyopencl
+    return arr
+
+
+def pick_opencl_cpu_platform(opencl_cpus):
+    """
+    Pick the best OpenCL implementation for the opencl cpu.
+    This function assume that there is only one opencl-enabled CPU on the
+    current machine, but there might be several OpenCL implementations/vendors.
+
+
+    Parameters
+    ----------
+    opencl_cpus: dict
+        Dictionary with the available opencl-enabled CPUs.
+        Usually obtained with collect_opencl_cpus().
+
+    Returns
+    -------
+    cpu: dict
+        A dictionary describing the CPU.
+    """
+    if len(opencl_cpus) == 0:
+        raise ValueError("No CPU to pick")
+    name2device = {}
+    for platform, devices in opencl_cpus.items():
+        for device_id, device_desc in devices.items():
+            name2device.setdefault(device_desc["name"], [])
+            name2device[device_desc["name"]].append(platform)
+    if len(name2device) > 1:
+        raise ValueError("Expected at most one CPU but got %d: %s" % (len(name2device), list(name2device.keys())))
+    cpu_name = list(name2device.keys())[0]
+    platforms = name2device[cpu_name]
+    # Several platforms for the same CPU
+    res = opencl_cpus[platforms[0]]
+    if len(platforms) > 1:
+        if "intel" in cpu_name.lower():
+            for platform in platforms:
+                if "intel" in platform.lower():
+                    res = opencl_cpus[platform]
+    #
+    return res[list(res.keys())[0]]
+
+
+def allocate_texture(ctx, shape, support_1D=False):
+    """
+    Allocate an OpenCL image ("texture").
+
+    Parameters
+    ----------
+    ctx: OpenCL context
+        OpenCL context
+    shape: tuple of int
+        Shape of the image. Note that pyopencl and OpenCL < 1.2
+        do not support 1D images, so 1D images are handled as 2D with one row
+    support_1D: bool, optional
+         force the image to be 1D if the shape has only one dim
+    """
+    if len(shape) == 1 and not (support_1D):
+        shape = (1,) + shape
+    return cl.Image(
+        ctx,
+        cl.mem_flags.READ_ONLY | cl.mem_flags.USE_HOST_PTR,
+        cl.ImageFormat(cl.channel_order.INTENSITY, cl.channel_type.FLOAT),
+        hostbuf=np.zeros(shape[::-1], dtype=np.float32),
+    )
+
+
+def check_textures_availability(ctx):
+    """
+    Check whether textures are supported on the current OpenCL context.
+    """
+    try:
+        dummy_texture = allocate_texture(ctx, (16, 16))
+        # Need to further access some attributes (pocl)
+        dummy_height = dummy_texture.height
+        textures_available = True
+        del dummy_texture, dummy_height
+    except (cl.RuntimeError, cl.LogicError):
+        textures_available = False
+    # Nvidia Fermi GPUs (compute capability 2.X) do not support opencl read_imagef
+    # There is no way to detect this until a kernel is compiled
+    try:
+        cc = ctx.devices[0].compute_capability_major_nv
+        textures_available &= cc >= 3
+    except (cl.LogicError, AttributeError):  # probably not a Nvidia GPU
+        pass
+    #
+    return textures_available
+
+
+def copy_to_texture(queue, dst_texture, src_array, dtype=np.float32):
+    shape = src_array.shape
+    if isinstance(src_array, parray.Array):
+        return cl.enqueue_copy(queue, dst_texture, src_array.data, offset=0, origin=(0, 0), region=shape[::-1])
+    elif isinstance(src_array, np.ndarray):
+        if not (src_array.flags["C_CONTIGUOUS"] and src_array.dtype == dtype):
+            src_array = np.ascontiguousarray(src_array, dtype=dtype)
+        return cl.enqueue_copy(queue, dst_texture, src_array, origin=(0, 0), region=shape[::-1])
+    else:
+        raise ValueError("Unknown source array type")