Distributed real fft & ifft

[sijieh23]

Hi team,

Just wondering, do you guys plan to implement distributed real fft & ifft? Current distributed fft only supports "C2C" transform, but not "R2C" and "C2R".

[juanjgalvez]

Hi @sijieh23 ,

Distributed R2C and C2R FFT implementation is part of the next nvmath-python release (v0.6.0), which is targeted for release in the next few days.

[samaid]

C2R and R2C FFTs are now live. Enjoy and let us know if you experience any issues!

[sijieh23]

I didn't find anything in the document about gathering the distributed data, something like mpi_allgather in the MPI language (for example, I want to collect data and write to disk). Is it automatically done inside the nvmath, or are there other tricks to do this?

[juanjgalvez]

nvmath-python is designed to interoperate with the Python ecosystem and the output of any nvmath-python operation will be a NumPy, CuPy or PyTorch array/tensor. You can use existing packages to gather the result on one process.

For example, you can use mpi4py to gather NumPy arrays and, if you have a CUDA-aware MPI implementation, you can also gather CuPy arrays directly without explicit copies to host memory. You can gather torch tensors with torch.distributed.

[sijieh23]

I'm using cupy. So I guess I can use cupy.distributed?

BTW, just wondering how nvmath's distributed fft handles numpy arrays? Moving to GPU memory, performing fft, and then transferring back to CPU memory?

[juanjgalvez]

I haven't used cupy.distributed, but you should be able to wrap your cupy array in a cupyx.distributed.array.distributed_array and then perform collective operations on it using NCCL. But if you don't care very much about performance for this step, then probably the easiest thing is to convert the cupy arrays to numpy with cupy.asnumpy(cp_array) and do the gather with mpi4py. Also, keep in mind that if every process doesn't hold the same amount of data, that you'll probably need to do a gatherv instead of gather (not sure if this is your case, it depends on the size of input and whether you're doing R2C, C2R, C2C).

BTW, just wondering how nvmath's distributed fft handles numpy arrays? Moving to GPU memory, performing fft, and then transferring back to CPU memory?

Exactly. The distributed FFT is executed on GPU with cuFFTMp library, so if you pass a numpy array, nvmath-python copies the operand to GPU memory, executes the transform and copies the result back to the original numpy array.