"""
Contains classes that abstract the MPI C API.
Classes in this module model the MPI node network and manage the relationship
of the complete domain to each
"""
from builtins import range
from builtins import object
import abc
from coffee.backend import backend as be
###############################################################################
# Abstract Base Class for MPI Interfaces
###############################################################################
[docs]class MPIInterface(object, metaclass=abc.ABCMeta):
def __init__(self, mpi_comm, *args, **kwds):
"""Initialisaer for MPIInterface objects.
Parameters
==========
mpi_comm: mpi4py.Comm
A wrapper to an MPI_COMM object via mpi4py.
"""
self.comm = mpi_comm
super(MPIInterface, self).__init__(*args, **kwds)
[docs] @abc.abstractmethod
def subdomain(self):
"""Returns a tuple of slice objects that represents the subdomain of the full
grid that this MPI node should perform computations over.
Returns
=======
tuple of slices
"""
[docs] @abc.abstractmethod
def communicate(self, data):
"""Returns a list of tuples. The first element of each tuple is a slice
representing the portion of data which corresponds to the second
element of the tuple. The second element of the tuple is an array of
data which are the communicated values that correspond to the
positions described by the slice."""
[docs] def barrier(self):
if self.comm is None:
return
return self.comm.barrier()
###############################################################################
# Concrete implementations
###############################################################################
[docs]class EvenCart(MPIInterface):
def __init__(self, domain, boundary_data, *args, **kwds):
"""Initialiser for EvenCart objects
An EvenCart object is an implementation of MPIInterface that
assumes that the global topology uses the MPI Cartesian methods.
Parameters
==========
domain: numpy.ndarray
A numpy array representing the values of points in the global
computational domain
boundary_data: grid.ABCBoundary
"""
super(EvenCart, self).__init__(*args, **kwds)
self.domain = domain
self.domain_mapping = self._make_domain_mappings(domain, boundary_data)
def _make_domain_mappings(self, domain, boundary_data):
"""A utility that constructs the tuples of slices that represent
the sub-grids for each MPI node based on that nodes rank.
The calculation assumes that each MPI node should handle roughly the
same load as all others and that the MPI topology is Cartesian.
Parameters
==========
domain: numpy.ndarray
A numpy array representing the values of points in the global
computational domain
boundary_data: grid.ABCBoundary
Returns
=======
list of tuples of slices:
Each tuple of slices in index i gives the sub grid for MPI node
with rank i.
"""
if self.comm is None:
return tuple([slice(None, None, None) for dim in domain])
r_map = []
for rank in range(self.comm.size):
coords = self.comm.Get_coords(rank)
r_map += [
tuple(
[
self._array_slice(
domain[i],
coord,
self.comm.dims[i],
boundary_data.ghost_points(i, -1),
boundary_data.ghost_points(i, 1),
)
for i, coord in enumerate(coords)
]
)
]
return r_map
def _neighbour_slices(self, shape, b_data):
"""A utility function that collects the source and destination
ranks for neighbouring MPI nodes along with the slices needed
to describe what data should be transferred.
Parameters
==========
shape: tuple of ints
The shape of the data that will be communicated to neighbouring MPI
nodes
b_data: grid.ABCBoundary
Returns
=======
list of tuples:
Each tuple has three members, first the source MPI node rank,
second the destination MPI node rank and third the list of tuples
of slices that describes the data to be communicated.
"""
if self.comm is None:
return []
dims = self.comm.Get_dim()
pos_neighbours = [
b_data.source_and_dest(d, 1) + b_data.internal_slice(shape, d, 1)
for d in range(dims)
]
neg_neighbours = [
b_data.source_and_dest(d, -1) + b_data.internal_slice(shape, d, -1)
for d in range(dims)
]
neighbours = pos_neighbours + neg_neighbours
return neighbours
def _array_slice(
self, array_length, rank, num_ranks, ghost_points_start, ghost_points_end
):
"""The utility method that calculates the slices that describe the
sub-grid managed by the MPI node with rank ``rank``.
The method assumes that it is doing the calculation for a single
dimension. To build the full list of tuples of slices, all
dimensions will need to be iterated over.
Parameters
==========
array_length: int
The number of grid points in this particular dimension of the
total computational domain.
rank: int
The MPI rank that the slices are computed for.
num_ranks: int
The number of MPI nodes that cover this particular dimension.
ghost_points_start: int
The number of ghost points at the start of the dimension.
ghost_points_end: int
The number of ghost points at the end of the dimension.
Returns
=======
slice:
A slice which gives the subdomain for the particular dimension and
the needed MPI node rank.
"""
# divide domain into appropriate parts
q, r = divmod(array_length, num_ranks)
# use the rank to details which part is relevant for this process
s = rank * q + min(rank, r)
e = s + q
# Adjust e to account for min(rank, r) term, which spreads the remainder
# over the appropriate number of processes
if rank < r:
e = e + 1
# add in ghost_points if we can
# this currently works for gp = 1 for an SAT boudnary method.
# The code will need to be checked for consistency if this changes.
if s - ghost_points_start > -1:
s = s - ghost_points_start
if e + ghost_points_end < array_length:
e = e + ghost_points_end
return slice(s, e, None)
@property
def subdomain(self):
"""Returns the slices which give the subdomain for the MPI node
that this method is called on.
Returns
=======
list of tuples of slices
"""
if self.comm is None:
return self.domain_mapping
return tuple(self.domain_mapping[self.comm.rank])
[docs] def communicate(self, data, b_data):
"""Communicates the given data to appropriate neighbouring nodes.
Parameters
==========
data: numpy.ndarray
An array representing the values of the computed function over
the sub-grid managed by the MPI node this code is being computed on.
b_data: grid.ABCBoundary
Returns
=======
list of tuples (list of slices, numpy.ndarray)
The returned data is a list of tuples. The first element of the tuple
describes which grid points the communicated data relates to.
The second element is the data received from whichever MPI node
manages the sub-grid which looks after the relevant grid points.
"""
# if self.comm.size == 1:
# return []
nslices = self._neighbour_slices(data.shape, b_data)
r_data = []
for source, dest, send_slice, recv_slice in nslices:
if dest < 0:
send_data = None
else:
# I really did not want to copy this view. For large data
# sets this feels like an unnecessary delay.
# However mpi4py did not seem to work with views and using
# self.comm.sendrecv to send and recv arbitrary objects.
# Plus self.comm.Sendrecv requires contiguous data
# so only in rare cases is it possible to send.
# Hence the be.copy statement is required.
send_data = be.array(data[send_slice], copy=True, order="C")
if source < 0:
recv_data = None
else:
# when testing the above assertion I suggest changing this
# line to be.ones_like, rather than be.empty_like
recv_data = be.empty_like(data[recv_slice])
self.comm.Sendrecv(send_data, dest=dest, recvbuf=recv_data, source=source)
if source >= 0:
r_data += [(recv_slice, recv_data)]
return r_data
[docs] def collect_data(self, data):
"""A method that collects all computed function values over all
sub-grids managed by all MPI nodes and returns the amalgamated data
if the rank of this process is 0.
This gives a way to get a single copy of all computed data so far. It does,
however, assume that all subsequence processing occurs on a single MPI
node. This is a bottle neck.
Parameters
==========
data: numpy.ndarray
The data to be communicated.
Returns
=======
None or numpy.ndarray
The collated data if MPI_Comm_Rank is 0, none otherwise.
"""
# Note that this method does not take account of ghost_points in
# the domains. This does not cause a problem. It just means
# that more data than necessary is passed.
# If there are no partitions of the domain we don't need to collect
if self.comm is None:
return data
# Gather the data
fields = self.comm.gather(data, root=0)
# If this process is root then collate the data and return
if self.comm.rank == 0:
rdata_edims_shape = (data.shape[0],) + data.shape[len(self.domain) + 1 :]
rdata_edims_slice = tuple(
[slice(None, None, None) for i in rdata_edims_shape]
)
rdata = be.zeros(
(rdata_edims_shape[0],) + self.domain + rdata_edims_shape[1:],
dtype=data.dtype,
)
for rank, field in enumerate(fields):
dslice = self.domain_mapping[rank]
rdata[(rdata_edims_slice[0],) + dslice + rdata_edims_slice[1:]] = field
return rdata
else:
return None