MAINT, ENH: Implement calling pocketfft via gufunc and allow out argument #25536

mhvk · 2024-01-04T02:53:29Z

EDITED as no longer a draft.

Inspired by #25399, where @serge-sans-paille aimed to add an out parameter to the various fft routines, but where it turned out to be very hard to get it right, I instead turned to calling the fft routines via gufuncs, where the iterator takes care of things more automatically.

It all works well (with one small exception; see below), and I could even remove some of the copying of the whole array done for strange n (now doing much smaller copies in the ufunc). I also added the tests for passing in out from #25399 and udated them.

Note on the commits: I can merge these later, but for now kept them separate so one can see the logic a bit more. In particular, the first commit separates out the numpy-specific code from what one gets from upstream pocketfft, with the second adding a few small changes (that we probably should push upstream, if it is really a problem to call malloc(0)).

All test pass except one where it is checked that longdouble input gets dealt with properly. @seberg - might you be able to advice? I thought that with casting='unsafe' (or casting='same_kind') I would be able to force the use of the double loops that I defined, but somehow I still a TypeError that no suitable loop is found for casting='safe' (i.e., what I pass in for casting seems to be ignored). Trying to debug that (see #25535 for the fix that was needed), I found that the error is raised at

numpy/numpy/_core/src/umath/ufunc_object.c

Lines 4713 to 4724 in ce3e462

    
               /* 
        
                * Note that part of the promotion is to the complete the signature 
        
                * (until here it only represents the fixed part and is usually NULLs). 
        
                * 
        
                * After promotion, we could push the following logic into the ArrayMethod 
        
                * in the future.  For now, we do it here.  The type resolution step can 
        
                * be shared between the ufunc and gufunc code. 
        
                */ 
        
               PyArrayMethodObject *ufuncimpl = promote_and_get_ufuncimpl(ufunc, 
        
                       operands, signature, 
        
                       operand_DTypes, force_legacy_promotion, allow_legacy_promotion, 
        
                       promoting_pyscalars, NPY_FALSE);

serge-sans-paille · 2024-01-04T14:21:40Z

@mhvk thanks for looking at this :-)

seberg · 2024-01-05T08:27:29Z

@mhvk sorry, I have been focusing on other work (and was lazy on the holidays). The default casting is (still) "same-kind", so if you force the double loop it should be accepted, however, you need to force it of course (sorry tautology :)).

There are two ways to do this in principle:

You install a promoter as here (and its use)
You use the new style loops (also that file) and just guirilla add it for longdouble->double explicitly. (although the first seems clearer)

But we want the first, I think. For this ufunc, it might also be interesting to use the new-style inner-loop. That way, you would be able to cache the fft plan over multiple calls (But I am not too picky about actually doing it).

mhvk · 2024-01-05T15:54:51Z

@seberg - that sounds good in principle. I'm still confused, though, why it is not good enough to just pass in casting='unsafe'. Shouldn't that work?

seberg · 2024-01-05T16:38:14Z

No, casting doesn't change how promotion works, and it shouldn't! (And it never did) It would be enough to pass dtype=np.double, casting="unsafe", thought.

seberg · 2024-01-05T16:46:35Z

(Sorry, to make that one step further: it is should also be enough to pass just dtype, as the default is already same-kind casting)

mhvk · 2024-01-05T16:59:03Z

Ah, that works indeed, though instead of dtype, I have to pass in signature=ufunc.types[0], so that inputs and outputs are both covered.

mhvk

I added some comments throughout to help review.

mhvk · 2024-01-07T01:57:12Z

numpy/fft/_pocketfft.py

@@ -698,22 +736,22 @@ def _cook_nd_args(a, s=None, axes=None, invreal=0):
    return s, axes


-def _raw_fftnd(a, s=None, axes=None, function=fft, norm=None):
+def _raw_fftnd(a, s=None, axes=None, function=fft, norm=None, out=None):


In this loop, out is either never used, or automatically updated in-place. In principle, one could use in-place any time the shape does not change, but that would need extra logic, so perhaps out of scope here.

mhvk · 2024-01-07T01:57:40Z

numpy/fft/_pocketfft.py

@@ -46,33 +46,36 @@
 # divided. This replaces the original, more intuitive 'fct` parameter to avoid
 # divisions by zero (or alternatively additional checks) in the case of
 # zero-length axes during its computation.
