WIP: Correlation matrix

dynamix/correlator/common.py

@@ -1,15 +1,29 @@
 import numpy as np
-from os import linesep
 
 import pyopencl.array as parray
 from pyopencl.tools import dtype_to_ctype
 from silx.opencl.common import pyopencl as cl
 from silx.opencl.processing import OpenclProcessing, KernelContainer
+from ..utils import _compile_kernels
 
 
-class BaseCorrelator(object):
+class BaseCorrelator:
     "Abstract base class for all Correlators"
 
+    default_extra_options = {
+        # dtype for indices offsets
+        # must be extended to unsigned long (np.uint64) if total_nnz > 4294967295
+        "offset_dtype": np.uint32,
+        # dtype for intermediate result (correlation in integer).
+        # must be extended to unsigned long for large nnz_per_frame and/or events counts
+        "res_dtype": np.uint32,
+        "sums_dtype": np.uint32,
+        # dtype for q-mask.
+        #  must be extended to int if number of q-bins > 127
+        "qmask_dtype": np.int8,
+    }
+
+
     def __init__(self):
         self.nframes = None
         self.shape = None
@@ -19,12 +33,30 @@ def __init__(self):
         self.weights = None
         self.scale_factors = None
 
-    def _set_parameters(self, shape, nframes, qmask, scale_factor, extra_options):
+
+    def _set_dtype(self, dtype):
+        # Configure data types - important for OpenCL kernels,
+        # as some operations are better performed on integer types (ex. atomic),
+        # but some overflow can occur for large/non-sparse data.
+
+        # data type for data. Usually the data values lie in a small range (less than 255)
+        self.dtype = dtype
+        # Other data types
+        self._offset_dtype = self.extra_options["offset_dtype"]
+        self._qmask_dtype = self.extra_options["qmask_dtype"]
+        self._res_dtype = self.extra_options["res_dtype"]
+        self._sums_dtype = self.extra_options["sums_dtype"]
+        self._pix_idx_dtype = np.uint32 # won't change (goes from 0 to N_x*N_y)
+        self._output_dtype = np.float32 # won't change
+
+
+    def _set_parameters(self, shape, nframes, qmask, scale_factor, extra_options, dtype):
+        self._configure_extra_options(extra_options)
+        self._set_dtype(dtype)
         self.nframes = nframes
         self._set_shape(shape)
         self._set_qmask(qmask=qmask)
         self._set_scale_factor(scale_factor=scale_factor)
-        self._configure_extra_options(extra_options)
 
     def _set_shape(self, shape):
         if np.isscalar(shape):
@@ -36,21 +68,21 @@ def _set_shape(self, shape):
     def _set_qmask(self, qmask=None):
         self.qmask = None
         if qmask is None:
-            self.bins = np.array([1], dtype=np.int32)
+            self.bins = np.array([1], dtype=self._qmask_dtype)
             self.n_bins = 1
-            self.qmask = np.ones(self.shape, dtype=np.int32)
+            self.qmask = np.ones(self.shape, dtype=self._qmask_dtype)
         else:
-            self.qmask = np.ascontiguousarray(qmask, dtype=np.int32)
+            self.qmask = np.ascontiguousarray(qmask, dtype=self._qmask_dtype)
             self.n_bins = self.qmask.max()
-            self.bins = np.arange(1, self.n_bins + 1, dtype=np.int32)
+            self.bins = np.arange(1, self.n_bins + 1, dtype=self._qmask_dtype)
         self.output_shape = (self.n_bins, self.nframes)
 
     def _set_weights(self, weights=None):
         if weights is None:
-            self.weights = np.ones(self.shape, dtype=self.output_dtype)
+            self.weights = np.ones(self.shape, dtype=self._output_dtype)
             return
         assert weights.shape == self.shape
-        self.weights = np.ascontiguousarray(weights, dtype=self.output_dtype)
+        self.weights = np.ascontiguousarray(weights, dtype=self._output_dtype)
         raise ValueError("Advanced weighting is not implemented yet")
 
     def _set_scale_factor(self, scale_factor=None):
@@ -75,15 +107,14 @@ def _configure_extra_options(self, extra_options):
         """
         :param extra_options: dict
         """
