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 "ccu_context_all_to_all_v_mesh2d.h"
12 : #include "ccu_instruction_all_to_all_v_mesh2d.h"
13 :
14 : namespace Hccl {
15 :
16 : constexpr int CKE_IDX_0 = 0;
17 : constexpr int CKE_IDX_1 = 1;
18 : constexpr int CKE_IDX_2 = 2;
19 : constexpr int CKE_IDX_3 = 3;
20 : constexpr int CKE_IDX_4 = 4;
21 : constexpr int FST_AXIS_ID = 0;
22 : constexpr int SEC_AXIS_ID = 1;
23 :
24 : constexpr int SEND_LOOP_UPDATE_FLAG = 1;
25 : constexpr int RECV_LOOP_UPDATE_FLAG = 2;
26 :
27 0 : CcuContextAllToAllVMesh2D::CcuContextAllToAllVMesh2D(const CcuCtxArg &arg, const std::vector<CcuTransport*> &transports,
28 0 : const CcuTransportGroup &group)
29 0 : : CcuContextAlgBase(arg, transports, group)
30 : {
31 0 : localAxisSignal_ = CreateMaskSignal();
32 :
33 0 : firstScratchBaseOffset_ = CreateVariable();
34 0 : secondScratchBaseOffset_ = CreateVariable();
35 0 : firstScratchSliceOffset_ = CreateVariable();
36 0 : firstScratchSliceStep_ = CreateVariable();
37 0 : secondScratchSliceOffset_ = CreateVariable();
38 0 : secondScratchSliceStep_ = CreateVariable();
39 :
40 0 : xnConst1_ = CreateVariable();
41 0 : completedRankCount_ = CreateVariable();
42 0 : xnHalfTransportSize_ = CreateVariable();
43 0 : xnMaxTransportSize_ = CreateVariable();
44 0 : curSendTailSize_ = CreateVariable();
45 0 : xnHalfTransportGoSize_ = CreateGroupOpSize();
46 0 : curSendTailGoSize_ = CreateGroupOpSize();
47 :
48 0 : if (transports.size() == 0) {
49 0 : THROW<InvalidParamsException>(StringFormat("CcuContextAllToAllVMesh2D transports is empty"));
50 : }
51 :
52 0 : const CcuCtxArgAllToAllVMesh2D *ctxArg = dynamic_cast<const CcuCtxArgAllToAllVMesh2D *>(&arg);
53 0 : if (ctxArg == nullptr) {
54 0 : THROW<NullPtrException>(StringFormat("CcuContextAllToAllVMesh2D::ctxArg ptr is null"));
55 : }
56 0 : rankId_ = ctxArg->rankId;
57 0 : axisId_ = ctxArg->axisId;
58 0 : dimSize_ = ctxArg->dimSize;
59 0 : if (dimSize_.size() != 2 || axisId_ > 1) { // dimSize不为2，或axisId超过1，则不为2D场景
60 0 : THROW<InvalidParamsException>(StringFormat("CcuContextAlltoAllVMesh2D::dimSize[%u] or axisId[%u] is invalid",
61 : dimSize_.size(), axisId_));
62 : }
63 0 : if (dimSize_[0] <= 1 || dimSize_[1] <= 1) {
64 0 : THROW<InvalidParamsException>(StringFormat("CcuContextAlltoAllVMesh2D::dimSize[0] is [%u], dimSize[1] is [%u] are invalid",
65 0 : dimSize_[0], dimSize_[1]));
66 : }
67 0 : dimId_.emplace_back(rankId_ % dimSize_[0]);
68 0 : dimId_.emplace_back(rankId_ / dimSize_[0]);
69 0 : localId_ = dimId_[axisId_];
70 0 : localSize_ = dimSize_[axisId_];
71 0 : anotherId_ = dimId_[1 - axisId_]; // 本rank在另一个轴上的Id
72 0 : anotherSize_ = dimSize_[1 - axisId_];
73 0 : rankSize_ = dimSize_[0] * dimSize_[1];
74 0 : HCCL_INFO("[CcuContextAlltoAllVMesh2D] RankId[%u], DimSize: D0[%u]--D1[%u], localId[%u], localSize[%u]",
75 : rankId_, dimSize_[0], dimSize_[1], localId_, localSize_);
76 :
77 0 : localAxisSignalName_ = "CcuContextAlltoAllVMesh2DAxisSync_" + std::to_string(axisId_);
78 0 : anotherAxisSignalName_ = "CcuContextAlltoAllVMesh2DAxisSync_" + std::to_string(1 - axisId_);
79 0 : }
80 :
81 0 : void CcuContextAllToAllVMesh2D::InitResources()
82 : {
83 0 : ExportMaskSignal(localAxisSignal_, localAxisSignalName_);
84 0 : anotherAxisSignal_ = ImportMaskSignal(anotherAxisSignalName_);
85 :
86 0 : uint32_t transportIdx = 0;
87 0 : u32 ckeNum = 2;
88 0 : input_ = CreateVariable();
89 :
90 0 : sendLoopNumRecorder_.resize(localSize_, std::vector<CcuRep::Variable>(anotherSize_));
91 0 : recvLoopNumRecorder_.resize(localSize_, std::vector<CcuRep::Variable>(anotherSize_));
92 0 : LocSendLoopNumRecorder_.resize(localSize_, std::vector<CcuRep::Variable>(anotherSize_));
93 0 : LocRecvLoopNumRecorder_.resize(localSize_, std::vector<CcuRep::Variable>(anotherSize_));
94 0 : for (uint32_t peerId = 0; peerId < localSize_; peerId++) {
95 0 : isPostFlag_.emplace_back(CreateVariable());
96 0 : sendRecorder_.emplace_back(CreateVariable());
97 0 : sendRecorder_[peerId] = 0;
98 0 : if (peerId == localId_) {
99 0 : scratch_.emplace_back(CreateVariable());
100 0 : output_.emplace_back(CreateVariable());
101 0 : token_.emplace_back(CreateVariable());
102 0 : for (uint16_t anotherId = 0; anotherId < anotherSize_; anotherId++) {
103 0 : sendLoopNumRecorder_[peerId][anotherId] = CreateVariable();
104 0 : recvLoopNumRecorder_[peerId][anotherId] = CreateVariable();
105 0 : LocSendLoopNumRecorder_[peerId][anotherId] = CreateVariable();
106 0 : LocRecvLoopNumRecorder_[peerId][anotherId] = CreateVariable();
107 : }
108 : } else {
109 0 : HCCL_INFO("[CcuContextAllToAllVMesh2D]Rank[%u], PeerId[%u], TransportId[%u]", rankId_, peerId, transportIdx);
110 0 : scratch_.emplace_back(CreateVariable(*(transports[transportIdx]), CKE_IDX_1));
111 0 : output_.emplace_back(CreateVariable(*(transports[transportIdx]), CKE_IDX_2));
112 0 : token_.emplace_back(CreateVariable(*(transports[transportIdx]), CKE_IDX_3));
113 0 : for (uint16_t anotherId = 0; anotherId < anotherSize_; anotherId++) {
114 0 : LocSendLoopNumRecorder_[peerId][anotherId] = CreateVariable();
115 0 : LocRecvLoopNumRecorder_[peerId][anotherId] = CreateVariable();
116 0 : sendLoopNumRecorder_[peerId][anotherId] = (CreateVariable(*(transports[transportIdx]), CKE_IDX_4 + anotherId * ckeNum));
117 0 : recvLoopNumRecorder_[peerId][anotherId] = (CreateVariable(*(transports[transportIdx]), CKE_IDX_4 + anotherId * ckeNum + 1));
118 : }
119 0 : transportIdx++;
120 : }
121 : }
122 :
123 0 : for (uint16_t i = 0; i < localSize_; i++) {
124 0 : inputAddrs_.emplace_back(CreateMemory());
125 0 : bufferAddrs_.emplace_back(CreateMemory());
126 0 : outputAddrs_.emplace_back(CreateMemory());
127 : }
128 :
129 0 : for (uint16_t sliceId = 0; sliceId < anotherSize_; sliceId++) {
130 0 : firstSignal_.emplace_back(CreateMaskSignal()); // 每个对端发anotherSize个分片，localSize个分片共用一个信号，共anotherSize个
131 0 : secondSignal_.emplace_back(CreateMaskSignal());
132 : }
133 :
134 0 : sendRecvInfo_.resize(rankSize_);
135 0 : for (uint64_t rankIdx = 0; rankIdx < rankSize_; rankIdx++) {
136 0 : sendRecvInfo_[rankIdx].sendOffset = CreateVariable();
137 0 : sendRecvInfo_[rankIdx].recvOffset = CreateVariable();
138 0 : sendRecvInfo_[rankIdx].sendTailSizeA = CreateVariable();
139 0 : sendRecvInfo_[rankIdx].sendTailSizeB = CreateVariable();
140 0 : sendRecvInfo_[rankIdx].sendTailGoSizeA = CreateGroupOpSize();
141 0 : sendRecvInfo_[rankIdx].sendTailGoSizeB = CreateGroupOpSize();
142 0 : sendRecvInfo_[rankIdx].sendTailSize = CreateVariable();
143 0 : sendRecvInfo_[rankIdx].recvTailSizeA = CreateVariable();
144 0 : sendRecvInfo_[rankIdx].recvTailSizeB = CreateVariable();
145 0 : sendRecvInfo_[rankIdx].sendLoopNum = CreateVariable();
146 0 : sendRecvInfo_[rankIdx].recvLoopNum = CreateVariable();
147 : }
148 :
149 0 : return;
150 : }
151 :
152 0 : void CcuContextAllToAllVMesh2D::LoadArgs()
153 : {
154 0 : Load(input_);
155 0 : Load(output_[localId_]);
156 0 : Load(token_[localId_]);
157 0 : Load(scratch_[localId_]);
158 :
159 0 : Load(firstScratchBaseOffset_);
160 0 : Load(secondScratchBaseOffset_);
161 0 : Load(firstScratchSliceOffset_);
162 0 : Load(firstScratchSliceStep_);
163 0 : Load(secondScratchSliceOffset_);
164 0 : Load(secondScratchSliceStep_);
165 0 : Load(xnHalfTransportSize_);
166 0 : Load(xnHalfTransportGoSize_);
167 :
168 0 : xnMaxTransportSize_ = xnHalfTransportSize_;
169 0 : xnMaxTransportSize_ += xnHalfTransportSize_;
170 :
171 0 : for (uint64_t rankIdx = 0; rankIdx < rankSize_; rankIdx++) {
172 0 : Load(sendRecvInfo_[rankIdx].sendOffset);
173 0 : Load(sendRecvInfo_[rankIdx].recvOffset);
174 0 : Load(sendRecvInfo_[rankIdx].sendTailSizeA);
175 0 : Load(sendRecvInfo_[rankIdx].sendTailSizeB);
176 0 : Load(sendRecvInfo_[rankIdx].sendTailGoSizeA);
177 0 : Load(sendRecvInfo_[rankIdx].sendTailGoSizeB);
178 0 : Load(sendRecvInfo_[rankIdx].sendTailSize);
179 0 : Load(sendRecvInfo_[rankIdx].recvTailSizeA);
180 0 : Load(sendRecvInfo_[rankIdx].recvTailSizeB);
181 0 : Load(sendRecvInfo_[rankIdx].sendLoopNum);
182 0 : Load(sendRecvInfo_[rankIdx].recvLoopNum);
183 : }
184 :
185 0 : return;
186 : }
187 :
188 0 : void CcuContextAllToAllVMesh2D::ExchangeInfoAndSync()
189 : {
190 : // 交换信息并做同步，前同步固定用1,2,3号信号
191 0 : uint16_t selfBit = 1 << localId_;
192 0 : uint16_t allBit = ((1 << localSize_) - 1) & (~(1 << localId_));
193 :
194 0 : CcuRep::Variable tempDst = CreateVariable();
195 0 : u32 transportId = 0;
196 0 : u32 ckeNum = 2;
197 0 : for (u32 id = 0; id < localSize_; id++) {
198 0 : if (id == localId_) {
199 0 : continue;
200 : }
201 0 : uint32_t dst = CalcDstRank(anotherId_, id);
202 0 : tempDst = output_[localId_];
203 0 : tempDst += sendRecvInfo_[dst].recvOffset;
204 :
205 0 : WriteVariableWithSignal(*transports[transportId], scratch_[localId_], CKE_IDX_1, CKE_IDX_1, selfBit);
206 0 : WriteVariableWithSignal(*transports[transportId], tempDst, CKE_IDX_2, CKE_IDX_2, selfBit);
207 0 : WriteVariableWithSignal(*transports[transportId], token_[localId_], CKE_IDX_3, CKE_IDX_3, selfBit);
208 :
209 0 : for (u32 anotherId = 0; anotherId < anotherSize_; anotherId++) {
210 0 : dst = CalcDstRank(anotherId, id);
211 0 : WriteVariableWithSignal(*transports[transportId], sendRecvInfo_[dst].sendLoopNum,
212 0 : CKE_IDX_4 + anotherId * ckeNum, CKE_IDX_4 + anotherId * ckeNum, selfBit);
213 0 : WriteVariableWithSignal(*transports[transportId], sendRecvInfo_[dst].recvLoopNum,
214 0 : CKE_IDX_4 + anotherId * ckeNum + 1, CKE_IDX_4 + anotherId * ckeNum + 1, selfBit);
215 : }
216 0 : transportId++;
217 : }
218 0 : GroupWait(*transportGroup, CKE_IDX_1, allBit);
219 0 : GroupWait(*transportGroup, CKE_IDX_2, allBit);
220 0 : GroupWait(*transportGroup, CKE_IDX_3, allBit);
221 0 : for (u32 anotherId = 0; anotherId < anotherSize_; anotherId++) {
222 0 : GroupWait(*transportGroup, CKE_IDX_4 + anotherId * ckeNum, allBit);
223 0 : GroupWait(*transportGroup, CKE_IDX_4 + anotherId * ckeNum + 1, allBit);
224 : }
225 :
226 0 : return;
227 0 : }
228 :
229 0 : void CcuContextAllToAllVMesh2D::RankSync(uint32_t signalIndex)
230 : {
231 : // 与远端做同步
232 0 : uint16_t selfBit = 1 << localId_;
233 0 : uint16_t waitBit = 0;
234 0 : uint16_t transportId = 0;
235 0 : for (u32 id = 0; id < localSize_; id++) {
236 0 : isPostFlag_[id] = 0;
237 0 : if (id == localId_) {
238 0 : continue;
239 : }
240 0 : for (uint16_t anotherId = 0; anotherId < anotherSize_; anotherId++) {
241 0 : u32 dstRank = CalcDstRank(anotherId, id);
242 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum != UINT64_MAX) {
243 0 : isPostFlag_[id] = 1;
244 0 : }
245 0 : CCU_IF(sendRecvInfo_[dstRank].recvLoopNum != UINT64_MAX) {
246 0 : isPostFlag_[id] = 1;
247 0 : }
248 0 : if (anotherId == anotherId_) {
249 0 : continue;
250 : }
251 0 : CCU_IF(LocSendLoopNumRecorder_[id][anotherId] != UINT64_MAX) {
252 0 : isPostFlag_[id] = 1;
253 0 : }
254 0 : CCU_IF(LocRecvLoopNumRecorder_[id][anotherId] != UINT64_MAX) {
255 0 : isPostFlag_[id] = 1;
256 0 : }
257 : }
258 0 : CCU_IF(isPostFlag_[id] == 1) {
259 0 : RemotePost(*transports[transportId], signalIndex, selfBit);
260 0 : }
261 0 : transportId++;
262 : }
263 0 : for (u32 id = 0; id < localSize_; id++) {
264 0 : if (id == localId_) {
265 0 : continue;
266 : }
267 0 : waitBit = 1 << id;
268 0 : CCU_IF(isPostFlag_[id] == 1) {
269 0 : GroupWait(*transportGroup, signalIndex, waitBit);
270 0 : }
271 : }
272 :
273 0 : return;
274 : }
275 :
276 0 : void CcuContextAllToAllVMesh2D::PostSync()
277 : {
278 0 : uint16_t selfBit = 1 << localId_;
279 0 : uint16_t allBit = ((1 << localSize_) - 1) & (~(1 << localId_));
280 :
281 0 : for (auto t : transports) {
282 0 : if (t == nullptr) {
283 0 : THROW<NullPtrException>(StringFormat("CcuContextAllToAllVMesh2D::Algorithm transport ptr is null"));
284 : }
285 0 : RemotePost(*t, CKE_IDX_0, selfBit);
286 : }
287 0 : GroupWait(*transportGroup, CKE_IDX_0, allBit);
288 0 : return;
289 : }
290 :
291 0 : void CcuContextAllToAllVMesh2D::UpdateLoopRecorder(uint16_t flag)
292 : {
293 0 : if (flag == SEND_LOOP_UPDATE_FLAG) {
294 0 : for (uint16_t peerId = 0; peerId < localSize_; peerId++) {
295 0 : for(uint16_t anotherId = 0; anotherId < anotherSize_; anotherId++) {
296 0 : u32 dstRank = CalcDstRank(anotherId, peerId);
297 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum != UINT64_MAX) {
298 0 : sendRecvInfo_[dstRank].sendLoopNum += xnConst1_;
299 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum == UINT64_MAX) {
300 0 : completedRankCount_ += xnConst1_;
301 0 : }
302 0 : }
303 0 : if (anotherId == anotherId_) {
304 0 : CCU_IF(sendRecvInfo_[dstRank].recvLoopNum != UINT64_MAX) {
305 0 : sendRecvInfo_[dstRank].recvLoopNum += xnConst1_;
306 0 : CCU_IF(sendRecvInfo_[dstRank].recvLoopNum == UINT64_MAX) {
307 0 : completedRankCount_ += xnConst1_;
308 0 : }
309 0 : }
310 : }
311 0 : if (anotherId == anotherId_ || peerId == localId_) {
312 0 : continue;
313 : }
314 0 : CCU_IF(LocSendLoopNumRecorder_[peerId][anotherId] != UINT64_MAX) {
315 0 : LocSendLoopNumRecorder_[peerId][anotherId] += xnConst1_;
316 0 : CCU_IF(LocSendLoopNumRecorder_[peerId][anotherId] == UINT64_MAX) {
317 0 : completedRankCount_ += xnConst1_;
318 0 : }
319 0 : }
320 : }
321 : }
322 0 : } else if (flag == RECV_LOOP_UPDATE_FLAG) {
323 0 : for (uint16_t peerId = 0; peerId < localSize_; peerId++) {
324 0 : for(uint16_t anotherId = 0; anotherId < anotherSize_; anotherId++) {
325 0 : u32 srcRank = CalcDstRank(anotherId, peerId);
326 0 : if (anotherId != anotherId_) {
327 0 : CCU_IF(sendRecvInfo_[srcRank].recvLoopNum != UINT64_MAX) {
328 0 : sendRecvInfo_[srcRank].recvLoopNum += xnConst1_;
329 0 : CCU_IF(sendRecvInfo_[srcRank].recvLoopNum == UINT64_MAX) {
330 0 : completedRankCount_ += xnConst1_;
331 0 : }
332 0 : }
333 : }
334 0 : if (anotherId == anotherId_ || peerId == localId_) {
335 0 : continue;
336 : }
337 0 : CCU_IF(LocRecvLoopNumRecorder_[peerId][anotherId] != UINT64_MAX) {
338 0 : LocRecvLoopNumRecorder_[peerId][anotherId] += xnConst1_;
339 0 : CCU_IF(LocRecvLoopNumRecorder_[peerId][anotherId] == UINT64_MAX) {
340 0 : completedRankCount_ += xnConst1_;
341 0 : }
342 0 : }
343 : }
344 : }
345 : }
346 :
347 0 : return;
348 : }
349 :
350 0 : void CcuContextAllToAllVMesh2D::AxisSync(uint32_t signalIndex)
351 : {
352 0 : const uint32_t DIE_NUM = 2; // 2个die
353 0 : if (signalIndex > 1) {
354 0 : THROW<InvalidParamsException>(StringFormat(
355 : "[CcuContextAllToAllVMesh2D] Unexpected SignalInex[%u]", signalIndex));
356 : }
357 0 : LocalCtxPost(anotherAxisSignal_, 1 << (axisId_ + signalIndex * DIE_NUM));
358 0 : LocalWait(localAxisSignal_, 1 << (1 - axisId_ + signalIndex * DIE_NUM));
359 0 : return;
360 : }
361 :
362 0 : uint32_t CcuContextAllToAllVMesh2D::CalcDstRank(uint32_t sliceId, uint32_t peerId) const
363 : {
364 : uint32_t dstRank;
365 0 : if (axisId_ == 0) {
366 0 : dstRank = sliceId * localSize_ + peerId;
367 : } else {
368 0 : dstRank = sliceId + anotherSize_ * peerId;
369 : }
370 0 : return dstRank;
371 : }
372 :
373 0 : uint32_t CcuContextAllToAllVMesh2D::CalcTransIdx(uint32_t peerId) const
374 : {
375 : uint32_t transIdx;
376 0 : if (peerId < localId_) {
377 0 : transIdx = peerId;
378 : } else {
379 0 : transIdx = peerId - 1;
380 : }
381 0 : return transIdx;
382 : }
383 :
384 0 : void CcuContextAllToAllVMesh2D::DoAll2AllVMultiLoop()
385 : {
386 : // 需要等待的次数：2 * rankSize_ + (localSize_ - 1) * (anotherSize_ - 1) * 2
387 0 : completedRankCount_ = 0;
388 0 : xnConst1_ = 1;
389 0 : uint64_t targetCount = 2 * rankSize_ + (localSize_ - 1) * (anotherSize_ - 1) * 2;
390 0 : CCU_WHILE(completedRankCount_ != targetCount) {
391 : // 第一轮，直连的rank间直接搬运数据。将需要中转的数据搬到中转rank的scratchBuf上
392 0 : FirstStep();
393 0 : RankSync(CKE_IDX_1);
394 0 : UpdateLoopRecorder(SEND_LOOP_UPDATE_FLAG);
395 0 : AxisSync(FST_AXIS_ID);
396 :
397 : // 第二轮，从input和buffer中将剩余的本端分片以及待转发分片发给对端；其中给每个对端发1个本端分片，localSize-1个转发分片
398 0 : HCCL_INFO("[CcuContextAlltoAllVMesh2D] Algorithm second step begins.");
399 0 : RankSync(CKE_IDX_2);
400 0 : SecondStep();
401 0 : RankSync(CKE_IDX_3);
402 0 : UpdateLoopRecorder(RECV_LOOP_UPDATE_FLAG);
403 0 : AxisSync(SEC_AXIS_ID);
404 0 : }
405 0 : }
406 :
407 0 : void CcuContextAllToAllVMesh2D::WriteToDstOutput(uint16_t sliceId, uint16_t peerId)
408 : {
409 0 : uint32_t dstRank = CalcDstRank(sliceId, peerId);
410 0 : uint32_t transIdx = CalcTransIdx(peerId);
411 :
412 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum == UINT64_MAX) { // 已经搬完了，仅同步
413 0 : LocalPost(firstSignal_[sliceId], (1 << peerId));
414 0 : }
415 :
416 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum != UINT64_MAX) { // 还没有搬完
417 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum == UINT64_MAX - 1)
418 : { // 最后一次搬运
419 0 : CCU_IF(sendRecvInfo_[dstRank].sendTailSize == 0)
420 : {
421 0 : LocalPost(firstSignal_[sliceId], (1 << peerId));
422 0 : }
423 0 : CCU_IF(sendRecvInfo_[dstRank].sendTailSize != 0)
424 : {
425 0 : Write(*(transports[transIdx]), outputAddrs_[peerId], inputAddrs_[peerId],
426 0 : sendRecvInfo_[dstRank].sendTailSize, firstSignal_[sliceId], (1 << peerId));
427 0 : }
428 0 : }
429 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum != UINT64_MAX - 1) { // 正常搬运
430 0 : Write(*(transports[transIdx]), outputAddrs_[peerId], inputAddrs_[peerId], xnMaxTransportSize_,
431 0 : firstSignal_[sliceId], (1 << peerId));
432 0 : sendRecvInfo_[dstRank].sendOffset += xnMaxTransportSize_;
433 0 : sendRecorder_[peerId] += xnMaxTransportSize_;
434 0 : }
435 0 : }
436 0 : return;
437 : }
438 :
439 0 : void CcuContextAllToAllVMesh2D::GroupCopyToDstOutput(uint16_t sliceId, uint16_t peerId)
440 : {
441 0 : HCCL_DEBUG("[CcuContextAlltoAllVMesh2D] GroupCopyToDstOutput Start.");
442 0 : uint32_t dstRank = CalcDstRank(sliceId, peerId);
443 :
444 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum == UINT64_MAX)
445 : { // 已经搬完了，仅同步
446 0 : LocalPost(firstSignal_[sliceId], (1 << peerId));
447 0 : }
448 :
449 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum != UINT64_MAX)
450 : { // 还没有完成，则继续循环
451 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum == UINT64_MAX - 1)
452 : { // 最后一轮循环, 发送尾块数据
453 0 : curSendTailSize_ = (axisId_ == 0) ? sendRecvInfo_[dstRank].sendTailSizeA :
454 0 : sendRecvInfo_[dstRank].sendTailSizeB;
455 0 : curSendTailGoSize_ = (axisId_ == 0) ? sendRecvInfo_[dstRank].sendTailGoSizeA :
456 0 : sendRecvInfo_[dstRank].sendTailGoSizeB;
457 0 : if (axisId_ == 1) {
458 0 : inputAddrs_[peerId].addr += sendRecvInfo_[dstRank].sendTailSizeA;
459 0 : outputAddrs_[peerId].addr += sendRecvInfo_[dstRank].sendTailSizeA;
460 : }
461 :
462 0 : CCU_IF(curSendTailSize_ == 0)
463 : {
464 0 : LocalPost(firstSignal_[sliceId], (1 << peerId));
465 0 : }
466 0 : CCU_IF(curSendTailSize_ != 0)
467 : {
468 0 : outputAddrs_[peerId].addr += sendRecvInfo_[dstRank].recvOffset;
469 0 : GroupCopy(outputAddrs_[peerId], inputAddrs_[peerId], curSendTailGoSize_);
470 0 : LocalPost(firstSignal_[sliceId], (1 << peerId));
471 0 : }
472 0 : }
473 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum != UINT64_MAX - 1)
474 : {
475 0 : outputAddrs_[peerId].addr += sendRecvInfo_[dstRank].recvOffset;
476 0 : if (axisId_ == 1) {
477 0 : inputAddrs_[peerId].addr += xnHalfTransportSize_;
478 0 : outputAddrs_[peerId].addr += xnHalfTransportSize_;
479 : }
480 0 : GroupCopy(outputAddrs_[peerId], inputAddrs_[peerId], xnHalfTransportGoSize_);
481 0 : LocalPost(firstSignal_[sliceId], (1 << peerId));
482 0 : sendRecvInfo_[dstRank].sendOffset += xnMaxTransportSize_;
483 0 : sendRecorder_[peerId] += xnMaxTransportSize_;
484 0 : }
485 0 : }
486 0 : HCCL_DEBUG("[CcuContextAlltoAllVMesh2D] GroupCopyToDstOutput end.");
487 0 : }
488 :
489 0 : void CcuContextAllToAllVMesh2D::WriteToDstScratch(uint16_t sliceId, uint16_t peerId)
490 : {
491 0 : uint32_t dstRank = CalcDstRank(sliceId, peerId);
492 0 : uint32_t transIdx = CalcTransIdx(peerId);
493 :
494 0 : if (peerId == localId_) {
495 0 : LocalPost(firstSignal_[sliceId], (1 << peerId));
496 : } else {
497 0 : if (axisId_ == 0) {
498 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum == UINT64_MAX) { // 已经搬完了，仅同步
499 0 : LocalPost(firstSignal_[sliceId], (1 << peerId));
500 0 : }
501 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum != UINT64_MAX) { // 还没有搬完
502 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum == UINT64_MAX - 1) { // 最后一次搬运
503 0 : CCU_IF(sendRecvInfo_[dstRank].sendTailSizeA == 0) {
504 0 : LocalPost(firstSignal_[sliceId], (1 << peerId));
505 0 : }
506 0 : CCU_IF(sendRecvInfo_[dstRank].sendTailSizeA != 0) {
507 0 : Write(*(transports[transIdx]), outputAddrs_[peerId], inputAddrs_[peerId],
508 0 : sendRecvInfo_[dstRank].sendTailSizeA, firstSignal_[sliceId], (1 << peerId));
509 0 : }
510 0 : }
511 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum != UINT64_MAX - 1) { // 正常搬运
512 0 : Write(*(transports[transIdx]), outputAddrs_[peerId], inputAddrs_[peerId],
513 0 : xnHalfTransportSize_, firstSignal_[sliceId], (1 << peerId));
514 0 : }
515 0 : }
516 : } else {
517 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum == UINT64_MAX) { // 已经搬完了，仅同步
518 0 : LocalPost(firstSignal_[sliceId], (1 << peerId));
519 0 : }
520 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum != UINT64_MAX) { // 还没有搬完
521 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum == UINT64_MAX - 1) { // 最后一次搬运
522 0 : CCU_IF(sendRecvInfo_[dstRank].sendTailSizeB == 0) {
523 0 : LocalPost(firstSignal_[sliceId], (1 << peerId));
524 0 : }
525 0 : CCU_IF(sendRecvInfo_[dstRank].sendTailSizeB != 0) {
526 0 : inputAddrs_[peerId].addr += sendRecvInfo_[dstRank].sendTailSizeA;
527 0 : Write(*(transports[transIdx]), outputAddrs_[peerId], inputAddrs_[peerId],
528 0 : sendRecvInfo_[dstRank].sendTailSizeB, firstSignal_[sliceId], (1 << peerId));
529 0 : }
530 0 : }
531 0 : CCU_IF(sendRecvInfo_[dstRank].sendLoopNum != UINT64_MAX - 1) { // 正常搬运
532 0 : inputAddrs_[peerId].addr += xnHalfTransportSize_;
533 0 : Write(*(transports[transIdx]), outputAddrs_[peerId], inputAddrs_[peerId],
534 0 : xnHalfTransportSize_, firstSignal_[sliceId], (1 << peerId));
535 0 : }
536 0 : }
537 : }
538 0 : sendRecvInfo_[dstRank].sendOffset += xnMaxTransportSize_;
539 : }
540 :
541 0 : return;
542 : }
543 :
544 0 : void CcuContextAllToAllVMesh2D::ReadFromSrc(uint16_t sliceId, uint16_t peerId)
545 : {
546 0 : uint32_t srcRank = CalcDstRank(sliceId, peerId);
547 0 : uint32_t transIdx = CalcTransIdx(peerId);
548 :
549 0 : CCU_IF(sendRecvInfo_[srcRank].recvLoopNum == UINT64_MAX) {
550 0 : LocalPost(secondSignal_[sliceId], (1 << peerId));
551 0 : }
552 0 : CCU_IF(sendRecvInfo_[srcRank].recvLoopNum != UINT64_MAX) {
553 0 : CCU_IF(sendRecvInfo_[srcRank].recvLoopNum == UINT64_MAX - 1) { // 最后一次搬运
554 0 : if (axisId_ == 0) {
555 0 : CCU_IF(sendRecvInfo_[srcRank].recvTailSizeB == 0) {
556 0 : LocalPost(secondSignal_[sliceId], (1 << peerId));
557 0 : }
558 0 : CCU_IF(sendRecvInfo_[srcRank].recvTailSizeB != 0) {
559 0 : outputAddrs_[peerId].addr += sendRecvInfo_[srcRank].recvTailSizeA;
560 0 : Read(*(transports[transIdx]), outputAddrs_[peerId], bufferAddrs_[peerId], sendRecvInfo_[srcRank].recvTailSizeB,
561 0 : secondSignal_[sliceId], (1 << peerId));
562 0 : }
563 : } else {
564 0 : CCU_IF(sendRecvInfo_[srcRank].recvTailSizeA == 0) {
565 0 : LocalPost(secondSignal_[sliceId], (1 << peerId));
566 0 : }
567 0 : CCU_IF(sendRecvInfo_[srcRank].recvTailSizeA != 0) {
568 0 : Read(*(transports[transIdx]), outputAddrs_[peerId], bufferAddrs_[peerId], sendRecvInfo_[srcRank].recvTailSizeA,
569 0 : secondSignal_[sliceId], (1 << peerId));
570 0 : }
571 : }
572 0 : }
573 0 : CCU_IF(sendRecvInfo_[srcRank].recvLoopNum != UINT64_MAX - 1) { // 正常搬运
574 0 : if (axisId_ == 0) {
575 0 : outputAddrs_[peerId].addr += xnHalfTransportSize_;
576 0 : Read(*(transports[transIdx]), outputAddrs_[peerId], bufferAddrs_[peerId], xnHalfTransportSize_,
577 0 : secondSignal_[sliceId], (1 << peerId));
578 : } else {
579 0 : Read(*(transports[transIdx]), outputAddrs_[peerId], bufferAddrs_[peerId], xnHalfTransportSize_,
580 0 : secondSignal_[sliceId], (1 << peerId));
581 : }
582 0 : sendRecvInfo_[srcRank].recvOffset += xnMaxTransportSize_;
583 0 : }
584 0 : }
585 0 : return;
586 : }
587 :
588 0 : void CcuContextAllToAllVMesh2D::FirstStep()
589 : {
590 : // 统一处理token，访问第i个对端需要使用对应的token
591 0 : for (uint16_t peerId = 0; peerId < localSize_; peerId++) {
592 0 : inputAddrs_[peerId].token = token_[peerId];
593 0 : bufferAddrs_[peerId].token = token_[peerId];
594 0 : outputAddrs_[peerId].token = token_[peerId];
595 : }
596 :
597 : // 统一处理bufferAddrs的初始值
598 0 : for (uint16_t peerId = 0; peerId < localSize_; peerId++) {
599 0 : bufferAddrs_[peerId].addr = scratch_[peerId];
600 0 : bufferAddrs_[peerId].addr += firstScratchBaseOffset_;
601 0 : if (peerId < localId_) {
602 0 : for (uint16_t i = 1; i < localId_; i++) {
603 0 : bufferAddrs_[peerId].addr += firstScratchSliceOffset_;
604 : }
605 : } else {
606 0 : for (uint16_t i = 0; i < localId_; i++) {
607 0 : bufferAddrs_[peerId].addr += firstScratchSliceOffset_;
608 : }
609 : }
610 : }
611 :
612 0 : for (uint16_t sliceId = 0; sliceId < anotherSize_; sliceId++) { // sliceId等于dstRank在另一个维度上的id
613 0 : for (uint32_t peerId = 0; peerId < localSize_; peerId++) {
614 0 : u32 dstRank = CalcDstRank(sliceId, peerId);
615 0 : if (peerId == localId_ && sliceId == anotherId_) {
616 0 : continue;
617 : }
618 0 : if (sliceId == anotherId_) {
619 0 : inputAddrs_[peerId].addr = input_;
620 0 : inputAddrs_[peerId].addr += sendRecvInfo_[dstRank].sendOffset;
621 0 : outputAddrs_[peerId].addr = output_[peerId];
622 0 : outputAddrs_[peerId].addr += sendRecorder_[peerId];
623 0 : WriteToDstOutput(sliceId, peerId);
624 : } else {
625 0 : inputAddrs_[peerId].addr = input_;
626 0 : inputAddrs_[peerId].addr += sendRecvInfo_[dstRank].sendOffset;
627 0 : outputAddrs_[peerId].addr = bufferAddrs_[peerId].addr;
628 0 : WriteToDstScratch(sliceId, peerId);
629 0 : bufferAddrs_[peerId].addr += firstScratchSliceStep_;
630 : }
631 : }
632 : }
633 0 : uint32_t dstRankForSelf = CalcDstRank(anotherId_, localId_);
634 0 : inputAddrs_[localId_].addr = input_;
635 0 : inputAddrs_[localId_].addr += sendRecvInfo_[dstRankForSelf].sendOffset;
636 0 : outputAddrs_[localId_].addr = output_[localId_];
637 0 : outputAddrs_[localId_].addr += sendRecorder_[localId_];
638 0 : GroupCopyToDstOutput(anotherId_, localId_);
639 :
640 : // 检查第一轮的数据是否已发完
641 0 : for (uint16_t sliceId = 0; sliceId < anotherSize_; sliceId++) {
642 0 : LocalWait(firstSignal_[sliceId], (1 << localSize_) - 1); // 等待第一轮所有分片都发完
643 : }
644 :
645 0 : return;
646 : }
647 :
648 0 : void CcuContextAllToAllVMesh2D::SecondStep()
649 : {
650 : // 统一处理bufferAddrs的初始值
651 0 : for (uint16_t peerId = 0; peerId < localSize_; peerId++) {
652 0 : bufferAddrs_[peerId].addr = scratch_[peerId];
653 0 : bufferAddrs_[peerId].addr += secondScratchBaseOffset_;
654 0 : if (peerId < localId_) {
655 0 : for (uint16_t i = 1; i < localId_; i++) {
656 0 : bufferAddrs_[peerId].addr += secondScratchSliceOffset_;
657 : }
658 : } else {
659 0 : for (uint16_t i = 0; i < localId_; i++) {
660 0 : bufferAddrs_[peerId].addr += secondScratchSliceOffset_;
661 : }
662 : }
663 : }
664 :
665 : // 本端从直连rank的scratchmem上读取数据
666 0 : for (uint16_t sliceId = 0; sliceId < anotherSize_; sliceId++) {
667 0 : for (uint32_t peerId = 0; peerId < localSize_; peerId++) {
668 0 : if (peerId == localId_ || sliceId == anotherId_) { // 直连链路之前已经搬过了
669 0 : LocalPost(secondSignal_[sliceId], (1 << peerId));
670 0 : continue;
671 : } else {
672 0 : u32 srcRank = CalcDstRank(sliceId, peerId);
673 0 : outputAddrs_[peerId].addr = output_[localId_];
674 0 : outputAddrs_[peerId].addr += sendRecvInfo_[srcRank].recvOffset;
675 0 : ReadFromSrc(sliceId, peerId);
676 : }
677 0 : bufferAddrs_[peerId].addr += secondScratchSliceStep_;
678 : }
679 : }
680 :
681 0 : for (uint16_t sliceId = 0; sliceId < anotherSize_; sliceId++) {
682 0 : LocalWait(secondSignal_[sliceId], (1 << localSize_) - 1); // 等待第二轮所有分片都发完
683 : }
684 :
685 0 : return;
686 : }
687 :
688 0 : void CcuContextAllToAllVMesh2D::CopyLoopNumRecorder()
689 : {
690 0 : for (uint16_t peerId = 0; peerId < localSize_; peerId++) {
691 0 : if (peerId == localId_) {
692 0 : continue;
693 : }
694 0 : for (uint16_t anotherId = 0; anotherId < anotherSize_; anotherId++) {
695 0 : LocSendLoopNumRecorder_[peerId][anotherId] = sendLoopNumRecorder_[peerId][anotherId];
696 0 : LocRecvLoopNumRecorder_[peerId][anotherId] = recvLoopNumRecorder_[peerId][anotherId];
697 : }
698 : }
699 0 : }
700 :
701 0 : void CcuContextAllToAllVMesh2D::Algorithm()
702 : {
703 : // 初始化寄存器资源 & 加载外部输入参数
704 0 : HCCL_INFO("[CcuContextAlltoAllVMesh2D] AllgatherMesh1D Algorithm Init Begins.");
705 0 : InitResources();
706 0 : LoadArgs();
707 :
708 : // 第一轮，X方向发a，Y方向发后b，到对端的块均放在output，要沿X转发的b块放在对端的bufferX，根据转发目的、自身locId两级偏移
709 0 : HCCL_INFO("[CcuContextAlltoAllVMesh2D] Algorithm first step begins.");
710 0 : ExchangeInfoAndSync();
711 0 : PostSync();
712 0 : AxisSync(SEC_AXIS_ID);
713 0 : CopyLoopNumRecorder();
714 :
715 0 : DoAll2AllVMultiLoop();
716 0 : PostSync();
717 0 : AxisSync(FST_AXIS_ID);
718 0 : HCCL_INFO("[CcuContextAlltoAllVMesh2D] Algorithm Ends.");
719 0 : return;
720 : }
721 :
722 0 : void CcuContextAllToAllVMesh2D::CalculateArgs()
723 : {
724 0 : if (axisId_ == 0) {
725 0 : firstScratchBaseOffset = 0;
726 0 : secondScratchBaseOffset = scratchSliceSize * (localSize_ - 1) * (anotherSize_ - 1);
727 :
728 0 : firstScratchSliceOffset = scratchSliceSize * (anotherSize_ - 1);
729 0 : firstScratchSliceStep = scratchSliceSize;
730 0 : secondScratchSliceOffset = scratchSliceSize;
731 0 : secondScratchSliceStep = scratchSliceSize * (localSize_ - 1);
732 : } else {
733 0 : firstScratchBaseOffset = scratchSliceSize * (localSize_ - 1) * (anotherSize_ - 1);
734 0 : secondScratchBaseOffset = 0;
735 :
736 0 : firstScratchSliceOffset = scratchSliceSize * (anotherSize_ - 1);
737 0 : firstScratchSliceStep = scratchSliceSize;
738 0 : secondScratchSliceOffset = scratchSliceSize;
739 0 : secondScratchSliceStep = scratchSliceSize * (localSize_ - 1);
740 : }
741 :
742 0 : return;
743 : }
744 :
745 0 : std::vector<uint64_t> CcuContextAllToAllVMesh2D::GeneArgs(const CcuTaskArg &arg)
746 : {
747 0 : const CcuTaskArgAllToAllVMesh2D *taskArg = dynamic_cast<const CcuTaskArgAllToAllVMesh2D *>(&arg);
748 0 : if (taskArg == nullptr) {
749 0 : THROW<NullPtrException>(StringFormat("CcuContextAllToAllVMesh2D::taskArg ptr is null"));
750 : }
751 :
752 0 : uint64_t inputAddr = taskArg->inputAddr;
753 0 : uint64_t outputAddr = taskArg->outputAddr;
754 0 : uint64_t scratchAddr = taskArg->scratchAddr;
755 0 : uint64_t tokenInfo = taskArg->token;
756 :
757 0 : scratchSliceSize = std::min(taskArg->scratchSliceSize, UB_MAX_TRANS_SIZE / MESH_2D_NUM); // 最小值
758 0 : CalculateArgs();
759 0 : auto scratchGoSliceSize = CalGoSize(scratchSliceSize);
760 :
761 0 : HCCL_INFO("[CcuContextAllToAllVMesh2D][GeneArgs] inputAddr[%llu], outputAddr[%llu], scratchAddr[%llu], " \
762 : "scratchSliceSize[%llu], firstScratchBaseOffset[%llu], secondScratchBaseOffset[%llu], " \
763 : "firstScratchSliceOffset[%llu], firstScratchSliceStep[%llu], secondScratchSliceOffset[%llu], " \
764 : "secondScratchSliceStep[%llu]", inputAddr, outputAddr, scratchAddr, scratchSliceSize,
765 : firstScratchBaseOffset, secondScratchBaseOffset, firstScratchSliceOffset, firstScratchSliceStep,
766 : secondScratchSliceOffset, secondScratchSliceStep);
767 :
768 0 : std::vector<uint64_t> processReturn = {inputAddr, outputAddr, tokenInfo, scratchAddr, firstScratchBaseOffset,
769 0 : secondScratchBaseOffset, firstScratchSliceOffset, firstScratchSliceStep, secondScratchSliceOffset,
770 0 : secondScratchSliceStep, scratchSliceSize};
771 :
772 0 : processReturn.insert(processReturn.end(), scratchGoSliceSize.begin(), scratchGoSliceSize.end());
773 :
774 0 : for (uint16_t i = 0; i < rankSize_; i++) {
775 0 : uint64_t perTranSize = scratchSliceSize * MESH_2D_NUM;
776 0 : uint64_t sendLoopNum = UINT64_MAX - 1 - taskArg->localSendRecvInfo.sendLength[i] / perTranSize;
777 0 : uint64_t recvLoopNum = UINT64_MAX - 1 - taskArg->localSendRecvInfo.recvLength[i] / perTranSize;
778 :
779 0 : uint64_t sendTailSize = taskArg->localSendRecvInfo.sendLength[i] - taskArg->localSendRecvInfo.sendLength[i] / perTranSize * perTranSize;
780 0 : uint64_t recvTailSize = taskArg->localSendRecvInfo.recvLength[i] - taskArg->localSendRecvInfo.recvLength[i] / perTranSize * perTranSize;
781 :
782 0 : uint64_t sendTailSizeA = sendTailSize / MESH_2D_NUM;
783 0 : uint64_t sendTailSizeB = sendTailSize - sendTailSizeA;
784 0 : auto sendTailGoSizeA = CalGoSize(sendTailSizeA);
785 0 : auto sendTailGoSizeB = CalGoSize(sendTailSizeB);
786 0 : uint64_t recvTailSizeA = recvTailSize / MESH_2D_NUM;
787 0 : uint64_t recvTailSizeB = recvTailSize - recvTailSizeA;
788 :
789 0 : uint64_t sendOffset = taskArg->localSendRecvInfo.sendOffset[i];
790 0 : uint64_t recvOffset = taskArg->localSendRecvInfo.recvOffset[i];
791 :
792 0 : processReturn.push_back(sendOffset);
793 0 : processReturn.push_back(recvOffset);
794 0 : processReturn.push_back(sendTailSizeA);
795 0 : processReturn.push_back(sendTailSizeB);
796 0 : processReturn.insert(processReturn.end(), sendTailGoSizeA.begin(), sendTailGoSizeA.end());
797 0 : processReturn.insert(processReturn.end(), sendTailGoSizeB.begin(), sendTailGoSizeB.end());
798 0 : processReturn.push_back(sendTailSize);
799 0 : processReturn.push_back(recvTailSizeA);
800 0 : processReturn.push_back(recvTailSizeB);
801 0 : processReturn.push_back(sendLoopNum);
802 0 : processReturn.push_back(recvLoopNum);
803 0 : HCCL_INFO("[CcuContextAllToAllVMesh2D][sliceInfo] curRankIdx[%u], dstrankIdx[%u]: sendOffset[%llu], "\
804 : "recvOffset[%llu], sendTailSizeA[%llu], sendTailSizeB[%llu], recvTailSizeA[%llu], recvTailSizeB[%llu],"\
805 : "sendLoopNum[%llu], recvLoopNum[%llu]", rankId_, i, sendOffset, recvOffset, sendTailSizeA, sendTailSizeB,
806 : recvTailSizeA, recvTailSizeB, sendLoopNum, recvLoopNum);
807 0 : }
808 :
809 0 : return processReturn;
810 0 : }
811 :
812 : }
|
