MPIProcessGroup#

class nvmath.distributed.MPIProcessGroup(mpi_comm)[source]#

ProcessGroup implemented on mpi4py.

Parameters:

mpi_comm – mpi4py communicator.

Attributes

nranks#
rank#

Methods

allreduce_buffer(
array: ndarray,
*,
op: ReductionOp,
) → None[source]#

Allreduce an array.

Parameters:

  • array – Input and output of the collective. The function operates in-place.

  • op – One of the values from ReduceOp enum. Specifies an operation for element-wise reductions.

allreduce_object(
obj: T,
*,
op: Callable[[T, T], T],
) → T[source]#

Reduces all Python objects contributed by members of the group. The result is a single reduced object which is returned on every process.

Parameters:

  • obj – object contributed by this process.

  • op – A Python function that takes two objects and returns a single (reduced) object.

broadcast_buffer(
array: ndarray,
*,
root: int = 0,
) → None[source]#

Broadcast an array from one process to every process.

Parameters:

  • array – input (on root) and output of the collective.

  • root – rank of sending process.