reuse fft plan1d for c2c convolve/correlate
+import contextlib
+
import cupy
import cupyx
import cupyx.scipy.fft
@@ -5,6 +7,7 @@ import cupyx.scipy.fft
from cupy import core
from cupy.core import _routines_math as _math
from cupy.core import fusion
+from cupy.cuda import cufft
from cupy.lib import stride_tricks
@@ -63,12 +66,25 @@ def _fft_convolve(a1, a2, mode):
if a1.dtype.kind == 'c' or a2.dtype.kind == 'c':
fft, ifft = cupy.fft.fft, cupy.fft.ifft
+ is_c2c = True
else:
fft, ifft = cupy.fft.rfft, cupy.fft.irfft
+ is_c2c = False
dtype = cupy.result_type(a1, a2)
n1, n2 = a1.size, a2.size
out_size = cupyx.scipy.fft.next_fast_len(n1 + n2 - 1)
+ if is_c2c:
+ # skip calling get_fft_plan() as we know the args exactly
+ fft_plan = cufft.Plan1d(
+ out_size,
+ cufft.CUFFT_C2C if dtype == cupy.complex64 else cufft.CUFFT_Z2Z,
+ 1)
+ else:
+ # this is a no-op context manager
+ # TODO(leofang): use contextlib.nullcontext() for PY37+?
+ fft_plan = contextlib.suppress()
+ with fft_plan:
fa1 = fft(a1, out_size)
fa2 = fft(a2, out_size)
out = ifft(fa1 * fa2, out_size)
