heat.communication

Module implementing the communication layer of HeAT

Module Contents

class MPIRequest(handle, sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any = None, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any = None, tensor: torch.Tensor = None, permutation: Tuple[int, Ellipsis] = None)

Represents a handle on a non-blocking operation

Parameters:
  • handle (MPI.Communicator) – Handle for the mpi4py Communicator

  • sendbuf (DNDarray or torch.Tensor or Any) – The buffer for the data to be send

  • recvbuf (DNDarray or torch.Tensor or Any) – The buffer to the receive data

  • tensor (torch.Tensor) – Internal Data

  • permutation (Tuple[int,...]) – Permutation of the tensor axes

Wait(status: mpi4py.MPI.Status = None)

Waits for an MPI request to complete

__getattr__(name: str) Callable

Default pass-through for the communicator methods.

Parameters:

name (str) – The name of the method to be called.

class Communication

Base class for Communications (inteded for other backends)

is_distributed() NotImplementedError

Whether or not the Communication is distributed

chunk(shape, split) NotImplementedError

Calculates the chunk of data that will be assigned to this compute node given a global data shape and a split axis. Returns (offset, local_shape, slices): the offset in the split dimension, the resulting local shape if the global input shape is chunked on the split axis and the chunk slices with respect to the given shape

Parameters:
  • shape (Tuple[int,...]) – The global shape of the data to be split

  • split (int) – The axis along which to chunk the data

class MPICommunication(handle=MPI.COMM_WORLD)

Bases: Communication

Class encapsulating all MPI Communication

Parameters:

handle (MPI.Communicator) – Handle for the mpi4py Communicator

__mpi_type_mappings
is_distributed() bool

Determines whether the communicator is distributed, i.e. handles more than one node.

chunk(shape: Tuple[int], split: int, rank: int = None, w_size: int = None, sparse: bool = False) Tuple[int, Tuple[int], Tuple[slice]]

Calculates the chunk of data that will be assigned to this compute node given a global data shape and a split axis. Returns (offset, local_shape, slices): the offset in the split dimension, the resulting local shape if the global input shape is chunked on the split axis and the chunk slices with respect to the given shape

Parameters:
  • shape (Tuple[int,...]) – The global shape of the data to be split

  • split (int) – The axis along which to chunk the data

  • rank (int, optional) – Process for which the chunking is calculated for, defaults to self.rank. Intended for creating chunk maps without communication

  • w_size (int, optional) – The MPI world size, defaults to self.size. Intended for creating chunk maps without communication

  • sparse (bool, optional) – Specifies whether the array is a sparse matrix

counts_displs_shape(shape: Tuple[int], axis: int) Tuple[Tuple[int], Tuple[int], Tuple[int]]

Calculates the item counts, displacements and output shape for a variable sized all-to-all MPI-call (e.g. MPI_Alltoallv). The passed shape is regularly chunk along the given axis and for all nodes.

Parameters:
  • shape (Tuple[int,...]) – The object for which to calculate the chunking.

  • axis (int) – The axis along which the chunking is performed.

mpi_type_and_elements_of(obj: heat.core.dndarray.DNDarray | torch.Tensor, counts: Tuple[int], displs: Tuple[int], is_contiguous: bool | None) Tuple[mpi4py.MPI.Datatype, Tuple[int, Ellipsis]]

Determines the MPI data type and number of respective elements for the given tensor (DNDarray or ``torch.Tensor). In case the tensor is contiguous in memory, a native MPI data type can be used. Otherwise, a derived data type is automatically constructed using the storage information of the passed object.

