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