Skip to content

fix: keep C order when zero-padding input that is both C- and F-contiguous - #397

Open
antonwolfy wants to merge 1 commit into
masterfrom
fix/pad-array-memory-order
Open

antonwolfy wants to merge 1 commit into
masterfrom
fix/pad-array-memory-order

Conversation

@antonwolfy

Copy link
Copy Markdown
Collaborator

Follow-up to #396.

_pad_array allocated the zero-padded copy of the input in F order whenever the input had the F-contiguous flag. An input with a single non-unit dimension, such as shape (m, 1), has both the C- and F-contiguous flags, so padding it along the unit axis (n larger than the input length) gave an F-contiguous copy of shape (m, n). That had two effects:

  • the result of fft, ifft, rfft and irfft (and of the N-D functions, which pad through the same helper) was F-contiguous, where NumPy returns a C-contiguous array;
  • a C-ordered out did not match the layout of the copy passed to MKL, so the result was computed into a temporary array and then copied into out. With perf: check x_arr, not x, when deciding whether irfft can write into out directly #396 this also applies to irfft, which before only got the direct path by checking the original input.

This PR uses F order only when the input is F- but not C-contiguous (PyArray_ISFORTRAN), so a both-contiguous input is padded into a C-ordered copy, as NumPy does. Input that is only F-contiguous is still padded in F order.

Result layout for padding to n=8, NumPy 2.5.4 vs mkl_fft on master:

input shape axis NumPy mkl_fft (master) mkl_fft (this PR)
(5, 1) 1 C F C
(1, 5) 0 C F C
(4, 1, 1) 2 C F C
(1, 1, 4) 0 C F C

The same holds for fft, ifft, rfft and irfft.

Timings for input of shape (4096, 1) padded to n=64 along axis 1 with a C-ordered out, on Lunar Lake, median in ms:

function master #396 this PR #396 + this PR
fft, C-ordered out 2.35 2.34 0.36 0.34
rfft, C-ordered out 0.41 0.41 0.17 0.17
irfft, C-ordered out 0.21 0.52 0.18 0.18
fft, no out 0.39 0.40 0.33 0.33

test_pad_both_contiguous_input checks the result layout and values with and without out for these shapes; it fails on master and passes with this change. The full test suite passes with this PR alone and merged with #396, and the two PRs merge without conflicts.

…guous

`_pad_array` allocated the zero-padded copy in F order whenever the input
had the F-contiguous flag. An input with a single non-unit dimension, such
as shape (m, 1), has both flags, so padding it along the unit axis gave an
F-contiguous copy of shape (m, n), and so an F-contiguous result of `fft`,
`ifft`, `rfft` and `irfft`, where NumPy returns a C-contiguous one.

It also made a C-ordered `out` incompatible with the copy passed to MKL,
so the result was computed into a temporary array and copied into `out`.

Use F order only when the input is F- but not C-contiguous, which is what
NumPy's PyArray_ISFORTRAN checks.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant