diff --git a/CHANGELOG.md b/CHANGELOG.md index 0fca2dc1..577e1cc2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added ### Changed +* `irfft` with `out=` now writes directly into `out` when the cast or zero-padded copy of a non-contiguous input has the same memory order as `out`, avoiding a temporary array and an extra copy [gh-396](https://github.com/IntelPython/mkl_fft/pull/396) ### Fixed diff --git a/mkl_fft/_pydfti.pyx b/mkl_fft/_pydfti.pyx index 76a30600..77936f31 100644 --- a/mkl_fft/_pydfti.pyx +++ b/mkl_fft/_pydfti.pyx @@ -750,8 +750,10 @@ def _c2r_fft1d_impl( # strides cannot 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 (e.g. when x is non-contiguous and gets cast to complex). + 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/tests/test_fft1d.py b/mkl_fft/tests/test_fft1d.py index 4336aef2..e63421e3 100644 --- a/mkl_fft/tests/test_fft1d.py +++ b/mkl_fft/tests/test_fft1d.py @@ -488,6 +488,31 @@ def test_irfft_out_strided(axis): assert np.all(base == -1.0) +@pytest.mark.parametrize("order", ["C", "F"]) +@pytest.mark.parametrize("axis", [0, 1, 2]) +@pytest.mark.parametrize("dt", ["i8", "f4", "f8", "c16"]) +def test_irfft_out_strided_input_copied(dt, axis, order): + # x is strided, but it is cast (i8, f4, f8) or padded (c16) into a + # contiguous copy, whose layout decides if MKL can write into out directly + x = rnd.randint(-50, 50, size=(20, 33, 54)).astype(dt) + x = np.asarray(x, order=order)[::2, ::3, ::4] + n = None + if dt == "c16": + # n // 2 + 1 > x.shape[axis] forces a zero-padded copy + n = 2 * x.shape[axis] + 4 + out_shape = list(x.shape) + out_shape[axis] = 2 * (x.shape[axis] - 1) if n is None else n + out_dtype = np.float32 if dt == "f4" else np.float64 + # sentinel: catches out being returned untouched + out = np.full(out_shape, -1.0, dtype=out_dtype, order=order) + result = mkl_fft.irfft(x, n=n, axis=axis, out=out) + expected = np.fft.irfft(x.astype(np.complex128), n=n, axis=axis) + + assert result is out + tol = 1e-4 if dt == "f4" else 1e-10 + assert_allclose(result, expected, rtol=tol, atol=tol) + + @requires_numpy_2 @pytest.mark.parametrize("dt", ["i4", "f4", "f8", "c8", "c16"]) def test_irfft_dtype(dt):