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 "alltoallv_staged_pairwise.h"
12 : #include "log.h"
13 : #include "alg_template_register.h"
14 :
15 : namespace hccl {
16 : using namespace std;
17 :
18 0 : AlltoAllVStagedPairwise::AlltoAllVStagedPairwise(const HcclDispatcher dispatcher)
19 0 : : AlltoAllVStagedBase(dispatcher)
20 : {
21 0 : }
22 :
23 0 : AlltoAllVStagedPairwise::~AlltoAllVStagedPairwise() {}
24 :
25 : // 图模式Prepare入口
26 0 : HcclResult AlltoAllVStagedPairwise::Prepare(DeviceMem &sendMem, DeviceMem &recvMem,
27 : StageAlltoAllVAddrInfo &sendAddrInfo, StageAlltoAllVAddrInfo &recvAddrInfo,
28 : bool isAlltoAllZCopyMode, Stream &mainStream)
29 : {
30 0 : DeviceMem scratchInputMem = DeviceMem();
31 0 : DeviceMem scratchOutputMem = DeviceMem();
32 0 : return AlltoAllVStagedPairwise::Prepare(sendMem, recvMem, scratchInputMem, scratchOutputMem,
33 0 : sendAddrInfo, recvAddrInfo, isAlltoAllZCopyMode, mainStream);
34 0 : }
35 :
36 : // 单算子Prepare入口
37 0 : HcclResult AlltoAllVStagedPairwise::Prepare(DeviceMem &sendMem, DeviceMem &recvMem, DeviceMem &scratchInputMem,
38 : DeviceMem &scratchOutputMem, StageAlltoAllVAddrInfo &sendAddrInfo, StageAlltoAllVAddrInfo &recvAddrInfo,
39 : bool isAlltoAllZCopyMode, Stream &mainStream)
40 : {
41 0 : CHK_RET(AlltoAllVStagedBase::Prepare(sendMem, recvMem, sendAddrInfo, recvAddrInfo,
42 : isAlltoAllZCopyMode, mainStream));
43 0 : if (GetWorkflowMode() == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE &&
44 0 : !isAlltoAllZCopyMode_) {
45 0 : CHK_PRT_RET((scratchInputMem.size() != scratchOutputMem.size()),
46 : HCCL_ERROR(
47 : "[AlltoAllVStagedPairwise][Prepare]scratchInputMem and scratchOutputMem should be the same size, "
48 : "ScratchInputMem[%llu] ScratchOutputMem[%llu]", scratchInputMem.size(), scratchOutputMem.size()),
49 : HCCL_E_MEMORY);
50 :
51 0 : CHK_PRT_RET(scratchInputMem.size() == 0,
52 : HCCL_ERROR("[AlltoAllVStagedPairwise][Prepare] invilad scratchMemSize[%llu]", scratchInputMem.size()),
53 : HCCL_E_PARA);
54 0 : scratchInputMem_ = scratchInputMem;
55 0 : scratchOutputMem_ = scratchOutputMem;
56 0 : scratchMemSize_ = scratchInputMem.size();
57 : }
58 :
59 0 : HCCL_DEBUG("[AlltoAllVStagedPairwise][Prepare] finished");
60 0 : return HCCL_SUCCESS;
61 : }
62 :
63 0 : HcclResult AlltoAllVStagedPairwise::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
64 : {
65 0 : HCCL_INFO("[AlltoAllVStagedPairwise][RunAsync]: rank[%u] transportSize[%llu]", rank, links.size());
66 0 : CHK_SMART_PTR_NULL(dispatcher_);
67 0 : CHK_PTR_NULL(mainStreamPtr_);
68 :
69 0 : CHK_PRT_RET(rankSize == 0, HCCL_ERROR("[AlltoAllVStagedPairwise][Prepare] invilad rankSize[%u]", rankSize),
70 : HCCL_E_PARA);
71 :
72 0 : CHK_PRT_RET(rankSize != links.size(),
73 : HCCL_ERROR("[AlltoAllVStagedPairwise][RunAsync]: rankSize[%u] and transport size[%llu] do not match", rankSize,
74 : links.size()),
75 : HCCL_E_PARA);
76 :
77 0 : bool sizeEqual = (sendAddrInfo_.size() == recvAddrInfo_.size() && sendAddrInfo_.size() == rankSize);
78 0 : CHK_PRT_RET(!sizeEqual,
79 : HCCL_ERROR("[AlltoAllVStagedPairwise][RunAsync] invilad params: "\
80 : "sendAddrInfo size[%u] recvAddrInfo size[%u] rankSize[%u]",
81 : sendAddrInfo_.size(), recvAddrInfo_.size(), rankSize),
82 : HCCL_E_PARA);
83 :
84 0 : CHK_RET(LocalCopy(rank));
85 :
86 0 : if (GetWorkflowMode() == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE &&
87 0 : !isAlltoAllZCopyMode_) { // 单算子 && BCopy模式
88 0 : CHK_RET(RunBCopyAlltoAll(rank, rankSize, links));
89 : } else {
90 0 : CHK_RET(RunZCopyAlltoAll(rank, rankSize, links));
91 : }
92 :
93 0 : return HCCL_SUCCESS;
94 : }
95 :
96 0 : HcclResult AlltoAllVStagedPairwise::RunZCopyAlltoAll(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
97 : {
98 0 : for (u32 i = 1; i < rankSize; i++) {
99 0 : u32 prevRank = (rank + rankSize - i) % rankSize;
100 0 : u32 nextRank = (rank + i) % rankSize;
101 0 : std::shared_ptr<Transport> prevTransport = links[prevRank];
102 0 : std::shared_ptr<Transport> nextTransport = links[nextRank];
103 :
104 0 : CHK_SMART_PTR_NULL(prevTransport);
105 0 : CHK_SMART_PTR_NULL(nextTransport);
106 :
107 0 : HCCL_DEBUG("[AlltoAllVStagedPairwise][RunZCopyAlltoAll]: prevRank[%u] nextRank[%u], step[%u]", prevRank,
108 : nextRank, i);
109 :
110 0 : CHK_RET(prevTransport->TxAck(*mainStreamPtr_)); // transport sync record
111 0 : CHK_RET(nextTransport->RxAck(*mainStreamPtr_)); // transport sync wait
112 :
113 0 : u32 sendDataNum = sendAddrInfo_[nextRank].size();
114 0 : vector<TxMemoryInfo> txMems(sendDataNum);
115 0 : u32 index = 0;
116 0 : for (auto &addrInfo : sendAddrInfo_[nextRank]) {
117 0 : txMems[index].dstMemType = UserMemType::OUTPUT_MEM;
118 0 : txMems[index].dstOffset = addrInfo.remoteOffset;
119 0 : txMems[index].src = static_cast<u8 *>(sendMem_.ptr()) + addrInfo.localOffset;
120 0 : txMems[index].len = addrInfo.localLength;
121 0 : index++;
122 : }
123 :
124 0 : u32 recvDataNum = recvAddrInfo_[prevRank].size();
125 0 : vector<RxMemoryInfo> rxMems(recvDataNum);
126 0 : index = 0;
127 0 : for (auto &addrInfo : recvAddrInfo_[prevRank]) {
128 0 : rxMems[index].srcMemType = UserMemType::INPUT_MEM;
129 0 : rxMems[index].srcOffset = addrInfo.remoteOffset;
130 0 : rxMems[index].dst = static_cast<u8 *>(recvMem_.ptr()) + addrInfo.localOffset;
131 0 : rxMems[index].len = addrInfo.localLength;
132 0 : index++;
133 : }
134 0 : CHK_RET(nextTransport->TxAsync(txMems, *mainStreamPtr_)); // send payload + data notify
135 0 : CHK_RET(prevTransport->RxAsync(rxMems, *mainStreamPtr_)); // wait data notify
136 0 : CHK_RET(ExecuteBarrier(prevTransport, nextTransport));
137 0 : }
138 :
139 0 : return HCCL_SUCCESS;
140 : }
141 :
142 0 : HcclResult AlltoAllVStagedPairwise::RunBCopyAlltoAll(const u32 rank, const u32 rankSize,
143 : const std::vector<LINK> &links)
144 : {
145 0 : for (u32 i = 1; i < rankSize; ++i) {
146 0 : u32 prevRank = (rank + rankSize - i) % rankSize;
147 0 : u32 nextRank = (rank + i) % rankSize;
148 0 : std::shared_ptr<Transport> prevTransport = links[prevRank];
149 0 : std::shared_ptr<Transport> nextTransport = links[nextRank];
150 :
151 0 : CHK_SMART_PTR_NULL(prevTransport);
152 0 : CHK_SMART_PTR_NULL(nextTransport);
153 :
154 0 : HCCL_DEBUG("[AlltoAllVStagedPairwise][RunBCopyAlltoAll]: prevRank[%u] nextRank[%u], step[%u]", prevRank,
155 : nextRank, i);
156 :
157 : // 计算本轮收发次数
158 0 : u64 sendTimes = 0;
159 0 : u64 recvTimes = 0;
160 :
161 0 : CalcSendRecvTimes(sendTimes, recvTimes, prevRank, nextRank);
162 0 : HCCL_DEBUG("sendTimes[%llu] recvTimes[%llu]", sendTimes, recvTimes);
163 :
164 0 : std::vector<std::list<OneSendRecvAddrInfo>> sendPolicies(sendTimes);
165 0 : std::vector<std::list<OneSendRecvAddrInfo>> recvPolicies(recvTimes);
166 0 : LoadPolicies(nextRank, sendAddrInfo_, sendPolicies);
167 0 : CHK_RET(CheckPolicies(sendTimes, sendPolicies));
168 :
169 0 : LoadPolicies(prevRank, recvAddrInfo_, recvPolicies);
170 0 : CHK_RET(CheckPolicies(recvTimes, recvPolicies));
171 :
172 0 : u64 curSendTime = 0;
173 0 : u64 curRecvTime = 0;
174 :
175 0 : while (curSendTime < sendTimes || curRecvTime < recvTimes) {
176 0 : CHK_RET(SendRecv(curSendTime, sendPolicies, curRecvTime, recvPolicies, prevTransport, nextTransport));
177 0 : curSendTime = curSendTime < sendTimes ? curSendTime + 1 : curSendTime;
178 0 : curRecvTime = curRecvTime < recvTimes ? curRecvTime + 1 : curRecvTime;
179 : }
180 0 : }
181 0 : return HCCL_SUCCESS;
182 : }
183 :
184 0 : void AlltoAllVStagedPairwise::CalcSendRecvTimes(u64 &sendTimes, u64 &recvTimes, const u32 prevRank, const u32 nextRank)
185 : {
186 0 : u64 sendBytes = 0;
187 0 : u64 recvBytes = 0;
188 0 : for (auto &addrInfo : sendAddrInfo_[nextRank]) {
189 0 : sendBytes += addrInfo.localLength;
190 : }
191 0 : for (auto &addrInfo : recvAddrInfo_[prevRank]) {
192 0 : recvBytes += addrInfo.localLength;
193 : }
194 0 : sendTimes = (sendBytes / scratchMemSize_) + ((sendBytes % scratchMemSize_) == 0 ? 0 : 1);
195 0 : recvTimes = (recvBytes / scratchMemSize_) + ((recvBytes % scratchMemSize_) == 0 ? 0 : 1);
196 0 : }
197 :
198 0 : void AlltoAllVStagedPairwise::LoadPolicies(const u32 rank, StageAlltoAllVAddrInfo &addrInfos,
199 : std::vector<std::list<OneSendRecvAddrInfo>> &policies)
200 : {
201 0 : std::list<OneSendRecvAddrInfo> tempPolicies;
202 0 : u64 curSendTime = 0;
203 0 : u64 curCCLBufSize = scratchMemSize_;
204 : // 当CCLbuf剩余空间不够发送、接收一整个task时，对task做拆分
205 : OneSendRecvAddrInfo curLastInfo;
206 0 : for (auto &addrInfo : addrInfos[rank]) {
207 : // 若当前task收发数据量为0，直接看下一个
208 0 : if (addrInfo.localLength == 0) {
209 0 : continue;
210 : }
211 0 : u64 curBytes = addrInfo.localLength;
212 0 : if (curBytes <= curCCLBufSize) {
213 0 : tempPolicies.push_back(addrInfo);
214 0 : curCCLBufSize -= curBytes;
215 0 : if (curCCLBufSize == 0) {
216 0 : curCCLBufSize = scratchMemSize_;
217 0 : policies[curSendTime] = tempPolicies;
218 0 : ++curSendTime;
219 0 : tempPolicies.clear();
220 : }
221 : } else {
222 0 : OneSendRecvAddrInfo tmpInfo = addrInfo;
223 0 : u64 tmpBytes = curBytes;
224 0 : while (tmpBytes > curCCLBufSize) {
225 0 : SplitSendRecvAddrInfo(curLastInfo, tmpInfo, curCCLBufSize);
226 0 : tempPolicies.push_back(curLastInfo);
227 0 : curCCLBufSize = scratchMemSize_;
228 0 : policies[curSendTime] = tempPolicies;
229 0 : ++curSendTime;
230 0 : tempPolicies.clear();
231 0 : tmpBytes = tmpInfo.localLength;
232 : }
233 0 : if (tmpBytes != 0) {
234 0 : tempPolicies.push_back(tmpInfo);
235 0 : curCCLBufSize -= tmpBytes;
236 : }
237 : }
238 : }
239 0 : if (curCCLBufSize != scratchMemSize_) {
240 0 : policies[curSendTime] = tempPolicies;
241 0 : ++curSendTime;
242 : }
243 0 : }
244 :
245 0 : HcclResult AlltoAllVStagedPairwise::CheckPolicies(const u64 times,
246 : const std::vector<std::list<OneSendRecvAddrInfo>> &policies) const
247 : {
248 0 : CHK_PRT_RET(times != policies.size(),
249 : HCCL_ERROR(
250 : "[AlltoAllVStagedPairwise][CheckPolicies] invilad params: times[%llu] policies size[%u]", times,
251 : policies.size()), HCCL_E_PARA);
252 :
253 0 : for (u32 i = 0; i < times; ++i) {
254 0 : u64 sum = 0;
255 0 : for (auto &addrInfo : policies[i]) {
256 0 : sum += addrInfo.localLength;
257 : }
258 0 : CHK_PRT_RET(sum > scratchMemSize_,
259 : HCCL_ERROR(
260 : "[AlltoAllVStagedPairwise][CheckPolicies] invilad params: curTime[%u] sum[%llu] scratchMemSize_[%u]", i,
261 : sum, scratchMemSize_), HCCL_E_PARA);
262 : }
263 0 : return HCCL_SUCCESS;
264 : }
265 :
266 0 : HcclResult AlltoAllVStagedPairwise::SendRecv(const u64 curSendTime,
267 : const std::vector<std::list<OneSendRecvAddrInfo>> &sendPolicies, const u64 curRecvTime,
268 : const std::vector<std::list<OneSendRecvAddrInfo>> &recvPolicies, std::shared_ptr<Transport> prevTransport,
269 : std::shared_ptr<Transport> nextTransport)
270 : {
271 0 : bool hasSend = curSendTime < sendPolicies.size();
272 0 : bool hasRecv = curRecvTime < recvPolicies.size();
273 0 : if (hasRecv) {
274 0 : CHK_RET(prevTransport->TxAck(*mainStreamPtr_)); // transport sync record
275 : }
276 0 : if (hasSend) {
277 0 : CHK_RET(nextTransport->RxAck(*mainStreamPtr_)); // transport sync wait
278 : }
279 0 : if (hasSend) {
280 : // 1、把对应内存块从sendbuf copy到CCLInputBuf
281 0 : u64 curCCLInputBufOffset = 0;
282 0 : for (auto &addrInfo : sendPolicies[curSendTime]) {
283 0 : DeviceMem dstMem = scratchInputMem_.range(curCCLInputBufOffset, addrInfo.localLength);
284 0 : DeviceMem srcMem = sendMem_.range(addrInfo.localOffset, addrInfo.localLength);
285 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstMem, srcMem, *mainStreamPtr_));
286 0 : curCCLInputBufOffset += addrInfo.localLength;
287 0 : }
288 : // 2、send CCLInputBuf to CCLOutPutBuf + record
289 0 : CHK_RET(nextTransport->TxAsync(UserMemType::OUTPUT_MEM, 0, scratchInputMem_.ptr(),
290 : curCCLInputBufOffset, *mainStreamPtr_));
291 : }
292 :
293 0 : if (hasRecv) {
294 : // 3、recv CCLOutPutBuf from CCLInputBuf
295 0 : u64 recvBytes = 0;
296 0 : for (auto &addrInfo : recvPolicies[curRecvTime]) {
297 0 : recvBytes += addrInfo.localLength;
298 : }
299 : // wait
300 0 : CHK_RET(prevTransport->RxAsync(UserMemType::INPUT_MEM, 0, scratchOutputMem_.ptr(), recvBytes, *mainStreamPtr_));
301 : // 4、把对应内存块从CCLOutputBuf copy到recvBuf
302 0 : u64 curCCLOutputBufOffset = 0;
303 0 : for (auto &addrInfo : recvPolicies[curRecvTime]) {
304 0 : DeviceMem srcMem = scratchOutputMem_.range(curCCLOutputBufOffset, addrInfo.localLength);
305 0 : DeviceMem dstMem = recvMem_.range(addrInfo.localOffset, addrInfo.localLength);
306 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstMem, srcMem, *mainStreamPtr_));
307 0 : curCCLOutputBufOffset += addrInfo.localLength;
308 0 : }
309 : }
310 :
311 0 : CHK_RET(ExecuteBarrier(hasSend, hasRecv, prevTransport, nextTransport));
312 :
313 0 : return HCCL_SUCCESS;
314 : }
315 :
316 0 : void AlltoAllVStagedPairwise::SplitSendRecvAddrInfo(OneSendRecvAddrInfo &curLastInfo, OneSendRecvAddrInfo &addrInfo,
317 : const u64 &curCCLBufSize) const
318 : {
319 : // 单算子收发暂不使用remote offset、len，考虑演进性，remote的数据也进行更新
320 0 : curLastInfo = addrInfo;
321 0 : curLastInfo.localLength = curCCLBufSize;
322 0 : curLastInfo.remoteLength = curCCLBufSize;
323 :
324 : // 切分后剩余部分，可能大于CCLbuf size，需要循环处理
325 0 : addrInfo.localOffset += curCCLBufSize;
326 0 : addrInfo.localLength -= curCCLBufSize;
327 0 : addrInfo.remoteOffset += curCCLBufSize;
328 0 : addrInfo.remoteLength -= curCCLBufSize;
329 0 : }
330 :
331 0 : HcclResult AlltoAllVStagedPairwise::ExecuteBarrier(std::shared_ptr<Transport> preLink,
332 : std::shared_ptr<Transport> aftLink)
333 : {
334 : // 同步与preLink保证数据收发已结束
335 0 : CHK_RET(preLink->TxAck(*mainStreamPtr_)); // record
336 :
337 0 : CHK_RET(aftLink->RxAck(*mainStreamPtr_)); // wait
338 :
339 : // 同步与aftLink保证数据收发已结束
340 0 : CHK_RET(aftLink->TxDataSignal(*mainStreamPtr_)); // record
341 :
342 0 : CHK_RET(preLink->RxDataSignal(*mainStreamPtr_)); // wait
343 :
344 0 : CHK_RET(preLink->RxWaitDone(*mainStreamPtr_));
345 0 : CHK_RET(aftLink->TxWaitDone(*mainStreamPtr_));
346 :
347 0 : return HCCL_SUCCESS;
348 : }
349 :
350 0 : HcclResult AlltoAllVStagedPairwise::ExecuteBarrier(bool hasSend, bool hasRecv,
351 : std::shared_ptr<Transport> preLink, std::shared_ptr<Transport> aftLink)
352 : {
353 : // 同步与preLink保证数据收发已结束
354 0 : if (hasRecv) {
355 0 : CHK_RET(preLink->TxAck(*mainStreamPtr_)); // record
356 : }
357 0 : if (hasSend) {
358 0 : CHK_RET(aftLink->RxAck(*mainStreamPtr_)); // wait
359 : }
360 :
361 : // 同步与aftLink保证数据收发已结束
362 0 : if (hasSend) {
363 0 : CHK_RET(aftLink->TxDataSignal(*mainStreamPtr_)); // record
364 : }
365 0 : if (hasRecv) {
366 0 : CHK_RET(preLink->RxDataSignal(*mainStreamPtr_)); // wait
367 : }
368 :
369 0 : if (hasRecv) {
370 0 : CHK_RET(preLink->RxWaitDone(*mainStreamPtr_));
371 : }
372 :
373 0 : if (hasSend) {
374 0 : CHK_RET(aftLink->TxWaitDone(*mainStreamPtr_));
375 : }
376 :
377 0 : return HCCL_SUCCESS;
378 : }
379 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_2_ALL_V_STAGED_PAIRWISE, AlltoAllVStagedPairwise);
380 : } // namespace hccl
|
