MPI segfaulting

[abid1214]

When I try using the 'gpumpi' simulator for running qaoa, I segfault when using more than 1 node.

Specifically, I run

mpiexec -n 2 python benchmark.py

where benchmark.py has the lines

f = get_qaoa_maxcut_objective(N, p, G, simulator='gpumpi')
f(theta)

when calling f(theta), it eventually calls furx_all in qokit/fur/mpi_nbcuda/fur.py. I get a segfault when running this, and specifically in the following line:

from ..lazy_import import MPI
from ..nbcuda.fur import furx_all as furx_local


def furx_all(x, theta: float, n_local_qubits: int, n_all_qubits: int, comm):
    assert n_local_qubits <= n_all_qubits
    assert n_all_qubits <= 2 * n_local_qubits, "n_all_qubits > 2*n_local_qubits is not yet implemented"

    for i in range(n_local_qubits):
        furx_local(x, theta, i)

    if n_all_qubits > n_local_qubits:
        comm.Alltoall(MPI.IN_PLACE, x)  #  <- SEGFAULTS AT THIS LINE ------ #
        for i in range(2 * n_local_qubits - n_all_qubits, n_local_qubits):
            furx_local(x, theta, i)
        comm.Alltoall(MPI.IN_PLACE, x)

I've attached benchmark.py to replicate this issue. The branch I'm working on is fix-mpi

benchmark.py.txt