Skip to content

Commit d71aa89

Browse files
MAINT: Fixed performance issue with rfft2 and irfft2
When fftn or ifftn are called on 1D arrays we should go though fft/ifft, instead of using _direct_fftnd to take advantage of caching. Unlike in rfftn_numpy, where out-of-place rfft_numpy is performed first, and the subsequent fftn calls can be in-place, for irfftn_numpy the ifftn need to be performance first, and one must check whether the first iteration can in fact be performed in place.
1 parent 8a01af5 commit d71aa89

File tree

1 file changed

+6
-3
lines changed

1 file changed

+6
-3
lines changed

mkl_fft/_pydfti.pyx

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -709,7 +709,7 @@ def _fftnd_impl(x, shape=None, axes=None, overwrite_x=False, direction=+1):
709709
raise ValueError("Direction of FFT should +1 or -1")
710710

711711
# _direct_fftnd requires complex type, and full-dimensional transform
712-
if isinstance(x, np.ndarray) and x.size != 0:
712+
if isinstance(x, np.ndarray) and x.size != 0 and x.ndim > 1:
713713
_direct = shape is None and axes is None
714714
if _direct:
715715
_direct = x.ndim <= 7 # Intel MKL only supports FFT up to 7D
@@ -841,6 +841,7 @@ def irfftn_numpy(x, s=None, axes=None):
841841
if len(s) > 1:
842842
if not no_trim:
843843
a = _fix_dimensions(a, s, axes)
844+
ovr_x = True if _datacopied(<cnp.ndarray> a, x) else False
844845
if len(set(axes)) == len(axes) and len(axes) == a.ndim:
845846
ss, aa = _remove_axis(s, axes, la)
846847
ind = [slice(None,None,1),] * len(s)
@@ -849,9 +850,11 @@ def irfftn_numpy(x, s=None, axes=None):
849850
tind = tuple(ind)
850851
a[tind] = _fftnd_impl(
851852
a[tind], shape=ss, axes=aa,
852-
overwrite_x=True, direction=-1)
853+
overwrite_x=ovr_x, direction=-1)
854+
ovr_x = True
853855
else:
854856
for ii in range(len(axes)-1):
855-
a = ifft(a, s[ii], axes[ii], overwrite_x=True)
857+
a = ifft(a, s[ii], axes[ii], overwrite_x=ovr_x)
858+
ovr_x = True
856859
a = irfft_numpy(a, n = s[-1], axis=la)
857860
return a

0 commit comments

Comments
 (0)