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 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
