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