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 communicationw_size (int, optional) – The MPI world size, defaults to
self.size
. Intended for creating chunk maps without communicationsparse (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.
- 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
- 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
- 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
- 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
- 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
- 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
- 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”
- 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
- __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 whichrecvbuf
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 whichrecvbuf
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 whichrecvbuf
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 whichrecvbuf
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 whichrecvbuf
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 thrownif
send_axis
orrecv_axis
areNone
, 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 thrownif
send_axis
orrecv_axis
areNone
, 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 thrownif
send_axis
orrecv_axis
areNone
, 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 thrownif
send_axis
orrecv_axis
areNone
, 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 thrownif
send_axis
orrecv_axis
areNone
, 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 packedrecv_axis (int) – The axis along which
recvbuf
is packedsend_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 packedrecv_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 packedrecv_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 packedrecv_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 packedrecv_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 packedrecv_axis (int) – The axis along which
recvbuf
is packedsend_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 packedrecv_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 packedrecv_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 packedrecv_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 packedrecv_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 properDevice
.- 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