Ah, OK, this switches to the C++ version, but uses only its 1D transforms!
Makes a lot of sense for a quick fix of the data type issue, but it misses the chance of massive speedups (through vectorization and/or multithreading) in multi-D transforms.
If numpy 2.0 is close that's probably the best option, though I would still urge to unlock the other features in later releases.
If only 1D FFTs are used, a lot of C++ code could be #if 0'd to reduce compile time. I can do this once the rest of the PR has converged.

but it misses the chance of massive speedups (through vectorization and/or multithreading) in multi-D transforms.

That is actually a (small or not?) issue with using the gufuncs, it may be tricky to expand to arbitrary dimensions in a nice way for those (which was mentioned to me, I admit; although I didn't realize it might be an issue any time soon).

And yes, so long as this looks all fine, it seems fine to not do the N-D versions.

Indeed, I went the easy route, based on the gufunc I had just introduced (mostly to be able to allow out). Is the threading/vectorization specific to the multi-dimensional iterator, or something that makes the 1D transform faster? From a quick look, it looks like it is the latter, so that should be relatively straightforward, by just using c2c, etc., in my calls. Indeed, that would mean no need for creating a buffer any more too...

Is the threading/vectorization specific to the multi-dimensional iterator, or something that makes the 1D transform faster?

No, threading/vectorization only affects 2D and higher transforms. I have played with applying this to 1D as well, but only in pocketfft's successor package.

Just to be clear, do you mean that vectorization helps for FTs over multiple axes, or FTs over a single axis in higher-D arrays? From the code, it looks like the latter, which might still work with the ufunc, as we get an extra axis passed in.

No, threading/vectorization only affects 2D and higher transforms. I have played with applying this to 1D as well, but only in pocketfft's successor package.

OK, that seems reasonable for follow-up, then. More generally, it may well make sense to use your iterator too, and/or copy what scipy does... I definitely went for what seemed the fastest route...

Just to be clear, do you mean that vectorization helps for FTs over multiple axes, or FTs over a single axis in higher-D arrays? From the code, it looks like the latter, which might still work with the ufunc, as we get an extra axis passed in.

Yes, it works for transforming just a single axis in a multi-D array too. The idea is to "bundle" a few of the 1D transforms into a single one of SIMD type.

The macos_arm64_test failure seems unrelated...

@mreineck - OK, I'll look into using c2c and friends...

Calling c2c, etc., using the outer dimension actually seems fairly straightforward and gives a nice speed-up. The only unfortunate thing is that pocketfft does not deal with the case that the input is shorter than the output, so in that case one still needs to buffer/copy. But the speed-up is worth the complexity, so I'll push a further refactoring tomorrow.

Nice!

Furthermore, I have not tried setting up a pocketfft cache, or done anything with threading. Scipy uses its meson.build to set a cache size of 16 and have compiler-dependent threading; happy to do the same here.

The cache size may help, assuming it's a useful setting in single-threaded mode (@mreineck?). I'd leave the threading along, we can't do much with that. In SciPy it's controlled by a workers= keyword to opt in to multi-threading. NumPy doesn't have that, and everything is supposed to be single-threaded (underlying BLAS/LAPACK is an exception for [reasons]).

Should the license be added to LICENSES.txt? It wasn't there for the C version, but perhaps that was an oversight.

Probably an oversight indeed, since the SciPy licenses list has it (here). So +1 for adding it.

With it, we can support all the dtypes that it supports: float, double, and long double.

Adding float32 support is great and obviously desirable. I'm less sure about longdouble support, since that's not really all that useful and something we may have more work to get rid of again later. We didn't get around to revisiting this topic for 2.0, but overall I'm still convinced that longdouble is overall more harmful than helpful for NumPy as a whole, so why add more support when I don't think anyone asked for it here?

Build time wise this PR in its current state looks fine - impact of C++ is okay overall, see quick profiling result below. Binary size is fine too, it's ~240 kb for a dev (debugoptimized) build. A note though that the SciPy pocketfft extension is more than 4x larger, so enabling all functionality that SciPy uses in the NumPy build is probably not the way to go (it'd be ~25% of the size of _multiarray_umath at that point).

The cache size may help, assuming it's a useful setting in single-threaded mode (@mreineck?)

It will work in single-threaded mode. However, my feeling is that the plan caching isn't really worth it: it may cause large memory consumption (see, e.g., scipy/scipy#12297) for very little performance gain (unless you call very short 1D FFTs repeatedly). Scipy maintainers tend towards keeping the cache, however.

OK, I now have a version that calls the vectorized pocketfft version if possible. It still has the direct loops for the case that nin < nout as well as for 1D transforms (where one can avoid copies in the probably fairly common case of very long 1D transforms with contiguous output).

Some remaining questions:

1. On the implementation: pocketfft vectorization properties are set by what is available at compile-time, not run-time. Is this a problem? Or can this be left for later optimization?
2. Would it make sense to include pocketfft as a submodule? That would automatically include the license (and not much more except for a small demo version). I think it is included as a submodule by scipy.

