qimpy.mpi.ProcessGrid

class ProcessGrid(comm, dim_names, shape=None)

Bases: object

Process grid of shape dimensions over communicator comm. Any -1 entries in shape are undetermined and will be resolved after the number of tasks split along that dimension are set using provide_n_tasks. Subsequently, use get_comm to get arbitrary hyperplane communicators that connect processes whose index only varies along specified subsets of dimensions.

Parameters:

  • comm (Comm) –

  • dim_names (str) –

  • shape (ndarray) –

__init__(comm, dim_names, shape=None)
Parameters:

  • comm (Comm) –

  • dim_names (str) –

  • shape (Sequence[int] | None) –

Return type:

None

Methods

__init__

get_comm

Get communicator for a hyper-plane spanning dim_names.

provide_n_tasks

Provide task count for a process grid dimension named dim_name.

Attributes

comm

Overall communicator within which grid is set-up.

n_procs

Total number of processes in grid.

i_proc

Overall rank of current process in grid.

dim_names

Each character (must be unique) names a dimension.

shape

Grid dimensions.
get_comm(dim_names)

Get communicator for a hyper-plane spanning dim_names. The resulting communicator will connect processes whose index in the process grid only varies along dimensions within dim_names. Dimensions before and including those in dim_names must be known, except when dim_names is a contiguous block of dimensions whose product can be determined now based on other dimensions.

Parameters:

dim_names (str) –

Return type:

Comm

provide_n_tasks(dim_name, n_tasks)

Provide task count for a process grid dimension named dim_name. If that dimension is undetermined (-1), set it to a suitable value that is compatible with the total processes, any other known dimensions, and with splitting n_tasks tasks with reasonable load balancing over this dimension.

Parameters:

  • dim_name (str) – Name of dimension (single charcater) to provide n_tasks for.

  • n_tasks (int) – Number of tasks available to split on this dimension of the process grid, used for setting dimension to ensure reasonable load balancing.

Return type:

None

comm: Comm

Overall communicator within which grid is set-up.

dim_names: str

Each character (must be unique) names a dimension.

i_proc: int

Overall rank of current process in grid.

n_procs: int

Total number of processes in grid.

shape: ndarray

Grid dimensions. Unresolved dimensions are -1.