# -*- coding: utf-8 -*- # PyVkFFT # (c) 2021- : ESRF-European Synchrotron Radiation Facility # authors: # Vincent Favre-Nicolin, favre@esrf.fr import warnings import ctypes import numpy as np import pyopencl as cl import pyopencl.array as cla from .base import load_library, VkFFTApp as VkFFTAppBase, check_vkfft_result, ctype_int_size_p from . import config from .tune import tune_vkfft try: _vkfft_opencl = load_library("_vkfft_opencl") class _types: """Aliases""" vkfft_config = ctypes.c_void_p vkfft_app = ctypes.c_void_p _vkfft_opencl.make_config.restype = ctypes.c_void_p _vkfft_opencl.make_config.argtypes = [ctype_int_size_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_size_t, ctype_int_size_p, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctype_int_size_p, ctypes.c_int] _vkfft_opencl.init_app.restype = ctypes.c_void_p _vkfft_opencl.init_app.argtypes = [_types.vkfft_config, ctypes.c_void_p, ctypes.POINTER(ctypes.c_int), ctypes.POINTER(ctypes.c_size_t), ctype_int_size_p, ctype_int_size_p] _vkfft_opencl.fft.restype = ctypes.c_int _vkfft_opencl.fft.argtypes = [_types.vkfft_app, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] _vkfft_opencl.ifft.restype = ctypes.c_int _vkfft_opencl.ifft.argtypes = [_types.vkfft_app, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] _vkfft_opencl.free_app.restype = None _vkfft_opencl.free_app.argtypes = [_types.vkfft_app] _vkfft_opencl.free_config.restype = None _vkfft_opencl.free_config.argtypes = [_types.vkfft_config] _vkfft_opencl.vkfft_version.restype = ctypes.c_uint32 _vkfft_opencl.vkfft_version.argtypes = None _vkfft_opencl.vkfft_max_fft_dimensions.restype = ctypes.c_uint32 _vkfft_opencl.vkfft_max_fft_dimensions.argtypes = None except OSError: # This is used for doc generation import sys if 'sphinx' in sys.modules: pass else: raise class VkFFTApp(VkFFTAppBase): """ VkFFT application interface implementing a FFT plan. """ def __init__(self, shape, dtype: type, queue: cl.CommandQueue, ndim=None, inplace=True, norm=1, r2c=False, dct=False, dst=False, axes=None, strides=None, tune_config=None, r2c_odd=False, verbose=False, **kwargs): """ Init function for the VkFFT application. :param shape: the shape of the array to be transformed. The number of dimensions of the array can be larger than the FFT dimensions. :param dtype: the numpy dtype of the source array (can be complex64 or complex128) :param queue: the pyopencl CommandQueue to use for the transform. :param ndim: the number of dimensions to use for the FFT. By default, uses the array dimensions. Can be smaller, e.g. ndim=2 for a 3D array to perform a batched 3D FFT on all the layers. The FFT is always performed along the last axes if the array's number of dimension is larger than ndim, i.e. on the x-axis for ndim=1, on the x and y axes for ndim=2, etc.. Unless axes are given. :param inplace: if True (the default), performs an inplace transform and the destination array should not be given in fft() and ifft(). :param norm: if 0 (unnormalised), every transform multiplies the L2 norm of the array by its size (or the size of the transformed array if ndimcomplex transform, where the complex destination is a half-hermitian array. For an inplace transform, if the transformed data shape is (...,nx), the input float array should have a shape of (..., nx+2), the last two columns being ignored in the input data, and the resulting complex array (using pycuda's GPUArray.view(dtype=np.complex64) to reinterpret the type) will have a shape (..., nx//2 + 1). For an out-of-place transform, if the input (real) shape is (..., nx), the output (complex) shape should be (..., nx//2+1). Note that for C2R transforms with ndim>=2, the source (complex) array is modified. For an inplace transform with an odd-sized x-axis, see the r2c_odd parameter. :param dct: used to perform a Direct Cosine Transform (DCT) aka a R2R transform. An integer can be given to specify the type of DCT (1, 2, 3 or 4). if dct=True, the DCT type 2 will be performed, following scipy's convention. :param dst: used to perform a Direct Sine Transform (DST) aka a R2R transform. An integer can be given to specify the type of DST (1, 2, 3 or 4). if dst=True, the DST type 2 will be performed, following scipy's convention. :param axes: a list or tuple of axes along which the transform should be made. if None, the transform is done along the ndim fastest axes, or all axes if ndim is None. For R2C transforms, the fast axis must be transformed. :param strides: the array strides - needed if not C-ordered. :param tune_config: this can be used to automatically generate an optimised set of VkFFT parameters by testing various configurations and measuring the FFT speed, in a manner similar to fftw's FFTW_MEASURE. This should be a dictionary including the backend used and the parameter values which will be tested. This is EXPERIMENTAL, as wrong parameters may lead to crashes. Note that this will allocate temporary GPU arrays, unless the arrays to used have been passed as parameters ('dest' and 'src'). Examples: tune={'backend':'cupy} - minimal example, will automatically test a small set of parameters (4 to 9 tests), depending on the GPU type. tune={'backend':'cupy, 'warpSize':[8,16,32,64,128]}: this will test 5 possible values for the warpSize. tune={'backend':'cupy, 'groupedBatch':[[-1,-1,-1],[8,8,8], [4,16,16}: this will test 3 possible values for groupedBatch. This one is more tricky to use. tune={'backend':'cupy, 'warpSize':[8,16,32,64,128], 'src':a}: this will test 5 possible values for the warpSize, with a given source GPU array. This would only be valid for an inplace transform as no destination array is given. :param r2c_odd: this should be set to True to perform an inplace r2c/c2r transform with an odd-sized fast (x) axis. Explanation: to perform a 1D inplace transform of an array with 100 elements, the input array should have a 100+2 size, resulting in a half-Hermitian array of size 51. If the input data has a size of 101, the input array should also be padded to 102 (101+1), and the resulting half-Hermitian array also has a size of 51. A flag is thus needed to differentiate the cases of 100+2 or 101+1. :param verbose: if True, print a 1-string info about this VkFFTApp. See __str__ for details. :raises RuntimeError: if the initialisation fails, e.g. if the GPU driver has not been properly initialised, or if the transform dimensions or data type are not allowed by VkFFT. """ if tune_config is not None: kwargs = tune_vkfft(tune_config, shape=shape, dtype=dtype, ndim=ndim, queue=queue, inplace=inplace, norm=norm, r2c=r2c, dct=dct, dst=dst, axes=axes, strides=strides, verbose=False, r2c_odd=r2c_odd, **kwargs)[0] super().__init__(shape, dtype, ndim=ndim, inplace=inplace, norm=norm, r2c=r2c, dct=dct, dst=dst, axes=axes, strides=strides, r2c_odd=r2c_odd, **kwargs) self.queue = queue if self.precision == 2 and 'cl_khr_fp16' not in self.queue.device.extensions: raise RuntimeError("Half precision required but cl_khr_fp16 extension is not available") if self.precision == 8 and 'cl_khr_fp64' not in self.queue.device.extensions: raise RuntimeError("Double precision required but cl_khr_fp64 extension is not available") self.config = self._make_config() if self.config is None: raise RuntimeError("Error creating VkFFTConfiguration. Was the OpenCL context properly initialised ?") res = ctypes.c_int(0) # Size of tmp buffer allocated by VkFFT - if any tmp_buffer_nbytes = ctypes.c_size_t(0) # 0 or 1 for each axis, only if the Bluestein algorithm is used use_bluestein_fft = np.zeros(vkfft_max_fft_dimensions(), dtype='long') # number of axis upload per dimension num_axis_upload = np.zeros(vkfft_max_fft_dimensions(), dtype='long') self.app = _vkfft_opencl.init_app(self.config, queue.int_ptr, ctypes.byref(res), ctypes.byref(tmp_buffer_nbytes), use_bluestein_fft, num_axis_upload) try: check_vkfft_result(res, shape=shape, dtype=dtype, ndim=ndim, inplace=inplace, norm=norm, r2c=r2c, dct=dct, dst=dst, axes=axes, backend="opencl:%s:%s" % (queue.device.platform.name, queue.device.name), strides=self.strides0, vkfft_shape=self.shape, vkfft_skip=self.skip_axis, vkfft_nbatch=self.n_batch) except Exception as ex: print("clVkFFTApp error:", shape, strides, dtype, ndim, axes, self.shape, self.skip_axis) raise ex if self.app is None: raise RuntimeError("Error creating VkFFTApplication. Was the OpenCL context properly initialised ?") self.tmp_buffer_nbytes = np.int64(tmp_buffer_nbytes) self.use_bluestein_fft = [bool(n) for n in use_bluestein_fft[:len(self.shape)]] self.nb_axis_upload = [int(num_axis_upload[i] * (self.skip_axis[i] is False)) for i in range(len(self.shape))] if verbose: print(self) def __del__(self): """ Takes care of deleting allocated memory in the underlying VkFFTApplication and VkFFTConfiguration. """ if self.app is not None: _vkfft_opencl.free_app(self.app) if self.config is not None: _vkfft_opencl.free_config(self.config) def _make_config(self): """ Create a vkfft configuration for a FFT transform""" if len(self.shape) > vkfft_max_fft_dimensions(): raise RuntimeError(f"Too many FFT dimensions after collapsing non-transform axes: " f"{len(self.shape)}>{vkfft_max_fft_dimensions()}") shape = np.ones(vkfft_max_fft_dimensions(), dtype='long') shape[:len(self.shape)] = self.shape skip = np.zeros(vkfft_max_fft_dimensions(), dtype='long') skip[:len(self.skip_axis)] = self.skip_axis grouped_batch = np.empty(vkfft_max_fft_dimensions(), dtype='long') grouped_batch.fill(-1) grouped_batch[:len(self.groupedBatch)] = self.groupedBatch if self.r2c and self.inplace: # the last one or two columns are ignored in the R array, and will be used # in the C array with a size nx//2+1 if self.r2c_odd: shape[0] -= 1 else: shape[0] -= 2 if self.norm == "ortho": norm = 0 else: norm = self.norm # We pass fake buffer pointer addresses to VkFFT. The real ones will be # given when performing the actual FFT. ctx = self.queue.context device = ctx.devices[0] platform = device.platform dest_gpudata = 2 if self.inplace: dest_gpudata = 0 return _vkfft_opencl.make_config(shape, self.ndim, 1, dest_gpudata, platform.int_ptr, device.int_ptr, ctx.int_ptr, norm, self.precision, int(self.r2c), int(self.dct), int(self.dst), int(self.disableReorderFourStep), int(self.registerBoost), int(self.use_lut), int(self.keepShaderCode), self.n_batch, skip, int(self.coalescedMemory), int(self.numSharedBanks), int(self.aimThreads), int(self.performBandwidthBoost), int(self.registerBoostNonPow2), int(self.registerBoost4Step), int(self.warpSize), grouped_batch, int(self.forceCallbackVersionRealTransforms)) def fft(self, src: cla.Array, dest: cla.Array = None, queue: cl.CommandQueue = None): """ Compute the forward FFT :param src: the source pyopencl Array :param dest: the destination pyopencl Array. Should be None for an inplace transform :param queue: the pyopencl CommandQueue to use for the transform. If not given, the queue of the source array is used. If the queue does not match the application's, a warning is emitted (see config.WARN_OPENCL_QUEUE_MISMATCH). If a queue is not supplied and the source and destination arrays do not have the same queue, then a RuntimeError is raised. :raises RuntimeError: in case of a GPU kernel launch error :return: the transformed array. For a R2C inplace transform, the complex view of the array is returned. """ if not queue: if dest is None: if src.queue is None: warnings.warn("queue is not given and the source array does not have a queue " "attached to it. Falling back to the queue given to the constructor " "of VkFFTApp. This is deprecated and will stop working in the future. " "Use src_array.with_queue(queue) to attach a queue to the array or " "pass a queue to this method", DeprecationWarning) queue = self.queue else: queue = src.queue elif dest.queue is None and src.queue is None: warnings.warn("queue is not given and the source/dest arrays do not have a queue " "attached to it. Falling back to the queue given to the constructor " "of VkFFTApp. This is deprecated and will stop working in the future. " "Use src_array.with_queue(queue) to attach a queue to the array or " "pass a queue to this method", DeprecationWarning) queue = self.queue elif dest.queue != src.queue: if dest.queue is None: queue = src.queue elif src.queue is None: queue = dest.queue else: raise RuntimeError("queue is not given and the source/dest arrays queues differ. " "Please supply a queue when calling fft(...), or " "use array.with_queue(queue)") else: queue = src.queue if config.WARN_OPENCL_QUEUE_MISMATCH and self.queue != queue: warnings.warn("Using the array queue for transform, which differs from the application " "queue. NB: this warning will be removed in a future version , and you can " "suppress it using config.WARN_OPENCL_QUEUE_MISMATCH", UserWarning) if self.inplace: if dest is not None: if src.data.int_ptr != dest.data.int_ptr: raise RuntimeError("VkFFTApp.fft: dest is not None but this is an inplace transform") res = _vkfft_opencl.fft(self.app, int(src.data.int_ptr), int(src.data.int_ptr), int(queue.int_ptr)) check_vkfft_result(res, src.shape, src.dtype, self.ndim, self.inplace, self.norm, self.r2c, self.dct, self.dst, backend="opencl:%s:%s" % (queue.device.platform.name, queue.device.name), strides=self.strides0, vkfft_shape=self.shape, vkfft_skip=self.skip_axis, vkfft_nbatch=self.n_batch) if self.norm == "ortho": src *= self._get_fft_scale(norm=0) if self.r2c: if src.dtype == np.float32: return src.view(dtype=np.complex64) elif src.dtype == np.float64: return src.view(dtype=np.complex128) return src else: if dest is None: raise RuntimeError("VkFFTApp.fft: dest is None but this is an out-of-place transform") elif src.data.int_ptr == dest.data.int_ptr: raise RuntimeError("VkFFTApp.fft: dest and src are identical but this is an out-of-place transform") if self.r2c: tmp = src.size // src.shape[self.fast_axis] * (src.shape[self.fast_axis] // 2 + 1) if dest.size != tmp: raise RuntimeError(f"VkFFTApp.fft the destination array size ({dest.size}) is not " f"equal to the expected one ({tmp}) for an R2C transform " f"[fast axis={self.fast_axis}, src shape={src.shape}, " f"axes={self.axes0}, strides={src.strides}]") res = _vkfft_opencl.fft(self.app, int(src.data.int_ptr), int(dest.data.int_ptr), int(queue.int_ptr)) check_vkfft_result(res, src.shape, src.dtype, self.ndim, self.inplace, self.norm, self.r2c, self.dct, self.dst, backend="opencl:%s:%s" % (queue.device.platform.name, queue.device.name), strides=self.strides0, vkfft_shape=self.shape, vkfft_skip=self.skip_axis, vkfft_nbatch=self.n_batch) if self.norm == "ortho": dest *= self._get_fft_scale(norm=0) return dest def ifft(self, src: cla.Array, dest: cla.Array = None, queue: cl.CommandQueue = None): """ Compute the backward FFT :param src: the source pyopencl.Array :param dest: the destination pyopencl.Array. Can be None for an inplace transform :param queue: the pyopencl CommandQueue to use for the transform. If not given, the queue of the source array is used. If the queue does not match the application, a warning is emitted (see config.WARN_OPENCL_QUEUE_MISMATCH). If a queue is not supplied and the source and destination arrays do not have the same queue, then a RuntimeError is raised. :raises RuntimeError: in case of a GPU kernel launch error :return: the transformed array. For a C2R inplace transform, the float view of the array is returned. """ if not queue: if dest is None: if src.queue is None: warnings.warn("queue is not given and the source array does not have a queue " "attached to it. Falling back to the queue given to the constructor " "of VkFFTApp. This is deprecated and will stop working in the future. " "Use src_array.with_queue(queue) to attach a queue to the array or " "pass a queue to this method", DeprecationWarning) queue = self.queue else: queue = src.queue elif dest.queue is None and src.queue is None: warnings.warn("queue is not given and the source/dest arrays do not have a queue " "attached to it. Falling back to the queue given to the constructor " "of VkFFTApp. This is deprecated and will stop working in the future. " "Use src_array.with_queue(queue) to attach a queue to the array or " "pass a queue to this method", DeprecationWarning) queue = self.queue elif dest.queue != src.queue: if dest.queue is None: queue = src.queue elif src.queue is None: queue = dest.queue else: raise RuntimeError("queue is not given and the source/dest arrays queues differ. " "Please supply a queue when calling fft(...), or " "use array.with_queue(queue)") else: queue = src.queue if config.WARN_OPENCL_QUEUE_MISMATCH and self.queue != queue: warnings.warn("Using the array queue for transform, which differs from the application " "queue. NB: this warning will be removed in a future version , and you can " "suppress it using config.WARN_OPENCL_QUEUE_MISMATCH", UserWarning) if self.inplace: if dest is not None: if src.data.int_ptr != dest.data.int_ptr: raise RuntimeError("VkFFTApp.fft: dest!=src but this is an inplace transform") res = _vkfft_opencl.ifft(self.app, int(src.data.int_ptr), int(src.data.int_ptr), int(queue.int_ptr)) check_vkfft_result(res, src.shape, src.dtype, self.ndim, self.inplace, self.norm, self.r2c, self.dct, self.dst, backend="opencl:%s:%s" % (queue.device.platform.name, queue.device.name), strides=self.strides0, vkfft_shape=self.shape, vkfft_skip=self.skip_axis, vkfft_nbatch=self.n_batch) if self.norm == "ortho": src *= self._get_ifft_scale(norm=0) if self.r2c: if src.dtype == np.complex64: return src.view(dtype=np.float32) elif src.dtype == np.complex128: return src.view(dtype=np.float64) return src else: if dest is None: raise RuntimeError("VkFFTApp.ifft: dest is None but this is an out-of-place transform") elif src.data.int_ptr == dest.data.int_ptr: raise RuntimeError("VkFFTApp.ifft: dest and src are identical but this is an out-of-place transform") if self.r2c: assert (src.size == dest.size // dest.shape[self.fast_axis] * (dest.shape[self.fast_axis] // 2 + 1)) # Special case, src and dest buffer sizes are different, # VkFFT is configured to go back to the source buffer res = _vkfft_opencl.ifft(self.app, int(dest.data.int_ptr), int(src.data.int_ptr), int(queue.int_ptr)) else: res = _vkfft_opencl.ifft(self.app, int(src.data.int_ptr), int(dest.data.int_ptr), int(queue.int_ptr)) check_vkfft_result(res, src.shape, src.dtype, self.ndim, self.inplace, self.norm, self.r2c, self.dct, self.dst, backend="opencl:%s:%s" % (queue.device.platform.name, queue.device.name), strides=self.strides0, vkfft_shape=self.shape, vkfft_skip=self.skip_axis, vkfft_nbatch=self.n_batch) if self.norm == "ortho": dest *= self._get_ifft_scale(norm=0) return dest def vkfft_version(): """ Get VkFFT version :return: version as X.Y.Z """ int_ver = _vkfft_opencl.vkfft_version() return "%d.%d.%d" % (int_ver // 10000, (int_ver % 10000) // 100, int_ver % 100) def vkfft_max_fft_dimensions(): """ Get the maximum number of dimensions VkFFT can handle. This is set at compile time. VkFFT default is 4, pyvkfft sets this to 8. Note that consecutive non-transformed are collapsed into a single axis, reducing the effective number of dimensions. :return: VKFFT_MAX_FFT_DIMENSIONS """ return _vkfft_opencl.vkfft_max_fft_dimensions()