Repository navigation
fix: respect out layout in real-input ifft and in rfft of cast input - #389
antonwolfy wants to merge 3 commits into
Conversation
Two pre-existing bugs with a user-provided `out`: - `ifft`/`ifftn` of real input compute a forward transform and then complex-conjugate the result with `v?Conj`, treating `out` as one contiguous block. When `out` was not contiguous (e.g. a strided or reversed view whose strides match the input), part of it was left unconjugated and memory outside of `out` was modified, up to heap corruption. Conjugate element-wise along the strides of `out` unless it is contiguous. - `_r2c_fft1d_impl` decided whether MKL can write into `out` directly from the contiguity of the original `x` rather than of `x_arr`, which for dtypes other than float32/float64 is a C-contiguous float64 copy. With Fortran-ordered integer/float16/longdouble input and `out`, `rfft`/`rfftn` gave wrong results.
The cast of input with dtypes other than float32/float64 to float64 in `_r2c_fft1d_impl` requested NPY_ARRAY_ENSURECOPY, which makes PyArray_FROM_OTF return a C-contiguous copy. For Fortran-ordered input this is a slow reordering copy, and with `out=` it also required a C-ordered temporary that is then copied into the Fortran-ordered `out`. The cast always copies anyway, so drop the flag: the copy keeps the memory order of the input, as the cast in `_c2r_fft1d_impl` already does. For Fortran-ordered int64 input of shape 128^3 this brings `rfft` from ~45 ms to ~5 ms, and the result layout now matches NumPy, which allocates the output with `empty_like` and so keeps the memory order of the input.
|
Performance results for this PR (head 4c71b70) against master (5bc05b5). Setup. Intel Core Ultra 5 236V (Lunar Lake: 4 P-cores and 4 LP E-cores), Ubuntu 25.10, Python 3.13, NumPy 2.5.3, oneMKL 2026.2. Both builds were compiled from source with gcc 15.2 and pass the full test suite.
1.
|
| 128³ input | master | this PR | speedup (1 / 4 threads) |
|---|---|---|---|
rfft int64, F-order, axis 0, out= F-order |
27.5 ms | 3.4 ms | 8.1x / 8.8x |
rfft int64, F-order, axis 0 |
20.9 ms | 3.4 ms | 6.1x / 7.6x |
rfft float16, F-order, axis 0, out= F-order |
20.2 ms | 3.3 ms | 6.1x / 5.0x |
rfftn int64, F-order |
22.3 ms | 13.1 ms | 1.7x / 2.6x |
rfft int64, strided C-order, out= C-order |
4.0 ms | 3.5 ms | 1.15x / 1.25x |
rfft float64, F-order (control, unchanged path) |
2.03 ms | 2.04 ms | 1.00x / 0.93x |
In the strided C-order case the peak temporary memory also drops from 32 to 16 MiB, because no extra temporary is allocated.
2. Real-input ifft / ifftn conjugation (C change)
The inverse transform of real input is a forward transform followed by an in-place conjugation. So the conjugation cost is measured as time(ifft) - time(fft) for the same input and out, timed alternately in the same process. Values are for 1 thread, with the total ifft time in parentheses.
| case | master | this PR |
|---|---|---|
contiguous out, 1-D 2^20 |
0.64 ms (6.03) | 0.65 ms (6.02) |
contiguous out, ifftn 64³ |
0.19 ms (1.14) | 0.09 ms (1.04) |
out[::2], 1-D 2^20 |
0.94 ms (6.61) | 2.35 ms (8.04) |
64³ strided view, axis=0 |
0.22 ms (4.75) | 1.25 ms (5.80) |
64³ strided view, axis=2 |
0.18 ms (3.11) | 1.10 ms (4.04) |
64³ strided view, ifftn |
0.25 ms (4.54) | 1.16 ms (5.45) |
- Contiguous
out: unchanged. It still uses the singlev?Conjcall. - Non-contiguous
out: the conjugation now costs 1.1–2.9 ms, which is 20–36% of theiffttime, with 1 or 4 threads. - Why master looks cheaper: master's numbers for non-contiguous
outaren't a valid baseline. It conjugated a contiguous span, which is the wrong elements. - Compared with NumPy:
np.conjugate(view, out=view)on the same views takes 0.9 ms (64³) and 1.6 ms (2^20[::2]). So the new loop matches NumPy for the 3-D views and is about 1.4x slower for the 1-D strided case.
3. Contiguous paths: no regression
These are 60 cases mirroring the ASV suite, all contiguous or with no out:
fft/ifft/rfft/irfft1-D at n = 1024, 16384 and 2^20;fftn/ifftn/rfftnat 64³ and 128³;- float32, float64, complex64 and complex128;
- contiguous
out=for real-inputifft/ifftn.
As a control, master was also compared with a byte-identical copy of itself:
| 1 thread, 60 cases | median ratio | range | cases > 10% off |
|---|---|---|---|
| master copy vs master (A/A) | 1.000 | 0.71–1.37 | 31 |
| this PR vs master | 0.995 | 0.74–1.33 | 34 |
| this PR vs master copy | 1.000 | 0.85–1.17 | 3 |
With 4 threads, the median ratio of this PR against master is 1.003.
The PR stays within the A/A noise. Individual cases vary by up to ±30% between processes even with identical binaries. This tracks the 64-byte alignment of the arrays: for example, fft of 1024 complex128 takes 1.9 µs with the output at offset 0 and 1.4 µs at offset 16 or 48. Single-case swings in ASV results for this PR should be read with that in mind.
Fixes two bugs with a user-provided
outthat are already on master, plus a related slowdown inrfft. Both were found while reviewing #387.1.
ifft/ifftnof real input with a non-contiguousoutThe inverse transform of real input is computed as a forward transform followed by an in-place complex conjugation, using
v?Conj(xout_size, xout_data, ...). That call treatsoutas one contiguous block. Cython passesoutto the backend directly whenever its element strides match those of the input, so for a strided or reversedout:outis left unconjugated, which gives wrong values;outis conjugated, which corrupts other data and can segfault.This affects the 1-D (
ifft) and N-D (ifftnover all axes) backends, and is reachable throughmkl_fft.interfaces.numpy_fft.Fix (
mklfft.c.src): a new{cfloat,cdouble}_conj_inplacehelper keeps thev?Conjfast path for contiguousout. Otherwise it conjugates element by element along the strides ofout, the same way the existing "copy conjugate even harmonics" loop walksout.2.
rfft/rfftnof cast input with Fortran-orderedout_r2c_fft1d_impldecides whether MKL can write straight intooutby checking that the input andoutare both C- or both F-contiguous. It checks the originalx, notx_arr. For dtypes other than float32/float64 (e.g. integer, float16, longdouble),x_arris a C-contiguous float64 copy. So with Fortran-ordered input andout, MKL was handed a C-ordered input and an F-ordered output:Fix (
_pydfti.pyx): check the contiguity ofx_arr._c2r_fft1d_implhas the same check, but its cast keeps the input's memory order, so it is left unchanged.3. Slow
rfftof cast Fortran-ordered inputAlso on master: the cast in
_r2c_fft1d_implrequestedNPY_ARRAY_ENSURECOPY, which makesPyArray_FROM_OTFreturn a C-contiguous copy. For Fortran-ordered input this is a slow reordering copy. With fix 2 alone,out=would also need a C-ordered temporary, copied into the Fortran-orderedout. The cast always copies anyway, since the dtype differs. So this PR drops the flag: the copy keeps the input's memory order, as the cast in_c2r_fft1d_implalready does.rfftof Fortran-ordered int64 input of shape 128³ along axis 0, on 1 thread:rfft(x, axis=0, out=out_F)rfft(x, axis=0)The result layout now also matches NumPy, which allocates the output with
empty_likeand so keeps the input's order. Over 18 combinations of dtype, input order and axis, master matched NumPy in 12; this PR matches in all 18. On master, Fortran-ordered int64/float16 input gave C-ordered results.Performance of fix 1
For contiguous
out, or noout, thev?Conjcall is unchanged and so is its cost. That is the only path the ASV benchmarks cover. For a non-contiguousout, the element-by-element conjugation costs about as much as NumPy'snp.conjugate(view, out=view)on the same view: about 1.4–3.4 ms for 2^20 or 64³ elements. Master was cheaper there only because it conjugated the wrong memory. Stridedv?ConjIcalls were tried and gave no gain for 3-D views, where memory traffic sets the cost.Tests
New tests are appended to
test_fft1d.pyandtest_fftnd.py:fft/ifftandfftn/ifftnof real input (float32 and float64) into strided or reversedout, checking that nothing outsideoutis written;rfft/rfftn/rfft2with Fortran-orderedoutfor int64, float16, longdouble and float64 input.The cases that pass on master are the forward transforms and the float64 / unaffected-axis controls.
There are no new compiler warnings. This PR is independent of #387 and #388; it may conflict with #387 only in
CHANGELOG.md.