Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2025 Huawei Technologies Co., Ltd.
       3              :  * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
       4              :  * CANN Open Software License Agreement Version 2.0 (the "License").
       5              :  * Please refer to the License for details. You may not use this file except in compliance with the License.
       6              :  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
       7              :  * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
       8              :  * See LICENSE in the root of the software repository for the full text of the License.
       9              :  */
      10              : 
      11              : #ifndef COLL_BATCH_SEND_RECV_GROUP_EXECUTOR_H
      12              : #define COLL_BATCH_SEND_RECV_GROUP_EXECUTOR_H
      13              : 
      14              : #include "coll_comm_executor.h"
      15              : #include "coll_batch_send_recv_executor.h"
      16              : #include <map>
      17              : 
      18              : namespace hccl {
      19              : class CollBatchSendRecvGroupExecutor : public CollBatchSendRecvExecutor {
      20              : public:
      21              :     CollBatchSendRecvGroupExecutor(const HcclDispatcher dispatcher, std::unique_ptr<TopoMatcher> &topoMatcher);
      22          128 :     ~CollBatchSendRecvGroupExecutor() override = default;
      23              :     HcclResult Orchestrate(OpParam& param, AlgResourceResponse& algResource) override;
      24              : protected:
      25              : 
      26              :     /* *************** 算法编排 *************** */
      27              :     u64 CalcSendLoopMaxCount(const u32 unitSize) const;
      28              :     u64 CalcRecvLoopMaxCount(const u32 unitSize) const;
      29              :     HcclResult CalcSendSlices();
      30              :     HcclResult CalcRecvSlices();
      31              :     HcclResult OrganizeSendItemByStream();
      32              :     HcclResult OrganizeRecvItemByStream();
      33              :     HcclResult CalcCommInfo(std::vector<LevelNSubCommTransport>& opTransport) override;
      34              :     struct SendRecvSlice {
      35              :         u8* addr;
      36              :         u64 size;
      37              :         u32 remoteRank;
      38              :         bool isRdma; // true=RDMA任务(用CCLOut拆半，无ping-pong)；false=SDMA任务(用CCLIn，ping-pong)
      39           37 :         SendRecvSlice(u8* addr, u64 size, u32 remoteRank, bool isRdma = false)
      40           37 :             : addr(addr), size(size), remoteRank(remoteRank), isRdma(isRdma) {}
      41              :     };
      42              : 
      43              :     // SDMA slice按remoteRank % streamNum分发到各SDMA从流(仅SDMA，ping-pong)。
      44              :     std::vector<std::deque<SendRecvSlice>> sendDataSlicesBySendStream_;
      45              :     std::vector<std::deque<SendRecvSlice>> recvDataSlicesByRecvStream_;
      46              : 
      47              :     // RDMA与SDMA分离：SDMA占sendStreamNum_+recvStreamNum_条从流(CCLIn ping-pong)；
      48              :     // RDMA单独占2条从流(1 send + 1 recv)，对端顺序按alltoallv_direct_fullmesh对称场景规则
      49              :     // (send前向递增、recv后向递减，均跳过本pod；对端无任务则跳过)。
      50              :     // RDMA使用CCLOut拆半：A(offset 0)=send scratch(单slot), B(offset rdmaDataBlockSize_)=recv scratch(单slot)。
      51              :     // RDMA不做ping-pong，逐slice处理。
      52              :     static constexpr u32 RDMA_CCLOUT_HALF_NUM = 2; // CCLOut拆成两半：A=send, B=recv
      53              :     static constexpr u32 RDMA_STREAM_NUM = 2;      // RDMA专用从流数：1 send + 1 recv
      54              : 
      55              :     // RDMA slice(已按对称场景规则排序)，分别由rdmaSendStreamIdx_/rdmaRecvStreamIdx_对应的从流处理。
      56              :     std::deque<SendRecvSlice> rdmaSendSlices_;
      57              :     std::deque<SendRecvSlice> rdmaRecvSlices_;
      58              : 
      59              : private:
      60              :     HcclResult RunLoop(OpParam& param);
      61              :     HcclResult RunTasks(OpParam& param);
      62              :     HcclResult ProcessPreloadedSendSlice(u32 streamIdx, u32& pendingSendCount, u32& nonEmptySendStream);
      63              :     HcclResult ProcessNewRankSendSlice(u32 streamIdx, u32& pendingSendCount, u32& nonEmptySendStream);
      64              :     HcclResult ProcessRecvSlice(u32 streamIdx, u32& nonEmptyRecvStream);
      65              :     // RDMA专用从流(单send/单recv)上处理一个slice，对端顺序由rdmaSendSlices_/rdmaRecvSlices_保证。
      66              :     HcclResult ProcessRdmaSendSlice();
      67              :     HcclResult ProcessRdmaRecvSlice();
      68              :     HcclResult CalcPodRange();
      69              :     bool IsRemoteRankRdma(u32 remoteRank) const;
      70              :     HcclResult SetNormalModeIfDeviceDirect();
      71              :     HcclResult CalcStreamNum(u32& streamNum) override;
      72              :     HcclResult CalcPingPongHalfSize();
      73              : 
      74              :     HcclResult MainPostSubWait(Stream& mainStream);
      75              :     HcclResult MainWaitSubPost(Stream& mainStream);
      76              :     // 统计各从流(SDMA send/recv + RDMA send/recv)是否有任务并记录到成员变量，
      77              :     // 同时返回SDMA非空流数供RunTasks循环使用。循环中不再更新。
      78              :     HcclResult CalcStreamTaskStatus(u32& nonEmptySendStream, u32& nonEmptyRecvStream);
      79              : 
      80              :     // 对称场景规则：send前向递增/recv后向递减遍历，跳过本pod。返回当前候选rank并推进游标。
      81              :     u32 GetNextDstRank(u32& curDstRank);
      82              :     u32 GetPreSrcRank(u32& curSrcRank);
      83              :     // 按对称场景规则对RDMA slice排序：isSend=true用前向(GetNextDstRank)，false用后向(GetPreSrcRank)。
      84              :     // 遍历所有跨pod对端，仅输出存在任务的对端的slice(起点初始化与每次更新均跳过无任务对端)。
      85              :     void OrderRdmaSlices(bool isSend, const std::map<u32, std::deque<SendRecvSlice>>& byRank,
      86              :                          std::deque<SendRecvSlice>& out);
      87              : 
      88              :     // RDMA专用从流在slaveStreams中的索引。
      89            8 :     u32 RdmaSendStreamIdx() const { return sendStreamNum_ + recvStreamNum_; }
      90            4 :     u32 RdmaRecvStreamIdx() const { return sendStreamNum_ + recvStreamNum_ + 1; }
      91              : 
      92              : private:
      93              : 
      94              :     std::vector<std::deque<HcclSendRecvItem*>> sendQueueBySendstream_;
      95              :     std::vector<std::deque<HcclSendRecvItem*>> recvQueueByRecvstream_;
      96              :     u32 sendStreamNum_ = 0;
      97              :     u32 recvStreamNum_ = 0;
      98              :     u64 bufferSliceSize_ = 0;
      99              :     u64 rdmaDataBlockSize_ = 0; // RDMA单流slot大小 = CCLOut/2(A=send半区, B=recv半区)
     100              :     u32 podStartRank_ = 0; // pod(rank)范围，用于判定isRdma：[podStartRank_, podEndRank_]内为SDMA，跨pod为RDMA
     101              :     u32 podEndRank_ = 0;
     102              :     u32 devNumInlocalPod_ = 0; // 本pod内rank数，用于对称场景规则起点的计算
     103              :     // Ping-pong state
     104              :     std::vector<u32> sendCurPhase_;         // which half has loaded data, ready to Record+Send (0=A, 1=B)
     105              :     std::vector<u64> sendLoadedSize_;       // size of data loaded in current sendCurPhase_ half (0 = nothing)
     106              :     std::vector<u32> sendLoadedRemoteRank_; // remote rank for the data loaded in current sendCurPhase_
     107              :     std::vector<u32> recvCurPhase_;         // which half to read from for recv (0=A, 1=B)
     108              :     std::vector<u32> recvCurRemoteRank_;    // current remote rank being received on this stream
     109              : 
     110              :     // 各从流是否有任务(在RunTasks起始时记录，循环中不更新)，用于头尾同步只唤醒有任务的从流。
     111              :     std::vector<bool> sendStreamHasTask_;   // send从流(共sendStreamNum_条)是否有任务
     112              :     std::vector<bool> recvStreamHasTask_;   // recv从流(共recvStreamNum_条)是否有任务
     113              :     bool rdmaSendHasTask_ = false;          // RDMA专用send从流是否有任务
     114              :     bool rdmaRecvHasTask_ = false;          // RDMA专用recv从流是否有任务
     115              : };
     116              : } // namespace hccl
     117              : 
     118              : #endif