Use mpi4py-fft with CuPy arrays?

[leofang]

@dalcinl I had a tiny bit of time over the holidays to pick up this old task that has been on my mind for a couple of years. I noticed DistArray is the workhorse behind mpi4py-fft, however it requires subclassing numpy.ndarray

mpi4py-fft/mpi4py_fft/distarray.py

Line 10 in ac510f8

class DistArray(np.ndarray):

which makes it not possible to use cupy.ndarray as a drop-in replacement, as CuPy currently does not support subclassing (cupy/cupy#3972).

Is it absolutely necessary to make DistArray a subclass? Can it be tweaked (hopefully mildly) to become an array container to support both NumPy and CuPy arrays (and potentially other array/tensor libraries)? Could this change break backward compatibility? From my first glimpse it seems possible but I'd like to discuss further and see if I miss anything.

[mikaem]

Hi @leofang
I think it's more correct to consider DistArray a convenience class to create distributed Numpy arrays of the right shape, whereas the PFFT class is the workhorse. So I don't really think DistArray is the correct place to look for a cuda drop-in replacement. Note that the Numpy work arrays of PFFT are created in the FFT class and perhaps that is a better place to start? With a serial cuda transform class?

[leofang]

Thanks, @mikaem! I am not sure I fully follow, maybe I still miss something...

Yes, DistArray is a convenient class to create an array with the right shape locally. But my understanding of the workflow is

• Create a PFFT object with the desired transform specified
• Use the PFFT object to create a local array using newDistArray
• Perform the transform on the local array

In terms of code, I am looking at the tutorial:

import numpy as np
from mpi4py import MPI
from mpi4py_fft import PFFT, newDistArray

N = np.array([128, 128, 128], dtype=int)
fft = PFFT(MPI.COMM_WORLD, N, axes=(0, 1, 2), dtype=np.float, grid=(-1,))
u = newDistArray(fft, False)
u[:] = np.random.random(u.shape).astype(u.dtype)
u_hat = fft.forward(u, normalize=True)
uj = np.zeros_like(u)
uj = fft.backward(u_hat, uj)
assert np.allclose(uj, u)
print(MPI.COMM_WORLD.Get_rank(), u.shape)

From a CuPy developer/user's perspective, I expect that this code should just work after replacing numpy by cupy, but it'd require newDistArray to return a CuPy array.

[mikaem]

My point is that there are Numpy arrays within the PFFT class, or more precisely the FFT class. The arrays you set up with DistArray are copied into these before the transforms start. So just changing DistArray does not change the underlying routines. Transforms are not in-place. Which is why I think you need to add a serial fft class for cupy just to get started. Is there a serial cupy fft library?

[brownbaerchen]

What's the status of this? Support for subclassing has since been added to cupy (see the issue mentioned above) and cupy has FFTs as well.
I am just starting to look into this, but I suppose it may be possible to replace np with cp in both DistArray and PFFT and be done with it now? Of course I have no idea about the details...
Are there any plans to add cupy support from your side?

[mikaem]

There are no plans as far as I know. But contributions are welcome😃

[leofang]

I didn't have bandwidth for this, so feel free to give it a shot and let us know how it goes!

[brownbaerchen]

I started out with CuPyfying the DistArray class. There were a few issues, but I think so far nothing catastrophic.
It is not as simple as replacing np with cp after all, as the signature for __new__ is different and the get function needs to stay on CPU, I believe.
I will continue with the PFFT class next week.

I would like to merge this back here eventually. I guess you would reasonably ask the changes to be tested. However, testing code on GPU is not straightforward on GitHub, at least without paying. So is there any chance of merging this here? Do you have a way of testing on GPUs? Otherwise, I guess I would not have to pay as much attention to cleaning up the code. That's why I ask already now.
We are currently working on testing CuPy code in pySDC by mirroring to a local gitlab and running the tests on the Jülich supercomputers, but it's unclear when this will be running and how the accounting for compute budget will work.

If you are interested, you can track my progress here.

[leofang]

I won't have time before the holidays, unfortunately, but after that I could help test it locally. I agree too that it should be merged if everything works locally.

For accessing GPUs in GitHub Actions, that's a longer discussion that won't happen over night, what I can share now is interested parties are working on this.

[brownbaerchen]

I made some progress, but there is still plenty of room for improvement and I was hoping you have some suggestions.
First of all, I made some scaling tests with the current implementation on the JUWELS booster cluster in Jülich which has 96 CPU threads per node and 4 A100 40GB GPUs per node. The following is 3D c2c single precision back and forth FFTs with slab decomposition:

You can see up to three lines:

• fftw+MPI: Current state of mpi4py-fft on CPUs
• cupy+customMPI: GPU version of mpi4py-fft with custom Alltoallw
• cuFFTMp+NVSHMEM: Distributed FFT library from NVIDIA that doesn't have Python bindings, but is included for comparison

The x-axis is per node to judge if the compute budged is better spent on CPUs or GPUs. Clearly, the GPUs perform pretty well and mpi4py-fft works on them pretty well with only minor changes.
The weak scaling plot includes one data point with one GPU (1/4 node) which is basically a single 3D transform with no communication.

However, cuFFTMp is faster by about a factor of two in weak scaling. I see two big issues with the GPU version of mpi4py-fft. First, the default Alltoallw from mpi4py is really slow for some reason. I did not check what is going on under the hood, but the custom Alltoallw is based on a series of ring sends and is faster by about a factor of 10. To conserve memory, I block between each ring send. This can definitely be improved. Matching the performance of NVSHMEM with MPI is, however, not possible afaik, especially for strong scaling. Weak scaling is a higher priority for me, though.
But a second big issue is memory operations which consume about 25% of compute time. First a 2D transform is performed on the slabs, then the data is transformed and a second 1D transform is applied. In the second transform, I can see that a kernel cupy_copy__complex64_complex64 is launched before and after the actual transform and the data is returned in 'F' ordering. I cannot see this kind of memory operation in cuFFTMp, but maybe I did miss it. Do you have suggestions as to how I could avoid these memory operations?

[leofang]

Thanks for working on this and sharing your finding, @brownbaerchen, this is very encouraging!

First, the default Alltoallw from mpi4py is really slow for some reason.

I'd let @dalcinl or @mikaem comment on this, but my hunch is it's related to the CUDA aware MPI implementation -- which MPI did you use, and how did you build it?

In the second transform, I can see that a kernel cupy_copy__complex64_complex64 is launched before and after the actual transform and the data is returned in 'F' ordering.

I think this roughly explains the perf difference between CuPy + mpi4py-fft and cuFFTMp + NVSHMEM. If you share your source code of the mpi4py-fft integration, I could take a look next week and try to see if there's any improvement I can suggest.

btw how did you use cuFFTMp in Python?

[brownbaerchen]

Regarding the slow Alltoallw, I used OpenMPI and ParaStationMPI (MPICH) modules as installed on the Jülich supercomputers. I trust they installed it well. Maybe the data is copied to host and back with this due to non-contiguous buffers? I should check that at some point.

I didn't actually run cuFFTMp in Python. I may try to build bindings for that at some point. I just ran this file from the NVIDIA samples repository.

My code is available on this branch.
I guess the main thing is the FFT backend, which has the following code:

def _Xfftn_plan_cupy(shape, axes, dtype, transforms, options):
    import cupy as cp

    transforms = {} if transforms is None else transforms
    if tuple(axes) in transforms:
        plan_fwd, plan_bck = transforms[tuple(axes)]
    else:
        if cp.issubdtype(dtype, cp.floating):
            plan_fwd = cp.fft.rfftn
            plan_bck = cp.fft.irfftn
        else:
            plan_fwd = cp.fft.fftn
            plan_bck = cp.fft.ifftn


    s = tuple(np.take(shape, axes))
    U = cp.array(fftw.aligned(shape, dtype=dtype))  # TODO: avoid going via CPU
    _V = plan_fwd(U, s=s, axes=axes)
    V = cp.array(fftw.aligned_like(_V.get()))  # TODO: avoid going via CPU
    M = np.prod(s)

    # CuPy has forward transform unscaled and backward scaled with 1/N
    return (
        _Yfftn_wrap(plan_fwd, U, V, 1, {'s': s, 'axes': axes}),
        _Yfftn_wrap(plan_bck, V, U, M, {'s': s, 'axes': axes}),
    )

It's basically copy-paste from the numpy backend. I also tried cupyx.scipy.fftpack.fftn with overwrite_x=True, but the result was the same.

I don't construct plans myself because they are cached during the first transform and that is fine by me. But maybe there is a plan that avoids going to F ordering?

[llukas]

@brownbaerchen would this https://github.com/NVIDIA/CUDALibrarySamples/tree/master/cuFFTMp/JAX_FFT help running cuFFTMp in Python?

[brownbaerchen]

@brownbaerchen would this https://github.com/NVIDIA/CUDALibrarySamples/tree/master/cuFFTMp/JAX_FFT help running cuFFTMp in Python?

Thanks for mentioning this! Since I haven't done any work with pybind11 or JAX before, this is a bit opaque to me. I am wondering why this was done with JAX and not with CuPy.
I want to do distributed FFTs in a library that relies on subclassing numpy arrays. Subclassing cupy arrays works just fine and transitioning to JAX would probably be feasible but I haven't tried. At this point, I don't know what the best course of action is.

I guess I have three options:

• Figure out how to call cuFFT from CuPy more efficiently in the transform after transposing and continue in mpi4py-fft
• Design a new datatype in my library based on JAX and use the bindings you mentioned
• Make bindings for cuFFTMp for CuPy analogous to the ones using JAX

Any recommendations?

[leofang]

Sorry for long delay. @brownbaerchen Your snippet above caught my attention. You shouldn't need the lines calling fftw.aligned() and then converting to CuPy arrays. It's a waste of memory bandwidth. fftw.aligned() is just a convenient helper function allocating aligned CPU buffers. Could you try removing it and redo the benchmarks?

[brownbaerchen]

Sorry for long delay. @brownbaerchen Your snippet above caught my attention. You shouldn't need the lines calling fftw.aligned() and then converting to CuPy arrays. It's a waste of memory bandwidth. fftw.aligned() is just a convenient helper function allocating aligned CPU buffers. Could you try removing it and redo the benchmarks?

You're right, thanks for your input! I ignored this so far because it is only done once during setup. Indeed, just removing the detour to fftw.aligned() significantly speeds up instantiation of the PFFT object. But once it is instantiated, it doesn't make any difference. Since I didn't include the setup in the plots, the results don't change. In particular, the expensive memory operations remain. Still, your comment is very useful in practice!