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