def recv(self, tensors, recv_options=RecvOptions()):
"""Receive a tensor from a source gpu in the group.
Args:
tensors (List): the received tensor.
recv_options: Receive options.
Returns:
None
"""
def p2p_fn(tensor, comm, stream, peer):
comm.recv(
nccl_util.get_tensor_ptr(tensor),
recv_options.n_elements
if recv_options.n_elements > 0
else nccl_util.get_tensor_n_elements(tensor),
nccl_util.get_nccl_tensor_dtype(tensor),
peer,
stream.ptr,
)
self._point2point(
tensors, p2p_fn, recv_options.src_rank, recv_options.src_gpu_index
)
def recv(self, tensors, recv_options=RecvOptions()):
"""Receive a tensor from a source rank in the group.
Args:
tensors (List): the received tensor.
recv_options: Receive options.
Returns:
None
"""
def p2p_fn(tensor, context, peer):
pygloo.recv(context, gloo_util.get_tensor_ptr(tensor),
gloo_util.get_tensor_n_elements(tensor),
gloo_util.get_gloo_tensor_dtype(tensor), peer)
self._point2point(tensors, p2p_fn, recv_options.src_rank)
