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 : #include "allltoall_pipeline_mesh_pairwise_ping_pong.h"
12 : #include <numeric>
13 : #include "alg_template_register.h"
14 :
15 : namespace hccl {
16 :
17 : // 需要将 ccl 切成两份，ping-pong 时也根据收发次数取模2决定使用 ping mem 还是 pong mem
18 : static const u32 PING_PONG_CONST_NUM = 2;
19 : static const u32 INTRA_STREAM_INFO_SENDLEN_INDEX = 0; // intraStreamInfo 中 sendLen 的下标
20 : static const u32 INTRA_STREAM_INFO_RECVLEN_INDEX = 1; // intraStreamInfo 中 recvLen 的下标
21 : static const u32 INTRA_STREAM_INFO_RECV_LOCAL_OFFSET_INDEX = 2; // intraStreamInfo 中 recvRemoteOffset 的下标
22 :
23 0 : AlltoallPipelineMeshPairwisePingPong::AlltoallPipelineMeshPairwisePingPong(const HcclDispatcher dispatcher)
24 0 : : AlltoallPipelineBase(dispatcher)
25 0 : {}
26 :
27 0 : AlltoallPipelineMeshPairwisePingPong::~AlltoallPipelineMeshPairwisePingPong() {}
28 :
29 0 : u32 AlltoallPipelineMeshPairwisePingPong::CalcInterNumSteps() { return interRankSize_ - 1; }
30 :
31 : // 适配新CollExecutor接口
32 0 : HcclResult AlltoallPipelineMeshPairwisePingPong::Prepare(
33 : u32 userRank, A2aPipelineMemory A2aPipelineMemory, const SubCommInfo& level0CommInfo,
34 : const SubCommInfo& level1CommInfo, Stream& mainStream, std::vector<Stream>& subStream,
35 : std::vector<std::shared_ptr<LocalNotify>>& notifyMain, std::vector<std::shared_ptr<LocalNotify>>& notifySub,
36 : std::vector<SendRecvInfo>& allMeshAggregationSendRecvInfo, HcclWorkflowMode workMode)
37 : {
38 0 : AlltoallPipelineBase::Prepare(
39 : userRank, A2aPipelineMemory, level0CommInfo, level1CommInfo, mainStream, subStream, notifyMain, notifySub,
40 : allMeshAggregationSendRecvInfo, workMode);
41 0 : if (workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
42 0 : pingPongMemSize_ = (cclIn_.size() / PING_PONG_CONST_NUM);
43 : } else {
44 0 : pingPongMemSize_ = (scratchMem_.size() / PING_PONG_CONST_NUM);
45 : }
46 0 : intraDataBlockSize_ = (pingPongMemSize_ / intraRankSize_);
47 0 : if (intraDataBlockSize_ > HCCL_MIN_SLICE_ALIGN_910B) {
48 0 : intraDataBlockSize_ = (intraDataBlockSize_ / HCCL_MIN_SLICE_ALIGN_910B) * HCCL_MIN_SLICE_ALIGN_910B;
49 : }
50 0 : memStatusInMesh_ = std::vector<bool>(intraRankSize_, false);
51 0 : CHK_RET(DeviceMemMapping());
52 0 : return HCCL_SUCCESS;
53 : }
54 :
55 0 : HcclResult AlltoallPipelineMeshPairwisePingPong::DeviceMemMapping()
56 : {
57 0 : if (workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
58 0 : interTransportSend_ = cclIn_;
59 0 : interTransportRecv_ = cclOut_;
60 0 : intraTransportSend_ = cclOut_;
61 0 : interSendPing_ = interTransportSend_.range(0, pingPongMemSize_);
62 0 : interSendPong_ = interTransportSend_.range(pingPongMemSize_, pingPongMemSize_);
63 0 : interRecvPing_ = interTransportRecv_.range(0, pingPongMemSize_);
64 0 : interRecvPong_ = interTransportRecv_.range(pingPongMemSize_, pingPongMemSize_);
65 : } else {
66 0 : interRecvPing_ = scratchMem_.range(0, pingPongMemSize_);
67 0 : interRecvPong_ = scratchMem_.range(pingPongMemSize_, pingPongMemSize_);
68 : }
69 0 : intraSendPing_ = interRecvPing_;
70 0 : intraSendPong_ = interRecvPong_;
71 :
72 0 : for (u32 intraRank = 0; intraRank < intraRankSize_; intraRank++) {
73 0 : if (intraRank == intraRankId_) {
74 0 : continue;
75 : }
76 0 : LINK& intraNeighboorTransport = intraLinks_[intraRank];
77 0 : void* remDMAMemPtr = nullptr;
78 0 : CHK_RET(intraNeighboorTransport->GetRemoteMem(UserMemType::INPUT_MEM, &remDMAMemPtr));
79 : DeviceMem remoteIntraSend = DeviceMem::create(
80 : static_cast<u8*>(remDMAMemPtr),
81 0 : workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE ? cclIn_.size() : scratchMem_.size());
82 0 : DeviceMem remoteIntraSendPing = remoteIntraSend.range(0, pingPongMemSize_);
83 0 : DeviceMem remoteIntraSendPong = remoteIntraSend.range(pingPongMemSize_, pingPongMemSize_);
84 0 : intraNeighBoorMemory_[intraRank] = {remoteIntraSendPing, remoteIntraSendPong};
85 0 : }
86 0 : return HCCL_SUCCESS;
87 0 : }
88 :
89 : // 将需要发送给其他 mesh 的数据准备好，并计算好 TxMemoryInfo
90 0 : HcclResult AlltoallPipelineMeshPairwisePingPong::PrepareInterSendData(u32 mainStep, u32 subStep)
91 : {
92 0 : nextInterSendData_.clear();
93 0 : u32 interSendRankStart = ((interRankId_ + 1 + mainStep) % interRankSize_) * intraRankSize_;
94 0 : DeviceMem interSendMem = (interSendUsePingMem_ ? interSendPing_ : interSendPong_);
95 0 : HCCL_DEBUG(
96 : "[AlltoallPipelineMeshPairwisePingPong][PrepareInterSendData] userRank %u, interRank %u, "
97 : "intraRank %u in main step %llu, sub step %llu send to remote %s",
98 : userRank_, interRankId_, intraRankId_, mainStep, subStep,
99 : sendToInterDstMemPing_ ? "interRecvPingMem" : "interRecvPongMem");
100 0 : u64 preStepMaxSend = intraDataBlockSize_ * subStep;
101 0 : for (u32 i = 0; i < intraRankSize_; i++) {
102 0 : u32 dataIndex = i + interSendRankStart;
103 0 : u64 totalSendLen = localSendRecvInfo_.sendLength[dataIndex];
104 0 : u64 sendLen = std::min(intraDataBlockSize_, std::max(totalSendLen, preStepMaxSend) - preStepMaxSend);
105 0 : if (sendLen == 0) {
106 0 : continue;
107 : }
108 0 : HCCL_DEBUG(
109 : "[AlltoallPipelineMeshPairwisePingPong][PrepareInterSendData] userRank %u, interRank %u, "
110 : "intraRank %u data index %llu move from userInput offset %llu length %llu to %s, total size %llu"
111 : "send to remote %s",
112 : userRank_, interRankId_, intraRankId_, dataIndex, localSendRecvInfo_.sendOffset[dataIndex] + preStepMaxSend,
113 : sendLen, interSendUsePingMem_ ? "localInterSendPingMem" : "localInterSendPongMem", totalSendLen,
114 : sendToInterDstMemPing_ ? "interRecvPingMem" : "interRecvPongMem");
115 0 : DeviceMem src = inputMem_.range(localSendRecvInfo_.sendOffset[dataIndex] + preStepMaxSend, sendLen);
116 0 : DeviceMem dst = interSendMem.range(i * intraDataBlockSize_, sendLen);
117 : // 单算子模式需要搬到 CCL，图模式省去这一步
118 0 : if (workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
119 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, mainStream_));
120 : }
121 0 : nextInterSendData_.emplace_back(TxMemoryInfo{
122 0 : UserMemType::OUTPUT_MEM, (sendToInterDstMemPing_ ? 0 : pingPongMemSize_) + i * intraDataBlockSize_,
123 0 : workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE ? dst.ptr() : src.ptr(), sendLen});
124 0 : }
125 0 : return HCCL_SUCCESS;
126 0 : }
127 :
128 : // 将需要发送给其他 mesh 的数据准备好，并准备好 TxMemoryInfo
129 0 : HcclResult AlltoallPipelineMeshPairwisePingPong::PrepareInterRecvData(u32 mainStep, u32 subStep)
130 : {
131 0 : nextInterRecvData_.clear();
132 0 : if (workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
133 : // 单算子模式本次要接收的数据都放在 CCL，直接整块接收
134 0 : nextInterRecvData_.emplace_back(RxMemoryInfo{
135 0 : UserMemType::INPUT_MEM, recvFromInterSrcMemPing_ ? 0u : pingPongMemSize_,
136 0 : (interRecvUsePingMem_ ? interRecvPing_ : interRecvPong_).ptr(), pingPongMemSize_});
137 : } else {
138 : // 图模式需要计算数据放在对端 userInput 的位置
139 0 : u32 recvFromRank = (userRank_ + groupRankSize_ - (mainStep + 1) * intraRankSize_) % groupRankSize_;
140 0 : const std::vector<u64>& remoteSendLength = (*allMeshAggregationSendRecvInfo_)[recvFromRank].sendLength;
141 0 : const std::vector<u64>& remoteSendOffset = (*allMeshAggregationSendRecvInfo_)[recvFromRank].sendOffset;
142 0 : u64 dataStartOffset = subStep * intraDataBlockSize_;
143 0 : for (u32 i = 0; i < intraRankSize_; i++) {
144 0 : u64 totalRecvDataLen = remoteSendLength[meshRankStart_ + i];
145 0 : u64 recvLen = std::min(std::max(totalRecvDataLen, dataStartOffset) - dataStartOffset, intraDataBlockSize_);
146 0 : if (recvLen == 0) {
147 0 : continue;
148 : }
149 0 : u64 recvRemoteOffset = remoteSendOffset[meshRankStart_ + i] + dataStartOffset;
150 0 : nextInterRecvData_.emplace_back(RxMemoryInfo{
151 : UserMemType::INPUT_MEM, recvRemoteOffset,
152 0 : (interRecvUsePingMem_ ? interRecvPing_ : interRecvPong_).range(i * intraDataBlockSize_, recvLen).ptr(),
153 : recvLen});
154 0 : HCCL_DEBUG(
155 : "[AlltoallPipelineMeshPairwisePingPong][PrepareInterRecvData] userRank %u, interRank %u, "
156 : "intraRank %u recv from remote userInput offset %llu length %llu to %s offset %llu",
157 : userRank_, interRankId_, intraRankId_, recvRemoteOffset, recvLen,
158 : interRecvUsePingMem_ ? "localInterRecvPingMem" : "localInterRecvPongMem", i * intraDataBlockSize_);
159 : }
160 : }
161 0 : return HCCL_SUCCESS;
162 : }
163 :
164 : // 准备下一次mesh间需要收发的数据，单算子模式需要从 userInput 搬到 CCLBuffer，图模式则仅需要准备好 TxMemoryInfo
165 0 : HcclResult AlltoallPipelineMeshPairwisePingPong::PrepareInterData(u32 mainStep, u32 subStep)
166 : {
167 0 : CHK_RET(PrepareInterSendData(mainStep, subStep));
168 0 : CHK_RET(PrepareInterRecvData(mainStep, subStep));
169 0 : return HCCL_SUCCESS;
170 : }
171 :
172 : // 将原先在userInput，且需要发到本mesh内其他卡的数据搬到CCL
173 0 : HcclResult AlltoallPipelineMeshPairwisePingPong::PrepareIntraData(u32 subStep)
174 : {
175 0 : u64 dataStartOffset = subStep * intraDataBlockSize_;
176 0 : for (u32 i = 0; i < intraRankSize_; i++) {
177 0 : u32 dataIndex = i + meshRankStart_;
178 0 : u64 totalSendDataLen = localSendRecvInfo_.sendLength[dataIndex];
179 0 : u64 sendLen = std::min(std::max(totalSendDataLen, dataStartOffset) - dataStartOffset, intraDataBlockSize_);
180 0 : if (i == intraRankId_ || sendLen == 0) {
181 0 : continue;
182 : }
183 0 : DeviceMem src = inputMem_.range(localSendRecvInfo_.sendOffset[dataIndex] + dataStartOffset, sendLen);
184 : DeviceMem dst
185 0 : = (intraSendUsePingMem_ ? intraSendPing_ : intraSendPong_).range(i * intraDataBlockSize_, sendLen);
186 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, mainStream_));
187 0 : HCCL_DEBUG(
188 : "[AlltoallPipelineMeshPairwisePingPong][PrepareIntraData] userRank %u, interRank %u, intraRank %u"
189 : "data index %u move from userInput offset %llu length %llu to %s, total size %llu ",
190 : userRank_, interRankId_, intraRankId_, dataIndex,
191 : localSendRecvInfo_.sendOffset[dataIndex] + dataStartOffset, sendLen,
192 : intraSendUsePingMem_ ? "IntraPingMem" : "IntraPongMem", totalSendDataLen);
193 0 : }
194 0 : return HCCL_SUCCESS;
195 : }
196 :
197 : // 计算mesh内其他卡此时RDMA接收到的数据是在 cclIn 还是 cclOut
198 0 : void AlltoallPipelineMeshPairwisePingPong::UpdateRemoteMemStatusIntra(u32 step)
199 : {
200 0 : for (u32 intraRank = 0; intraRank < intraRankSize_; intraRank++) {
201 0 : if (intraRank == intraRankId_)
202 0 : continue;
203 0 : u32 intraRankHaveRecv = 0;
204 0 : for (u32 i = 1; i <= step; i++) {
205 : const std::vector<u64>& intraRankRecvFrom
206 0 : = (*allMeshAggregationSendRecvInfo_)
207 0 : [(meshRankStart_ + groupRankSize_ + intraRank - i * intraRankSize_) % groupRankSize_]
208 0 : .sendLength;
209 0 : u64 maxRecvLen = std::accumulate(
210 0 : intraRankRecvFrom.begin() + meshRankStart_, intraRankRecvFrom.begin() + meshRankStart_ + intraRankSize_,
211 0 : 0ULL, [](u64 a, u64 b) {
212 0 : return a > b ? a : b;
213 : });
214 0 : intraRankHaveRecv += ((maxRecvLen + intraDataBlockSize_ - 1) / intraDataBlockSize_);
215 : }
216 0 : memStatusInMesh_[intraRank] = ((intraRankHaveRecv % PING_PONG_CONST_NUM) == 0);
217 : }
218 0 : }
219 :
220 : // 计算本卡接收数据的那张卡和本卡将要发数据的那张卡在这个大步骤中的
221 : // 第一个小步骤从哪块ccl收以及发到哪块ccl, 每次切换
222 0 : void AlltoallPipelineMeshPairwisePingPong::UpdateRemoteMemStatusInter(u32 step)
223 : {
224 0 : u32 recvGlobalRank = (userRank_ + groupRankSize_ - (step + 1) * intraRankSize_) % groupRankSize_;
225 0 : u32 recvInterRank = ((interRankId_ + interRankSize_ - (step + 1)) % interRankSize_);
226 0 : u32 numRecvRankHaveSend = 0;
227 0 : u32 numSendRankHaveRecv = 0;
228 0 : const std::vector<u64>& recvRankSendLen = (*allMeshAggregationSendRecvInfo_)[recvGlobalRank].sendLength;
229 0 : for (u32 i = 1; i <= step; i++) {
230 0 : u32 firstBlockIndex = (((recvInterRank + i) % interRankSize_) * intraRankSize_);
231 0 : u64 maxSendLen = std::accumulate(
232 0 : recvRankSendLen.begin() + firstBlockIndex, recvRankSendLen.begin() + firstBlockIndex + intraRankSize_, 0ULL,
233 0 : [](u64 a, u64 b) {
234 0 : return a > b ? a : b;
235 : });
236 0 : numRecvRankHaveSend += ((maxSendLen + intraDataBlockSize_ - 1) / intraDataBlockSize_);
237 : const std::vector<u64>& sendRankRecvFrom
238 0 : = (*allMeshAggregationSendRecvInfo_)[(userRank_ + i * intraRankSize_) % groupRankSize_].sendLength;
239 0 : u64 maxRecvLen = std::accumulate(
240 0 : sendRankRecvFrom.begin() + meshRankStart_, sendRankRecvFrom.begin() + meshRankStart_ + intraRankSize_, 0ULL,
241 0 : [](u64 a, u64 b) {
242 0 : return a > b ? a : b;
243 : });
244 0 : numSendRankHaveRecv += ((maxRecvLen + intraDataBlockSize_ - 1ULL) / intraDataBlockSize_);
245 : }
246 : // 首次默认都从对端pingMem读，本卡接收数据来源的那张卡每发一次数据切换一次
247 0 : recvFromInterSrcMemPing_ = ((numRecvRankHaveSend % PING_PONG_CONST_NUM) == 0);
248 : // 首次默认发到对端pingMem，本卡发送数据目的地的那张卡每接收一次数据切换一次
249 0 : sendToInterDstMemPing_ = ((numSendRankHaveRecv % PING_PONG_CONST_NUM) == 0);
250 0 : }
251 :
252 : // 收集本次 SDMA 子步骤每条流需要收发的长度，偏移地址，内存状态信息避免重复计算影响性能
253 0 : void AlltoallPipelineMeshPairwisePingPong::UpdateIntraStreamInfo(u32 interRankDistance, u32 subStep)
254 : {
255 0 : intraStreamInfo_.clear();
256 0 : u32 firstDataBlockIndex = (meshRankStart_ + groupRankSize_ - interRankDistance * intraRankSize_) % groupRankSize_;
257 : const std::vector<u64>& sendInfo
258 0 : = (*allMeshAggregationSendRecvInfo_)[firstDataBlockIndex + intraRankId_].sendLength;
259 0 : u64 dataStartOffset = subStep * intraDataBlockSize_;
260 0 : HCCL_DEBUG(
261 : "[AlltoallPipelineMeshPairwisePingPong][UpdateSDMAStreamInfo] userRank %u, "
262 : "interRank %u, intraRank %u, interRankDistance %llu, sub step %llu",
263 : userRank_, interRankId_, intraRankId_, interRankDistance, subStep);
264 0 : for (u32 i = 0; i < intraRankSize_; i++) {
265 0 : u64 totalSendDataLen = sendInfo[meshRankStart_ + i];
266 0 : u64 totalRecvDataLen = localSendRecvInfo_.recvLength[i + firstDataBlockIndex];
267 0 : u64 sendLen = std::min(std::max(totalSendDataLen, dataStartOffset) - dataStartOffset, intraDataBlockSize_);
268 0 : u64 recvLen = std::min(std::max(totalRecvDataLen, dataStartOffset) - dataStartOffset, intraDataBlockSize_);
269 0 : u64 localOffset = localSendRecvInfo_.recvOffset[i + firstDataBlockIndex] + subStep * intraDataBlockSize_;
270 0 : if (i != intraRankId_) {
271 0 : intraStreamInfo_[i] = {sendLen, recvLen, localOffset};
272 0 : HCCL_DEBUG(
273 : "[AlltoallPipelineMeshPairwisePingPong][UpdateSDMAStreamInfo] userRank %u, interRank %u, "
274 : "intraRank %u, sdma stream %llu need send %llu and read length %llu to local offset %llu",
275 : userRank_, interRankId_, intraRankId_, i, sendLen, recvLen, localOffset);
276 : }
277 : }
278 0 : }
279 :
280 0 : HcclResult AlltoallPipelineMeshPairwisePingPong::SendRecvDataIntraMesh()
281 : {
282 0 : HCCL_DEBUG(
283 : "[AlltoallPipelineMeshPairwisePingPong][ReadDataInMesh] userRank %u, "
284 : "interRank %u, intraRank %u, sdma stream %s wait main stream",
285 : userRank_, interRankId_, intraRankId_, GetStreamIndexString().c_str());
286 0 : bool anySend = false;
287 0 : for (auto& sdmaInfo : intraStreamInfo_) {
288 0 : u32 streamIndex = sdmaInfo.first;
289 0 : u64 recvLen = sdmaInfo.second[INTRA_STREAM_INFO_RECVLEN_INDEX];
290 0 : u64 recvOffset = sdmaInfo.second[INTRA_STREAM_INFO_RECV_LOCAL_OFFSET_INDEX];
291 0 : Stream& currStream = subStream_[streamIndex];
292 0 : LINK& readTransport = intraLinks_[streamIndex];
293 0 : CHK_RET(readTransport->TxAck(currStream));
294 0 : CHK_RET(readTransport->RxAck(currStream));
295 0 : if (recvLen > 0) {
296 0 : DeviceMem src = intraNeighBoorMemory_[streamIndex][(memStatusInMesh_[streamIndex] ? 0 : 1)].range(
297 0 : intraRankId_ * intraDataBlockSize_, recvLen);
298 0 : DeviceMem dst = outputMem_.range(recvOffset, recvLen);
299 0 : CHK_RET(HcclD2DMemcpyAsync(
300 : dispatcher_, dst, src, currStream, readTransport->GetRemoteRank(), readTransport->GetLinkType()));
301 0 : }
302 0 : CHK_RET(readTransport->TxDataSignal(currStream));
303 0 : HCCL_DEBUG(
304 : "[AlltoallPipelineMeshPairwisePingPong][ReadDataInMesh] userRank %u, interRank %u, "
305 : "intraRank %u, sdma stream %llu read data from remote %s offset %llu len %llu to local %llu",
306 : userRank_, interRankId_, intraRankId_, streamIndex,
307 : memStatusInMesh_[streamIndex] ? "IntraSendPingMem" : "IntraSendPongMem", intraRankId_ * intraDataBlockSize_,
308 : recvLen, recvOffset);
309 0 : memStatusInMesh_[streamIndex] = (!memStatusInMesh_[streamIndex]);
310 0 : CHK_RET(readTransport->RxDataSignal(currStream));
311 0 : anySend = true;
312 : }
313 0 : HCCL_DEBUG(
314 : "[AlltoallPipelineMeshPairwisePingPong][ReadDataInMesh] userRank %u, "
315 : "interRank %u, intraRank %u, sdma stream %s notify main stream",
316 : userRank_, interRankId_, intraRankId_, GetStreamIndexString().c_str());
317 0 : intraSendUsePingMem_ ^= anySend;
318 0 : return HCCL_SUCCESS;
319 : }
320 :
321 0 : HcclResult AlltoallPipelineMeshPairwisePingPong::SendRecvDataInterMesh(u32 step, bool doSend, bool doRecv)
322 : {
323 0 : Stream& interStream = subStream_[intraRankId_];
324 0 : LINK& interRecvTransport = interLinks_[(interRankId_ + interRankSize_ - 1 - step) % interRankSize_];
325 0 : LINK& interSendTransport = interLinks_[(interRankId_ + 1 + step) % interRankSize_];
326 0 : if (doRecv) {
327 0 : CHK_RET(interRecvTransport->TxAck(interStream));
328 : }
329 0 : if (doSend) {
330 0 : CHK_RET(interSendTransport->RxAck(interStream));
331 0 : CHK_RET(interSendTransport->TxAsync(
332 : UserMemType::OUTPUT_MEM, (sendToInterDstMemPing_ ? 0u : pingPongMemSize_),
333 : (interSendUsePingMem_ ? interSendPing_ : interSendPong_).ptr(), pingPongMemSize_, interStream));
334 0 : interSendUsePingMem_ ^= true;
335 0 : sendToInterDstMemPing_ ^= true;
336 : }
337 0 : if (doRecv) {
338 0 : CHK_RET(interRecvTransport->RxAsync(
339 : UserMemType::INPUT_MEM, (recvFromInterSrcMemPing_ ? 0u : pingPongMemSize_),
340 : (interRecvUsePingMem_ ? interRecvPing_ : interRecvPong_).ptr(), pingPongMemSize_, interStream));
341 0 : CHK_RET(interRecvTransport->PostFinAck(interStream));
342 0 : interRecvUsePingMem_ ^= true;
343 0 : recvFromInterSrcMemPing_ ^= true;
344 : }
345 0 : if (doSend) {
346 0 : CHK_RET(interSendTransport->WaitFinAck(interStream));
347 : }
348 0 : CHK_RET(ExecuteBarrier(interRecvTransport, interSendTransport, interStream));
349 0 : return HCCL_SUCCESS;
350 : }
351 :
352 0 : HcclResult AlltoallPipelineMeshPairwisePingPong::LocalCopyDataRecvFromInter(u32 mainStep, u32 subStep)
353 : {
354 : u64 localRecvLen = localSendRecvInfo_
355 0 : .recvLength[(userRank_ + groupRankSize_ - (mainStep + 1) * intraRankSize_) % groupRankSize_];
356 : u64 localRecvOff = localSendRecvInfo_
357 0 : .recvOffset[(userRank_ + groupRankSize_ - (mainStep + 1) * intraRankSize_) % groupRankSize_];
358 0 : u64 currStepRecvLen = std::min(localRecvLen - subStep * intraDataBlockSize_, intraDataBlockSize_);
359 0 : DeviceMem src = (interRecvUsePingMem_ ? interRecvPong_ : interRecvPing_)
360 0 : .range(intraRankId_ * intraDataBlockSize_, currStepRecvLen);
361 0 : DeviceMem dst = outputMem_.range(localRecvOff + subStep * intraDataBlockSize_, currStepRecvLen);
362 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, mainStream_));
363 0 : return HCCL_SUCCESS;
364 0 : }
365 :
366 0 : HcclResult AlltoallPipelineMeshPairwisePingPong::PreProcess()
367 : {
368 0 : HCCL_DEBUG(
369 : "[AlltoallPipelineMeshPairwisePingPong][PreProcess] userRank %u, interRank %u, intraRank %u, "
370 : "main stream notify RDMA stream %llu start send",
371 : userRank_, interRankId_, intraRankId_, intraRankId_);
372 : // 搬下次要做 Server 间收发的数据到 ccl buffer
373 0 : CHK_RET(PrepareInterData(0u, 0u));
374 : // 主流notify RDMA流
375 0 : ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_);
376 0 : CHK_RET(NotifyInterStreamStart());
377 :
378 : // 先做一部分 mesh 内 SDMA 操作，剩下的数据待到整体 RDMA 做完之后再补
379 0 : CHK_RET(PrepareIntraData(0u));
380 0 : UpdateIntraStreamInfo(0u, 0u);
381 0 : ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_);
382 0 : CHK_RET(NotifyIntraStreamStart());
383 0 : CHK_RET(SendRecvDataIntraMesh());
384 0 : ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_);
385 : // 主流搬本地那块数据
386 0 : DeviceMem src = inputMem_.range(localSendRecvInfo_.sendOffset[userRank_], localSendRecvInfo_.sendLength[userRank_]);
387 : DeviceMem dst
388 0 : = outputMem_.range(localSendRecvInfo_.recvOffset[userRank_], localSendRecvInfo_.recvLength[userRank_]);
389 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, mainStream_));
390 0 : return HCCL_SUCCESS;
391 0 : }
392 :
393 : // 分别计算当前大步骤需要做几次mesh间收和发，和mesh内收和发（mesh间收的次数和mesh内发的次数相同）
394 0 : void AlltoallPipelineMeshPairwisePingPong::GetNumSubStep(
395 : u32 step, u32& interSendSubStep, u32& interRecvSubStep, u32& intraSubStep)
396 : {
397 0 : u32 sendRankStart = ((interRankId_ + 1 + step) % interRankSize_) * intraRankSize_;
398 0 : u32 recvRankStart = ((interRankId_ + interRankSize_ - 1 - step) % interRankSize_) * intraRankSize_;
399 0 : const std::vector<u64>& sendInfo = (*allMeshAggregationSendRecvInfo_)
400 0 : [(userRank_ + groupRankSize_ - intraRankSize_ * (step + 1)) % intraRankSize_]
401 0 : .sendLength;
402 0 : u64 maxInterSendLen = std::accumulate(
403 0 : localSendRecvInfo_.sendLength.begin() + sendRankStart,
404 0 : localSendRecvInfo_.sendLength.begin() + sendRankStart + intraRankSize_, 0ULL, [](u64 a, u64 b) {
405 0 : return a > b ? a : b;
406 : });
407 0 : u64 maxInterRecvLen = std::accumulate(
408 0 : sendInfo.begin() + meshRankStart_, sendInfo.begin() + meshRankStart_ + intraRankSize_, 0ULL, [](u64 a, u64 b) {
409 0 : return a > b ? a : b;
410 0 : });
411 0 : u64 maxIntraRecvLen = std::accumulate(
412 0 : localSendRecvInfo_.recvLength.begin() + recvRankStart,
413 0 : localSendRecvInfo_.recvLength.begin() + recvRankStart + intraRankSize_, 0ULL, [](u64 a, u64 b) {
414 0 : return a > b ? a : b;
415 0 : });
416 0 : interSendSubStep = (maxInterSendLen + intraDataBlockSize_ - 1) / intraDataBlockSize_;
417 0 : interRecvSubStep = (maxInterRecvLen + intraDataBlockSize_ - 1) / intraDataBlockSize_;
418 : // mesh 的收发步数取决于本卡从其它mesh收到的需要转发到mesh内其他卡的数据以及本卡需要做mesh内读的其他卡数据
419 0 : intraSubStep = (std::max(maxInterRecvLen, maxIntraRecvLen) + intraDataBlockSize_ - 1) / intraDataBlockSize_;
420 0 : }
421 :
422 0 : HcclResult AlltoallPipelineMeshPairwisePingPong::PipelineSend(u32 step, bool isLastStep)
423 : {
424 0 : CHK_RET(ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_));
425 0 : u64 maxDataBlock = 0;
426 0 : for (const SendRecvInfo& info : (*allMeshAggregationSendRecvInfo_)) {
427 0 : for (u64 sendLen : info.sendLength) {
428 0 : maxDataBlock = std::max(maxDataBlock, sendLen);
429 : }
430 : }
431 0 : u32 totalSubStep = (maxDataBlock + intraDataBlockSize_ - 1) / intraDataBlockSize_;
432 : // 计算需要从源端哪块内存收数据和发到哪块目的内存
433 : u64 localRecvLen
434 0 : = localSendRecvInfo_.recvLength[(userRank_ + groupRankSize_ - (step + 1) * intraRankSize_) % groupRankSize_];
435 0 : for (u32 subStep = 0; subStep < totalSubStep; subStep++) {
436 : // RDMA 收发数据
437 0 : SendRecvDataInterMesh(step, true, true);
438 0 : if ((subStep + 1u) == totalSubStep) {
439 0 : CHK_RET(PrepareInterData(step + 1, 0u));
440 : } else {
441 0 : CHK_RET(PrepareInterData(step, subStep + 1u));
442 : }
443 0 : ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_);
444 0 : CHK_RET(WaitIntraStreamFinish());
445 0 : CHK_RET(WaitInterStreamFinish());
446 0 : ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_);
447 0 : UpdateIntraStreamInfo(step + 1u, subStep);
448 0 : CHK_RET(NotifyIntraStreamStart());
449 0 : CHK_RET(SendRecvDataIntraMesh());
450 0 : ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_);
451 0 : if ((!isLastStep && subStep < totalSubStep) || (isLastStep && subStep < (totalSubStep - 1))) {
452 0 : CHK_RET(NotifyInterStreamStart());
453 : }
454 0 : if (localRecvLen > subStep * intraDataBlockSize_) {
455 0 : CHK_RET(LocalCopyDataRecvFromInter(step, subStep));
456 : }
457 : }
458 0 : return HCCL_SUCCESS;
459 : }
460 :
461 0 : HcclResult AlltoallPipelineMeshPairwisePingPong::PostProcess()
462 : {
463 0 : CHK_RET(ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_));
464 0 : CHK_RET(WaitIntraStreamFinish());
465 0 : ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_);
466 0 : u64 maxDataBlock = 0;
467 0 : for (const SendRecvInfo& info : (*allMeshAggregationSendRecvInfo_)) {
468 0 : for (u64 sendLen : info.sendLength) {
469 0 : maxDataBlock = std::max(maxDataBlock, sendLen);
470 : }
471 : }
472 0 : u64 stepLast = (maxDataBlock + intraDataBlockSize_ - 1) / intraDataBlockSize_;
473 0 : for (u64 i = 1; i < stepLast; i++) {
474 0 : UpdateIntraStreamInfo(0, i);
475 0 : CHK_RET(PrepareIntraData(i));
476 0 : ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_);
477 0 : CHK_RET(NotifyIntraStreamStart());
478 0 : ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_);
479 0 : CHK_RET(SendRecvDataIntraMesh());
480 0 : CHK_RET(WaitIntraStreamFinish());
481 0 : CHK_RET(ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_));
482 : }
483 0 : return HCCL_SUCCESS;
484 : }
485 : REGISTER_TEMPLATE(
486 : TemplateType::TEMPLATE_ALL_2_ALL_PIPELINE_MESH_PAIRWISE_PING_PONG, AlltoallPipelineMeshPairwisePingPong);
487 : } // namespace hccl
|
