From 36ccaa8f73545ab2e0f40f6010e5d7b8eb99b088 Mon Sep 17 00:00:00 2001 From: Anton Volkov Date: Tue, 6 Oct 2026 14:46:57 +0200 Subject: [PATCH 1/2] fix: respect `out` layout in real-input ifft and in rfft of cast input 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. --- CHANGELOG.md | 2 ++ mkl_fft/_pydfti.pyx | 5 ++-- mkl_fft/src/mklfft.c.src | 55 ++++++++++++++++++++++++++++++++++--- mkl_fft/tests/test_fft1d.py | 52 +++++++++++++++++++++++++++++++++++ mkl_fft/tests/test_fftnd.py | 33 ++++++++++++++++++++++ 5 files changed, 141 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3d1795ff..c428a1aa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * Fixed possible memory leaks when `PyMem_Malloc` fails, and raise `MemoryError` [gh-373](https://github.com/IntelPython/mkl_fft/pull/373) * Fixed possible memory leaks when multi-iterator constructors fail [gh-373](https://github.com/IntelPython/mkl_fft/pull/373) * Fixed N-D transforms returning a success status when a scratch allocation fails [gh-373](https://github.com/IntelPython/mkl_fft/pull/373) +* Fixed wrong results and writes outside of `out` in `ifft` and `ifftn` of real input when `out` is not contiguous: the final complex conjugation treated `out` as one contiguous block [gh-389](https://github.com/IntelPython/mkl_fft/pull/389) +* Fixed wrong results in `rfft` and `rfftn` with `out` for input that is cast to `float64` (e.g. integer, `float16`) when the input and `out` are Fortran-ordered [gh-389](https://github.com/IntelPython/mkl_fft/pull/389) ## [2.3.2] - 2026-08-04 diff --git a/mkl_fft/_pydfti.pyx b/mkl_fft/_pydfti.pyx index 36e152fb..c529e12f 100644 --- a/mkl_fft/_pydfti.pyx +++ b/mkl_fft/_pydfti.pyx @@ -642,8 +642,9 @@ def _r2c_fft1d_impl( # be compared directly. # TODO: currently instead of this condition, we check both input # and output to be c_contig or f_contig, relax this condition - c_contig = x.flags.c_contiguous and out.flags.c_contiguous - f_contig = x.flags.f_contiguous and out.flags.f_contiguous + # Check x_arr, which may be a cast copy of x with a different layout. + c_contig = x_arr.flags.c_contiguous and out.flags.c_contiguous + f_contig = x_arr.flags.f_contiguous and out.flags.f_contiguous if c_contig or f_contig: f_arr = out else: diff --git a/mkl_fft/src/mklfft.c.src b/mkl_fft/src/mklfft.c.src index 18410480..755535dc 100644 --- a/mkl_fft/src/mklfft.c.src +++ b/mkl_fft/src/mklfft.c.src @@ -603,6 +603,54 @@ int @name@_mkl_@mode@_in(PyArrayObject* x_inout, npy_intp n, int axis, double fs { ((mkl_complex_p_dest)->real) = ((mkl_complex_p_src)->real); \ ((mkl_complex_p_dest)->imag) = -((mkl_complex_p_src)->imag); } +/**begin repeat +* #name=cfloat,cdouble# +* #MKL_TYPE=MKL_Complex8,MKL_Complex16# +* #vml_conj_func=vmcConj,vmzConj# +*/ +/* Complex conjugate x in place, honoring its strides: x may be a + non-contiguous out array provided by the user */ +static int +@name@_conj_inplace(PyArrayObject *x) +{ + multi_iter_t mit; + npy_intp *x_shape = 0, *x_strides = 0; + npy_intp x_itemsize, x_size; + int x_rank; + char *x_data = (char *) PyArray_DATA(x); + + get_basic_array_data(x, &x_rank, &x_shape, + &x_strides, &x_itemsize, &x_size); + + if (PyArray_ISONESEGMENT(x)) { + @vml_conj_func@(x_size, (@MKL_TYPE@ *) x_data, + (@MKL_TYPE@ *) x_data, VML_HA); + return 0; + } + + if (multi_iter_new(&mit, x_shape, x_rank)) + return DFTI_MEMORY_ERROR; + + while(!MultiIter_Done(mit)) { + char *tmp = x_data; + @MKL_TYPE@ *elem; + int i; + + for(i = 0; i < x_rank; i++) + tmp += x_strides[i] * MultiIter_IndexElem(mit, i); + + elem = (@MKL_TYPE@ *) tmp; + SET_CONJ(elem, elem); + + if (multi_iter_next(&mit)) + break; + } + + multi_iter_free(&mit); + return 0; +} +/**end repeat**/ + /**begin repeat * #REALIN=(float,double)*2# * #COMPLEXOUT=(cfloat,cdouble)*2# @@ -611,7 +659,6 @@ int @name@_mkl_@mode@_in(PyArrayObject* x_inout, npy_intp n, int axis, double fs * #DFTI_PRECISION=(DFTI_SINGLE,DFTI_DOUBLE)*2# * #mode=(fft1d)*2,(ifft1d)*2# * #POST_CONJUGATE=(FALSE)*2,(TRUE)*2# -* #vml_conj_func=(vmcConj,vmzConj)*2# */ int @REALIN@_@COMPLEXOUT@_mkl_@mode@_out( PyArrayObject *x_in, npy_intp n, int axis, PyArrayObject *x_out, @@ -879,7 +926,8 @@ int @REALIN@_@COMPLEXOUT@_mkl_@mode@_out( } if (@POST_CONJUGATE@) { - @vml_conj_func@(xout_size, xout_data, xout_data, VML_HA); + status = @COMPLEXOUT@_conj_inplace(x_out); + if (status != 0) goto failed; } status = OK; @@ -1544,7 +1592,6 @@ int * #op_name=fftnd*2,ifftnd*2# * #DftiCompute_MODE=DftiComputeForward*4# * #POST_CONJUGATE=FALSE*2,TRUE*2# -* #vml_conj_func=(vmcConj,vmzConj)*2# */ int @inp_type_name@_@out_type_name@_mkl_@op_name@_out( @@ -1727,7 +1774,7 @@ int if (@POST_CONJUGATE@) { Py_BEGIN_ALLOW_THREADS - @vml_conj_func@(xout_size, xout_data, xout_data, VML_HA); + status = @out_type_name@_conj_inplace(x_out); Py_END_ALLOW_THREADS } diff --git a/mkl_fft/tests/test_fft1d.py b/mkl_fft/tests/test_fft1d.py index 16464e75..5f178d2c 100644 --- a/mkl_fft/tests/test_fft1d.py +++ b/mkl_fft/tests/test_fft1d.py @@ -416,3 +416,55 @@ def test_irfft_dtype(dt): result = mkl_fft.irfft(x) expected = np.fft.irfft(x) assert_allclose(result, expected, rtol=1e-7, atol=1e-7, strict=True) + + +@pytest.mark.parametrize("step", [2, -1, -3]) +@pytest.mark.parametrize("func", ["fft", "ifft"]) +def test_fft_real_input_out_strided_1d(func, step): + # ifft of real input conjugates the result in place in out, + # which must honor the strides of out + x = rnd.random(48)[::step] + base = np.full(48, -1 - 1j) + out = base[::step] + result = getattr(mkl_fft, func)(x, out=out) + expected = getattr(np.fft, func)(x) + + assert result is out + assert_allclose(result, expected) + # nothing outside of out was written to + out[...] = -1 - 1j + assert np.all(base == -1 - 1j) + + +@pytest.mark.parametrize("axis", [0, 1, 2]) +@pytest.mark.parametrize("dt", [np.float32, np.float64]) +@pytest.mark.parametrize("func", ["fft", "ifft"]) +def test_fft_real_input_out_strided(func, dt, axis): + shape = (20, 33, 54) + x = rnd.random(shape).astype(dt)[::2, ::3, ::4] + base = np.full(shape, -1 - 1j, dtype=np.result_type(dt, np.complex64)) + out = base[::2, ::3, ::4] + result = getattr(mkl_fft, func)(x, axis=axis, out=out) + expected = getattr(np.fft, func)(x.astype(np.float64), axis=axis) + + assert result is out + tol = 1e-5 if dt == np.float32 else 1e-10 + assert_allclose(result, expected, rtol=tol, atol=tol) + # nothing outside of out was written to + out[...] = -1 - 1j + assert np.all(base == -1 - 1j) + + +@pytest.mark.parametrize("axis", [0, 1, 2]) +@pytest.mark.parametrize("dt", ["i8", "f2", "f8", "g"]) +def test_rfft_out_fortran(dt, axis): + # dtypes other than f4/f8 are cast to a C-contiguous float64 copy + x = np.asfortranarray(rnd.randint(-50, 50, size=(6, 7, 8)).astype(dt)) + out_shape = list(x.shape) + out_shape[axis] = x.shape[axis] // 2 + 1 + out = np.full(out_shape, -1 - 1j, order="F") + result = mkl_fft.rfft(x, axis=axis, out=out) + expected = np.fft.rfft(x.astype(np.float64), axis=axis) + + assert result is out + assert_allclose(result, expected, atol=1e-10) diff --git a/mkl_fft/tests/test_fftnd.py b/mkl_fft/tests/test_fftnd.py index 4aa6f159..0a6b4f5a 100644 --- a/mkl_fft/tests/test_fftnd.py +++ b/mkl_fft/tests/test_fftnd.py @@ -373,3 +373,36 @@ def test_empty_axes_returns_same_object(dtype, func): assert ( result is x ), f"{func} with axes=() should return the same object, not a copy" + + +@pytest.mark.parametrize("axes", [None, (0, 1), (1, 2)]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("func", ["fftn", "ifftn"]) +def test_out_strided_real_input(func, dtype, axes): + # ifftn of real input conjugates the result in place in out, + # which must honor the strides of out + shape = (20, 30, 40) + x = rnd.random(shape).astype(dtype)[::2, ::3, ::4] + base = np.full(shape, -1 - 1j, dtype=np.result_type(dtype, np.complex64)) + out = base[::2, ::3, ::4] + result = getattr(mkl_fft, func)(x, axes=axes, out=out) + expected = getattr(np.fft, func)(x.astype(np.float64), axes=axes) + + assert result is out + rtol, atol = _get_rtol_atol(result) + assert_allclose(result, expected, rtol=rtol, atol=atol) + # nothing outside of out was written to + out[...] = -1 - 1j + assert np.all(base == -1 - 1j) + + +@pytest.mark.parametrize("dt", ["i8", "f8"]) +@pytest.mark.parametrize("func", ["rfftn", "rfft2"]) +def test_rfftn_out_fortran(func, dt): + x = np.asfortranarray(rnd.randint(-50, 50, size=(6, 7, 8)).astype(dt)) + expected = getattr(np.fft, func)(x.astype(np.float64)) + out = np.full(expected.shape, -1 - 1j, order="F") + result = getattr(mkl_fft, func)(x, out=out) + + assert result is out + assert_allclose(result, expected, atol=1e-9) From 4c71b7097c2348da3d52b37f3b15448dd334c392 Mon Sep 17 00:00:00 2001 From: Anton Volkov Date: Tue, 6 Oct 2026 15:23:24 +0200 Subject: [PATCH 2/2] perf: keep memory order when casting input in rfft 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. --- CHANGELOG.md | 1 + mkl_fft/_pydfti.pyx | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c428a1aa..69102586 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * `_direct_fftnd` now also checks the status returned by the backend instead of discarding it [gh-373](https://github.com/IntelPython/mkl_fft/pull/373) * Pinned Cython in the Coverity Scan workflow so generated code stays stable between scans, and added `coverity/README.md` documenting the known false-positive families and the scan review checklist [gh-374](https://github.com/IntelPython/mkl_fft/pull/374) * Reduced Python overhead in `norm="forward"`/`"ortho"` scaling by computing the scale factor without `numpy.prod` [gh-384](https://github.com/IntelPython/mkl_fft/pull/384) +* `rfft` now casts input of dtypes other than `float32`/`float64` to `float64` keeping its memory order instead of making it C-contiguous, avoiding a slow reordering copy for Fortran-ordered input [gh-389](https://github.com/IntelPython/mkl_fft/pull/389) ### Fixed * Fixed `norm="forward"`/`"ortho"` scaling in `fftn`, `ifftn`, `rfftn`, `irfftn` and the `fft2` family when only a subset of axes is transformed: the scale used the full array shape instead of the transformed axes [gh-336](https://github.com/IntelPython/mkl_fft/issues/336), [gh-370](https://github.com/IntelPython/mkl_fft/pull/370) diff --git a/mkl_fft/_pydfti.pyx b/mkl_fft/_pydfti.pyx index c529e12f..10fd49b8 100644 --- a/mkl_fft/_pydfti.pyx +++ b/mkl_fft/_pydfti.pyx @@ -616,8 +616,10 @@ def _r2c_fft1d_impl( pass else: # we must cast the input to doubles and allocate the output, + # the cast always copies; without NPY_ARRAY_ENSURECOPY the copy keeps + # the memory order of x instead of being made C-contiguous try: - requirement = cnp.NPY_ARRAY_BEHAVED | cnp.NPY_ARRAY_ENSURECOPY + requirement = cnp.NPY_ARRAY_BEHAVED if x_type is cnp.NPY_LONGDOUBLE: requirement = requirement | cnp.NPY_ARRAY_FORCECAST x_arr = cnp.PyArray_FROM_OTF(