Update vendored finufft and add GPU support

[lgarrison]

In recent months, cufinufft has been merged into the primary finufft codebase (thanks to @blackwer), which itself has matured a lot over the last few years. The vendored finufft hasn't been updated in a long time, so this PR does that, fixes some minor compatibility issues, and removes the deprecated vendored cufinufft. It also adds GPU support.

The GPU support comes from #4 but with an important update, which is that we can now pass the GPU stream that JAX gives us to the cufinufft library via flatironinstitute/finufft#330. This is probably required for multi-GPU support (although I haven't tested it) and may help with performance too. The only other changes were various updates to match the new API, build fixes, etc.

Some notes about the CMake build: cufinufft no longer requires building two static libraries with different macro definitions (one for float and one for double). The multiple precisions are supported via C++ templates, so the library can be built all at once. Also, now that finufft itself uses CMake, we might prefer to include it as a CMake sub-package rather than itemizing source files. But I left the itemization approach for now, since it required fewer changes.

For the moment, the vendored finufft is maintained in my own fork while we wait for some important PRs (including flatironinstitute/finufft#330 and flatironinstitute/finufft#354) to be merged. It should be easy to re-target the primary repo later.

The tests pass on the GPU and CPU, but more help testing the GPU in particular would be great! For anybody looking to run this on the Flatiron clusters, this is my build environment:

env.sh

ml modules/2.2
ml gcc
ml python/3.11
ml fftw
ml cmake
ml cuda/12
ml cudnn

[dfm]

Do we need this here? Perhaps we could adjust the tolerance for the allclose calls below to test both? I think the jax._src.public_test_util.check_close might do what we want.

[dfm]

This is awesome! I wonder if it's worth getting the Flatiron Jenkins CI set up to test the GPU ops? I've been doing that with the JAX CUDA ops in my exoplanet-core package: https://github.com/exoplanet-dev/exoplanet-core/blob/main/ci/Jenkinsfile

I'd like to update all the custom call stuff because the best practices have changed a little in the meantime, but that should be a separate PR.

For now, I'm keen to merge this with the only question being about running a GPU-enabled CI.

[lgarrison]

Yeah, I think it definitely makes sense to set this up on Jenkins! I actually started looking into this, but @Matematija reported an non-deterministic crash that I'm looking into first, might be when using the latest JAX. I'll keep you posted.

[lgarrison]

The problems appear at JAX 0.4.9, which is also the version where JAX starts using a non-blocking CUDA stream, according to jax-ml/jax#16580. I've gone through and fixed some cufinufft stream race conditions for kernel launches, cudaMemcpyAsync, Thrust, and (maybe) cufft, which resolves the particular crash that @Matematija reported, but there are still problems, because the tests don't pass. I'll need to dig into why (I'll probably be delayed in working on this by jury duty).

[lgarrison]

I think this is fixed now! The relevant patch is in the finufft submodule, so don't forget to update submodules when you pull. I'll work on Jenkins next.

[lgarrison]

@dfm Can you add the Jenkins webhook to this project, or give me permissions to do so? I think I would need to be a repo admin.