def send(self, tensors, send_options=SendOptions()):
"""Send a tensor to a destination gpu in the group.
Args:
tensors (List): the tensor to send.
send_options: send options.
Returns:
None
"""
def p2p_fn(tensor, comm, stream, peer):
comm.send(
nccl_util.get_tensor_ptr(tensor),
send_options.n_elements
if send_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, send_options.dst_rank, send_options.dst_gpu_index
)
def send(self, tensors, send_options=SendOptions()):
"""Send a tensor to a destination rank in the group.
Args:
tensors (List): the tensor to send.
send_options: send options.
Returns:
None
"""
def p2p_fn(tensor, context, peer):
pygloo.send(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, send_options.dst_rank)
