Add GPU support

[tyler-a-cox]

Both finufft and matvis have GPU support for NUFFT, beam interpolation, and coordinate rotation. It shouldn't be too difficult to add GPU support for fftvis, but it may require some reorganizing of code.

[steven-murray]

Let me know if you want to bounce ideas on how to do this!

[kartikmandar]

Hey there, I saw that the cuFINUFFT was merged into FINUFFT itself, and we can use the already present cuFINUFFT’s API (now in FINUFFT) to make a simulate_gpu.py module.

We can abstract away the CPU implementation from simulate.py to simulate_cpu.py, and simulate.py can act as a high-level wrapper with parameters to choose CPU or GPU implementation.
In the simulate_gpu.py, each ray task could use a specific GPU using a round-robin assignment.

However, since cuFINUFFT expects GPUArray objects (via PyCUDA) instead of NumPy arrays, we can make helper functions to convert them. I don't think it would be much overhead.

As far as beam interpolation is concerned, it can be directly taken from matvis: GPUBeamInterpolator.

We will, of course, need some more tests, though.

I could take this up if nobody else is working on it.

[steven-murray]

Thanks @kartikmandar, that sounds like a reasonable way forward. Overall, I think we should stick to a similar API as in matvis.

[tyler-a-cox]

Thanks @kartikmandar, I’m definitely open to contributions if you’re interested in taking this. I agree with Steven that we should probably try to the API as similar to matvis as possible. Let me know if you’d like to meet to discuss!

[kartikmandar]

Sure I would follow the matvis API closely.
If it’s not extremely urgent can we have a discussion on Friday and I can work on the draft over the weekend.

[kartikmandar]

Hello @tyler-a-cox and @steven-murray,
So I went through both matvis and fftvis API structure, and this is what I am proposing for the change of structure:

fftvis/
├── pyproject.toml
├── README.md
├── LICENSE
├── .coveragerc
├── src/fftvis/
│ ├── init.py
│ ├── wrapper.py # High-level wrapper similar to matvis
│ ├── logutils.py

│ ├── core/ # Abstract/shared functionality
│ ├── init.py
│ ├── cli.py # Update for GPU profiling
│ ├── utils.py # Common utilities for both CPU/GPU

│ ├── cpu/
│ ├── init.py
│ ├── simulate.py
│ ├── beams.py

│ ├── gpu/
│ ├── init.py
│ ├── simulate.py
│ ├── beams.py

│ ├── tests/ # Add GPU tests
├── ci/
├── docs/ # Have to update for changes in API structure and GPU support
├── .github/

I will first initialise the change in structure, with a PR (struc-change) and check that everything for CPU simulation is working as it is.

If this feels okay let me know, I will go ahead with this?

[kartikmandar]

Should we be closing this?