-        self.extra_options = {}
-        if extra_options is not None:
-            self.extra_options.update(extra_options)
+        self.extra_options = self.default_extra_options.copy()
+        self.extra_options.update(extra_options or {})
 
 
 class OpenclCorrelator(BaseCorrelator, OpenclProcessing):
 
     def __init__(
-        self, shape, nframes, qmask=None, dtype=np.int8, weights=None,
+        self, shape, nframes, qmask=None, dtype=np.uint8, weights=None,
         scale_factor=None, extra_options={},
         ctx=None, devicetype="all", platformid=None, deviceid=None,
         block_size=None, memory=None, profile=False
@@ -129,8 +160,11 @@ def __init__(
             the same normalization is used for all bins), or an array which
             must have the same length as the number of bins.
         extra_options: dict, optional
-            Advanced extra options. None available yet.
-
+            Advanced extra options. Available options are:
+                - "offset_dtype": np.uint32
+                - "res_dtype": np.uint32,
+                - "sums_dtype": np.uint32,
+                - "qmask_dtype": np.int8,
 
         Other parameters
         -----------------
@@ -147,24 +181,35 @@ def __init__(
         self._allocate_memory()
 
     def _set_parameters(self, shape, nframes, dtype, qmask, weights, scale_factor, extra_options):
-        BaseCorrelator._set_parameters(self, shape, nframes, qmask, scale_factor, extra_options)
-        self._set_dtype(dtype)
+        BaseCorrelator._set_parameters(self, shape, nframes, qmask, scale_factor, extra_options, dtype)
+        self._set_cl_dtypes()
         self._set_weights(weights)
         self.is_cpu = (self.device.type == "CPU")
+        self.extra_options = self.default_extra_options.copy().update((extra_options or {}))
 
-    def _set_dtype(self, dtype="f"):
-        # add checks ?
-        self.dtype = dtype
-        self.output_dtype = np.float32  # TODO custom ?
-        self.sums_dtype = np.uint32  # TODO custom ?
+    def _set_cl_dtypes(self):
         self.c_dtype = dtype_to_ctype(self.dtype)
-        self.c_sums_dtype = dtype_to_ctype(self.sums_dtype)
+        self.c_sums_dtype = dtype_to_ctype(self._sums_dtype)
         self.idx_c_dtype = "int"  # TODO custom ?
+        self._res_c_dtype = dtype_to_ctype(self._res_dtype)
+        self._dtype_compilation_flags = [
+            "-DDTYPE=%s" % self.c_dtype,
+            "-DOFFSET_DTYPE=%s" % (dtype_to_ctype(self._offset_dtype)),
+            "-DQMASK_DTYPE=%s" % (dtype_to_ctype(self._qmask_dtype)),
+            "-DRES_DTYPE=%s" % self._res_c_dtype
+        ]
+
 
     def _allocate_memory(self):
         # self.d_norm_mask = parray.to_device(self.queue, self.weights)
         if self.qmask is not None:
             self.d_qmask = parray.to_device(self.queue, self.qmask)
+        self._allocated = {}
+        # send scale_factors to device
+        self.allocate_array("d_scale_factors", (self.n_bins,), dtype=np.float64)
+        for bin_idx, scale_factor in enumerate(self.scale_factors.values()):
+            self.d_scale_factors[bin_idx] = scale_factor
+
 
     def _set_data(self, arrays):
         """
@@ -194,22 +239,68 @@ def _reset_arrays(self, arrays_names):
                 setattr(self, array_name, old_array)
                 setattr(self, old_array_name, None)
 
+
+    # -----------
+
+    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 = parray.zeros(self.queue, 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):
+        """
+        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'.
+        """
+        if isinstance(array_ref, parray.Array):
+            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=array_ref.dtype)
+            getattr(self, array_name).set(array_ref)
+        else:
+            raise ValueError("Expected numpy array or pyopencl array")
+        return getattr(self, array_name)
+
+    def get_array(self, array_name):
+        return getattr(self, array_name, None)
+
+
+    # -----------
+
+
     # Overwrite OpenclProcessing.compile_kernel, as it does not support
     # kernels outside silx/opencl/resources
     def compile_kernels(self, kernel_files=None, compile_options=None):
-        kernel_files = kernel_files or self.kernel_files
-
-        allkernels_src = []
-        for kernel_file in kernel_files:
-            with open(kernel_file) as fid:
-                kernel_src = fid.read()
-            allkernels_src.append(kernel_src)
-        allkernels_src = linesep.join(allkernels_src)
-
-        compile_options = compile_options or self.get_compiler_options()
-        try:
-            self.program = cl.Program(self.ctx, allkernels_src).build(options=compile_options)
-        except (cl.MemoryError, cl.LogicError) as error:
-            raise MemoryError(error)
-        else:
-            self.kernels = KernelContainer(self.program)
+        _compile_kernels(self, kernel_files=kernel_files, compile_options=compile_options)
+
+
+    def get_timings(self):
+        if not(self.profile):
+            raise RuntimeError("Need to instantiate this class with profile=True")
+        evd = lambda e: (e.stop - e.start)/1e6
+        return {e.name: evd(e) for e in self.events}
+
+

dynamix/correlator/dense.py

@@ -199,12 +199,12 @@ def _allocate_arrays(self):
         self.d_sums = parray.zeros(
             self.queue,
             self.output_shape,
-            self.sums_dtype
+            self._sums_dtype
         )
         self.d_sums_f = parray.zeros(
             self.queue,
             self.output_shape,
-            self.output_dtype,
+            self._output_dtype,
         )
         self.d_output = parray.zeros(
             self.queue,
@@ -217,7 +217,7 @@ def _normalize_sums(self):
             self.d_sums_f[:] *= self.scale_factors[0]
         else:
             for i, factor in enumerate(self.scale_factors.values()):
-                self.d_sums_f[i] /= np.array([factor], dtype=self.output_dtype)[0]
+                self.d_sums_f[i] /= np.array([factor], dtype=self._output_dtype)[0]
         self.d_sums_f.finish()
 
     def correlate(self, frames):