All-to-all collective communication for numpy & cupy arrays with ChainerMN (demo)

alltoall_demo.py

import chainermn
from chainermn.communicators._memory_utility import array_to_buffer_object
import mpi4py.MPI

comm = chainermn.create_communicator('naive')


def alltoall_demo(xp):
    mpi_comm = comm.mpi_comm

    n = comm.size
    a = xp.arange(n * 3).reshape((n, -1)).astype(xp.float32)

    print('BEFORE', comm.rank, ':', a.ravel())
    mpi_comm.Alltoall(mpi4py.MPI.IN_PLACE, array_to_buffer_object(a))
    print('AFTER', comm.rank, ':', a.ravel())


import cupy as cp
import numpy as np
alltoall_demo(np)
alltoall_demo(cp)