Updating Healpix CUDA primitive

[ASKabalan]

Adding a few updates

• Updating to the newest custom call API (API 4) using FFI
• implementing a grad rule for healpix cuda FFT
• Implementing a Batching rule

A batching rule seems to be very important for two things
Being able to jacrev/ jacfwd
and because in most cases .. the size of a healpix map can fit on a single GPU but sometimes we want to batch the spherical transform

I will be doing that next

[ASKabalan]

Hello @matt-graham @jasonmcewen @CosmoMatt

Just a quick PR to wrap up a few stuff

1. Updated the binding API to the newest FFI
2. Added a vmap implementation of the cuda primitive
3. Added a transpose rule which allows jacfwd and jacrev (consequently grad aswell)
4. added more tests https://github.com/astro-informatics/s2fft/blob/ASKabalan/tests/test_healpix_ffts.py#L100
5. Removed two files which are now no longer needed with the FFI API (kernel helpers) (so maybe they should be removed from the license section)
6. Constrained nanobind to be nanobind >=2.0,<2.6" because of a regression [BUG]: Regression when using scikit build tools and nanobind wjakob/nanobind#982

And finally I added cudastreamhandler which is used to split the XLA provided stream for the VMAP lowering (this header is my own work)

There is an issue with building pyssht not sure that this is my fault

I will check the failing worflows when I get the chance, but in the meantime a review is appreciated

[matt-graham]

Hello @matt-graham @jasonmcewen @CosmoMatt

Just a quick PR to wrap up a few stuff

1. Updated the binding API to the newest [FFI](https://docs.jax.dev/en/latest/ffi.html)

2. Added a vmap implementation of the cuda primitive

3. Added a transpose rule which allows jacfwd and jacrev (consequently grad aswell)

4. added more tests https://github.com/astro-informatics/s2fft/blob/ASKabalan/tests/test_healpix_ffts.py#L100

5. Removed two files which are now no longer needed with the FFI API ([kernel helpers](https://github.com/astro-informatics/s2fft/blob/main/lib/include/kernel_helpers.h)) (so maybe they should be removed from the license section)

6. Constrained nanobind to be nanobind >=2.0,<2.6" because of a regression [[BUG]: Regression when using scikit build tools and nanobind wjakob/nanobind#982](https://github.com/wjakob/nanobind/issues/982)

And finally I added cudastreamhandler which is used to split the XLA provided stream for the VMAP lowering (this header is my own work)

There is an issue with building pyssht not sure that this is my fault

I will check the failing worflows when I get the chance, but in the meantime a review is appreciated

Hi @ASKabalan, sorry for the delay in getting back to you.

This all sounds great - thanks for picking up #237 in particular and for the updates to use the newer FFI interface.

With regards to the failing workflows - this was probably due to #292 which was fixed in #293. If you merge in latest main here that should hopefully resolve the upstream dependency build problems that were causing the test workflows to fail.

I've added some initial review comments below. Will have a closer look next week and try testing this out, but don't have access to GPU machine atm.

[matt-graham]

I think we could use s2fft.inverse(flm, L=L, reality=False, method="jax", sampling="healpix") here instead of going via healpy? Rationale being that I would have a slight preference for minimising the number of additional tests that depend on healpy as it we are no longer requiring it as direct dependency for package and in the long run it might be possible to also remove it as a test dependency.

[matt-graham]

I've tried building, installing and running this on a system with CUDA 12.6 + a NVIDIA A100, and running the HEALPix FFT tests with

pytest tests/test_healpix_ffts.py

consistently the tests hang when trying to run the first test_healpix_fft_cuda instance.

Running just the IFFT tests with

pytest tests/test_healpix_ffts.py::test_healpix_ifft_cuda

the tests for both set of test parameters pass.

Trying to dig into this a bit, running the following locally

import healpy
import jax
import s2fft
import numpy

jax.config.update("jax_enable_x64", True)

seed = 20250416
nside = 4
L = 2 * nside
reality = False

rng = numpy.random.default_rng(seed)
flm = s2fft.utils.signal_generator.generate_flm(rng=rng, L=L, reality=reality)
flm_hp = s2fft.sampling.s2_samples.flm_2d_to_hp(flm, L)
f = healpy.sphtfunc.alm2map(flm_hp, nside, lmax=L - 1)
flm_cuda = s2fft.utils.healpix_ffts.healpix_fft_cuda(f=f, L=L, nside=nside, reality=reality).block_until_ready()

raises an error

jaxlib.xla_extension.XlaRuntimeError: INTERNAL: CUDA error: : CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered

so it looks like there is some memory addressing issue somewhere in the healpix_fft_cuda implementation?

[ASKabalan]

Thank you

I was able to reproduce with 12.4.1 but not locally with 12.4

I will take a look

[codecov]

Codecov Report

Attention: Patch coverage is 75.00000% with 4 lines in your changes missing coverage. Please review.

Project coverage is 96.07%. Comparing base (0de6f11) to head (d29af9b).

Files with missing lines	Patch %	Lines
s2fft/utils/healpix_ffts.py	66.66%	3 Missing ⚠️
s2fft/utils/jax_primitive.py	75.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #290      +/-   ##
==========================================
- Coverage   96.55%   96.07%   -0.48%     
==========================================
  Files          32       32              
  Lines        3450     3469      +19     
==========================================
+ Hits         3331     3333       +2     
- Misses        119      136      +17     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ASKabalan]

@matt-graham Hey
I am picking were i left off
So it seems that there is an error when building with python3.8
Doesn't seem to be coming from my code.
It seems to be because of some compile error when compiling sht

I would suggest dropping python3.8 from the test suite since JAX no longer supports it anyway

[matt-graham]

@matt-graham Hey I am picking were i left off So it seems that there is an error when building with python3.8 Doesn't seem to be coming from my code. It seems to be because of some compile error when compiling sht

Hi @ASKabalan. Do you mean so3 rather than (py)ssht? From a quick look at the logs of the failing Actions workflow job on Python 3.8 / ubuntu-latest it appears like it's an error with building so3 (ERROR: Failed building wheel for so3). If so this is likely the same issue as described in #308. I've opened a PR to try to fix this upstream in so3 (astro-informatics/so3#31).

I would suggest dropping python3.8 from the test suite since JAX no longer supports it anyway

Yes agreed we should drop Python 3.8 from test matrix - we have an open pull request #305 to update to only supporting Python 3.11+ but this is partially blocked by #212 as the tests currently exit with fatal errors when running on MacOS / Python 3.9+ due to an incompatibility between the OpenMP runtime's the MacOS wheels for healpy and PyTorch are built for (healpy/healpy#1012)