o "j@s`ddlZddlmZddlmZddlmZmZGdddZ ej ddZ dd Z d d Z dS) N) framework)_get_global_group_warn_cur_rank_not_in_groupc@seZdZdZdddZdS)P2POpa A class that makes point-to-point operations for "batch_isend_irecv". This class creates the type of P2P operation, communication buffer, peer rank, Group. Instances of this class will be passed to ``paddle.distributed.batch_isend_irecv`` for point-to-point communication. Args: op (callable): A function to send data to or receive data from a peer process. The type of ``op`` is either ``paddle.distributed.isend`` or ``paddle.distributed.irecv``. tensor (Tensor): Tensor to send or receive. peer (int): The destination or source rank. group (Group, optional): The group instance return by new_group or None for global default group. Default: None. Examples: .. code-block:: python >>> # doctest: +REQUIRES(env: DISTRIBUTED) >>> import paddle >>> import paddle.distributed as dist >>> dist.init_parallel_env() >>> rank = dist.get_rank() >>> world_size = dist.get_world_size() >>> send_t = paddle.arange(2) + rank >>> # paddle.tensor([0, 1]) # Rank-0 >>> # paddle.tensor([1, 2]) # Rank-1 >>> recv_t = paddle.empty(shape=[2], dtype=send_t.dtype) >>> send_op = dist.P2POp(dist.isend, send_t, (rank + 1) % world_size) >>> recv_op = dist.P2POp(dist.irecv, recv_t, (rank - 1 + world_size) % world_size) NcCsH|tjtjfvr td||_||_||_|durt|_dS||_dS)NztInvalid ``op`` function. Expected ``op`` to be of type ``paddle.distributed.isend`` or ``paddle.distributed.irecv``.) distZisendZirecv RuntimeErroroptensorpeerrgroup)selfrr r r r s/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/paddle/distributed/communication/batch_isend_irecv.py__init__@szP2POp.__init__N)__name__ __module__ __qualname____doc__rr r r rrs&rccsT|dkr tjjzdVW|dkrtjjdSdS|dkr)tjjww)NZNCCL)rcoreZProcessGroupNCCLZ group_startZ group_endbackendr r r_with_batch_p2p_guardNs  rcsRt|trtdd|Dstd|djjtfdd|Ds'tddS)zu Helper to check that the ``p2p_op_list`` is a list of P2POp instances and all ops use the same backend. css|]}t|tVqdSr) isinstancer.0p2p_opr r r ^s  z%_check_p2p_op_list..z[Invalid ``p2p_op_list``. Each op is expected to to be of type ``paddle.distributed.P2POp``.rc3s|] }|jjkVqdSr)r rrrr rrgsz(All groups need to use the same backend.N)rlistallrr r) p2p_op_listr rr_check_p2p_op_listYs r!c Cst||dj}t|rdStrX|durtn|}|j}g}t|)|D]}|j}|j }|j }|j}||||} | durE| | q(Wd|S1sQwY|St d)a Send or Receive a batch of tensors asynchronously and return a list of requests. Process each of the point-to-point operations in ``p2p_op_list`` and return the corresponding tasks. NCCL are currently supported. Args: p2p_op_list (List[P2POp]): A list of point-to-point operations(type of each operator is ``paddle.distributed.P2POp``). The order of the isend/irecv in the list matters and it needs to match with corresponding isend/irecv on the remote end. Returns: A list of distributed tasks returned by calling the corresponding op in the op_list. Warning: This API only supports the dygraph mode. Examples: .. code-block:: python >>> # doctest: +REQUIRES(env: DISTRIBUTED) >>> import paddle >>> import paddle.distributed as dist >>> dist.init_parallel_env() >>> rank = dist.get_rank() >>> world_size = dist.get_world_size() >>> send_t = paddle.arange(2) + rank >>> # paddle.tensor([0, 1]) # Rank-0 >>> # paddle.tensor([1, 2]) # Rank-1 >>> recv_t = paddle.empty(shape=[2], dtype=send_t.dtype) >>> send_op = dist.P2POp(dist.isend, send_t, (rank + 1) % world_size) >>> recv_op = dist.P2POp(dist.irecv, recv_t, (rank - 1 + world_size) % world_size) >>> tasks = dist.batch_isend_irecv([send_op, recv_op]) >>> for task in tasks: ... task.wait() >>> print(recv_t) >>> # paddle.tensor([1, 2]) # Rank-0 >>> # paddle.tensor([0, 1]) # Rank-1 rNz*Don't support static graph mode currently.) r!r rrZin_dynamic_moderrrrr r appendr) r r rtasksrrr r Z comm_grouptaskr r rbatch_isend_irecvks02       r%) contextlibZpaddle.distributed distributedrZpaddlerZ&paddle.distributed.communication.grouprrrcontextmanagerrr!r%r r r rs  5