Adding float32 support is great and obviously desirable. I'm less sure about longdouble support, since that's not really all that useful and something we may have more work to get rid of again later.

@rgommers - I would love to have a proper float128 instead of longdouble, but on the other hand, the present behaviour is not that longdouble is not supported (in which case it would be fine to continue with that), but rather that longdouble is just converted to double without any warning. Given that, my sense is that supporting it correctly is the least disruptive change.

the present behaviour is not that longdouble is not supported (in which case it would be fine to continue with that), but rather that longdouble is just converted to double without any warning. Given that, my sense is that supporting it correctly is the least disruptive change.

I'm not really opposed, but it is a change, with technically only increased binary size as a penalty and no upsides for Windows and macOS arm64 users. Doing nothing and continuing to convert to float64 either internally like before or, as an improvement, use the float64 loops internally but cast back to longdouble before returning the result, are therefore also options (even the former is zero change from now, so hard to label it as more disruptive).

I'm not sure how serious the plan was to include all the pocketfft functionality that SciPy needs into NumPy - but binary size does seem to become relevant then. If's it's <0.1 MB I'm not too worried, but I wouldn't really want to pay 0.25 MB or more for long double loops that no one has asked for.

ls -lh build/numpy/fft/_pocketfft_umath.cpython-311-x86_64-linux-gnu.so
# 5.1M present PR
# 4.1M without longdouble
# 2.9M w/ longdouble but with POCKETFFT_NO_VECTORS
# (which uses low-level routines only, but halves speed;
# I checked that high-level routines are not compiled for longdouble)

For comparison, multiarray_umath is 73 MB, the simd stuff 15MB, and everything in random adds up to 43 MB.

Thanks for the comparison, that is helpful. Note that those are debugoptimized numbers; what I had in mind was wheel sizes (I should have specified). Our wheel sizes are in the 13 - 21 MB range, and the new macOS arm64 ones even as small as 5 MB.

I did a quick test building a wheel from this branch; the pocketfft extension is 508 kb uncompressed, 208 kb compressed in it (total wheel size 7.8 MB; _multiarray_umath.so 2.6 MB compressed). So that should still be fine. If dst/dct transforms would be enabled, it may get a bit questionable. But that could still be taken care of mostly by disabling the longdouble loops for platforms where it's an alias for double. Anyway, it doesn't seem too likely to become an issue at this moment. Let's not worry about it further.

OK, great! The other transforms may not be as bad as one would think, since they mostly use the same code. But anyway, not sure those are needed, and definitely not relevant here.

The relevant flags are here:

Its only limited to core only, I thought it was a global flag

Are we sure those flags are seen in fft -- since they don't appear in my build log... I'm not sure how meson exactly works, but does what happens in subdir(_core) influence `subdir(fft)? See

OK, I just empirically confirmed that that flag does not get passed on. And with that, that I can catch C++ exceptions. I'll push that up in a bit.

And with that, that I can catch C++ exceptions

If its about memory allocations, I think we should use our custom allocations over malloc or new.

@seiko2plus - the problem is that the allocations (and possible other errors) happen inside pocketfft and ideally we do not start rewriting that...

My current solution is just to put everything in a try/catch using macros and convert C++ exceptions to python errors. I don't really know how to test memory errors, but I checked that if I purposely pass in an invalid axis number, the corresponding exception properly leads to a RuntimeError. See last commit, ab16ef4

If its about memory allocations, I think we should use our custom allocations over malloc or new.

I'd really recommend not doing that. Most likely these custom operations are C-style and therefore will silently return NULL on error, which goes against fundamental assumptions C++ codes are making. You'd have to patch every place in pocketfft where memory is allocated and add error handling code.

p.s. Note that I now on purpose use pocketfft allocators, since those throw an exception, hence I could remove my custom code.

My current solution is just to put everything in a try/catch using macros and convert C++ exceptions to python errors.

To me this looks exactly like the right solution. It's located directly at the interface, does not require any big patching and can be extended in case more detailed error handling is required.
If memory allocations behave in the standard C++ way, they will be caught here as well.

"-fno-exceptions" should certainly not be used when compiling pocketfft and the wrapper code.
To be honest, I'd consider removing -fno-exceptions -fno-rtti from your configuration altogether. This was a good idea 20 years ago, but today the overhead (both memory and runtime) should be pretty negligible.

@seberg - thanks! Both suggestions were really nice, and make the code a lot cleaner. I pushed a revised version that uses them and that also replaces the addition of the C++ pocketfft header with a new submodule (which includes the license).

p.s. In the checklist above, this just leaves the question of whether this is tested thoroughly enough. I could adjust more tests to try different dtype, but I'm a bit wary of just increasing test run time for little benefit. Currently, the tests that I did change exercise the options to do longer or shorter FTs than the input (i.e., which set n), which has been the main pain point so far.

p.p.s. Well the above about tests was not actually true, but now it is, so I ticked my "more tests" item...

Is there anything more to do for this? (C++ pocketfft in np.fft)