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.h"
12 : #include "ccu_instruction_all_reduce_mesh2d_two_shot.h"
13 :
14 : namespace Hccl {
15 : constexpr uint32_t AXIS_NUM = 2;
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 CKE_IDX_5 = 5;
22 : constexpr int CKE_IDX_6 = 6;
23 : constexpr int INPUT_XN_ID = 0;
24 : constexpr int OUTPUT_XN_ID = 1;
25 : constexpr int TOKEN_XN_ID = 2;
26 :
27 0 : CcuContextAllReduceMesh2DTwoShot::CcuContextAllReduceMesh2DTwoShot(const CcuCtxArg &arg,
28 : const std::vector<CcuTransport *> &transports,
29 0 : const CcuTransportGroup &group)
30 0 : : CcuContextAlgBase(arg, transports, group)
31 : {
32 0 : const CcuCtxArgAllReduceMesh2DTwoShot *ctxArg = dynamic_cast<const CcuCtxArgAllReduceMesh2DTwoShot *>(&arg);
33 0 : if (ctxArg == nullptr) {
34 0 : THROW<NullPtrException>(StringFormat("CcuContextAllReduceMesh2DTwoShot::ctxArg ptr is null"));
35 : }
36 0 : dimSize_ = ctxArg->dimSize_;
37 0 : axisId_ = ctxArg->axisId_;
38 0 : rankId_ = ctxArg->rankId_;
39 0 : dataType_ = ctxArg->op_.dataType;
40 0 : outputDataType_ = ctxArg->op_.outputDataType;
41 0 : reduceOp_ = ctxArg->op_.reduceOp;
42 0 : if (outputDataType_ == DataType::INVALID) {
43 0 : outputDataType_ = dataType_;
44 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] 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("[CcuContextAllReduceMesh2DTwoShot] 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("[CcuContextAllReduceMesh2DTwoShot] 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 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] Init, CtxArgs are rankSize[%llu], dimSize0[%llu], dimSize1[%llu], axisId[%u], "
60 : "rankId[%llu], dataType[%s], outputDataType[%s], reduceOp[%s]",
61 : rankSize_, dimSize_[0], dimSize_[1], axisId_, rankId_, dataType_.Describe().c_str(),
62 : outputDataType_.Describe().c_str(), reduceOp_.Describe().c_str());
63 :
64 0 : CHK_PRT_THROW(dimSize_[0] == 0 || dimSize_[1] == 0,
65 : HCCL_ERROR("[CcuContextAllReduceMesh2DTwoShot] dimSize0[%llu] or dimSize1[%llu] is zero",
66 : dimSize_[0], dimSize_[1]),
67 : InvalidParamsException, "dimSize[0] or dimSize[1] is invalid");
68 :
69 0 : myRankIdxInAxis_.push_back(rankId_ % dimSize_[0]); // 本 rank 在第 0 维上的 index
70 0 : myRankIdxInAxis_.push_back(rankId_ / dimSize_[0]); // 本 rank 在第 1 维上的 index
71 :
72 0 : myRankIdxInCurrentAxis_ = myRankIdxInAxis_[axisId_];
73 0 : currentAxisRankSize_ = dimSize_[axisId_];
74 :
75 0 : otherAxisId_ = 1 - axisId_;
76 0 : myRankIdxInOtherAxis_ = myRankIdxInAxis_[otherAxisId_];
77 0 : otherAxisRankSize_ = dimSize_[otherAxisId_];
78 :
79 : // 同步信号初始化
80 0 : currAxisSignalName_ = "CcuContextAllReduceMesh2DTwoShotAxisSync_" + std::to_string(axisId_);
81 0 : otherAxisSignalName_ = "CcuContextAllReduceMesh2DTwoShotAxisSync_" + std::to_string(otherAxisId_);
82 0 : currAxisSignal_ = CreateMaskSignal();
83 0 : ExportMaskSignal(currAxisSignal_, currAxisSignalName_);
84 0 : otherAxisSignal_ = ImportMaskSignal(otherAxisSignalName_);
85 :
86 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] Init, myRankIdx0[%llu], myRankIdx1[%llu], "
87 : "myRankIdxInCurrentAxis[%llu], currentAxisRankSize[%llu]",
88 : myRankIdxInAxis_[0], myRankIdxInAxis_[1], myRankIdxInCurrentAxis_, currentAxisRankSize_);
89 0 : }
90 :
91 0 : void CcuContextAllReduceMesh2DTwoShot::Algorithm()
92 : {
93 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] AllReduceMesh2DTwoShot run.");
94 0 : selfBit_ = 1 << myRankIdxInCurrentAxis_;
95 0 : allBit_ = ((1 << currentAxisRankSize_) - 1) & (~(1 << myRankIdxInCurrentAxis_));
96 :
97 0 : InitVariables();
98 0 : LoadArgs();
99 0 : PreSync();
100 :
101 0 : CcuRep::Variable currOffset = CreateVariable();
102 0 : GroupOpSize currGoSize = CreateGroupOpSize();
103 :
104 : // ==== TwoShot Step1 Reduce Scatter (GroupReduce) ====
105 : // 第1步reduce的第一个数据片：本轴 MyRank * 对轴 RankSize
106 0 : uint64_t currStepStartingSliceRankIdx = myRankIdxInCurrentAxis_ * otherAxisRankSize_;
107 0 : uint64_t currStepSliceNumber = otherAxisRankSize_; // 总片数为：对轴 RankSize
108 0 : uint64_t currStepSliceType = axisId_; // 数据片为：本轴数据片
109 0 : HCCL_INFO("[Algorithm] Step1: currStepStartingSliceRankIdx[%llu], currStepSliceNumber[%llu], "
110 : "currStepSliceType[%llu]",
111 : currStepStartingSliceRankIdx, currStepSliceNumber, currStepSliceType);
112 :
113 0 : for (uint64_t currentSliceRankIdx = currStepStartingSliceRankIdx;
114 0 : currentSliceRankIdx < currStepStartingSliceRankIdx + currStepSliceNumber; currentSliceRankIdx++) {
115 0 : GetSliceOffsetAndGoSize(currentSliceRankIdx, currStepSliceType, currOffset, currGoSize);
116 0 : DoGroupReduce(inputAddr_, inputAddr_[myRankIdxInCurrentAxis_], currOffset, currGoSize);
117 : }
118 0 : SyncAll(CKE_IDX_4);
119 :
120 : // ==== TwoShot Step2 Reduce Scatter (GroupReduce) ====
121 : // reduce数据片：对轴 MyRank * 本轴 RankSize + 对轴 MyRank
122 0 : currStepStartingSliceRankIdx = myRankIdxInOtherAxis_ * currentAxisRankSize_ + myRankIdxInCurrentAxis_;
123 0 : currStepSliceNumber = 1; // 总片数为：1
124 0 : currStepSliceType = otherAxisId_; // 数据片为：对轴数据片
125 0 : HCCL_INFO("[Algorithm] Step2: currStepStartingSliceRankIdx[%llu], currStepSliceNumber[%llu], "
126 : "currStepSliceType[%llu]",
127 : currStepStartingSliceRankIdx, currStepSliceNumber, currStepSliceType);
128 :
129 0 : for (uint64_t currentSliceRankIdx = currStepStartingSliceRankIdx;
130 0 : currentSliceRankIdx < currStepStartingSliceRankIdx + currStepSliceNumber; currentSliceRankIdx++) {
131 0 : GetSliceOffsetAndGoSize(currentSliceRankIdx, currStepSliceType, currOffset, currGoSize);
132 0 : DoGroupReduce(inputAddr_, inputAddr_[myRankIdxInCurrentAxis_], currOffset, currGoSize);
133 : }
134 0 : SyncAll(CKE_IDX_5);
135 :
136 : // ==== TwoShot Step3 All Gather (GroupBroadcast) ====
137 : // Broadcast 的第一个数据片：对轴 MyRank * 本轴 RankSize + 对轴 MyRank
138 0 : currStepStartingSliceRankIdx = myRankIdxInOtherAxis_ * currentAxisRankSize_ + myRankIdxInCurrentAxis_;
139 0 : currStepSliceNumber = 1; // 总片数为：1
140 0 : currStepSliceType = otherAxisId_; // 数据片为：对轴数据片
141 0 : HCCL_INFO("[Algorithm] Step3: currStepStartingSliceRankIdx[%llu], currStepSliceNumber[%llu], "
142 : "currStepSliceType[%llu]",
143 : currStepStartingSliceRankIdx, currStepSliceNumber, currStepSliceType);
144 :
145 0 : for (uint64_t currentSliceRankIdx = currStepStartingSliceRankIdx;
146 0 : currentSliceRankIdx < currStepStartingSliceRankIdx + currStepSliceNumber; currentSliceRankIdx++) {
147 0 : GetSliceOffsetAndGoSize(currentSliceRankIdx, currStepSliceType, currOffset, currGoSize);
148 0 : DoGroupBroadcast(inputAddr_[myRankIdxInCurrentAxis_], outputAddr_, currOffset, currGoSize);
149 : }
150 0 : SyncAll(CKE_IDX_6);
151 :
152 : // ==== TwoShot Step4 All Gather (GroupBroadcast) ====
153 : // Broadcast 的第一个数据片：本轴 MyRank * 对轴 RankSize
154 0 : currStepStartingSliceRankIdx = myRankIdxInCurrentAxis_ * otherAxisRankSize_;
155 0 : currStepSliceNumber = otherAxisRankSize_; // 总片数为：对轴 RankSize
156 0 : currStepSliceType = axisId_; // 数据片为：本轴数据片
157 0 : HCCL_INFO("[Algorithm] Step4: currStepStartingSliceRankIdx[%llu], currStepSliceNumber[%llu], "
158 : "currStepSliceType[%llu]",
159 : currStepStartingSliceRankIdx, currStepSliceNumber, currStepSliceType);
160 :
161 0 : for (uint64_t currentSliceRankIdx = currStepStartingSliceRankIdx;
162 0 : currentSliceRankIdx < currStepStartingSliceRankIdx + currStepSliceNumber; currentSliceRankIdx++) {
163 0 : GetSliceOffsetAndGoSize(currentSliceRankIdx, currStepSliceType, currOffset, currGoSize);
164 0 : DoGroupBroadcast(outputAddr_[myRankIdxInCurrentAxis_], outputAddr_, currOffset, currGoSize);
165 : }
166 0 : SyncAll(CKE_IDX_0);
167 :
168 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] AllReduceMesh2DTwoShot end.");
169 0 : return;
170 0 : }
171 :
172 0 : void CcuContextAllReduceMesh2DTwoShot::GetSliceOffsetAndGoSize(uint64_t currentSliceRankIdx, uint64_t currStepSliceType,
173 : CcuRep::Variable &currOffset, GroupOpSize &currGoSize)
174 : {
175 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] GetSliceOffsetAndGoSize Starts, currentSliceRankIdx[%llu], "
176 : "currStepSliceType[%llu]",
177 : currentSliceRankIdx, currStepSliceType);
178 0 : currOffset = 0;
179 :
180 0 : CcuRep::Variable normalSliceSize = CreateVariable();
181 0 : normalSliceSize = normalRankXSliceSize_;
182 0 : normalSliceSize += normalRankYSliceSize_;
183 : // currentSliceRankIdx * normalSliceSize 是每个 rank 的 slice 的起始位置
184 0 : for (uint64_t i = 0; i < currentSliceRankIdx; i++) {
185 0 : currOffset += normalSliceSize;
186 : }
187 :
188 0 : if(currentSliceRankIdx == rankSize_ - 1) {
189 : // 最后一个rank的数据量可能会大过 normalSliceSize，因为要额外处理尾块
190 0 : if(currStepSliceType == 0) {
191 0 : HCCL_INFO("[GetSliceOffsetAndGoSize] Last Rank X Slice");
192 0 : currGoSize = lastRankXGoSize_;
193 : } else {
194 0 : HCCL_INFO("[GetSliceOffsetAndGoSize] Last Rank Y Slice");
195 : // Y 轴上需要额外添加 X 轴数据块大小的偏移
196 0 : currOffset += lastRankXSliceSize_;
197 0 : currGoSize = lastRankYGoSize_;
198 : }
199 : } else {
200 0 : if(currStepSliceType == 0) {
201 0 : HCCL_INFO("[GetSliceOffsetAndGoSize] Normal Rank X Slice");
202 0 : currGoSize = normalRankXGoSize_;
203 : } else {
204 0 : HCCL_INFO("[GetSliceOffsetAndGoSize] Normal Rank Y Slice");
205 : // Y 轴上需要额外添加 X 轴数据块大小的偏移
206 0 : currOffset += normalRankXSliceSize_;
207 0 : currGoSize = normalRankYGoSize_;
208 : }
209 : }
210 :
211 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] GetSliceOffsetAndGoSize Ends");
212 0 : return;
213 0 : }
214 :
215 0 : void CcuContextAllReduceMesh2DTwoShot::InitVariables()
216 : {
217 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] InitVariables Starts");
218 : // 初始化资源
219 0 : uint16_t transportIdx = 0;
220 0 : if (transports.size() == 0) {
221 0 : THROW<NullPtrException>(StringFormat("CcuContextAllReduceMesh2DTwoShot transports is empty"));
222 : }
223 : // 按照rank号从小到大遍历transports，遇到本rank就填充本地资源，否则依次取远端资源，要求给框架返回的Link同样是按顺序排列的
224 0 : for (uint64_t peerId = 0; peerId < currentAxisRankSize_; peerId++) {
225 0 : if (peerId == myRankIdxInCurrentAxis_) {
226 0 : inputAddr_.push_back(CreateVariable());
227 0 : outputAddr_.push_back(CreateVariable());
228 0 : token_.push_back(CreateVariable());
229 : } else {
230 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] MyRank[%u], PeerId[%llu], TransportId[%u]",
231 : myRankIdxInCurrentAxis_, peerId, transportIdx);
232 0 : CHK_PRT_RET(transports[transportIdx] == nullptr || transportIdx >= transports.size(),
233 : HCCL_ERROR("[CcuContextAllReduceMesh2DTwoShot] Algorithm transport ptr is null or transportIdx is out of bounds"), );
234 0 : inputAddr_.push_back(CreateVariable((*transports[transportIdx]), INPUT_XN_ID));
235 0 : outputAddr_.push_back(CreateVariable((*transports[transportIdx]), OUTPUT_XN_ID));
236 0 : token_.push_back(CreateVariable((*transports[transportIdx]), TOKEN_XN_ID));
237 0 : transportIdx++;
238 : }
239 : }
240 :
241 0 : lastRankXSliceSize_ = CreateVariable();
242 0 : lastRankYSliceSize_ = CreateVariable();
243 0 : normalRankXSliceSize_ = CreateVariable();
244 0 : normalRankYSliceSize_ = CreateVariable();
245 :
246 0 : normalRankXGoSize_ = CreateGroupOpSize();
247 0 : normalRankYGoSize_ = CreateGroupOpSize();
248 0 : lastRankXGoSize_ = CreateGroupOpSize();
249 0 : lastRankYGoSize_ = CreateGroupOpSize();
250 :
251 0 : for (uint64_t rankIdx = 0; rankIdx < currentAxisRankSize_; rankIdx++) {
252 0 : tmpAddrList_.push_back(CreateMemory());
253 : }
254 0 : tmpAddr_ = CreateMemory();
255 :
256 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] InitVariables Ends");
257 0 : return;
258 : }
259 :
260 0 : void CcuContextAllReduceMesh2DTwoShot::LoadArgs()
261 : {
262 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] LoadArgs Starts");
263 0 : Load(inputAddr_[myRankIdxInCurrentAxis_]);
264 0 : Load(outputAddr_[myRankIdxInCurrentAxis_]);
265 0 : Load(token_[myRankIdxInCurrentAxis_]);
266 0 : Load(normalRankXSliceSize_);
267 0 : Load(normalRankYSliceSize_);
268 0 : Load(lastRankXSliceSize_);
269 0 : Load(lastRankYSliceSize_);
270 0 : Load(normalRankXGoSize_);
271 0 : Load(normalRankYGoSize_);
272 0 : Load(lastRankXGoSize_);
273 0 : Load(lastRankYGoSize_);
274 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] LoadArgs Ends");
275 0 : return;
276 : }
277 :
278 0 : void CcuContextAllReduceMesh2DTwoShot::PreSync()
279 : {
280 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] PreSync Starts");
281 : // 前同步
282 0 : for (auto t : transports) {
283 0 : WriteVariableWithSignal(*t, inputAddr_[myRankIdxInCurrentAxis_], INPUT_XN_ID, CKE_IDX_1, selfBit_);
284 0 : WriteVariableWithSignal(*t, outputAddr_[myRankIdxInCurrentAxis_], OUTPUT_XN_ID, CKE_IDX_2, selfBit_);
285 0 : WriteVariableWithSignal(*t, token_[myRankIdxInCurrentAxis_], TOKEN_XN_ID, CKE_IDX_3, selfBit_);
286 : }
287 :
288 0 : GroupWait(*transportGroup, CKE_IDX_1, allBit_);
289 0 : GroupWait(*transportGroup, CKE_IDX_2, allBit_);
290 0 : GroupWait(*transportGroup, CKE_IDX_3, allBit_);
291 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] PreSync Ends");
292 0 : }
293 :
294 0 : void CcuContextAllReduceMesh2DTwoShot::SyncAll(int ckeIdx)
295 : {
296 0 : DoAxisSync(0);
297 0 : DoGroupSync(ckeIdx, selfBit_, allBit_);
298 0 : DoAxisSync(1);
299 0 : }
300 :
301 0 : void CcuContextAllReduceMesh2DTwoShot::DoAxisSync(uint32_t signalIdx)
302 : {
303 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] DoAxisSync Starts, signalIdx[%u]", signalIdx);
304 0 : uint32_t sendBit = 1 << axisId_;
305 0 : uint32_t waitBit = 1 << (1 - axisId_);
306 0 : sendBit = sendBit << (AXIS_NUM * signalIdx);
307 0 : waitBit = waitBit << (AXIS_NUM * signalIdx);
308 0 : LocalCtxPost(otherAxisSignal_, sendBit);
309 0 : LocalWait(currAxisSignal_, waitBit);
310 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] DoAxisSync Ends");
311 0 : return;
312 : }
313 :
314 0 : void CcuContextAllReduceMesh2DTwoShot::DoGroupSync(int ckeIdx, uint16_t selfBit, uint16_t allBit)
315 : {
316 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] DoGroupSync Starts, ckeIdx[%d], selfBit[%u], allBit[%u]", ckeIdx,
317 : selfBit, allBit);
318 0 : for (auto t : transports) {
319 0 : RemotePost(*t, ckeIdx, selfBit);
320 : }
321 0 : GroupWait(*transportGroup, ckeIdx, allBit);
322 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] DoGroupSync Ends");
323 0 : return;
324 : }
325 :
326 0 : void CcuContextAllReduceMesh2DTwoShot::DoGroupReduce(std::vector<CcuRep::Variable> &srcBase, CcuRep::Variable &dstBase,
327 : CcuRep::Variable &offset, GroupOpSize &goSize)
328 : {
329 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] DoGroupReduce Starts");
330 : // 从轴上所有的对端读取数据
331 0 : std::vector<CcuRep::Memory> &srcAddrs = tmpAddrList_;
332 0 : uint32_t curId = 0;
333 0 : uint32_t rmtId = 0;
334 0 : for (uint64_t rankIdx = 0; rankIdx < currentAxisRankSize_; rankIdx++) {
335 0 : if (rankIdx != myRankIdxInCurrentAxis_) {
336 0 : curId = rmtId;
337 0 : rmtId++;
338 : } else {
339 0 : curId = currentAxisRankSize_ - 1;
340 : }
341 0 : srcAddrs[curId].addr = srcBase[rankIdx];
342 0 : srcAddrs[curId].token = token_[rankIdx];
343 0 : srcAddrs[curId].addr += offset;
344 : }
345 : // Reduce 到本端
346 0 : CcuRep::Memory &dstAddr = tmpAddr_;
347 0 : dstAddr.addr = dstBase;
348 0 : dstAddr.addr += offset;
349 0 : dstAddr.token = token_[myRankIdxInCurrentAxis_];
350 0 : GroupReduce(transports, dstAddr, srcAddrs, goSize, dataType_, outputDataType_, reduceOp_);
351 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] DoGroupReduce Ends");
352 0 : return;
353 : }
354 :
355 0 : void CcuContextAllReduceMesh2DTwoShot::DoGroupBroadcast(CcuRep::Variable &srcBase,
356 : std::vector<CcuRep::Variable> &dstBase,
357 : CcuRep::Variable &offset, GroupOpSize &goSize)
358 : {
359 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] DoGroupBroadcast Starts");
360 : // 从轴上所有的对端读取数据
361 0 : std::vector<CcuRep::Memory> &dstAddrs = tmpAddrList_;
362 0 : uint32_t rmtId = 0;
363 0 : uint32_t curId = 0;
364 0 : for (uint64_t rankIdx = 0; rankIdx < currentAxisRankSize_; rankIdx++) {
365 0 : if (rankIdx != myRankIdxInCurrentAxis_) {
366 0 : curId = rmtId;
367 0 : rmtId++;
368 : } else {
369 0 : curId = currentAxisRankSize_ - 1;
370 : }
371 0 : dstAddrs[curId].addr = dstBase[rankIdx];
372 0 : dstAddrs[curId].addr += offset;
373 0 : dstAddrs[curId].token = token_[rankIdx];
374 : }
375 : // Reduce 到本端
376 0 : CcuRep::Memory &srcAddr = tmpAddr_;
377 0 : srcAddr.addr = srcBase;
378 0 : srcAddr.addr += offset;
379 0 : srcAddr.token = token_[myRankIdxInCurrentAxis_];
380 0 : GroupBroadcast(transports, dstAddrs, srcAddr, goSize);
381 :
382 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] DoGroupBroadcast Ends");
383 0 : return;
384 : }
385 :
386 0 : std::vector<uint64_t> CcuContextAllReduceMesh2DTwoShot::GeneArgs(const CcuTaskArg &arg)
387 : {
388 0 : HCCL_INFO("[CcuContextReduceScatterMesh2D] GeneArgs Starts");
389 0 : const CcuTaskArgAllReduceMesh2DTwoShot *taskArg = dynamic_cast<const CcuTaskArgAllReduceMesh2DTwoShot *>(&arg);
390 0 : if (taskArg == nullptr) {
391 0 : THROW<NullPtrException>(StringFormat("CcuContextAllReduceMesh2DTwoShot::taskArg ptr is null"));
392 : }
393 0 : uint64_t tokenInfo = taskArg->token_;
394 0 : uint64_t inputAddr = taskArg->inputAddr_;
395 0 : uint64_t outputAddr = taskArg->outputAddr_;
396 :
397 0 : uint64_t normalRankXSliceSize = taskArg->normalRankXSliceSize_;
398 0 : uint64_t normalRankYSliceSize = taskArg->normalRankYSliceSize_;
399 0 : uint64_t lastRankXSliceSize = taskArg->lastRankXSliceSize_;
400 0 : uint64_t lastRankYSliceSize = taskArg->lastRankYSliceSize_;
401 :
402 0 : auto normalRankXGoSize = CalGoSize(normalRankXSliceSize);
403 0 : auto normalRankYGoSize = CalGoSize(normalRankYSliceSize);
404 0 : auto lastRankXGoSize = CalGoSize(lastRankXSliceSize);
405 0 : auto lastRankYGoSize = CalGoSize(lastRankYSliceSize);
406 :
407 0 : HCCL_INFO("[CcuContextAllReduceMesh2DTwoShot] GeneArgs, TaskArgs are inputAddr[%llu], "
408 : "outputAddr[%llu], normalRankXSliceSize[%llu], normalRankYSliceSize[%llu], lastRankXSliceSize[%llu], "
409 : "lastRankYSliceSize[%llu]",
410 : inputAddr, outputAddr, normalRankXSliceSize, normalRankYSliceSize, lastRankXSliceSize,
411 : lastRankYSliceSize);
412 :
413 : std::vector<uint64_t> taskArgList{
414 : inputAddr, outputAddr, tokenInfo, normalRankXSliceSize, normalRankYSliceSize,
415 0 : lastRankXSliceSize, lastRankYSliceSize};
416 :
417 : // push goSize
418 0 : for (auto goSize : {normalRankXGoSize, normalRankYGoSize, lastRankXGoSize, lastRankYGoSize}) {
419 0 : for (auto val : goSize) {
420 0 : taskArgList.push_back(val);
421 : }
422 0 : }
423 0 : return taskArgList;
424 0 : }
425 : } // namespace Hccl
|
