def get_fft_plan(a, shape=None, axes=None, value_type='C2C'): """ Generate a CUDA FFT plan for transforming up to three axes. Args: a (cupy.ndarray): Array to be transform, assumed to be either C- or F- contiguous. shape (None or tuple of ints): Shape of the transformed axes of the output. If ``shape`` is not given, the lengths of the input along the axes specified by ``axes`` are used. axes (None or int or tuple of int): The axes of the array to transform. If `None`, it is assumed that all axes are transformed. Currently, for performing N-D transform these must be a set of up to three adjacent axes, and must include either the first or the last axis of the array. value_type (str): The FFT type to perform. Acceptable values are: * 'C2C': complex-to-complex transform (default) * 'R2C': real-to-complex transform * 'C2R': complex-to-real transform Returns: a cuFFT plan for either 1D transform (``cupy.cuda.cufft.Plan1d``) or N-D transform (``cupy.cuda.cufft.PlanNd``). .. note:: The returned plan can not only be passed as one of the arguments of the functions in ``cupyx.scipy.fftpack``, but also be used as a context manager for both ``cupy.fft`` and ``cupyx.scipy.fftpack`` functions: .. code-block:: python x = cupy.random.random(16).reshape(4, 4).astype(cupy.complex) plan = cupyx.scipy.fftpack.get_fft_plan(x) with plan: y = cupy.fft.fftn(x) # alternatively: y = cupyx.scipy.fftpack.fftn(x) # no explicit plan is given! # alternatively: y = cupyx.scipy.fftpack.fftn(x, plan=plan) # pass plan explicitly In the first case, no cuFFT plan will be generated automatically, even if ``cupy.fft.config.enable_nd_planning = True`` is set. .. note:: If this function is called under the context of :func:`~cupy.fft.config.set_cufft_callbacks`, the generated plan will have callbacks enabled. .. warning:: This API is a deviation from SciPy's, is currently experimental, and may be changed in the future version. """ # check input array if a.flags.c_contiguous: order = 'C' elif a.flags.f_contiguous: order = 'F' else: raise ValueError('Input array a must be contiguous') if isinstance(shape, int): shape = (shape, ) if isinstance(axes, int): axes = (axes, ) if (shape is not None) and (axes is not None) and len(shape) != len(axes): raise ValueError('Shape and axes have different lengths.') # check axes # n=1: 1d (need axis1D); n>1: Nd if axes is None: n = a.ndim if shape is None else len(shape) axes = tuple(i for i in range(-n, 0)) if n == 1: axis1D = 0 else: # axes is a tuple n = len(axes) if n == 1: axis1D = axes[0] if axis1D >= a.ndim or axis1D < -a.ndim: err = 'The chosen axis ({0}) exceeds the number of '\ 'dimensions of a ({1})'.format(axis1D, a.ndim) raise ValueError(err) elif n > 3: raise ValueError('Only up to three axes is supported') # Note that "shape" here refers to the shape along trasformed axes, not # the shape of the output array, and we need to convert it to the latter. # The result is as if "a=_cook_shape(a); return a.shape" is called. # Because of this, we need to use (possibly unsorted) axes. transformed_shape = shape shape = list(a.shape) if transformed_shape is not None: for s, axis in zip(transformed_shape, axes): if s is not None: if axis == axes[-1] and value_type == 'C2R': s = s // 2 + 1 shape[axis] = s shape = tuple(shape) # check value_type out_dtype = _output_dtype(a.dtype, value_type) fft_type = _convert_fft_type(out_dtype, value_type) # TODO(leofang): figure out if we really have to skip F-order? if n > 1 and value_type != 'C2C' and a.flags.f_contiguous: raise ValueError('C2R/R2C PlanNd for F-order arrays is not supported') # generate plan # (load from cache if it exists, otherwise create one but don't cache it) if n > 1: # ND transform if cupy.cuda.runtime.is_hip and value_type == 'C2R': raise RuntimeError("hipFFT's C2R PlanNd is buggy and unsupported") out_size = _get_fftn_out_size(shape, transformed_shape, axes[-1], value_type) # _get_cufft_plan_nd interacts with plan cache and callback plan = _get_cufft_plan_nd(shape, fft_type, axes=axes, order=order, out_size=out_size, to_cache=False) else: # 1D transform # prepare plan arguments if value_type != 'C2R': out_size = shape[axis1D] else: out_size = _get_fftn_out_size(shape, transformed_shape, axis1D, value_type) batch = prod(shape) // shape[axis1D] devices = None if not config.use_multi_gpus else config._devices keys = (out_size, fft_type, batch, devices) mgr = config.get_current_callback_manager() if mgr is not None: # to avoid a weird segfault, we generate and cache distinct plans # for every possible (load_aux, store_aux) pairs; the plans are # still generated from the same external Python module load_aux = mgr.cb_load_aux_arr store_aux = mgr.cb_store_aux_arr keys += (mgr.cb_load, mgr.cb_store, 0 if load_aux is None else load_aux.data.ptr, 0 if store_aux is None else store_aux.data.ptr) cache = get_plan_cache() cached_plan = cache.get(keys) if cached_plan is not None: plan = cached_plan elif mgr is None: plan = cufft.Plan1d(out_size, fft_type, batch, devices=devices) else: # has callback # TODO(leofang): support multi-GPU callback (devices is ignored) if devices: raise NotImplementedError('multi-GPU cuFFT callbacks are not ' 'yet supported') plan = mgr.create_plan(('Plan1d', keys[:-3])) mgr.set_callbacks(plan) return plan
def _get_cufft_plan_nd(shape, fft_type, axes=None, order='C', out_size=None, to_cache=True): """Generate a CUDA FFT plan for transforming up to three axes. Args: shape (tuple of int): The shape of the array to transform fft_type (int): The FFT type to perform. Supported values are: `cufft.CUFFT_C2C`, `cufft.CUFFT_C2R`, `cufft.CUFFT_R2C`, `cufft.CUFFT_Z2Z`, `cufft.CUFFT_Z2D`, and `cufft.CUFFT_D2Z`. axes (None or int or tuple of int): The axes of the array to transform. Currently, these must be a set of up to three adjacent axes and must include either the first or the last axis of the array. If `None`, it is assumed that all axes are transformed. order ({'C', 'F'}): Specify whether the data to be transformed has C or Fortran ordered data layout. out_size (int): The output length along the last axis for R2C/C2R FFTs. For C2C FFT, this is ignored (and set to `None`). to_cache (bool): Whether to cache the generated plan. Default is ``True``. Returns: plan (cufft.PlanNd): A cuFFT Plan for the chosen `fft_type`. """ ndim = len(shape) if fft_type in (cufft.CUFFT_C2C, cufft.CUFFT_Z2Z): value_type = 'C2C' elif fft_type in (cufft.CUFFT_C2R, cufft.CUFFT_Z2D): value_type = 'C2R' else: # CUFFT_R2C or CUFFT_D2Z value_type = 'R2C' if axes is None: # transform over all axes fft_axes = tuple(range(ndim)) else: _, fft_axes = _prep_fftn_axes(ndim, s=None, axes=axes, value_type=value_type) if not _nd_plan_is_possible(fft_axes, ndim): raise ValueError( "An n-dimensional cuFFT plan could not be created. The axes must " "be contiguous and non-repeating. Between one and three axes can " "be transformed and either the first or last axis must be " "included in axes.") if order not in ['C', 'F']: raise ValueError('order must be \'C\' or \'F\'') """ For full details on idist, istride, iembed, etc. see: http://docs.nvidia.com/cuda/cufft/index.html#advanced-data-layout in 1D: input[b * idist + x * istride] output[b * odist + x * ostride] in 2D: input[b * idist + (x * inembed[1] + y) * istride] output[b * odist + (x * onembed[1] + y) * ostride] in 3D: input[b * idist + ((x * inembed[1] + y) * inembed[2] + z) * istride] output[b * odist + ((x * onembed[1] + y) * onembed[2] + z) * ostride] """ # At this point, _default_fft_func() guarantees that for F-order arrays # we only need to consider C2C, and not C2R or R2C. # TODO(leofang): figure out if we really have to skip F-order? in_dimensions = [shape[d] for d in fft_axes] if order == 'F': in_dimensions = in_dimensions[::-1] in_dimensions = tuple(in_dimensions) if fft_type in (cufft.CUFFT_C2C, cufft.CUFFT_Z2Z): out_dimensions = in_dimensions plan_dimensions = in_dimensions else: out_dimensions = list(in_dimensions) if out_size is not None: # for C2R & R2C out_dimensions[-1] = out_size # only valid for C order! out_dimensions = tuple(out_dimensions) if fft_type in (cufft.CUFFT_R2C, cufft.CUFFT_D2Z): plan_dimensions = in_dimensions else: # CUFFT_C2R or CUFFT_Z2D plan_dimensions = out_dimensions inembed = in_dimensions onembed = out_dimensions if fft_axes == tuple(range(ndim)): # tranfsorm over all axes nbatch = 1 idist = odist = 1 # doesn't matter since nbatch = 1 istride = ostride = 1 else: # batch along the first or the last axis if 0 not in fft_axes: # don't FFT along the first min_axis_fft axes min_axis_fft = _reduce(min, fft_axes) nbatch = _prod(shape[:min_axis_fft]) if order == 'C': # C-ordered GPU array with batch along first dim idist = _prod(in_dimensions) odist = _prod(out_dimensions) istride = 1 ostride = 1 elif order == 'F': # F-ordered GPU array with batch along first dim idist = 1 odist = 1 istride = nbatch ostride = nbatch elif (ndim - 1) not in fft_axes: # don't FFT along the last axis num_axes_batch = ndim - len(fft_axes) nbatch = _prod(shape[-num_axes_batch:]) if order == 'C': # C-ordered GPU array with batch along last dim idist = 1 odist = 1 istride = nbatch ostride = nbatch elif order == 'F': # F-ordered GPU array with batch along last dim idist = _prod(in_dimensions) odist = _prod(out_dimensions) istride = 1 ostride = 1 else: raise ValueError( 'General subsets of FFT axes not currently supported for ' 'GPU case (Can only batch FFT over the first or last ' 'spatial axes).') for n in plan_dimensions: if n < 1: raise ValueError('Invalid number of FFT data points specified.') keys = (plan_dimensions, inembed, istride, idist, onembed, ostride, odist, fft_type, nbatch, order, fft_axes[-1], out_size) mgr = config.get_current_callback_manager() if mgr is not None: # to avoid a weird segfault, we generate and cache distinct plans # for every possible (load_aux, store_aux) pairs; the plans are # still generated from the same external Python module load_aux = mgr.cb_load_aux_arr store_aux = mgr.cb_store_aux_arr keys += (mgr.cb_load, mgr.cb_store, 0 if load_aux is None else load_aux.data.ptr, 0 if store_aux is None else store_aux.data.ptr) cache = get_plan_cache() cached_plan = cache.get(keys) if cached_plan is not None: plan = cached_plan elif mgr is None: plan = cufft.PlanNd(*keys) if to_cache: cache[keys] = plan else: # has callback plan = mgr.create_plan(('PlanNd', keys[:-4])) mgr.set_callbacks(plan) if to_cache: cache[keys] = plan return plan
def _exec_fft(a, direction, value_type, norm, axis, overwrite_x, out_size=None, out=None, plan=None): fft_type = _convert_fft_type(a.dtype, value_type) if axis % a.ndim != a.ndim - 1: a = a.swapaxes(axis, -1) if a.base is not None or not a.flags.c_contiguous: a = a.copy() elif (value_type == 'C2R' and not overwrite_x and 10010 <= cupy.cuda.runtime.runtimeGetVersion()): # The input array may be modified in CUDA 10.1 and above. # See #3763 for the discussion. a = a.copy() elif cupy.cuda.runtime.is_hip and value_type != 'C2C': # hipFFT's R2C would overwrite input # hipFFT's C2R needs a workaround (see below) a = a.copy() n = a.shape[-1] if n < 1: raise ValueError('Invalid number of FFT data points (%d) specified.' % n) # Workaround for hipFFT/rocFFT: # Both cuFFT and hipFFT/rocFFT have this requirement that 0-th and # N/2-th element must be real, but cuFFT internally simply ignores it # while hipFFT handles it badly in both Plan1d and PlanNd, so we must # do the correction ourselves to ensure the condition is met. if cupy.cuda.runtime.is_hip and value_type == 'C2R': a[..., 0] = a[..., 0].real + 0j if out_size is None: a[..., -1] = a[..., -1].real + 0j elif out_size % 2 == 0: a[..., out_size // 2] = a[..., out_size // 2].real + 0j if out_size is None: out_size = n batch = a.size // n # plan search precedence: # 1. plan passed in as an argument # 2. plan as context manager # 3. cached plan # 4. create a new one curr_plan = cufft.get_current_plan() if curr_plan is not None: if plan is None: plan = curr_plan else: raise RuntimeError('Use the cuFFT plan either as a context manager' ' or as an argument.') if plan is None: devices = None if not config.use_multi_gpus else config._devices # TODO(leofang): do we need to add the current stream to keys? keys = (out_size, fft_type, batch, devices) mgr = config.get_current_callback_manager() if mgr is not None: # to avoid a weird segfault, we generate and cache distinct plans # for every possible (load_aux, store_aux) pairs; the plans are # still generated from the same external Python module load_aux = mgr.cb_load_aux_arr store_aux = mgr.cb_store_aux_arr keys += (mgr.cb_load, mgr.cb_store, 0 if load_aux is None else load_aux.data.ptr, 0 if store_aux is None else store_aux.data.ptr) cache = get_plan_cache() cached_plan = cache.get(keys) if cached_plan is not None: plan = cached_plan elif mgr is None: plan = cufft.Plan1d(out_size, fft_type, batch, devices=devices) cache[keys] = plan else: # has callback # TODO(leofang): support multi-GPU callback (devices is ignored) if devices: raise NotImplementedError('multi-GPU cuFFT callbacks are not ' 'yet supported') plan = mgr.create_plan(('Plan1d', keys[:-5])) mgr.set_callbacks(plan) cache[keys] = plan else: # check plan validity if not isinstance(plan, cufft.Plan1d): raise ValueError('expected plan to have type cufft.Plan1d') if fft_type != plan.fft_type: raise ValueError('cuFFT plan dtype mismatch.') if out_size != plan.nx: raise ValueError('Target array size does not match the plan.', out_size, plan.nx) if batch != plan.batch: raise ValueError('Batch size does not match the plan.') if config.use_multi_gpus != (plan.gpus is not None): raise ValueError('Unclear if multiple GPUs are to be used or not.') if overwrite_x and value_type == 'C2C': out = a elif out is not None: # verify that out has the expected shape and dtype plan.check_output_array(a, out) else: out = plan.get_output_array(a) if batch != 0: plan.fft(a, out, direction) sz = out.shape[-1] if fft_type == cufft.CUFFT_R2C or fft_type == cufft.CUFFT_D2Z: sz = n if norm is None: if direction == cufft.CUFFT_INVERSE: out /= sz else: out /= math.sqrt(sz) if axis % a.ndim != a.ndim - 1: out = out.swapaxes(axis, -1) return out
def _exec_fft(a, direction, value_type, norm, axis, overwrite_x, out_size=None, out=None, plan=None): fft_type = _convert_fft_type(a.dtype, value_type) if axis % a.ndim != a.ndim - 1: a = a.swapaxes(axis, -1) if a.base is not None or not a.flags.c_contiguous: a = a.copy() elif (value_type == 'C2R' and not overwrite_x and 10010 <= cupy.cuda.runtime.runtimeGetVersion()): # The input array may be modified in CUDA 10.1 and above. # See #3763 for the discussion. a = a.copy() n = a.shape[-1] if n < 1: raise ValueError('Invalid number of FFT data points (%d) specified.' % n) if out_size is None: out_size = n batch = a.size // n # plan search precedence: # 1. plan passed in as an argument # 2. plan as context manager # 3. cached plan # 4. create a new one curr_plan = cufft.get_current_plan() if curr_plan is not None: if plan is None: plan = curr_plan else: raise RuntimeError('Use the cuFFT plan either as a context manager' ' or as an argument.') if plan is None: devices = None if not config.use_multi_gpus else config._devices # TODO(leofang): do we need to add the current stream to keys? keys = (out_size, fft_type, batch, devices) cache = get_plan_cache() cached_plan = cache.get(keys) if cached_plan is not None: plan = cached_plan else: plan = cufft.Plan1d(out_size, fft_type, batch, devices=devices) cache[keys] = plan else: # check plan validity if not isinstance(plan, cufft.Plan1d): raise ValueError('expected plan to have type cufft.Plan1d') if fft_type != plan.fft_type: raise ValueError('cuFFT plan dtype mismatch.') if out_size != plan.nx: raise ValueError('Target array size does not match the plan.', out_size, plan.nx) if batch != plan.batch: raise ValueError('Batch size does not match the plan.') if config.use_multi_gpus != plan._use_multi_gpus: raise ValueError('Unclear if multiple GPUs are to be used or not.') if overwrite_x and value_type == 'C2C': out = a elif out is not None: # verify that out has the expected shape and dtype plan.check_output_array(a, out) else: out = plan.get_output_array(a) if batch != 0: plan.fft(a, out, direction) sz = out.shape[-1] if fft_type == cufft.CUFFT_R2C or fft_type == cufft.CUFFT_D2Z: sz = n if norm is None: if direction == cufft.CUFFT_INVERSE: out /= sz else: out /= math.sqrt(sz) if axis % a.ndim != a.ndim - 1: out = out.swapaxes(axis, -1) return out