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_reduce_mesh2d_two_shot_mem2mem.h"
12 : #include "ccu_instruction_all_reduce_mesh2d_two_shot_mem2mem.h"
13 :
14 : namespace Hccl {
15 : constexpr int INPUT_XN_ID = 0;
16 : constexpr int OUTPUT_XN_ID = 1;
17 : constexpr int TOKEN_XN_ID = 2;
18 : constexpr int CKE_IDX_0 = 0;
19 : constexpr int CKE_IDX_1 = 1;
20 : constexpr int CKE_IDX_2 = 2;
21 : constexpr int CKE_IDX_3 = 3;
22 : constexpr int CKE_IDX_4 = 4;
23 : constexpr int CKE_IDX_5 = 5;
24 : constexpr int CKE_IDX_6 = 6;
25 : constexpr uint32_t AXIS_NUM = 2;
26 0 : CcuContextAllReduceMeshTwoShotMem2Mem2D::CcuContextAllReduceMeshTwoShotMem2Mem2D(
27 0 : const CcuCtxArg &arg, const std::vector<CcuTransport *> &transports, const CcuTransportGroup &group)
28 0 : : CcuContextAlgBase(arg, transports, group)
29 : {
30 0 : const CcuCtxArgAllReduceMeshTwoShotMem2Mem2D *ctxArg
31 0 : = dynamic_cast<const CcuCtxArgAllReduceMeshTwoShotMem2Mem2D *>(&arg);
32 0 : if (ctxArg == nullptr) {
33 0 : THROW<NullPtrException>(StringFormat("CcuContextAllReduceMeshTwoShotMem2Mem2D::ctxArg ptr is null"));
34 : }
35 0 : dimSize_ = ctxArg->dimSize_;
36 0 : axisId_ = ctxArg->axisId_;
37 0 : rankId_ = ctxArg->rankId_;
38 0 : dataType_ = ctxArg->op_.dataType;
39 0 : outputDataType_ = ctxArg->op_.outputDataType;
40 0 : reduceOp_ = ctxArg->op_.reduceOp;
41 0 : if (outputDataType_ == DataType::INVALID) {
42 0 : outputDataType_ = dataType_;
43 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] outputDataType is [INVALID], set outputDataType to[%s]",
44 : outputDataType_.Describe().c_str());
45 : }
46 :
47 0 : uint32_t max_dimSize = 2;
48 0 : if (dimSize_.size() != max_dimSize or axisId_ > 1) {
49 0 : THROW<NullPtrException>(
50 0 : StringFormat("[CcuContextAllReduceMeshTwoShotMem2Mem2D] dimSize[%u] or axisId[%u] is invalid",
51 : dimSize_.size(), axisId_));
52 : }
53 0 : CHK_PRT_THROW(dimSize_[0] == 0 || dimSize_[1] == 0,
54 : HCCL_ERROR("[CcuContextAllReduceMeshTwoShotMem2Mem2D] dimSize0[%llu] or dimSize1[%llu] is zero",
55 : dimSize_[0], dimSize_[1]),
56 : InvalidParamsException, "dimSize[0] or dimSize[1] is invalid");
57 :
58 0 : rankSize_ = dimSize_[0] * dimSize_[1];
59 0 : myRankIdxInAxis_.push_back(rankId_ % dimSize_[0]); // 本 rank 在第 0 维上的 index
60 0 : myRankIdxInAxis_.push_back(rankId_ / dimSize_[0]); // 本 rank 在第 1 维上的 index
61 :
62 0 : myRankIdxInCurrentAxis_ = myRankIdxInAxis_[axisId_];
63 0 : currentAxisRankSize_ = dimSize_[axisId_];
64 :
65 0 : otherAxisId_ = 1 - axisId_;
66 0 : myRankIdxInOtherAxis_ = myRankIdxInAxis_[otherAxisId_];
67 0 : otherAxisRankSize_ = dimSize_[otherAxisId_];
68 :
69 : // 同步信号初始化
70 0 : currAxisSignalName_ = "CcuContextAllReduceMeshTwoShotMem2Mem2DAxisSync_" + std::to_string(axisId_);
71 0 : otherAxisSignalName_ = "CcuContextAllReduceMeshTwoShotMem2Mem2DAxisSync_" + std::to_string(otherAxisId_);
72 0 : currAxisSignal_ = CreateMaskSignal();
73 0 : ExportMaskSignal(currAxisSignal_, currAxisSignalName_);
74 0 : otherAxisSignal_ = ImportMaskSignal(otherAxisSignalName_);
75 0 : }
76 :
77 0 : void CcuContextAllReduceMeshTwoShotMem2Mem2D::Algorithm()
78 : {
79 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] AllReduceMeshMem2Mem2D run.");
80 0 : selfBit_ = 1 << myRankIdxInCurrentAxis_;
81 0 : allBit_ = ((1 << currentAxisRankSize_) - 1) & (~(1 << myRankIdxInCurrentAxis_));
82 :
83 0 : InitVariables();
84 0 : LoadArgs();
85 0 : PreSync();
86 :
87 : // ==== TwoShot Step1 Reduce Scatter (GroupReduce) ====
88 0 : uint64_t currStepStartingSliceRankIdx = myRankIdxInCurrentAxis_ * otherAxisRankSize_;
89 0 : uint64_t currStepSliceNumber = otherAxisRankSize_; // 总片数为：对轴 RankSize
90 0 : uint64_t currStepSliceType = axisId_; // 数据片为：本轴数据片
91 0 : for (uint64_t currentSliceRankIdx = currStepStartingSliceRankIdx;
92 0 : currentSliceRankIdx < currStepStartingSliceRankIdx + currStepSliceNumber; currentSliceRankIdx++) {
93 0 : GetSliceOffsetAndGoSize(currentSliceRankIdx, currStepSliceType);
94 0 : DoGroupReduce(inputAddr_, inputAddr_[myRankIdxInCurrentAxis_]);
95 : }
96 0 : SyncAll(CKE_IDX_4);
97 :
98 : // ==== TwoShot Step2 Reduce Scatter (GroupReduce) ====
99 0 : currStepStartingSliceRankIdx = myRankIdxInOtherAxis_ * currentAxisRankSize_ + myRankIdxInCurrentAxis_;
100 0 : currStepSliceNumber = 1; // 总片数为：1
101 0 : currStepSliceType = otherAxisId_; // 数据片为：对轴数据片
102 0 : for (uint64_t currentSliceRankIdx = currStepStartingSliceRankIdx;
103 0 : currentSliceRankIdx < currStepStartingSliceRankIdx + currStepSliceNumber; currentSliceRankIdx++) {
104 0 : GetSliceOffsetAndGoSize(currentSliceRankIdx, currStepSliceType);
105 0 : DoGroupReduce(inputAddr_, inputAddr_[myRankIdxInCurrentAxis_]);
106 : }
107 0 : SyncAll(CKE_IDX_5);
108 :
109 : // ==== TwoShot Step3 All Gather (allGatherStep) ====
110 0 : currStepStartingSliceRankIdx = myRankIdxInOtherAxis_ * currentAxisRankSize_ + myRankIdxInCurrentAxis_;
111 0 : currStepSliceNumber = 1; // 总片数为：1
112 0 : currStepSliceType = otherAxisId_; // 数据片为：对轴数据片
113 0 : for (uint64_t currentSliceRankIdx = currStepStartingSliceRankIdx;
114 0 : currentSliceRankIdx < currStepStartingSliceRankIdx + currStepSliceNumber; currentSliceRankIdx++) {
115 0 : GetSliceOffsetAndGoSize(currentSliceRankIdx, currStepSliceType);
116 0 : AllGatherStep(inputAddr_[myRankIdxInCurrentAxis_], outputAddr_);
117 : }
118 0 : SyncAll(CKE_IDX_6);
119 :
120 : // ==== TwoShot Step4 All Gather (allGatherStep) ====
121 0 : currStepStartingSliceRankIdx = myRankIdxInCurrentAxis_ * otherAxisRankSize_;
122 0 : currStepSliceNumber = otherAxisRankSize_; // 总片数为：对轴 RankSize
123 0 : currStepSliceType = axisId_; // 数据片为：本轴数据片
124 0 : for (uint64_t currentSliceRankIdx = currStepStartingSliceRankIdx;
125 0 : currentSliceRankIdx < currStepStartingSliceRankIdx + currStepSliceNumber; currentSliceRankIdx++) {
126 0 : GetSliceOffsetAndGoSize(currentSliceRankIdx, currStepSliceType);
127 0 : AllGatherStep(outputAddr_[myRankIdxInCurrentAxis_], outputAddr_);
128 : }
129 0 : SyncAll(CKE_IDX_0);
130 :
131 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] AllReduceMeshMem2Mem2D end.");
132 0 : return;
133 : }
134 :
135 0 : void CcuContextAllReduceMeshTwoShotMem2Mem2D::GetSliceOffsetAndGoSize(uint64_t currentSliceRankIdx,
136 : uint64_t currStepSliceType)
137 : {
138 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] GetSliceOffsetAndGoSize Starts, currentSliceRankIdx[%llu], "
139 : "currStepSliceType[%llu], myRankIdxInAxisX[%llu], myRankIdxInAxisY[%llu], axisId[%u]",
140 : currentSliceRankIdx, currStepSliceType, myRankIdxInAxis_[0], myRankIdxInAxis_[1], axisId_);
141 0 : curOffset_ = 0;
142 0 : CcuRep::Variable normalSliceSize = CreateVariable();
143 0 : normalSliceSize = normalRankXSliceSize_;
144 0 : normalSliceSize += normalRankYSliceSize_;
145 : // currentSliceRankIdx * normalSliceSize 是每个 rank 的 slice 的起始位置
146 0 : for (uint64_t i = 0; i < currentSliceRankIdx; i++) {
147 0 : curOffset_ += normalSliceSize;
148 : }
149 :
150 0 : if (currentSliceRankIdx == rankSize_ - 1) {
151 : // 最后一个rank的数据量可能会大过 normalSliceSize，因为要额外处理尾块
152 0 : if (currStepSliceType == 0) {
153 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D][GetSliceOffsetAndGoSize] Last Rank X Slice");
154 0 : currGoSize_ = lastRankXGoSize_;
155 0 : curSliceVec_ = lastXSlices_;
156 0 : curOffsetVec_ = lastXOffsets_;
157 0 : curSliceSize_ = lastRankXSliceSize_;
158 : } else {
159 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D][GetSliceOffsetAndGoSize] Last Rank Y Slice");
160 : // Y 轴上需要额外添加 X 轴数据块大小的偏移
161 0 : curOffset_ += lastRankXSliceSize_;
162 0 : currGoSize_ = lastRankYGoSize_;
163 0 : curSliceVec_ = lastYSlices_;
164 0 : curOffsetVec_ = lastYOffsets_;
165 0 : curSliceSize_ = lastRankYSliceSize_;
166 : }
167 : } else {
168 0 : if (currStepSliceType == 0) {
169 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D][GetSliceOffsetAndGoSize] Normal Rank X Slice");
170 0 : currGoSize_ = normalRankXGoSize_;
171 0 : curSliceVec_ = normalXSlices_;
172 0 : curOffsetVec_ = normalXOffsets_;
173 0 : curSliceSize_ = normalRankXSliceSize_;
174 : } else {
175 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D][GetSliceOffsetAndGoSize] Normal Rank Y Slice");
176 : // Y 轴上需要额外添加 X 轴数据块大小的偏移
177 0 : curOffset_ += normalRankXSliceSize_;
178 0 : currGoSize_ = normalRankYGoSize_;
179 0 : curSliceVec_ = normalYSlices_;
180 0 : curOffsetVec_ = normalYOffsets_;
181 0 : curSliceSize_ = normalRankYSliceSize_;
182 : }
183 : }
184 :
185 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] GetSliceOffsetAndGoSize Ends");
186 0 : return;
187 0 : }
188 :
189 0 : void CcuContextAllReduceMeshTwoShotMem2Mem2D::InitVariables()
190 : {
191 0 : uint16_t transportIdx = 0;
192 0 : if (transports.size() == 0) {
193 0 : THROW<NullPtrException>(StringFormat("CcuContextAllReduceMeshTwoShotMem2Mem2D transports is empty"));
194 : }
195 :
196 0 : for (uint64_t peerId = 0; peerId < currentAxisRankSize_; peerId++) {
197 0 : if (peerId == myRankIdxInCurrentAxis_) {
198 0 : inputAddr_.push_back(CreateVariable());
199 0 : outputAddr_.push_back(CreateVariable());
200 0 : token_.push_back(CreateVariable());
201 : } else {
202 0 : CHK_PRT_RET(transports[transportIdx] == nullptr || transportIdx >= transports.size(),
203 : HCCL_ERROR("[CcuContextAllReduceMeshTwoShotMem2Mem2D] Algorithm transport ptr is null or transportIdx is out of bounds"),);
204 0 : inputAddr_.push_back(CreateVariable((*transports[transportIdx]), INPUT_XN_ID));
205 0 : outputAddr_.push_back(CreateVariable((*transports[transportIdx]), OUTPUT_XN_ID));
206 0 : token_.push_back(CreateVariable((*transports[transportIdx]), TOKEN_XN_ID));
207 0 : transportIdx++;
208 : }
209 : }
210 0 : normalRankXSliceSize_ = CreateVariable();
211 0 : normalRankYSliceSize_ = CreateVariable();
212 0 : lastRankXSliceSize_ = CreateVariable();
213 0 : lastRankYSliceSize_ = CreateVariable();
214 0 : curOffset_ = CreateVariable();
215 0 : curSliceSize_ = CreateVariable();
216 0 : for (uint64_t i = 0; i < currentAxisRankSize_ - 1; i++) {
217 0 : normalXSlices_.push_back(CreateVariable());
218 0 : normalXOffsets_.push_back(CreateVariable());
219 0 : lastXSlices_.push_back(CreateVariable());
220 0 : lastXOffsets_.push_back(CreateVariable());
221 0 : normalYSlices_.push_back(CreateVariable());
222 0 : normalYOffsets_.push_back(CreateVariable());
223 0 : lastYSlices_.push_back(CreateVariable());
224 0 : lastYOffsets_.push_back(CreateVariable());
225 : }
226 :
227 0 : for (uint64_t i = 0; i < currentAxisRankSize_ - 1; i++) {
228 0 : curOffsetVec_.push_back(CreateVariable());
229 0 : curSliceVec_.push_back(CreateVariable());
230 : }
231 0 : normalRankXGoSize_ = CreateGroupOpSize();
232 0 : normalRankYGoSize_ = CreateGroupOpSize();
233 0 : lastRankXGoSize_ = CreateGroupOpSize();
234 0 : lastRankYGoSize_ = CreateGroupOpSize();
235 0 : currGoSize_ = CreateGroupOpSize();
236 0 : for (uint64_t rankIdx = 0; rankIdx < currentAxisRankSize_; rankIdx++) {
237 0 : tmpAddrList_.push_back(CreateMemory());
238 : }
239 0 : tmpAddr_ = CreateMemory();
240 0 : return;
241 : }
242 :
243 0 : void CcuContextAllReduceMeshTwoShotMem2Mem2D::LoadArgs()
244 : {
245 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] LoadArgs Starts");
246 0 : Load(inputAddr_[myRankIdxInCurrentAxis_]);
247 0 : Load(outputAddr_[myRankIdxInCurrentAxis_]);
248 0 : Load(token_[myRankIdxInCurrentAxis_]);
249 0 : Load(normalRankXSliceSize_);
250 0 : Load(normalRankYSliceSize_);
251 0 : Load(lastRankXSliceSize_);
252 0 : Load(lastRankYSliceSize_);
253 :
254 0 : for (uint64_t i = 0; i < currentAxisRankSize_ - 1; i++)
255 0 : Load(normalXSlices_[i]);
256 0 : for (uint64_t i = 0; i < currentAxisRankSize_ - 1; i++)
257 0 : Load(normalXOffsets_[i]);
258 0 : for (uint64_t i = 0; i < currentAxisRankSize_ - 1; i++)
259 0 : Load(normalYSlices_[i]);
260 0 : for (uint64_t i = 0; i < currentAxisRankSize_ - 1; i++)
261 0 : Load(normalYOffsets_[i]);
262 :
263 0 : for (uint64_t i = 0; i < currentAxisRankSize_ - 1; i++)
264 0 : Load(lastXSlices_[i]);
265 0 : for (uint64_t i = 0; i < currentAxisRankSize_ - 1; i++)
266 0 : Load(lastXOffsets_[i]);
267 0 : for (uint64_t i = 0; i < currentAxisRankSize_ - 1; i++)
268 0 : Load(lastYSlices_[i]);
269 0 : for (uint64_t i = 0; i < currentAxisRankSize_ - 1; i++)
270 0 : Load(lastYOffsets_[i]);
271 :
272 0 : Load(normalRankXGoSize_);
273 0 : Load(normalRankYGoSize_);
274 0 : Load(lastRankXGoSize_);
275 0 : Load(lastRankYGoSize_);
276 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] LoadArgs Ends");
277 0 : return;
278 : }
279 :
280 0 : void CcuContextAllReduceMeshTwoShotMem2Mem2D::PreSync()
281 : {
282 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] PreSync Starts");
283 : // 前同步
284 0 : for (auto t : transports) {
285 0 : WriteVariableWithSignal(*t, inputAddr_[myRankIdxInCurrentAxis_], INPUT_XN_ID, CKE_IDX_1, selfBit_);
286 0 : WriteVariableWithSignal(*t, outputAddr_[myRankIdxInCurrentAxis_], OUTPUT_XN_ID, CKE_IDX_2, selfBit_);
287 0 : WriteVariableWithSignal(*t, token_[myRankIdxInCurrentAxis_], TOKEN_XN_ID, CKE_IDX_3, selfBit_);
288 : }
289 :
290 0 : GroupWait(*transportGroup, CKE_IDX_1, allBit_);
291 0 : GroupWait(*transportGroup, CKE_IDX_2, allBit_);
292 0 : GroupWait(*transportGroup, CKE_IDX_3, allBit_);
293 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] PreSync Ends");
294 0 : }
295 :
296 0 : void CcuContextAllReduceMeshTwoShotMem2Mem2D::SyncAll(int ckeIdx)
297 : {
298 0 : DoAxisSync(0);
299 0 : DoGroupSync(ckeIdx, selfBit_, allBit_);
300 0 : DoAxisSync(1);
301 0 : }
302 :
303 0 : void CcuContextAllReduceMeshTwoShotMem2Mem2D::DoAxisSync(uint32_t signalIdx)
304 : {
305 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] DoAxisSync Starts, signalIdx[%u]", signalIdx);
306 0 : uint32_t sendBit = 1 << axisId_;
307 0 : uint32_t waitBit = 1 << (1 - axisId_);
308 0 : sendBit = sendBit << (AXIS_NUM * signalIdx);
309 0 : waitBit = waitBit << (AXIS_NUM * signalIdx);
310 0 : LocalCtxPost(otherAxisSignal_, sendBit);
311 0 : LocalWait(currAxisSignal_, waitBit);
312 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] DoAxisSync Ends");
313 0 : return;
314 : }
315 :
316 0 : void CcuContextAllReduceMeshTwoShotMem2Mem2D::DoGroupSync(int ckeIdx, uint16_t selfBit, uint16_t allBit)
317 : {
318 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] DoGroupSync Starts, ckeIdx[%d], selfBit[%u], allBit[%u]",
319 : ckeIdx, selfBit, allBit);
320 0 : for (auto t : transports) {
321 0 : RemotePost(*t, ckeIdx, selfBit);
322 : }
323 0 : GroupWait(*transportGroup, ckeIdx, allBit);
324 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] DoGroupSync Ends");
325 0 : return;
326 : }
327 :
328 0 : void CcuContextAllReduceMeshTwoShotMem2Mem2D::DoGroupReduce(std::vector<CcuRep::Variable> &srcAddr,
329 : CcuRep::Variable &dstAddr)
330 : {
331 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] DoGroupReduce starts");
332 0 : uint16_t allBit = ((1 << currentAxisRankSize_) - 1) & (~(1 << myRankIdxInCurrentAxis_));
333 0 : std::vector<CcuRep::Memory> &src = tmpAddrList_;
334 0 : CcuRep::Memory &dst = tmpAddr_;
335 :
336 0 : dst.token = token_[myRankIdxInCurrentAxis_];
337 0 : for (uint64_t rankIdx = 0; rankIdx < currentAxisRankSize_; rankIdx++) {
338 0 : src[rankIdx].token = token_[rankIdx];
339 : }
340 :
341 0 : CcuRep::MaskSignal locMask = CreateMaskSignal();
342 0 : for (uint64_t i = 0; i < (currentAxisRankSize_ - 1); i++) {
343 0 : for (uint64_t j = 0; j < (currentAxisRankSize_ - 1); j++) {
344 0 : uint16_t nextNum = i + j + 1;
345 0 : if (nextNum >= currentAxisRankSize_) {
346 0 : nextNum += 1;
347 : }
348 0 : uint16_t rmtRank = (myRankIdxInCurrentAxis_ + nextNum) % currentAxisRankSize_;
349 : uint16_t rmtTransport;
350 0 : if (rmtRank < myRankIdxInCurrentAxis_) {
351 0 : rmtTransport = rmtRank;
352 : } else {
353 0 : rmtTransport = rmtRank - 1;
354 : }
355 :
356 0 : dst.addr = dstAddr;
357 0 : src[rmtRank].addr = srcAddr[rmtRank];
358 0 : dst.addr += curOffset_;
359 0 : src[rmtRank].addr += curOffset_;
360 0 : dst.addr += curOffsetVec_[j];
361 0 : src[rmtRank].addr += curOffsetVec_[j];
362 0 : CCU_IF(curSliceVec_[j] == 0)
363 : {
364 0 : LocalPost(locMask, (1 << rmtRank));
365 0 : }
366 0 : CCU_IF(curSliceVec_[j] != 0)
367 : {
368 0 : ReadReduce(*transports[rmtTransport], dst, src[rmtRank], curSliceVec_[j], dataType_, reduceOp_, locMask,
369 0 : 1 << rmtRank);
370 0 : }
371 : }
372 0 : LocalWait(locMask, allBit);
373 : }
374 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] DoGroupReduce end");
375 0 : }
376 :
377 0 : void CcuContextAllReduceMeshTwoShotMem2Mem2D::AllGatherStep(CcuRep::Variable &srcAddr,
378 : std::vector<CcuRep::Variable> &dstAddr)
379 : {
380 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] AllGatherStep Starts");
381 0 : CcuRep::Memory &src = tmpAddr_;
382 0 : std::vector<CcuRep::Memory> &dst = tmpAddrList_;
383 0 : src.addr = srcAddr;
384 0 : src.addr += curOffset_;
385 0 : src.token = token_[myRankIdxInCurrentAxis_];
386 0 : CCU_IF(curSliceSize_ != 0)
387 : {
388 0 : uint32_t transportId = 0;
389 0 : CcuRep::MaskSignal locMask = CreateMaskSignal();
390 0 : for (uint64_t rankIdx = 0; rankIdx < currentAxisRankSize_; rankIdx++) {
391 0 : dst[rankIdx].addr = dstAddr[rankIdx];
392 0 : dst[rankIdx].addr += curOffset_;
393 0 : dst[rankIdx].token = token_[rankIdx];
394 :
395 0 : if (rankIdx == myRankIdxInCurrentAxis_) {
396 0 : LocalPost(locMask, (1 << rankIdx));
397 : } else {
398 0 : Write(*transports[transportId], dst[rankIdx], src, curSliceSize_, locMask, 1 << rankIdx);
399 0 : transportId++;
400 : }
401 : }
402 0 : GroupCopy(dst[myRankIdxInCurrentAxis_], src, currGoSize_);
403 0 : LocalWait(locMask, (1 << currentAxisRankSize_) - 1);
404 0 : }
405 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] AllGatherStep end");
406 0 : }
407 :
408 0 : void CcuContextAllReduceMeshTwoShotMem2Mem2D::CalMeshChunkSlices(uint64_t totalSize, uint64_t sliceNum,
409 : std::vector<uint64_t> &slices,
410 : std::vector<uint64_t> &offsets)
411 : {
412 0 : if (sliceNum == 0) {
413 0 : THROW<InvalidParamsException>(StringFormat(
414 : "[CcuContextAllReduceMeshTwoShotMem2Mem2D][CalMeshChunkSlices] Invalid sliceNum [%u] .", sliceNum));
415 : }
416 0 : uint64_t totalCount = totalSize / DataTypeSizeGet(dataType_);
417 0 : uint64_t bigNum = totalCount % sliceNum;
418 0 : uint64_t bigSize = (totalCount / sliceNum + 1) * DataTypeSizeGet(dataType_);
419 0 : uint64_t smallSize = (totalCount / sliceNum) * DataTypeSizeGet(dataType_);
420 :
421 : // 计算每个分片的大小和偏移量
422 0 : uint64_t currentOffset = 0;
423 0 : for (uint64_t i = 0; i < sliceNum; ++i) {
424 0 : uint64_t chunkSize = 0;
425 0 : if (i < bigNum) {
426 0 : chunkSize = bigSize;
427 : } else {
428 0 : chunkSize = smallSize;
429 : }
430 0 : slices.push_back(chunkSize);
431 0 : offsets.push_back(currentOffset);
432 0 : currentOffset += chunkSize;
433 : }
434 0 : }
435 :
436 0 : std::vector<uint64_t> CcuContextAllReduceMeshTwoShotMem2Mem2D::GeneArgs(const CcuTaskArg &arg)
437 : {
438 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] GeneArgs Starts");
439 0 : const CcuTaskArgAllReduceMeshTwoShotMem2Mem2D *taskArg
440 0 : = dynamic_cast<const CcuTaskArgAllReduceMeshTwoShotMem2Mem2D *>(&arg);
441 0 : if (taskArg == nullptr) {
442 0 : THROW<NullPtrException>(StringFormat("CcuContextAllReduceMeshTwoShotMem2Mem2D::taskArg ptr is null"));
443 : }
444 0 : uint64_t tokenInfo = taskArg->token_;
445 0 : uint64_t inputAddr = taskArg->inputAddr_;
446 0 : uint64_t outputAddr = taskArg->outputAddr_;
447 :
448 0 : uint64_t normalRankXSliceSize = taskArg->normalRankXSliceSize_;
449 0 : uint64_t normalRankYSliceSize = taskArg->normalRankYSliceSize_;
450 0 : std::vector<uint64_t> normalXSlices{};
451 0 : std::vector<uint64_t> normalXOffsets{};
452 0 : std::vector<uint64_t> normalYSlices{};
453 0 : std::vector<uint64_t> normalYOffsets{};
454 0 : CalMeshChunkSlices(normalRankXSliceSize, currentAxisRankSize_ - 1, normalXSlices, normalXOffsets);
455 0 : CalMeshChunkSlices(normalRankYSliceSize, currentAxisRankSize_ - 1, normalYSlices, normalYOffsets);
456 :
457 0 : uint64_t lastRankXSliceSize = taskArg->lastRankXSliceSize_;
458 0 : uint64_t lastRankYSliceSize = taskArg->lastRankYSliceSize_;
459 0 : std::vector<uint64_t> lastXSlices{};
460 0 : std::vector<uint64_t> lastXOffsets{};
461 0 : std::vector<uint64_t> lastYSlices{};
462 0 : std::vector<uint64_t> lastYOffsets{};
463 0 : CalMeshChunkSlices(lastRankXSliceSize, currentAxisRankSize_ - 1, lastXSlices, lastXOffsets);
464 0 : CalMeshChunkSlices(lastRankYSliceSize, currentAxisRankSize_ - 1, lastYSlices, lastYOffsets);
465 :
466 0 : auto normalRankXGoSize = CalGoSize(normalRankXSliceSize);
467 0 : auto normalRankYGoSize = CalGoSize(normalRankYSliceSize);
468 0 : auto lastRankXGoSize = CalGoSize(lastRankXSliceSize);
469 0 : auto lastRankYGoSize = CalGoSize(lastRankYSliceSize);
470 :
471 0 : HCCL_INFO("[CcuContextAllReduceMeshTwoShotMem2Mem2D] GeneArgs, TaskArgs are inputAddr[%llu], "
472 : "outputAddr[%llu], normalRankXSliceSize[%llu], normalRankYSliceSize[%llu], lastRankXSliceSize[%llu], "
473 : "lastRankYSliceSize[%llu]", inputAddr, outputAddr, normalRankXSliceSize, normalRankYSliceSize,
474 : lastRankXSliceSize, lastRankYSliceSize);
475 :
476 : std::vector<uint64_t> taskArgList{
477 : inputAddr, outputAddr, tokenInfo, normalRankXSliceSize, normalRankYSliceSize,
478 0 : lastRankXSliceSize, lastRankYSliceSize};
479 :
480 0 : for (const auto &vec : {normalXSlices, normalXOffsets, normalYSlices, normalYOffsets, lastXSlices, lastXOffsets,
481 0 : lastYSlices, lastYOffsets}) {
482 0 : for (auto val : vec) {
483 0 : taskArgList.push_back(val);
484 : }
485 0 : }
486 :
487 : // push goSize
488 0 : for (auto goSize : {normalRankXGoSize, normalRankYGoSize, lastRankXGoSize, lastRankYGoSize}) {
489 0 : for (auto val : goSize) {
490 0 : taskArgList.push_back(val);
491 : }
492 0 : }
493 0 : return taskArgList;
494 0 : }
495 : } // namespace Hccl
|
