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_reduce_scatter_mesh2d_mem2mem.h"
12 : #include "ccu_instruction_reduce_scatter_mesh2d_mem2mem.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 INPUT_XN_ID = 1;
22 : constexpr int TOKEN_XN_ID = 2;
23 : constexpr int FST_AXIS_ID = 0;
24 : constexpr int SEC_AXIS_ID = 1;
25 : constexpr int X_AXIS_ID = 0;
26 : constexpr int Y_AXIS_ID = 1;
27 : constexpr uint64_t CCU_MS_SIZE = 4096;
28 : constexpr uint64_t LOOP_COUNT = 8;
29 : constexpr uint64_t LOCAL_COPY_MS = 8;
30 : constexpr uint32_t max_dimSize = 2;
31 :
32 0 : CcuContextReduceScatterMeshMem2Mem2D::CcuContextReduceScatterMeshMem2Mem2D(const CcuCtxArg& arg,
33 0 : const std::vector<CcuTransport*>& transports, const CcuTransportGroup& group)
34 0 : : CcuContextAlgBase(arg, transports, group)
35 : {
36 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] Enter Constructor");
37 :
38 0 : const CcuCtxArgReduceScatterMeshMem2Mem2D* ctxArg = dynamic_cast<const CcuCtxArgReduceScatterMeshMem2Mem2D*>(&arg);
39 0 : if (ctxArg == nullptr) {
40 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMeshMem2Mem2D::ctxArg ptr is null"));
41 : }
42 0 : rankId_ = ctxArg->rankId_;
43 0 : dimSize_ = ctxArg->dimSize_;
44 0 : axisId_ = ctxArg->axisId_;
45 0 : if (dimSize_.size() != max_dimSize || axisId_ > 1 || dimSize_[0] == 0) { // 2D 拓扑校验
46 0 : THROW<NullPtrException>(
47 0 : StringFormat("[CcuContextReduceScatterMeshMem2Mem2D] dimSize[%zu] or axisId[%u] or dimSize[0] [%u] is invalid",
48 0 : dimSize_.size(), axisId_, dimSize_[0]));
49 : }
50 0 : dimId_.emplace_back(rankId_ % dimSize_[0]); // 当前 rank 所在列编号, 亦即 x 方向的 localId_
51 0 : dimId_.emplace_back(rankId_ / dimSize_[0]); // 当前 rank 所在行编号, 亦即 y 方向的 localId_
52 0 : localId_ = dimId_[axisId_]; // 当前 rank 在 axisId_ 轴上的编号
53 0 : localSize_ = dimSize_[axisId_]; // mesh2d 拓扑在 axisId_ 轴上的维度
54 0 : oppsiteSize_ = dimSize_[1 - axisId_];
55 :
56 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] RankId[%u], DimSize0[%u], "
57 : "DimSize1[%u], localId[%u], lcoalSize[%u], oppsiteSize[%u]",
58 : rankId_, dimSize_[0], dimSize_[1], localId_, localSize_, oppsiteSize_);
59 :
60 0 : dataType_ = ctxArg->op_.dataType;
61 0 : outputDataType_ = ctxArg->op_.outputDataType;
62 :
63 0 : if (outputDataType_ == DataType::INVALID) {
64 0 : outputDataType_ = dataType_;
65 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] outputDataType is [INVALID], set outputDataType to[%s]",
66 : outputDataType_.Describe().c_str());
67 : }
68 :
69 0 : reduceOp_ = ctxArg->op_.reduceOp;
70 :
71 0 : localAxisSignalName_ = "CcuContextReduceScatterMeshMem2Mem2DAxisSync_" + std::to_string(axisId_);
72 0 : anotherAxisSignalName_ = "CcuContextReduceScatterMeshMem2Mem2DAxisSync_" + std::to_string(1 - axisId_);
73 :
74 0 : moConfig.loopCount = LOOP_COUNT; // loop 展开 8 次 or 16 次
75 0 : moConfig.msInterleave = LOCAL_COPY_MS; // 一个 loop 8 个 MS
76 0 : moConfig.memSlice = LOCAL_COPY_MS * CCU_MS_SIZE; // 32k
77 0 : if (moRes.executor.size() == 0) {
78 0 : moRes.executor = CreateBlockExecutor(moConfig.loopCount);
79 0 : moRes.maskSignal = CreateBlockMaskSignal(moConfig.loopCount);
80 0 : moRes.ccuBuffer = CreateBlockCcuBuffer(moConfig.loopCount * moConfig.msInterleave);
81 : }
82 0 : }
83 :
84 0 : void CcuContextReduceScatterMeshMem2Mem2D::InitResources()
85 : {
86 0 : step0BaseOffset_ = CreateVariable();
87 0 : step0AddOffset_ = CreateVariable();
88 0 : step1AddOffset_ = CreateVariable();
89 0 : yAxisOffset_ = CreateVariable();
90 0 : xAxisSize_ = CreateVariable();
91 0 : yAxisSize_ = CreateVariable();
92 0 : localAxisSignal_ = CreateMaskSignal();
93 0 : anotherAxisSignal_ = CreateMaskSignal();
94 0 : xAxisGroupOpSize_ = CreateGroupOpSize();
95 0 : yAxisGroupOpSize_ = CreateGroupOpSize();
96 0 : curGoSize_ = CreateGroupOpSize();
97 :
98 0 : ExportMaskSignal(localAxisSignal_, localAxisSignalName_);
99 0 : anotherAxisSignal_ = ImportMaskSignal(anotherAxisSignalName_);
100 :
101 0 : output_.push_back(CreateVariable());
102 :
103 0 : if (transports.size() == 0) {
104 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMeshMem2Mem2D transports is empty"));
105 : }
106 :
107 : // 从小到大遍历 transports, 遇到本 rank 就填充本地资源
108 : // 否则依次取远端资源, 要求给框架返回的 Link 同样是按顺序排列的
109 0 : uint32_t transportIdx = 0;
110 0 : for (uint64_t peerId = 0; peerId < localSize_; peerId++) {
111 0 : if (peerId == localId_) {
112 0 : input_.push_back(CreateVariable());
113 0 : token_.push_back(CreateVariable());
114 : } else {
115 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] MyRankLocalId[%u], PeerId[%llu], TransportId[%u]",
116 : localId_, peerId, transportIdx);
117 0 : CHK_PRT_RET(transports[transportIdx] == nullptr,
118 : HCCL_ERROR("[CcuContextReduceScatterMeshMem2Mem2D] Algorithm transport ptr is null"), );
119 0 : input_.push_back(CreateVariable((*transports[transportIdx]), INPUT_XN_ID));
120 0 : token_.push_back(CreateVariable((*transports[transportIdx]), TOKEN_XN_ID));
121 0 : transportIdx++;
122 : }
123 : }
124 :
125 0 : strideSize_ = CreateVariable();
126 0 : for (uint64_t i = 0; i < localSize_ - 1; i++) {
127 0 : xlocalSlice_.push_back(CreateVariable());
128 0 : ylocalSlice_.push_back(CreateVariable());
129 : }
130 :
131 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] InitResources finished");
132 : }
133 :
134 0 : void CcuContextReduceScatterMeshMem2Mem2D::AxisSync(uint32_t signalIndex)
135 : {
136 0 : const uint32_t DIE_NUM = 2;
137 0 : if (signalIndex > 1) {
138 0 : THROW<InvalidParamsException>(StringFormat("[CcuContextReduceScatterMeshMem2Mem2D] Unexpected SignalInex[%u]",
139 : signalIndex));
140 : }
141 0 : LocalCtxPost(anotherAxisSignal_, 1 << (axisId_ + signalIndex * DIE_NUM));
142 0 : LocalWait(localAxisSignal_, 1 << (1 - axisId_ + signalIndex * DIE_NUM));
143 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] AxisSync run finished");
144 0 : }
145 :
146 0 : void CcuContextReduceScatterMeshMem2Mem2D::PreSync()
147 : {
148 0 : uint16_t selfBit = 1 << localId_; // selfBit = 1*2^{localId_}
149 0 : uint16_t allBit = ((1 << localSize_) - 1) & (~(1 << localId_));
150 :
151 0 : for (auto t : transports) {
152 0 : WriteVariableWithSignal(*t, input_[localId_], INPUT_XN_ID, CKE_IDX_1, selfBit); // index = 1，传递output信息
153 0 : WriteVariableWithSignal(*t, token_[localId_], TOKEN_XN_ID, CKE_IDX_2, selfBit); // index = 2，传递token信息
154 : }
155 0 : GroupWait(*transportGroup, CKE_IDX_1, allBit); // index = 1，传递output信息
156 0 : GroupWait(*transportGroup, CKE_IDX_2, allBit); // index = 2，传递token信息
157 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] PreSync run finished");
158 0 : }
159 :
160 0 : void CcuContextReduceScatterMeshMem2Mem2D::PostSync(uint32_t signalIndex)
161 : {
162 0 : uint16_t selfBit = 1 << localId_;
163 0 : uint16_t allBit = ((1 << localSize_) - 1) & (~(1 << localId_));
164 :
165 0 : for (auto t : transports) {
166 0 : RemotePost(*t, signalIndex, selfBit);
167 : }
168 0 : GroupWait(*transportGroup, signalIndex, allBit);
169 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] PostSync run finished");
170 0 : }
171 :
172 0 : void CcuContextReduceScatterMeshMem2Mem2D::LoadArgs()
173 : {
174 0 : Load(input_[localId_]);
175 0 : Load(output_[0]);
176 0 : Load(token_[localId_]);
177 :
178 0 : Load(step0BaseOffset_);
179 0 : Load(step0AddOffset_);
180 0 : Load(step1AddOffset_);
181 :
182 0 : Load(yAxisOffset_);
183 0 : Load(xAxisSize_);
184 0 : Load(yAxisSize_);
185 :
186 0 : for (uint64_t i = 0; i < localSize_ - 1; i++) {
187 0 : Load(xlocalSlice_[i]);
188 : }
189 0 : for (uint64_t i = 0; i < localSize_ - 1; i++) {
190 0 : Load(ylocalSlice_[i]);
191 : }
192 :
193 0 : Load(xAxisGroupOpSize_);
194 0 : Load(yAxisGroupOpSize_);
195 :
196 0 : curGoSize_ = (axisId_ == 0) ? yAxisGroupOpSize_ : xAxisGroupOpSize_;
197 :
198 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] LoadArgs run finished");
199 0 : }
200 :
201 0 : void CcuContextReduceScatterMeshMem2Mem2D::CreateLocalCopyLoop()
202 : {
203 0 : std::string loopType = "reducescatter";
204 0 : if (registeredLoop.find(loopType) != registeredLoop.end()) {
205 0 : return;
206 : }
207 :
208 0 : for (uint32_t index = 0; index < 2; index++) { // 需要 2 个 Loop
209 0 : CcuRep::Memory src = CreateMemory();
210 0 : CcuRep::Memory dst = CreateMemory();
211 0 : CcuRep::Variable len = CreateVariable();
212 0 : CcuRep::LoopBlock lb(this, loopType + "_localcopy_loop_" + std::to_string(index));
213 0 : lb(src, dst, len);
214 :
215 0 : CcuRep::MaskSignal sem = moRes.maskSignal[index];
216 0 : std::vector<CcuRep::CcuBuffer> bufs;
217 0 : for (uint32_t i = 0; i < LOCAL_COPY_MS; i++) {
218 0 : bufs.push_back(moRes.ccuBuffer[i]);
219 : }
220 :
221 0 : LocalCopy(bufs[0], src, len, sem);
222 0 : LocalWait(sem);
223 0 : LocalCopy(dst, bufs[0], len, sem);
224 0 : LocalWait(sem);
225 0 : }
226 0 : registeredLoop.insert(loopType);
227 0 : }
228 :
229 0 : void CcuContextReduceScatterMeshMem2Mem2D::LocalCopyByLoopGroup(CcuRep::Memory dst, CcuRep::Memory src)
230 : {
231 0 : CreateLocalCopyLoop();
232 :
233 0 : CCU_IF(curGoSize_.addrOffset != 0) {
234 0 : CcuRep::Variable loopParam = CreateVariable();
235 0 : loopParam = CcuRep::GetLoopParam(0, moConfig.memSlice * moConfig.loopCount, 0);
236 0 : loopParam += curGoSize_.loopParam;
237 :
238 0 : CcuRep::Variable sliceSize = CreateVariable();
239 0 : sliceSize = moConfig.memSlice;
240 0 : auto lc = Loop("reducescatter_localcopy_loop_0")(src, dst, sliceSize);
241 :
242 0 : CcuRep::Variable paraCfg = CreateVariable();
243 0 : paraCfg = CcuRep::GetParallelParam(moConfig.loopCount - 1, 0, 1);
244 0 : CcuRep::Variable offsetCfg = CreateVariable();
245 0 : offsetCfg = CcuRep::GetOffsetParam(moConfig.memSlice, moConfig.msInterleave, 1);
246 0 : LoopGroup({lc}, {loopParam}, paraCfg, offsetCfg);
247 0 : }
248 :
249 0 : CCU_IF(curGoSize_.parallelParam != 0) {
250 0 : CcuRep::Condition cond(this, curGoSize_.parallelParam != 0);
251 :
252 0 : src.addr += curGoSize_.addrOffset;
253 0 : dst.addr += curGoSize_.addrOffset;
254 0 : auto lc0 = Loop("reducescatter_localcopy_loop_0")(src, dst, curGoSize_.residual);
255 :
256 0 : src.addr += curGoSize_.residual;
257 0 : dst.addr += curGoSize_.residual;
258 0 : CcuRep::Variable sliceSize = CreateVariable();
259 0 : sliceSize = moConfig.memSlice;
260 0 : auto lc1 = Loop("reducescatter_localcopy_loop_1")(src, dst, sliceSize);
261 :
262 0 : CcuRep::Variable loopCfg0 = CreateVariable();
263 0 : loopCfg0 = CcuRep::GetLoopParam(0, 0, 1);
264 0 : CcuRep::Variable loopCfg1 = CreateVariable();
265 0 : loopCfg1 = CcuRep::GetLoopParam(0, 0, 1);
266 0 : CcuRep::Variable offsetCfg = CreateVariable();
267 0 : offsetCfg = CcuRep::GetOffsetParam(moConfig.memSlice, moConfig.msInterleave, 1);
268 0 : LoopGroup({lc0, lc1}, {loopCfg0, loopCfg1}, curGoSize_.parallelParam, offsetCfg);
269 0 : }
270 0 : }
271 :
272 0 : std::vector<uint64_t> CcuContextReduceScatterMeshMem2Mem2D::CalMeshChunkSlice(uint64_t dataSize, uint64_t sliceNum)
273 : {
274 0 : uint64_t dataCount = dataSize / DataTypeSizeGet(dataType_);
275 :
276 0 : uint64_t bigDataSliceNum = dataCount % sliceNum;
277 0 : uint64_t bigDataSliceSize = (dataCount / sliceNum + 1) * DataTypeSizeGet(dataType_);
278 0 : uint64_t smallDataSliceNum = sliceNum - dataCount % sliceNum;
279 0 : uint64_t smallDataSliceSize = dataCount / sliceNum * DataTypeSizeGet(dataType_);
280 :
281 0 : return {bigDataSliceNum, bigDataSliceSize, smallDataSliceNum, smallDataSliceSize};
282 : }
283 :
284 0 : void CcuContextReduceScatterMeshMem2Mem2D::Step1Reduce()
285 : {
286 0 : std::vector<CcuRep::Memory> src;
287 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
288 0 : src.push_back(CreateMemory());
289 : }
290 0 : for (uint32_t localIdx = 0; localIdx < localSize_; localIdx++) {
291 0 : src[localIdx].addr = input_[localIdx];
292 0 : src[localIdx].token = token_[localIdx];
293 : }
294 :
295 0 : CcuRep::Memory dst = CreateMemory();
296 0 : dst.addr = input_[localId_];
297 0 : dst.token = token_[localId_];
298 :
299 0 : CcuRep::Memory tempDst = CreateMemory();
300 0 : CcuRep::Memory tempSrc = CreateMemory();
301 :
302 0 : for (uint32_t oppsiteIdx = 0; oppsiteIdx < oppsiteSize_; oppsiteIdx++) {
303 : // 由 oppsiteIdx 造成的地址偏移
304 0 : for (uint32_t localIdx = 0; localIdx < localSize_; localIdx++) {
305 0 : src[localIdx].addr += (oppsiteIdx == 0) ? step0BaseOffset_ : step0AddOffset_;
306 : }
307 0 : dst.addr += (oppsiteIdx == 0) ? step0BaseOffset_ : step0AddOffset_;
308 :
309 0 : bool isXAxis = (axisId_ == X_AXIS_ID);
310 0 : CcuRep::Variable len = isXAxis ? xAxisSize_ : yAxisSize_;
311 0 : std::vector<CcuRep::Variable> sliceSize = isXAxis ? xlocalSlice_: ylocalSlice_;
312 0 : uint16_t allBit = ((1 << localSize_) - 1) & (~(1 << localId_));
313 :
314 0 : CcuRep::MaskSignal localMask = CreateMaskSignal();
315 0 : for (uint32_t i = 0; i < localSize_ - 1; i++) {
316 0 : CcuRep::Variable sliceOffset = CreateVariable();
317 0 : CcuRep::Variable strideOffset = CreateVariable();
318 0 : sliceOffset = 0;
319 0 : strideOffset = 0;
320 0 : for (uint64_t j = 0; j < localSize_ - 1; j++) { // 遍历 rmt rank
321 0 : tempDst.addr = dst.addr; // local tempDst 初始化
322 0 : tempDst.token = dst.token;
323 0 : tempDst.addr += sliceOffset; // local tempDst 地址按照 chunk 大小偏移
324 0 : uint16_t nextNum = i + j + 1;
325 0 : if (nextNum >= localSize_) {
326 0 : nextNum += 1;
327 : }
328 0 : uint16_t rmtRank = (localId_ + nextNum) % localSize_;
329 0 : tempSrc.addr = src[rmtRank].addr;
330 0 : tempSrc.token = src[rmtRank].token;
331 0 : tempSrc.addr += sliceOffset; // rmt tempSrc 地址按照 chunk 大小偏移
332 :
333 0 : CCU_IF(sliceSize[j] == 0) {
334 0 : LocalPost(localMask, 1 << rmtRank);
335 0 : continue;
336 0 : }
337 : uint16_t rmtTransport;
338 0 : if (rmtRank < localId_) {
339 0 : rmtTransport = rmtRank;
340 : } else {
341 0 : rmtTransport = rmtRank - 1;
342 : }
343 0 : ReadReduce(*transports[rmtTransport], tempDst, tempSrc, sliceSize[j],
344 0 : dataType_, reduceOp_, localMask, 1 << rmtRank);
345 0 : sliceOffset += sliceSize[j];
346 : }
347 0 : LocalWait(localMask, allBit);
348 0 : }
349 0 : }
350 :
351 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] Step1Reduce run finished");
352 0 : }
353 :
354 0 : void CcuContextReduceScatterMeshMem2Mem2D::Step2Reduce()
355 : {
356 0 : std::vector<CcuRep::Memory> src;
357 0 : for (uint32_t rankIdx = 0; rankIdx < localSize_; rankIdx++) {
358 0 : src.push_back(CreateMemory());
359 : }
360 0 : for (uint32_t localIdx = 0; localIdx < localSize_; localIdx++) {
361 0 : src[localIdx].addr = input_[localIdx];
362 0 : src[localIdx].addr += step1AddOffset_;
363 0 : src[localIdx].token = token_[localIdx];
364 : }
365 :
366 0 : CcuRep::Memory dst = CreateMemory();
367 0 : dst.addr = output_[0];
368 0 : dst.token = token_[localId_];
369 0 : if (axisId_ == X_AXIS_ID) {
370 0 : dst.addr += yAxisOffset_;
371 : }
372 :
373 0 : CcuRep::Memory tempDst = CreateMemory();
374 0 : CcuRep::Memory tempSrc = CreateMemory();
375 :
376 0 : bool isXAxis = (axisId_ == X_AXIS_ID);
377 0 : CcuRep::Variable len = isXAxis ? yAxisSize_ : xAxisSize_;
378 0 : std::vector<CcuRep::Variable> sliceSize = isXAxis ? ylocalSlice_: xlocalSlice_;
379 0 : uint16_t allBit = ((1 << localSize_) - 1) & (~(1 << localId_));
380 :
381 0 : CcuRep::MaskSignal localMask = CreateMaskSignal();
382 0 : for (uint32_t i = 0; i < localSize_ - 1; i++) {
383 0 : CcuRep::Variable sliceOffset = CreateVariable();
384 0 : sliceOffset = 0;
385 0 : for (uint64_t j = 0; j < localSize_ - 1; j++) { // 遍历 rmt rank
386 0 : tempDst.addr = src[localId_].addr;
387 0 : tempDst.token = src[localId_].token;
388 0 : tempDst.addr += sliceOffset; // dst 地址偏移
389 0 : uint16_t nextNum = i + j + 1;
390 0 : if (nextNum >= localSize_) {
391 0 : nextNum += 1;
392 : }
393 0 : uint16_t rmtRank = (localId_ + nextNum) % localSize_;
394 0 : tempSrc.addr = src[rmtRank].addr;
395 0 : tempSrc.token = src[rmtRank].token;
396 0 : tempSrc.addr += sliceOffset; // src 地址偏移
397 :
398 0 : CCU_IF(sliceSize[j] == 0) {
399 0 : LocalPost(localMask, 1 << rmtRank);
400 0 : continue;
401 0 : }
402 : uint16_t rmtTransport;
403 0 : if (rmtRank < localId_) {
404 0 : rmtTransport = rmtRank;
405 : } else {
406 0 : rmtTransport = rmtRank - 1;
407 : }
408 0 : ReadReduce(*transports[rmtTransport], tempDst, tempSrc, sliceSize[j],
409 0 : dataType_, reduceOp_, localMask, 1 << rmtRank);
410 0 : sliceOffset += sliceSize[j];
411 : }
412 0 : LocalWait(localMask, allBit);
413 0 : }
414 :
415 0 : LocalCopyByLoopGroup(dst, src[localId_]); // 将计算结果 copy 到 output_
416 :
417 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] Step2Reduce run finished");
418 0 : }
419 :
420 0 : void CcuContextReduceScatterMeshMem2Mem2D::Algorithm()
421 : {
422 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] ReduceScatterMeshMem2Mem2D run");
423 :
424 0 : InitResources(); // 读取数据
425 0 : LoadArgs();
426 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] Algorithm first step begins");
427 0 : PreSync();
428 :
429 0 : Step1Reduce();
430 0 : PostSync(CKE_IDX_3);
431 0 : AxisSync(FST_AXIS_ID);
432 :
433 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] Algorithm second step begins");
434 0 : PostSync(CKE_IDX_4);
435 :
436 0 : Step2Reduce();
437 0 : PostSync(CKE_IDX_0);
438 0 : AxisSync(SEC_AXIS_ID);
439 :
440 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] ReduceScatterMeshMem2Mem2D end");
441 0 : return;
442 : }
443 :
444 0 : std::vector<uint64_t> CcuContextReduceScatterMeshMem2Mem2D::GeneArgs(const CcuTaskArg &arg)
445 : {
446 0 : const CcuTaskArgReduceScatterMeshMem2Mem2D* taskArg = dynamic_cast<const CcuTaskArgReduceScatterMeshMem2Mem2D*>(&arg);
447 0 : if (taskArg == nullptr) {
448 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMeshMem2Mem2D::taskArg ptr is null"));
449 : }
450 0 : uint64_t inputAddr = taskArg->inputAddr_;
451 0 : uint64_t outputAddr = taskArg->outputAddr_;
452 0 : uint64_t tokenInfo = taskArg->token_;
453 0 : uint64_t outputSize = taskArg->outputSize_;
454 0 : uint64_t offset = taskArg->offset_;
455 0 : uint64_t yAxisOffset = taskArg->xAxisSize_;
456 0 : uint64_t xAxisSize = taskArg->xAxisSize_;
457 0 : uint64_t yAxisSize = taskArg->yAxisSize_;
458 :
459 : /* @brief 计算 mesh-chunk 分块数量和数据大小
460 : * @param 总是分成 localSize_ - 1 块,
461 : * @param xlocalSlice: 把 xAxisSize 切分成 localSize_ - 1 份
462 : * @param ylocalSlice: 把 yAxisSize 切分成 localSize_ - 1 份
463 : * @vector bigDataSliceNum, bigDataSliceSize, smallDataSliceNum, smallDataSliceSize
464 : */
465 0 : std::vector<uint64_t> xlocalSlice = CalMeshChunkSlice(xAxisSize, localSize_ - 1);
466 0 : std::vector<uint64_t> ylocalSlice = CalMeshChunkSlice(yAxisSize, localSize_ - 1);
467 :
468 : // 计算不同die的数据
469 0 : uint64_t step0BaseOffset = (axisId_ == 0) ? dimId_[0] * outputSize + offset :
470 0 : dimId_[1] * dimSize_[0] * outputSize + offset + xAxisSize;
471 0 : uint64_t step0AddOffset = (axisId_ == 0) ? dimSize_[0] * outputSize : outputSize;
472 0 : uint64_t step1AddOffset = rankId_ * outputSize + offset + (axisId_ == 0 ? xAxisSize : 0);
473 0 : auto xAxisGoSize = CalGoSize(xAxisSize);
474 0 : auto yAxisGoSize = CalGoSize(yAxisSize);
475 :
476 : std::vector<uint64_t> processReturn =
477 : {inputAddr, outputAddr, tokenInfo,
478 : step0BaseOffset, step0AddOffset, step1AddOffset,
479 0 : yAxisOffset, xAxisSize, yAxisSize};
480 :
481 0 : for (uint64_t i = 0; i < xlocalSlice[0]; i++) { // bigData 块
482 0 : processReturn.push_back(xlocalSlice[1]);
483 : }
484 0 : for (uint64_t i = 0; i < xlocalSlice[2]; i++) { // smallData 块
485 0 : processReturn.push_back(xlocalSlice[3]);
486 : }
487 :
488 0 : for (uint64_t i = 0; i < ylocalSlice[0]; i++) { // bigData 块
489 0 : processReturn.push_back(ylocalSlice[1]);
490 : }
491 0 : for (uint64_t i = 0; i < ylocalSlice[2]; i++) { // smallData 块
492 0 : processReturn.push_back(ylocalSlice[3]);
493 : }
494 :
495 0 : for (auto goSize : {xAxisGoSize, yAxisGoSize}) {
496 0 : for (auto val : goSize) {
497 0 : processReturn.push_back(val);
498 : }
499 0 : }
500 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem2D] GeneArgs: inputAddr[%llu], outputAddr[%llu],"
501 : "step0BaseOffset[%llu], step0AddOffset[%llu], step1AddOffset[%llu]",
502 : inputAddr, outputAddr, step0BaseOffset, step0AddOffset, step1AddOffset);
503 :
504 0 : return processReturn;
505 0 : }
506 : } // namespace Hccl
|
