Error using s2fft with healpix in GPU

[mj-gomes]

Hi, I am using s2fft in a script where I use healpix maps. Everything works perfectly when using a CPU. However, I now need to use the GPU and it does not seem to work.

I believe the problem is in the combination CUDA + Jax + Healpy, because if I use CUDA and Jax but without using 'healpix' as the method for the fft calculations (I tested for example sampling='mw', method='jax_cuda') it works. And using sampling='healpix' and method 'jax_healpy', without CUDA (in a CPU) also works.

I installed s2fft from source to have the CUDA extension support as detailed in here and I am using CUDA 12.3, jax 0.5.3, jaxlib 0.5.3, and Healpy 1.16.6. Also CMake 3.26.5. All these versions should be compatible with s2fft, I believe.

This is the error that I get, it looks like callbacks from s2fft :

ERROR: CUFFT call "cufftXtSetCallback(forwardPlan, getfftShiftFloat(equator), CUFFT_CB_ST_COMPLEX, (void **)&params_dev)" in line 284 of file /tmp/pip-req-build-trgn26kr/lib/src/s2fft_callbacks.cu failed with code (5).
ERROR: CUFFT call "cufftXtSetCallback(backwardPlan, getfftNormFloat(equator, false), CUFFT_CB_ST_COMPLEX, (void **)&params_dev)" in line 287 of file /tmp/pip-req-build-trgn26kr/lib/src/s2fft_callbacks.cu failed with code (5).

And then it keeps repeating the two error messages. Does anyone know what is happening ? Thank you in advance.

[matt-graham]

Hi @mj-gomes. Thanks for the detailed report. It looks like the error code (5) reported by cuFFT corresponds to

CUFFT_INTERNAL_ERROR     =  5, // Driver or internal cuFFT library error

which unfortunately doesn't narrow things down much.

I have managed to reproduce this in a fresh Python 3.12 virtual environment on a system with a NVIDIA GPU with following package versions

Environment details

astropy                  7.1.0
astropy-iers-data        0.2025.9.15.0.37.0
coverage                 7.10.6
ducc0                    0.38.0
filelock                 3.19.1
fsspec                   2025.9.0
healpy                   1.18.1
iniconfig                2.1.0
jax                      0.5.3
jax-cuda12-pjrt          0.5.3
jax-cuda12-plugin        0.5.3
jaxlib                   0.5.3
jinja2                   3.1.6
markupsafe               3.0.2
ml-dtypes                0.5.3
mpmath                   1.3.0
networkx                 3.5
numpy                    1.26.4
nvidia-cublas-cu12       12.9.1.4
nvidia-cuda-cupti-cu12   12.9.79
nvidia-cuda-nvcc-cu12    12.9.86
nvidia-cuda-nvrtc-cu12   12.9.86
nvidia-cuda-runtime-cu12 12.9.79
nvidia-cudnn-cu12        9.13.0.50
nvidia-cufft-cu12        11.4.1.4
nvidia-cusolver-cu12     11.7.5.82
nvidia-cusparse-cu12     12.5.10.65
nvidia-nccl-cu12         2.28.3
nvidia-nvjitlink-cu12    12.9.86
nvidia-nvshmem-cu12      3.4.5
opt-einsum               3.4.0
packaging                25.0
pluggy                   1.6.0
pyerfa                   2.0.1.5
pygments                 2.19.2
pyssht                   1.5.3
pytest                   8.4.2
pytest-cov               7.0.0
pyyaml                   6.0.2
s2fft                    1.3.1.dev12        /home/ccaemgr/projects/s2fft
scipy                    1.16.2
setuptools               80.9.0
so3                      1.3.6
sympy                    1.14.0
torch                    2.8.0
typing-extensions        4.15.0

by running

import jax
jax.config.update("jax_enable_x64", True)
import s2fft
import numpy as np

L = 128
nside = L // 2
reality = True
sampling = "healpix"
method = "jax_cuda"
iter = 3

rng = np.random.default_rng(1234)
flm = s2fft.utils.signal_generator.generate_flm(rng, L=L, reality=reality)
f = s2fft.inverse(flm, L=L, nside=nside, reality=reality, method=method , sampling=sampling)
flm_recov = s2fft.forward(f, L=L, nside=nside, reality=reality, method=method, sampling=sampling, iter=iter)

which gives a similar series of errors to what you report. The initial s2fft.inverse call runs without errors as does the s2fft.forward call if iter=0 but using iter=1 or above or calling s2fft.forward more than once results in an error. I'm not sure what is causing the error and unfortunately won't have anytime to investigate for a while.

There is an open PR to refactor the CUDA extensions used by the method="jax_cuda" implementations (#290) which may fix this issue but there is a race condition in the implementation there which needs fixing before it can be merged.

If your bandlimit L / nside parameter value is not too large, the best solution for now may just be to use the native JAX implementation method="jax" which will incur a compilation time overhead but should have similar runtime performance to method="jax_cuda". If L is large though for your application then the compilation time may be too much of a blocker.

[mj-gomes]

Hi @matt-graham, thank you for your answer !

If your bandlimit L / nside parameter value is not too large, the best solution for now may just be to use the native JAX implementation method="jax" which will incur a compilation time overhead but should have similar runtime performance to method="jax_cuda". If L is large though for your application then the compilation time may be too much of a blocker.

Yes, that is what I ended up doing basically (using "jax_healpy" instead of "jax_cuda"). The compilation time is indeed somewhat long but it is no big blocker for my use case.