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_pairwise.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : AlltoAllVPairWise::AlltoAllVPairWise(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
16 :
17 0 : AlltoAllVPairWise::~AlltoAllVPairWise() {}
18 :
19 0 : HcclResult AlltoAllVPairWise::Prepare(
20 : AlltoAllVBufferInfo& sendBuffer, AlltoAllVBufferInfo& recvBuffer, bool isAlltoAllZCopyMode, const Stream& stream,
21 : HcclWorkflowMode workMode, std::map<u32, std::vector<u64>>& rankSendDisplsMap,
22 : std::map<u32, std::vector<u64>>& rankRecvDisplsMap)
23 : {
24 0 : DeviceMem scratchInputMem = DeviceMem();
25 0 : DeviceMem scratchOutputMem = DeviceMem();
26 0 : CHK_RET(AlltoAllVPairWise::Prepare(
27 : sendBuffer, recvBuffer, scratchInputMem, scratchOutputMem, isAlltoAllZCopyMode, stream, workMode,
28 : rankSendDisplsMap, rankRecvDisplsMap));
29 0 : return HCCL_SUCCESS;
30 0 : }
31 :
32 0 : HcclResult AlltoAllVPairWise::Prepare(
33 : AlltoAllVBufferInfo& sendBuffer, AlltoAllVBufferInfo& recvBuffer, DeviceMem& scratchInputMem,
34 : DeviceMem& scratchOutputMem, bool isAlltoAllZCopyMode, const Stream& stream, HcclWorkflowMode workMode,
35 : std::map<u32, std::vector<u64>>& rankSendDisplsMap, std::map<u32, std::vector<u64>>& rankRecvDisplsMap)
36 : {
37 0 : HCCL_INFO("[AlltoAllVPairWise][Prepare] Begin");
38 0 : scratchMemSize_ = 0;
39 0 : sendDataUnitBytes_ = 0;
40 0 : recvDataUnitBytes_ = 0;
41 0 : isAlltoAllZCopyMode_ = isAlltoAllZCopyMode;
42 0 : workMode_ = workMode;
43 0 : rankSendDisplsMapPtr_ = &rankSendDisplsMap;
44 0 : rankRecvDisplsMapPtr_ = &rankRecvDisplsMap;
45 :
46 0 : if (workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
47 0 : CHK_PRT_RET(
48 : (!isAlltoAllZCopyMode_ && scratchInputMem.size() != scratchOutputMem.size()),
49 : HCCL_ERROR(
50 : "[AlltoAllVPairWise][Prepare]scratchInputMem and scratchOutputMem should be the same size, "
51 : "ScratchInputMem[%llu] ScratchOutputMem[%llu]",
52 : scratchInputMem.size(), scratchOutputMem.size()),
53 : HCCL_E_MEMORY);
54 :
55 0 : CHK_PRT_RET(
56 : scratchInputMem.size() == 0 || scratchOutputMem.size() == 0,
57 : HCCL_ERROR("[AlltoAllVPairWise][Prepare] invilad scratchMemSize[%llu]", scratchInputMem.size()),
58 : HCCL_E_PARA);
59 0 : scratchInputMem_ = scratchInputMem;
60 0 : scratchOutputMem_ = scratchOutputMem;
61 0 : scratchMemSize_ = scratchInputMem.size();
62 : }
63 :
64 0 : sendBuffer_ = sendBuffer;
65 0 : recvBuffer_ = recvBuffer;
66 0 : stream_ = stream;
67 :
68 0 : CHK_RET(SalGetDataTypeSize(sendBuffer_.dataType, sendDataUnitBytes_));
69 0 : CHK_RET(SalGetDataTypeSize(recvBuffer_.dataType, recvDataUnitBytes_));
70 :
71 0 : return HCCL_SUCCESS;
72 : }
73 :
74 0 : HcclResult AlltoAllVPairWise::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
75 : {
76 0 : HCCL_INFO("[AlltoAllVPairWise][RunAsync]: rank[%u] transportSize[%llu]", rank, links.size());
77 0 : CHK_SMART_PTR_NULL(dispatcher_);
78 0 : CHK_PTR_NULL(stream_.ptr());
79 :
80 0 : CHK_PRT_RET(rankSize == 0, HCCL_ERROR("[AlltoAllVPairWise][Prepare] invilad rankSize[%u]", rankSize), HCCL_E_PARA);
81 :
82 0 : CHK_PRT_RET(
83 : rankSize != links.size(),
84 : HCCL_ERROR(
85 : "[AlltoAllVPairWise][RunAsync]: rankSize[%u] and transport size[%llu] do not match", rankSize,
86 : links.size()),
87 : HCCL_E_PARA);
88 :
89 0 : CHK_RET(LocalCopy(rank));
90 0 : if (workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE && !isAlltoAllZCopyMode_) { // 单算子 && BCopy模式
91 0 : CHK_RET(RunBCopyAlltoAll(rank, rankSize, links));
92 0 : } else {
93 0 : CHK_RET(RunZCopyAlltoAll(rank, rankSize, links));
94 : }
95 0 : return HCCL_SUCCESS;
96 : }
97 :
98 : // 从本rank的sendbuffer拷贝到本rank的recvbuffer
99 0 : HcclResult AlltoAllVPairWise::LocalCopy(const u32 rank)
100 : {
101 : DeviceMem dstMem = recvBuffer_.mem.range(
102 0 : recvDataUnitBytes_ * recvBuffer_.displs[rank], recvBuffer_.counts[rank] * recvDataUnitBytes_);
103 : DeviceMem srcMem = sendBuffer_.mem.range(
104 0 : sendDataUnitBytes_ * sendBuffer_.displs[rank], sendBuffer_.counts[rank] * sendDataUnitBytes_);
105 0 : HCCL_DEBUG(
106 : "[AlltoAllVPairWise][LocalCopy]: Rank[%u] destAddr[%p], destMax[%llu], srcAddr[%p], size[%llu]", rank,
107 : dstMem.ptr(), dstMem.size(), srcMem.ptr(), srcMem.size());
108 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstMem, srcMem, stream_));
109 :
110 0 : return HCCL_SUCCESS;
111 0 : }
112 :
113 0 : HcclResult AlltoAllVPairWise::RunBCopyAlltoAll(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
114 : {
115 0 : for (u32 i = 1; i < rankSize; i++) {
116 0 : u32 prevRank = (rank + rankSize - i) % rankSize;
117 0 : u32 nextRank = (rank + i) % rankSize;
118 0 : std::shared_ptr<Transport> prevTransport = links[prevRank];
119 0 : std::shared_ptr<Transport> nextTransport = links[nextRank];
120 :
121 0 : CHK_SMART_PTR_NULL(prevTransport);
122 0 : CHK_SMART_PTR_NULL(nextTransport);
123 :
124 0 : HCCL_DEBUG("[AlltoAllVPairWise][RunBCopyAlltoAll]: prevRank[%u] nextRank[%u], step[%u]", prevRank, nextRank, i);
125 :
126 0 : u64 sendBytes = sendBuffer_.counts[nextRank] * sendDataUnitBytes_;
127 0 : u64 recvBytes = recvBuffer_.counts[prevRank] * recvDataUnitBytes_;
128 :
129 0 : u64 sendDispBytes = sendBuffer_.displs[nextRank] * sendDataUnitBytes_;
130 0 : u64 recvDispBytes = recvBuffer_.displs[prevRank] * recvDataUnitBytes_;
131 :
132 : // scratchMemSize_ 的合法性已经在 Prepare 函数中校验
133 0 : u32 sendTimes = (sendBytes / scratchMemSize_) + ((sendBytes % scratchMemSize_) == 0 ? 0 : 1);
134 0 : u32 recvTimes = (recvBytes / scratchMemSize_) + ((recvBytes % scratchMemSize_) == 0 ? 0 : 1);
135 :
136 0 : HCCL_DEBUG(
137 : "[AlltoAllVPairWise][RunBCopyAlltoAll]: rank[%u] "
138 : "sendTimes[%u] recvTimes[%u] sendBytes[%llu] recvBytes[%llu] scratchMemSize_[%llu]",
139 : rank, sendTimes, recvTimes, sendBytes, recvBytes, scratchMemSize_);
140 :
141 0 : u32 curSendTime = 0;
142 0 : u32 curRecvTime = 0;
143 0 : while (sendTimes != 0 || recvTimes != 0) {
144 0 : u8* sendAddr = reinterpret_cast<u8*>(sendBuffer_.mem.ptr()) + sendDispBytes + curSendTime * scratchMemSize_;
145 0 : u8* recvAddr = reinterpret_cast<u8*>(recvBuffer_.mem.ptr()) + recvDispBytes + curRecvTime * scratchMemSize_;
146 0 : u64 curSendBytes = 0;
147 0 : u64 curRecvBytes = 0;
148 0 : CHK_RET(CalcSendRecvCounts(sendTimes, curSendTime, sendBytes, curSendBytes));
149 0 : CHK_RET(CalcSendRecvCounts(recvTimes, curRecvTime, recvBytes, curRecvBytes));
150 :
151 0 : HCCL_DEBUG(
152 : "[AlltoAllVPairWise][RunBCopyAlltoAll]: "
153 : "curSendTime[%llu] curRecvTime[%llu] curSendBytes[%llu] curRecvBytes[%llu]",
154 : curSendTime, curRecvTime, curSendBytes, curRecvBytes);
155 :
156 0 : HcclResult ret = SendRecv(curSendBytes, curRecvBytes, sendAddr, recvAddr, prevTransport, nextTransport);
157 0 : CHK_PRT_RET(
158 : ret != HCCL_SUCCESS,
159 : HCCL_ERROR(
160 : "[AlltoAllVPairWise][RunBCopyAlltoAll]: errNo[0x%016llx] "
161 : "curSendBytes[%llu] curRecvBytes[%llu] sendAddr[%p] recvAddr[%p]",
162 : HCCL_ERROR_CODE(ret), curSendBytes, curRecvBytes, sendAddr, recvAddr),
163 : ret);
164 :
165 0 : curSendTime = curSendBytes != 0 ? curSendTime + 1 : curSendTime;
166 0 : curRecvTime = curRecvBytes != 0 ? curRecvTime + 1 : curRecvTime;
167 0 : if (curSendTime == sendTimes && curRecvTime == recvTimes) {
168 0 : break;
169 : }
170 : }
171 0 : }
172 :
173 0 : return HCCL_SUCCESS;
174 : }
175 :
176 0 : HcclResult AlltoAllVPairWise::CalcSendRecvCounts(u32 times, u32 curTime, u64 totalBytes, u64& curBytes) const
177 : {
178 0 : if (times == 0) { // 不需要发送
179 0 : curBytes = 0;
180 0 : } else if (times == 1 && curTime == times - 1) { // 只发一次
181 0 : curBytes = totalBytes;
182 0 : } else if (times > 1 && totalBytes % scratchMemSize_ == 0 && curTime < times) {
183 0 : curBytes = scratchMemSize_;
184 0 : } else if (times > 1 && totalBytes % scratchMemSize_ != 0 && curTime < times - 1) {
185 0 : curBytes = scratchMemSize_;
186 0 : } else if (times > 1 && totalBytes % scratchMemSize_ != 0 && curTime == times - 1) {
187 0 : curBytes = totalBytes % scratchMemSize_;
188 : } else {
189 0 : curBytes = 0;
190 : }
191 0 : return HCCL_SUCCESS;
192 : }
193 :
194 0 : HcclResult AlltoAllVPairWise::SendRecv(
195 : u64 curSendBytes, u64 curRecvBytes, u8* sendAddr, u8* recvAddr, std::shared_ptr<Transport> prevTransport,
196 : std::shared_ptr<Transport> nextTransport)
197 : {
198 0 : if (curRecvBytes > 0) {
199 0 : CHK_RET(prevTransport->TxAck(stream_)); // transport sync record
200 : }
201 0 : if (curSendBytes > 0) {
202 0 : CHK_RET(nextTransport->RxAck(stream_)); // transport sync wait
203 0 : DeviceMem srcMem1 = DeviceMem::create(sendAddr, curSendBytes);
204 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, scratchInputMem_, srcMem1, stream_));
205 : // send payload + notify
206 0 : CHK_RET(nextTransport->TxAsync(UserMemType::OUTPUT_MEM, 0, scratchInputMem_.ptr(), curSendBytes, stream_));
207 0 : }
208 0 : if (curRecvBytes > 0) {
209 0 : CHK_RET(prevTransport->RxAsync(UserMemType::INPUT_MEM, 0, scratchOutputMem_.ptr(), curRecvBytes, stream_));
210 0 : DeviceMem dstMem = DeviceMem::create(recvAddr, curRecvBytes);
211 0 : DeviceMem srcMem = scratchOutputMem_.range(0, curRecvBytes);
212 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstMem, srcMem, stream_));
213 0 : CHK_RET(prevTransport->TxAck(stream_)); // record
214 0 : }
215 0 : if (curSendBytes > 0) {
216 0 : CHK_RET(nextTransport->RxAck(stream_)); // wait
217 0 : CHK_RET(nextTransport->TxDataSignal(stream_)); // record
218 : }
219 0 : if (curRecvBytes > 0) {
220 0 : CHK_RET(prevTransport->RxDataSignal(stream_)); // wait
221 0 : CHK_RET(prevTransport->RxWaitDone(stream_));
222 : }
223 0 : if (curSendBytes > 0) {
224 0 : CHK_RET(nextTransport->TxWaitDone(stream_));
225 : }
226 0 : return HCCL_SUCCESS;
227 : }
228 :
229 0 : HcclResult AlltoAllVPairWise::SendRecv(
230 : TxMemoryInfo txMemoryInfo, RxMemoryInfo rxMemoryInfo, std::shared_ptr<Transport> prevTransport,
231 : std::shared_ptr<Transport> nextTransport)
232 : {
233 : // send payload + notify
234 0 : CHK_RET(nextTransport->TxAsync(
235 : txMemoryInfo.dstMemType, txMemoryInfo.dstOffset, txMemoryInfo.src, txMemoryInfo.len, stream_));
236 0 : CHK_RET(prevTransport->RxAsync(
237 : rxMemoryInfo.srcMemType, rxMemoryInfo.srcOffset, rxMemoryInfo.dst, rxMemoryInfo.len, stream_));
238 0 : CHK_RET(prevTransport->TxAck(stream_)); // record
239 0 : CHK_RET(nextTransport->RxAck(stream_)); // wait
240 0 : CHK_RET(nextTransport->TxDataSignal(stream_)); // record
241 0 : CHK_RET(prevTransport->RxDataSignal(stream_)); // wait
242 0 : CHK_RET(prevTransport->RxWaitDone(stream_));
243 0 : CHK_RET(nextTransport->TxWaitDone(stream_));
244 0 : return HCCL_SUCCESS;
245 : }
246 :
247 0 : HcclResult AlltoAllVPairWise::RunZCopyAlltoAll(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
248 : {
249 0 : for (u32 i = 1; i < rankSize; i++) {
250 0 : u32 prevRank = (rank + rankSize - i) % rankSize;
251 0 : u32 nextRank = (rank + i) % rankSize;
252 0 : std::shared_ptr<Transport> prevTransport = links[prevRank];
253 0 : std::shared_ptr<Transport> nextTransport = links[nextRank];
254 :
255 0 : CHK_SMART_PTR_NULL(prevTransport);
256 0 : CHK_SMART_PTR_NULL(nextTransport);
257 :
258 0 : HCCL_DEBUG("[AlltoAllVPairWise][RunZCopyAlltoAll]: prevRank[%u] nextRank[%u], step[%u]", prevRank, nextRank, i);
259 :
260 0 : CHK_RET(prevTransport->TxAck(stream_)); // transport sync record
261 0 : CHK_RET(nextTransport->RxAck(stream_)); // transport sync wait
262 :
263 0 : u64 sendBytes = sendBuffer_.counts[nextRank] * sendDataUnitBytes_;
264 0 : u64 recvBytes = recvBuffer_.counts[prevRank] * recvDataUnitBytes_;
265 0 : u64 sendDispBytes = sendBuffer_.displs[nextRank] * sendDataUnitBytes_;
266 0 : u64 recvDispBytes = recvBuffer_.displs[prevRank] * recvDataUnitBytes_;
267 0 : u8* sendAddr = reinterpret_cast<u8*>(sendBuffer_.mem.ptr()) + sendDispBytes;
268 0 : u8* recvAddr = reinterpret_cast<u8*>(recvBuffer_.mem.ptr()) + recvDispBytes;
269 :
270 0 : u64 dstOffset = rankRecvDisplsMapPtr_->at(nextRank)[rank];
271 0 : u64 srcOffset = rankSendDisplsMapPtr_->at(prevRank)[rank];
272 :
273 0 : TxMemoryInfo txMemoryInfo{UserMemType::OUTPUT_MEM, dstOffset, sendAddr, sendBytes};
274 0 : RxMemoryInfo rxMemoryInfo{UserMemType::INPUT_MEM, srcOffset, recvAddr, recvBytes};
275 :
276 0 : HCCL_DEBUG(
277 : "[AlltoAllVPairWise][RunZCopyAlltoAll]: sendBytes[%llu] recvBytes[%llu] sendDispBytes[%llu]"
278 : " dstOffset[%llu]",
279 : sendBytes, recvBytes, sendDispBytes, dstOffset);
280 0 : HcclResult ret = SendRecv(txMemoryInfo, rxMemoryInfo, prevTransport, nextTransport);
281 0 : CHK_PRT_RET(
282 : ret != HCCL_SUCCESS,
283 : HCCL_ERROR(
284 : "[AlltoAllVPairWise][RunZCopyAlltoAll]errNo[0x%016llx] "
285 : "sendBytes[%llu] recvBytes[%llu] sendAddr[%p] dstOffset[%llu]",
286 : HCCL_ERROR_CODE(ret), sendBytes, recvBytes, sendAddr, dstOffset),
287 : ret);
288 0 : }
289 :
290 0 : return HCCL_SUCCESS;
291 : }
292 0 : HcclResult AlltoAllVPairWise::GetNslbAdjInfo(
293 : const u32 rank, const u32 rankSize, const std::vector<LINK>& links, AdjInfo& nslbAdjInfo)
294 : {
295 : (void)links;
296 0 : for (u32 i = 1; i < rankSize; i++) {
297 0 : u32 nextRank = (rank + i) % rankSize;
298 0 : if (i < NSLBDP_PAIRWISE_MAXPHASE) {
299 0 : NslbDpAdjInfo adjInfoStep = {};
300 0 : adjInfoStep.dstLocalRankId = nextRank;
301 0 : adjInfoStep.phaseId = i;
302 0 : adjInfoStep.rev = 0;
303 0 : HCCL_INFO(
304 : "AlltoAllVPairWise-nslb: adjInfoStep.phaseId[%u], remoteuserRank[%u]", adjInfoStep.phaseId, nextRank);
305 0 : nslbAdjInfo.nsAdjInfo.push_back(adjInfoStep);
306 : }
307 : }
308 0 : nslbAdjInfo.dstRankNum = nslbAdjInfo.nsAdjInfo.size();
309 :
310 0 : return HCCL_SUCCESS;
311 : }
312 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_2_ALL_V_PAIRWISE, AlltoAllVPairWise);
313 : } // namespace hccl
|