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