PAOFLOW.utils.communication#

Attributes#

`comm`
`rank`
`size`

Functions#

`load_balancing`(size, rank, n)	Compute start and stop indices for a simple 1-D load-balanced partition.
`load_sizes`(size, n, dim)	Compute scatter/gather layout for all MPI ranks.
`scatter_array`(arr[, sroot])	Scatter the first dimension of an array to all MPI ranks.
`gather_array`(arr, arraux[, sroot])	Gather local sub-arrays into a full array on `sroot`.
`scatter_full`(arr, npool[, sroot])	Scatter the first dimension of an array using a chunked pool strategy.
`gather_full`(arr, npool[, sroot])	Gather local sub-arrays from all ranks using a chunked pool strategy.
`gather_scatter`(arr, scatter_axis, npool)	Redistribute an array so that `scatter_axis` is the distributed axis.
`gen_window`(array[, root])	Create a shared-memory window containing a copy of `array`.

Module Contents#

PAOFLOW.utils.communication.comm[source]#

PAOFLOW.utils.communication.rank[source]#

PAOFLOW.utils.communication.size[source]#

PAOFLOW.utils.communication.load_balancing(size, rank, n)[source]#

Compute start and stop indices for a simple 1-D load-balanced partition.

Parameters:

size (int) – Total number of MPI ranks.
rank (int) – Current MPI rank.
n (int) – Total number of items to distribute.

Returns:

start (int) – First item index assigned to this rank.
stop (int) – One past the last item index assigned to this rank.

PAOFLOW.utils.communication.load_sizes(size, n, dim)[source]#

Compute scatter/gather layout for all MPI ranks.

Parameters:

size (int) – Total number of MPI ranks.
n (int) – Number of elements in the first dimension to distribute.
dim (int) – Product of all remaining dimensions (i.e. the stride per element).

Returns:

Row i contains:

sizes[i, 0] — number of scalar elements on rank i (first-dimension count times dim).
sizes[i, 1] — scalar offset in the flat array where rank i’s sub-array begins.
sizes[i, 2] — first-dimension count on rank i.

Return type:

np.ndarray, shape (size, 3), int

PAOFLOW.utils.communication.scatter_array(arr, sroot=0)[source]#

Scatter the first dimension of an array to all MPI ranks.

Parameters:

arr (np.ndarray or None) – Array to scatter (only significant on sroot). The first dimension is distributed; all remaining dimensions are replicated.
sroot (int, optional) – Rank of the process that holds the data to scatter (default 0).

Returns:

Local sub-array for this rank. The first dimension has size determined by load_sizes(); remaining dimensions are identical to those of arr.

Return type:

np.ndarray

Notes

Uses MPI.Scatterv for variable-count scatter so that the array does not need to be exactly divisible by the number of ranks. Shape and dtype are broadcast from sroot before the transfer.

PAOFLOW.utils.communication.gather_array(arr, arraux, sroot=0)[source]#

Gather local sub-arrays into a full array on sroot.

Parameters:

arr (np.ndarray or None) – Destination array on rank sroot (only written on sroot). Its shape must match the result of gathering all arraux sub-arrays along the first dimension.
arraux (np.ndarray) – Local sub-array on this rank.
sroot (int, optional) – Rank that collects the gathered result (default 0).

Returns:

The result is written directly into arr on sroot.

Return type:

None

Notes

Uses MPI.Gatherv with offsets computed from load_sizes(). The layout is consistent with scatter_array().

PAOFLOW.utils.communication.scatter_full(arr, npool, sroot=0)[source]#

Scatter the first dimension of an array using a chunked pool strategy.

Parameters:

arr (np.ndarray or None) – Array to scatter (only significant on sroot).
npool (int) – Number of scatter pools; the array is chunked into npool groups to avoid large single MPI messages.
sroot (int, optional) – Root rank (default 0).

Returns:

Local sub-array on this rank, sized according to load_balancing().

Return type:

np.ndarray

Notes

Internally delegates to scatter_array() for each pool chunk to keep individual MPI message sizes manageable.

PAOFLOW.utils.communication.gather_full(arr, npool, sroot=0)[source]#

Gather local sub-arrays from all ranks using a chunked pool strategy.

Parameters:

arr (np.ndarray) – Local sub-array on this rank.
npool (int) – Number of gather pools.
sroot (int, optional) – Rank that collects the result (default 0).

Returns:

Assembled array on sroot; None on all other ranks.

Return type:

np.ndarray or None

Notes

Internally delegates to gather_array() for each pool chunk. The inverse operation of scatter_full().

PAOFLOW.utils.communication.gather_scatter(arr, scatter_axis, npool)[source]#

Redistribute an array so that scatter_axis is the distributed axis.

Parameters:

arr (np.ndarray) – Array to redistribute. All ranks hold an identical copy on entry.
scatter_axis (int) – Axis along which the output will be scattered across ranks.
npool (int) – Number of pools passed to the underlying scatter_full() / gather_full() calls.

Returns:

Sub-array for this rank, with scatter_axis distributed.

Return type:

np.ndarray

Notes

This is useful when a computation requires a different distribution axis than the one provided by an earlier scatter_full() call. Each rank first gathers the slices it needs, then the process with rank == r calls gather_full() with sroot=r in a loop over all ranks.

PAOFLOW.utils.communication.gen_window(array, root=0)[source]#

Create a shared-memory window containing a copy of array.

Parameters:

array (np.ndarray or None) – Array to place in shared memory (only significant on root).
root (int, optional) – Rank that owns the source data (default 0).

Returns:

View into the shared-memory buffer. All ranks on the same compute node can read the data without additional communication after the initial barrier.

Return type:

np.ndarray

Notes

Uses MPI.Win.Allocate_shared to create a one-sided memory window. Non-root ranks allocate zero bytes; the data pointer is obtained on all ranks via Shared_query(0).