Parameters:
  • obj (DNDarray or torch.Tensor) – The object for which to construct the MPI data type and number of elements

  • counts (Tuple[ints,...], optional) – Optional counts arguments for variable MPI-calls (e.g. Alltoallv)

  • displs (Tuple[ints,...], optional) – Optional displacements arguments for variable MPI-calls (e.g. Alltoallv)

  • is_contiguous (bool) – Information on global contiguity of the memory-distributed object. If None, it will be set to local contiguity via torch.Tensor.is_contiguous().

  • ToDo (#)

as_mpi_memory(obj) mpi4py.MPI.memory

Converts the passed torch.Tensor into an MPI compatible memory view.

Parameters:

obj (torch.Tensor) – The tensor to be converted into a MPI memory view.

as_buffer(obj: torch.Tensor, counts: Tuple[int] = None, displs: Tuple[int] = None, is_contiguous: bool | None = None) List[mpi4py.MPI.memory | Tuple[int, int] | mpi4py.MPI.Datatype]

Converts a passed torch.Tensor into a memory buffer object with associated number of elements and MPI data type.

Parameters:
  • obj (torch.Tensor) – The object to be converted into a buffer representation.

  • counts (Tuple[int,...], optional) – Optional counts arguments for variable MPI-calls (e.g. Alltoallv)

  • displs (Tuple[int,...], optional) – Optional displacements arguments for variable MPI-calls (e.g. Alltoallv)

  • is_contiguous (bool, optional) – Optional information on global contiguity of the memory-distributed object.

alltoall_sendbuffer(obj: torch.Tensor) List[mpi4py.MPI.memory | Tuple[int, int] | mpi4py.MPI.Datatype]

Converts a passed torch.Tensor into a memory buffer object with associated number of elements and MPI data type. XXX: might not work for all MPI stacks. Might require multiple type commits or so

Parameters:

obj (torch.Tensor) – The object to be transformed into a custom MPI datatype

alltoall_recvbuffer(obj: torch.Tensor) List[mpi4py.MPI.memory | Tuple[int, int] | mpi4py.MPI.Datatype]

Converts a passed torch.Tensor into a memory buffer object with associated number of elements and MPI data type. XXX: might not work for all MPI stacks. Might require multiple type commits or so

Parameters:

obj (torch.Tensor) – The object to be transformed into a custom MPI datatype

Free() None

Free a communicator.

Split(color: int = 0, key: int = 0) MPICommunication

Split communicator by color and key.

Parameters:
  • color (int, optional) – Determines the new communicator for a process.

  • key (int, optional) – Ordering within the new communicator.

Irecv(buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, source: int = MPI.ANY_SOURCE, tag: int = MPI.ANY_TAG) MPIRequest

Nonblocking receive

Parameters:
  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to place the received message

  • source (int, optional) – Rank of source process, that send the message

  • tag (int, optional) – A Tag to identify the message

Recv(buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, source: int = MPI.ANY_SOURCE, tag: int = MPI.ANY_TAG, status: mpi4py.MPI.Status = None)

Blocking receive

Parameters:
  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to place the received message

  • source (int, optional) – Rank of the source process, that send the message

  • tag (int, optional) – A Tag to identify the message

  • status (MPI.Status, optional) – Details on the communication

__send_like(func: Callable, buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, dest: int, tag: int) Tuple[heat.core.dndarray.DNDarray | torch.Tensor | None]

Generic function for sending a message to process with rank “dest”

Parameters:
  • func (Callable) – The respective MPI sending function

  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the message to be send

  • dest (int, optional) – Rank of the destination process, that receives the message

  • tag (int, optional) – A Tag to identify the message

Bsend(buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, dest: int, tag: int = 0)

Blocking buffered send

Parameters:
  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the message to be send

  • dest (int, optional) – Index of the destination process, that receives the message

  • tag (int, optional) – A Tag to identify the message

Ibsend(buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, dest: int, tag: int = 0) MPIRequest

Nonblocking buffered send

Parameters:
  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the message to be send

  • dest (int, optional) – Rank of the destination process, that receives the message

  • tag (int, optional) – A Tag to identify the message

Irsend(buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, dest: int, tag: int = 0) MPIRequest

Nonblocking ready send

Parameters:
  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the message to be send

  • dest (int, optional) – Rank of the destination process, that receives the message

  • tag (int, optional) – A Tag to identify the message

Isend(buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, dest: int, tag: int = 0) MPIRequest

Nonblocking send

Parameters:
  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the message to be send

  • dest (int, optional) – Rank of the destination process, that receives the message

  • tag (int, optional) – A Tag to identify the message

Issend(buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, dest: int, tag: int = 0) MPIRequest

Nonblocking synchronous send

Parameters:
  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the message to be send

  • dest (int, optional) – Rank of the destination process, that receives the message

  • tag (int, optional) – A Tag to identify the message

Rsend(buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, dest: int, tag: int = 0)

Blocking ready send

Parameters:
  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the message to be send

  • dest (int, optional) – Rank of the destination process, that receives the message

  • tag (int, optional) – A Tag to identify the message

Ssend(buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, dest: int, tag: int = 0)

Blocking synchronous send

Parameters:
  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the message to be send

  • dest (int, optional) – Rank of the destination process, that receives the message

  • tag (int, optional) – A Tag to identify the message

Send(buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, dest: int, tag: int = 0)

Blocking send

Parameters:
  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the message to be send

  • dest (int, optional) – Rank of the destination process, that receives the message

  • tag (int, optional) – A Tag to identify the message

__broadcast_like(func: Callable, buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, root: int) Tuple[heat.core.dndarray.DNDarray | torch.Tensor | None]

Generic function for broadcasting a message from the process with rank “root” to all other processes of the communicator

Parameters:
  • func (Callable) – The respective MPI broadcast function

  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the message to be broadcasted

  • root (int) – Rank of the root process, that broadcasts the message

Bcast(buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, root: int = 0) None

Blocking Broadcast

Parameters:
  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the message to be broadcasted

  • root (int) – Rank of the root process, that broadcasts the message

Ibcast(buf: heat.core.dndarray.DNDarray | torch.Tensor | Any, root: int = 0) MPIRequest

Nonblocking Broadcast

Parameters:
  • buf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the message to be broadcasted

  • root (int) – Rank of the root process, that broadcasts the message

__reduce_like(func: Callable, sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, *args, **kwargs) Tuple[heat.core.dndarray.DNDarray | torch.Tensor | None]

Generic function for reduction operations.

Parameters:
  • func (Callable) – The respective MPI reduction operation

  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result of the reduction

Allreduce(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, op: mpi4py.MPI.Op = MPI.SUM)

Combines values from all processes and distributes the result back to all processes

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result of the reduction

  • op (MPI.Op) – The operation to perform upon reduction

Exscan(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, op: mpi4py.MPI.Op = MPI.SUM)

Computes the exclusive scan (partial reductions) of data on a collection of processes

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result of the reduction

  • op (MPI.Op) – The operation to perform upon reduction

Iallreduce(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, op: mpi4py.MPI.Op = MPI.SUM) MPIRequest

Nonblocking allreduce reducing values on all processes to a single value

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result of the reduction

  • op (MPI.Op) – The operation to perform upon reduction

Iexscan(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, op: mpi4py.MPI.Op = MPI.SUM) MPIRequest

Nonblocking Exscan

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result of the reduction

  • op (MPI.Op) – The operation to perform upon reduction

Iscan(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, op: mpi4py.MPI.Op = MPI.SUM) MPIRequest

Nonblocking Scan

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result of the reduction

  • op (MPI.Op) – The operation to perform upon reduction

Ireduce(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, op: mpi4py.MPI.Op = MPI.SUM, root: int = 0) MPIRequest

Nonblocking reduction operation

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result of the reduction

  • op (MPI.Op) – The operation to perform upon reduction

  • root (int) – Rank of the root process

Reduce(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, op: mpi4py.MPI.Op = MPI.SUM, root: int = 0)

Reduce values from all processes to a single value on process “root”

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result of the reduction

  • op (MPI.Op) – The operation to perform upon reduction

  • root (int) – Rank of the root process

Scan(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, op: mpi4py.MPI.Op = MPI.SUM)

Computes the scan (partial reductions) of data on a collection of processes in a nonblocking way

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result of the reduction

  • op (MPI.Op) – The operation to perform upon reduction

__allgather_like(func: Callable, sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, axis: int, **kwargs)

Generic function for allgather operations.

Parameters:
  • func (Callable) – Type of MPI Allgather function (i.e. allgather, allgatherv, iallgather)

  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • axis (int) – Concatenation axis: The axis along which sendbuf is packed and along which recvbuf puts together individual chunks

Allgather(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recv_axis: int = 0)

Gathers data from all tasks and distribute the combined data to all tasks

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • recv_axis (int) – Concatenation axis: The axis along which sendbuf is packed and along which recvbuf puts together individual chunks

Allgatherv(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recv_axis: int = 0)

v-call of Allgather: Each process may contribute a different amount of data.

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • recv_axis (int) – Concatenation axis: The axis along which sendbuf is packed and along which recvbuf puts together individual chunks

Iallgather(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recv_axis: int = 0) MPIRequest

Nonblocking Allgather.

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • recv_axis (int) – Concatenation axis: The axis along which sendbuf is packed and along which recvbuf puts together individual chunks

Iallgatherv(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recv_axis: int = 0)

Nonblocking v-call of Allgather: Each process may contribute a different amount of data.

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • recv_axis (int) – Concatenation axis: The axis along which sendbuf is packed and along which recvbuf puts together individual chunks

__alltoall_like(func: Callable, sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, send_axis: int, recv_axis: int, **kwargs)

Generic function for alltoall operations.

Parameters:
  • func (Callable) – Specific alltoall function

  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • send_axis (int) –

    Future split axis, along which data blocks will be created that will be send to individual ranks

    • if send_axis==recv_axis, an error will be thrown

    • if send_axis or recv_axis are None, an error will be thrown

  • recv_axis (int) – Prior split axis, along which blocks are received from the individual ranks

Alltoall(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, send_axis: int = 0, recv_axis: int = None)

All processes send data to all processes: The jth block sent from process i is received by process j and is placed in the ith block of recvbuf.

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • send_axis (int) –

    Future split axis, along which data blocks will be created that will be send to individual ranks

    • if send_axis==recv_axis, an error will be thrown

    • if send_axis or recv_axis are None, an error will be thrown

  • recv_axis (int) – Prior split axis, along which blocks are received from the individual ranks

Alltoallv(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, send_axis: int = 0, recv_axis: int = None)

v-call of Alltoall: All processes send different amount of data to, and receive different amount of data from, all processes

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • send_axis (int) –

    Future split axis, along which data blocks will be created that will be send to individual ranks

    • if send_axis==recv_axis, an error will be thrown

    • if send_axis or recv_axis are None, an error will be thrown

  • recv_axis (int) – Prior split axis, along which blocks are received from the individual ranks

Ialltoall(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, send_axis: int = 0, recv_axis: int = None) MPIRequest

Nonblocking Alltoall

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • send_axis (int) –

    Future split axis, along which data blocks will be created that will be send to individual ranks

    • if send_axis==recv_axis, an error will be thrown

    • if send_axis or recv_axis are None, an error will be thrown

  • recv_axis (int) – Prior split axis, along which blocks are received from the individual ranks

Ialltoallv(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, send_axis: int = 0, recv_axis: int = None) MPIRequest

Nonblocking v-call of Alltoall: All processes send different amount of data to, and receive different amount of data from, all processes

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • send_axis (int) –

    Future split axis, along which data blocks will be created that will be send to individual ranks

    • if send_axis==recv_axis, an error will be thrown

    • if send_axis or recv_axis are None, an error will be thrown

  • recv_axis (int) – Prior split axis, along which blocks are received from the individual ranks

__gather_like(func: Callable, sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, send_axis: int, recv_axis: int, send_factor: int = 1, recv_factor: int = 1, **kwargs)

Generic function for gather operations.

Parameters:
  • func (Callable) – Type of MPI Scatter/Gather function

  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • send_axis (int) – The axis along which sendbuf is packed

  • recv_axis (int) – The axis along which recvbuf is packed

  • send_factor (int) – Number of elements to be scattered (vor non-v-calls)

  • recv_factor (int) – Number of elements to be gathered (vor non-v-calls)

Gather(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, root: int = 0, axis: int = 0, recv_axis: int = None)

Gathers together values from a group of processes

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • root (int) – Rank of receiving process

  • axis (int) – The axis along which sendbuf is packed

  • recv_axis (int) – The axis along which recvbuf is packed

Gatherv(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, root: int = 0, axis: int = 0, recv_axis: int = None)

v-call for Gather: All processes send different amount of data

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • root (int) – Rank of receiving process

  • axis (int) – The axis along which sendbuf is packed

  • recv_axis (int) – The axis along which recvbuf is packed

Igather(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, root: int = 0, axis: int = 0, recv_axis: int = None) MPIRequest

Non-blocking Gather

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • root (int) – Rank of receiving process

  • axis (int) – The axis along which sendbuf is packed

  • recv_axis (int) – The axis along which recvbuf is packed

Igatherv(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, root: int = 0, axis: int = 0, recv_axis: int = None) MPIRequest

Non-blocking v-call for Gather: All processes send different amount of data

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • root (int) – Rank of receiving process

  • axis (int) – The axis along which sendbuf is packed

  • recv_axis (int) – The axis along which recvbuf is packed

__scatter_like(func: Callable, sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, send_axis: int, recv_axis: int, send_factor: int = 1, recv_factor: int = 1, **kwargs)

Generic function for scatter operations.

Parameters:
  • func (Callable) – Type of MPI Scatter/Gather function

  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • send_axis (int) – The axis along which sendbuf is packed

  • recv_axis (int) – The axis along which recvbuf is packed

  • send_factor (int) – Number of elements to be scattered (vor non-v-calls)

  • recv_factor (int) – Number of elements to be gathered (vor non-v-calls)

Iscatter(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, root: int = 0, axis: int = 0, recv_axis: int = None) MPIRequest

Non-blocking Scatter

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • root (int) – Rank of sending process

  • axis (int) – The axis along which sendbuf is packed

  • recv_axis (int) – The axis along which recvbuf is packed

Iscatterv(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, root: int = 0, axis: int = 0, recv_axis: int = None) MPIRequest

Non-blocking v-call for Scatter: Sends different amounts of data to different processes

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • root (int) – Rank of sending process

  • axis (int) – The axis along which sendbuf is packed

  • recv_axis (int) – The axis along which recvbuf is packed

Scatter(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, root: int = 0, axis: int = 0, recv_axis: int = None)

Sends data parts from one process to all other processes in a communicator

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • root (int) – Rank of sending process

  • axis (int) – The axis along which sendbuf is packed

  • recv_axis (int) – The axis along which recvbuf is packed

Scatterv(sendbuf: heat.core.dndarray.DNDarray | torch.Tensor | Any, recvbuf: int, root: int = 0, axis: int = 0, recv_axis: int = None)

v-call for Scatter: Sends different amounts of data to different processes

Parameters:
  • sendbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address of the send message

  • recvbuf (Union[DNDarray, torch.Tensor, Any]) – Buffer address where to store the result

  • root (int) – Rank of sending process

  • axis (int) – The axis along which sendbuf is packed

  • recv_axis (int) – The axis along which recvbuf is packed

__getattr__(name: str)

Default pass-through for the communicator methods.

Parameters:

name (str) – The name of the method to be called.

get_comm() Communication

Retrieves the currently globally set default communication.

sanitize_comm(comm: Communication | None) Communication

Sanitizes a device or device identifier, i.e. checks whether it is already an instance of heat.core.devices.Device or a string with known device identifier and maps it to a proper Device.

Parameters:

comm (Communication) – The comm to be sanitized

Raises:

TypeError – If the given communication is not the proper type

use_comm(comm: Communication = None)

Sets the globally used default communicator.

Parameters:

comm (Communication or None) – The communication to be set