-def _raw_fft(a, n, axis, is_real, is_forward, inv_norm):
+def _raw_fft(a, n, axis, is_real, is_forward, inv_norm, out=None):


Note the simplification because axis can just get passed on to the gufunc.

mhvk · 2024-01-07T02:03:02Z

numpy/fft/_pocketfft_umath.c

+}
+
+/*
+ * For the forward real, we cannot know what the requested number of points is


Note that the python side assumes one can pass in npts larger or smaller than nin - only the output size is related to npts.

mhvk · 2024-01-07T02:04:24Z

numpy/fft/pocketfft/pocketfft.c

@@ -9,18 +9,12 @@
 *  Copyright (C) 2004-2018 Max-Planck-Society
 *  \author Martin Reinecke
 */
-#define NPY_NO_DEPRECATED_API NPY_API_VERSION


The changes in this file are just stripping out the previous numpy call code, making the file as close as possible to the original (e.g., removing the added static).

mhvk · 2024-01-07T02:05:57Z

numpy/fft/pocketfft/pocketfft.h

@@ -0,0 +1,34 @@
+/*


The header file is now needed, since compilation is done separately.

seberg

Using templating might make some things a bit nicer (the copy function avoiding memcpy() and also the static void data passing). But both seems so small that it's not worth to get out C++ for.

Overall looks good, a couple of small comments, need to have another closer look (especially at the Python side), but I doubt I will notice anything.

numpy/fft/_pocketfft.py

numpy/fft/_pocketfft_umath.c

mhvk · 2024-01-07T22:48:32Z

@seberg - on the templating, I may still try it, but the new API seemed not entirely trivial without an easy example of a gufunc to work from. That said, perhaps these gufuncs can be the good example...

ngoldbaum · 2024-01-08T15:32:42Z

I've had a bunch of experience writing new-style ufuncs for stringdtype, although never with the gufunc machinery. Happy to help out in north american time zones if needed.

seberg · 2024-01-08T15:36:59Z

I don't think gufuncs change anything at all: The only difference between the two is that the input strides differ.
That said, the static void *data field is not passed in currently, so if that isn't solved with templating, that might be annoying.

Might have a look at it, overall, it might be easiest to just get started/push a start. But: I am totally fine with not doing it!

mhvk · 2024-01-09T14:44:10Z

@seberg, @ngoldbaum - I still like the idea of using the new API, but don't really see myself having time for it this week -- term has started again... I think this PR is in pretty good shape, though it needs a squash of most of the commits (I'd prefer not to squash all, but keep the changes to pocketfft separate).

mhvk · 2024-01-09T23:37:32Z

I went ahead and rebased, squashing to 3 commits: separate out to upstream pocketfft, apply the small number of numpy patches, and the addition/use of the gufuncs.

seberg

I'll just approve this to be clear about thinking that it is a good idea. I still didn't read the tests super carefully, but that should be OK.

There is the small symbol issue I noticed looking through, I wouldn't mind reading things once more more carefully, but overall that time would probably better spend making it C++ ;).

Let's get this in, if nobody beats me to it: next time I come across, I may push a fix for the symbols and just merge.

seberg · 2024-01-18T19:56:58Z

numpy/fft/pocketfft/pocketfft.h

+void destroy_rfft_plan (rfft_plan plan);
+int rfft_backward(rfft_plan plan, double c[], double fct);
+int rfft_forward(rfft_plan plan, double c[], double fct);
+size_t rfft_length(rfft_plan plan);


I think these should all have a NPY_NO_EXPORT ideally. Small thing, but we do it everywhere.

My logic was that we're vendoring the pocketfft files, so ideally we really change as little as possible. Does it matter if it is only included in the ufunc generation code? If you think it does matter, even a little, then I'll change it. But unlike the other changes to pocketfft, we could not upstream these.

I am not super sure what the advantags are beyond making them more private (i.e. if it can avoid clashs or has other advantages).

So all else being equal, it is the right thing to add it I think. But it probably doesn't matter a whole lot, so 🤷.

OK, I just added NPY_VISIBILITY_HIDDEN (which is the same as NPY_NO_EXPORT, but defined in numpy/numpyconfig.h rather than somewhere in multiarray, so a bit more logical for a vendored routine).

This commit ensures the pockefft.c file is identical to upstream.

These numpy changes could in principle be added to upstream, as they add robustness.

These changes make no sense for upstream.

The ufuncs taken npts from the output, so that using less or more than the full input remains possible.

seberg · 2024-01-19T16:50:32Z

Thanks Marten, let's get this in. It isn't like we are going to find substantive improvements at this point (beyond the discussed refactors), at least I hope ;).

@mhvk if you want to add a small release note fragment, just make a new PR.

seberg · 2024-04-10T06:45:16Z

It was pointed out to me that we made rfft(imaginary_array) fail in this PR. TBH, I think that is probably just as well (it is a bit annoying that the error doesn't report the dtypes to make it more obvious, though).

So I will lean to consider it a good 2.0 change that shouldn't really affect many (you currently rightly get the annoying complex warning). (I.e. the only real improvement here would be to improve that error message, IMO.)

mhvk · 2024-04-10T13:22:17Z

I agree that if people call rfft on a complex array it is almost certainly a mistake, so an error is fine. I wouldn't mind a better error message, though I think it might be the ufunc message one could improve, by listing both the input types and the accepted types, which would help all ufuncs. But perhaps relatively low priority?

p.s. Right now,

np.fft.rfft(np.arange(10)+1j)
TypeError: ufunc 'rfft_n_even' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''

seberg · 2024-04-10T14:16:41Z

Ah, that particular error is really more a "no loop found" path but in an old code path. So it should be pretty straight forward to the input dtype(s) here, but you can't really say which cast "failed", because we may still be in the promotion step rather than casting step (and at least I think the code path can be used for either step).

mhvk · 2024-04-10T14:34:01Z

So, maybe the TODO here would be to upgrade to the new machinery, which would then give a clearer error message?

seberg · 2024-04-12T07:25:46Z

Yeah, I guess it doesn't matter enough to fix this quickly. Upgrading to the new machinery is a big chunk of work and impossible to back-port (we have to remove the "disable nep 50" option first), but will be a big churn that would be nice to at least get started on in the next months.

mhvk added 01 - Enhancement component: numpy.fft 03 - Maintenance component: numpy.ufunc labels Jan 4, 2024

mhvk added this to the 2.0.0 release milestone Jan 4, 2024

mhvk mentioned this pull request Jan 4, 2024

Provide an 'out' parameter for numpy.fft.fft #25399

Closed

mhvk force-pushed the fft-as-gufunc branch 2 times, most recently from e9844e2 to 8271cfc Compare January 4, 2024 21:39

mhvk mentioned this pull request Jan 4, 2024

MAINT: fix ufunc debug tracing #25535

Merged

mhvk force-pushed the fft-as-gufunc branch from 8271cfc to 4ae8b88 Compare January 4, 2024 22:38

mhvk marked this pull request as ready for review January 4, 2024 23:51

mhvk changed the title ~~DRAFT Implement calling pocketfft via gufunc~~ DRAFT Implement calling pocketfft via gufunc and allow out argument Jan 4, 2024

mhvk force-pushed the fft-as-gufunc branch from d7dfe35 to fb6ba45 Compare January 5, 2024 16:57

mhvk force-pushed the fft-as-gufunc branch from fb6ba45 to 73e300a Compare January 5, 2024 17:00

mhvk changed the title ~~DRAFT Implement calling pocketfft via gufunc and allow out argument~~ MAINT, ENH: Implement calling pocketfft via gufunc and allow out argument Jan 5, 2024

mhvk commented Jan 7, 2024

View reviewed changes

mhvk force-pushed the fft-as-gufunc branch from 73e300a to 6ef606c Compare January 7, 2024 02:06

seberg reviewed Jan 7, 2024

View reviewed changes

mhvk force-pushed the fft-as-gufunc branch 3 times, most recently from 87a911b to e907f3a Compare January 9, 2024 01:28

mhvk force-pushed the fft-as-gufunc branch from e907f3a to 819053c Compare January 9, 2024 23:36

asmeurer mentioned this pull request Jan 10, 2024

ENH: Add fft optional extension submodule to numpy.array_api #25317

Merged

This was referenced Jan 11, 2024

NEP: Initial draft for NEP 43 for extensible ufuncs #16723

Merged

MAINT: Return size_t from num_codepoints in string ufuncs Buffer class #25571

Merged

rgommers mentioned this pull request Jan 15, 2024

NEP: add NEP 56 on array API standard support in main namespace #25542

Merged

seberg reviewed Jan 18, 2024

View reviewed changes

seberg approved these changes Jan 18, 2024

View reviewed changes

mhvk added 4 commits January 19, 2024 11:10

MAINT: Separate out pocketfft from the numpy additions

a8813cc

This commit ensures the pockefft.c file is identical to upstream.

MAINT: put back guards against allocating zero bytes

64aa994

These numpy changes could in principle be added to upstream, as they add robustness.

MAINT: mark pocketfft interface visibility hidden

0ac6bde

These changes make no sense for upstream.

MAINT,ENH: Add FFT gufuncs, use them, and add option to provide out.

570f272

The ufuncs taken npts from the output, so that using less or more than the full input remains possible.

mhvk force-pushed the fft-as-gufunc branch from 819053c to 570f272 Compare January 19, 2024 16:11

seberg merged commit 539dafa into numpy:main Jan 19, 2024

seberg mentioned this pull request Jan 19, 2024

ENH,MAINT: FFT Gufuncs could be improved with C++ and new API #25637

Open

mhvk deleted the fft-as-gufunc branch January 19, 2024 17:50

jakevdp mentioned this pull request Jan 22, 2024

BUG: np.fft.hfft on main segfaults when n=1 #25661

Closed

mhvk mentioned this pull request Jan 23, 2024

BUG: correct irfft with n=1 on larger input #25668

Merged

jakevdp mentioned this pull request Jan 24, 2024

BUG: np.fft.hfft results differ from previous release when n > len(arr) #25679

Closed

asmeurer mentioned this pull request Jan 26, 2024

np.fft.fft always returns np.complex128 regardless of input type #17801

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

MAINT, ENH: Implement calling pocketfft via gufunc and allow out argument #25536

MAINT, ENH: Implement calling pocketfft via gufunc and allow out argument #25536

mhvk commented Jan 4, 2024 •

edited

Loading

serge-sans-paille commented Jan 4, 2024

seberg commented Jan 5, 2024

mhvk commented Jan 5, 2024

seberg commented Jan 5, 2024

seberg commented Jan 5, 2024

mhvk commented Jan 5, 2024

mhvk left a comment

mhvk Jan 7, 2024

mhvk Jan 7, 2024

mhvk Jan 7, 2024 •

edited

Loading

mhvk Jan 7, 2024

mhvk Jan 7, 2024

seberg left a comment

mhvk commented Jan 7, 2024

ngoldbaum commented Jan 8, 2024

seberg commented Jan 8, 2024

mhvk commented Jan 9, 2024

mhvk commented Jan 9, 2024

seberg left a comment

seberg Jan 18, 2024

mhvk Jan 18, 2024

seberg Jan 18, 2024

mhvk Jan 19, 2024

seberg commented Jan 19, 2024

seberg commented Apr 10, 2024

mhvk commented Apr 10, 2024

seberg commented Apr 10, 2024

mhvk commented Apr 10, 2024

seberg commented Apr 12, 2024

	/*
	* Note that part of the promotion is to the complete the signature
	* (until here it only represents the fixed part and is usually NULLs).
	*
	* After promotion, we could push the following logic into the ArrayMethod
	* in the future. For now, we do it here. The type resolution step can
	* be shared between the ufunc and gufunc code.
	*/
	PyArrayMethodObject *ufuncimpl = promote_and_get_ufuncimpl(ufunc,
	operands, signature,
	operand_DTypes, force_legacy_promotion, allow_legacy_promotion,
	promoting_pyscalars, NPY_FALSE);

MAINT, ENH: Implement calling pocketfft via gufunc and allow out argument #25536

MAINT, ENH: Implement calling pocketfft via gufunc and allow out argument #25536

Conversation

mhvk commented Jan 4, 2024 • edited Loading

serge-sans-paille commented Jan 4, 2024

seberg commented Jan 5, 2024

mhvk commented Jan 5, 2024

seberg commented Jan 5, 2024

seberg commented Jan 5, 2024

mhvk commented Jan 5, 2024

mhvk left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mhvk Jan 7, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

seberg left a comment

Choose a reason for hiding this comment

mhvk commented Jan 7, 2024

ngoldbaum commented Jan 8, 2024

seberg commented Jan 8, 2024

mhvk commented Jan 9, 2024

mhvk commented Jan 9, 2024

seberg left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

seberg commented Jan 19, 2024

seberg commented Apr 10, 2024

mhvk commented Apr 10, 2024

seberg commented Apr 10, 2024

mhvk commented Apr 10, 2024

seberg commented Apr 12, 2024

mhvk commented Jan 4, 2024 •

edited

Loading

mhvk Jan 7, 2024 •

edited

Loading