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