[FR] MPI.gatherv

[PhilipVinc]

Some time ago i tried adding support to mpi.gatherv but i was blocked by the following issue:

gatherv in C requires passing an array with the length of the various partitions on each node.
This, from the point of view of jax, is static information.
However i do not know how to create a static tuple in python and pass it to the cython module

[dionhaefner]

I think the key point to figure out is whether the length array should live on host or device in GPU mode. Can MPI handle it being on device memory? If so, you can just treat it like any other (non-static) input.

[PhilipVinc]

host

[dionhaefner]

OK. Then I see these options:

• Copy the offsets to host before the MPI call, then you can just pass it along with the input on both CPU and GPU.
• JAX does support passing constant arrays, so we can use this on CPU, and add it to opaque on GPU. But that would mean that it has to be static. Option 1 would allow for dynamic values as well.

[PhilipVinc]

this array comes from shapes, so it is static information.
I did not know how to pass the array but that seems to do the trick. I'll eventually try to resuscitate this.

[dionhaefner]

Actually, I think you will have to pass it as a int* or so on GPU, because the opaque struct needs to have a fixed length. So you might as well do that on CPU, too.

[YigitElma]

Hi! Are there any near-term plans to implement MPI.Gatherv? For this operation, I currently use mpi4py, but it would be perfect to be able to use mpi4jax and jit my computation 🙏

[dionhaefner]

No plans to implement this (or any other ops) to mpi4jax, but happy to review a PR that does it.

[YigitElma]

No plans to implement this (or any other ops) to mpi4jax, but happy to review a PR that does it.

Thanks for the reply! I might give it a shot later.