Line data Source code
1 : /**
2 : * Copyright (c) 2026 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_mesh1d_2die.h"
12 : #include "ccu_instruction_reduce_scatter_mesh1d_2die.h"
13 : #include "ccu_assist.h"
14 :
15 : namespace Hccl {
16 :
17 : constexpr int INPUT_XN_ID = 0;
18 : constexpr int TOKEN_XN_ID = 1;
19 : constexpr int CKE_IDX_0 = 0;
20 : constexpr int CKE_IDX_1 = 1;
21 : constexpr int CKE_IDX_2 = 2;
22 : constexpr int CKE_IDX_3 = 3;
23 : constexpr int LOOP_NUM = 128;
24 :
25 : constexpr int MISSION_NUM = 2;
26 : const std::string LOCAL_REDUCE_LOOP_BLOCK_TAG{"_local_reduce_loop_"};
27 :
28 0 : CcuContextReduceScatterMesh1D2Die::CcuContextReduceScatterMesh1D2Die(const CcuCtxArg &arg,
29 : const std::vector<CcuTransport *> &transports,
30 0 : const CcuTransportGroup &group)
31 0 : : CcuContextAlgBase(arg, transports, group)
32 : {
33 0 : const CcuCtxArgReduceScatterMesh1D2Die *ctxArg = dynamic_cast<const CcuCtxArgReduceScatterMesh1D2Die *>(&arg);
34 0 : if (ctxArg == nullptr) {
35 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMesh1D2Die::ctxArg ptr is null"));
36 : }
37 0 : moConfig.loopCount = LOOP_NUM;
38 :
39 0 : rmtReduceWithMyRank_ = ctxArg->rmtReduceWithMyRank_;
40 0 : myRankId_ = ctxArg->rankId_;
41 0 : rankSize_ = ctxArg->dimSize_[0];
42 :
43 0 : rmtReduceRankNum_ = transports.size() + (rmtReduceWithMyRank_ == true ? 1 : 0);
44 :
45 0 : rmtSyncMyBit_ = 1 << (myRankId_ % rmtReduceRankNum_);
46 : rmtSyncWaitBit_
47 0 : = rmtReduceWithMyRank_ ? ((1 << rmtReduceRankNum_) - 1) & (~rmtSyncMyBit_) : (1 << rmtReduceRankNum_) - 1;
48 :
49 0 : ctxName_ = ctxArg->GetCtxSignature().Describe();
50 0 : myMissionSignalName_ = ctxName_ + (rmtReduceWithMyRank_ ? "_withMyRank" : "_withoutMyRank");
51 0 : otherMissionSignalName_ = ctxName_ + (!rmtReduceWithMyRank_ ? "_withMyRank" : "_withoutMyRank");
52 :
53 0 : missionSyncMybit_ = 1 << (rmtReduceWithMyRank_ ? 1 : 0);
54 0 : missionSyncWaitBit_ = 1 << (!rmtReduceWithMyRank_ ? 1 : 0);
55 :
56 : // 数据类型处理
57 0 : dataType_ = ctxArg->op_.dataType;
58 0 : outputDataType_ = ctxArg->op_.outputDataType;
59 0 : if (outputDataType_ == DataType::INVALID) {
60 0 : outputDataType_ = dataType_;
61 0 : HCCL_INFO("[CcuContextReduceScatterMesh1D2Die] outputDataType is [INVALID], set outputDataType to[%s]",
62 : outputDataType_.Describe().c_str());
63 : }
64 0 : reduceOp_ = ctxArg->op_.reduceOp;
65 0 : }
66 :
67 0 : void CcuContextReduceScatterMesh1D2Die::InitResources()
68 : {
69 0 : moConfig.loopCount = LOOP_NUM;
70 :
71 0 : myInput_ = CreateVariable();
72 0 : myOutput_ = CreateVariable();
73 0 : myScratch_ = CreateVariable();
74 0 : myToken_ = CreateVariable();
75 :
76 0 : for (auto &t : transports) {
77 0 : peerInput_.push_back(CreateVariable(*t, INPUT_XN_ID));
78 0 : peerToken_.push_back(CreateVariable(*t, TOKEN_XN_ID));
79 : }
80 :
81 0 : sliceSize_ = CreateVariable();
82 :
83 0 : rmtReduceSliceOffset_ = CreateVariable();
84 :
85 0 : rmtReduceGoSize_ = CreateGroupOpSize();
86 :
87 0 : AllocGoResource(LOOP_NUM);
88 :
89 0 : myMissionSignal_ = CreateMaskSignal();
90 0 : ExportMaskSignal(myMissionSignal_, myMissionSignalName_);
91 0 : otherMissionSignal_ = ImportMaskSignal(otherMissionSignalName_);
92 0 : }
93 :
94 0 : void CcuContextReduceScatterMesh1D2Die::LoadArgs()
95 : {
96 0 : Load(myInput_);
97 0 : Load(myOutput_);
98 0 : Load(myToken_);
99 0 : Load(myScratch_);
100 0 : Load(sliceSize_);
101 0 : Load(rmtReduceSliceOffset_);
102 0 : Load(rmtReduceGoSize_);
103 0 : }
104 :
105 0 : void CcuContextReduceScatterMesh1D2Die::PreSync()
106 : {
107 0 : for (auto &t : transports) {
108 0 : WriteVariableWithSignal(*t, myInput_, INPUT_XN_ID, CKE_IDX_1, rmtSyncMyBit_);
109 0 : WriteVariableWithSignal(*t, myToken_, TOKEN_XN_ID, CKE_IDX_2, rmtSyncMyBit_);
110 : }
111 0 : GroupWait(*transportGroup, CKE_IDX_1, rmtSyncWaitBit_);
112 0 : GroupWait(*transportGroup, CKE_IDX_2, rmtSyncWaitBit_);
113 0 : }
114 :
115 0 : void CcuContextReduceScatterMesh1D2Die::PostSync(uint32_t signalIndex)
116 : {
117 0 : for (auto &t : transports) {
118 0 : RemotePost(*t, signalIndex, rmtSyncMyBit_);
119 : }
120 0 : GroupWait(*transportGroup, signalIndex, rmtSyncWaitBit_);
121 0 : }
122 :
123 0 : void CcuContextReduceScatterMesh1D2Die::MissionSync(uint32_t signalIndex)
124 : {
125 0 : HCCL_INFO("[CcuContextReduceScatterMesh1D2Die] MissionSync, missionSyncMybit_[%u], missionSyncWaitBit_[%u]",
126 : missionSyncMybit_, missionSyncWaitBit_);
127 0 : LocalCtxPost(otherMissionSignal_, missionSyncMybit_ << (signalIndex * MISSION_NUM));
128 0 : LocalWait(myMissionSignal_, missionSyncWaitBit_ << (signalIndex * MISSION_NUM));
129 0 : return;
130 : }
131 :
132 0 : void CcuContextReduceScatterMesh1D2Die::RmtReduce()
133 : {
134 0 : std::vector<CcuRep::Memory> src;
135 0 : src.reserve(rmtReduceRankNum_);
136 0 : for (uint32_t peerIdx = 0; peerIdx < transports.size(); peerIdx++) {
137 0 : src.push_back(CreateMemory());
138 0 : src.back().token = peerToken_[peerIdx];
139 0 : src.back().addr = peerInput_[peerIdx];
140 0 : src.back().addr += rmtReduceSliceOffset_;
141 : }
142 0 : if (rmtReduceWithMyRank_) {
143 0 : src.push_back(CreateMemory());
144 0 : src.back().token = myToken_;
145 0 : src.back().addr = myInput_;
146 0 : src.back().addr += rmtReduceSliceOffset_;
147 : }
148 :
149 0 : CcuRep::Memory dst = CreateMemory();
150 0 : dst.token = myToken_;
151 0 : dst.addr = rmtReduceWithMyRank_ ? myOutput_ : myScratch_;
152 :
153 0 : if (rmtReduceWithMyRank_) {
154 0 : GroupReduce(transports, dst, src, rmtReduceGoSize_, dataType_, outputDataType_, reduceOp_);
155 : } else {
156 0 : GroupReduceWithoutMyRank(transports, dst, src, rmtReduceGoSize_, dataType_, outputDataType_, reduceOp_);
157 : }
158 0 : }
159 :
160 0 : std::string CcuContextReduceScatterMesh1D2Die::GetLoopBlockTag(std::string loopType, int32_t index) const
161 : {
162 0 : return loopType + LOCAL_REDUCE_LOOP_BLOCK_TAG + std::to_string(index);
163 : }
164 :
165 0 : void CcuContextReduceScatterMesh1D2Die::CreateReduceLoop(uint32_t size, DataType dataType, DataType outputDataType,
166 : ReduceOp opType)
167 : {
168 0 : std::string loopType = CcuRep::GetReduceTypeStr(dataType, opType);
169 0 : loopType = "local_reduce_" + loopType;
170 0 : if (registeredLoop.find(loopType) != registeredLoop.end()) {
171 0 : return;
172 : }
173 :
174 0 : uint32_t expansionNum = CcuRep::GetReduceExpansionNum(opType, dataType, outputDataType);
175 0 : uint32_t usedBufNum = size > expansionNum ? size : expansionNum;
176 :
177 0 : for (int32_t index = 0; index < 2; index++) { // 需要实例化2个Loop
178 0 : CcuRep::Memory dst = CreateMemory();
179 0 : std::vector<CcuRep::Memory> src;
180 0 : for (uint32_t i = 0; i < size; i++) {
181 0 : src.emplace_back(CreateMemory());
182 : }
183 0 : CcuRep::Variable len = CreateVariable();
184 0 : CcuRep::Variable lenForExpansion = CreateVariable();
185 0 : CcuRep::LoopBlock lb(this, GetLoopBlockTag(loopType, index));
186 0 : lb(dst, src, len, lenForExpansion);
187 :
188 0 : std::vector<CcuRep::CcuBuffer> bufs = {moRes.ccuBuffer.begin() + index * moConfig.msInterleave,
189 0 : moRes.ccuBuffer.begin() + index * moConfig.msInterleave + usedBufNum};
190 0 : CcuRep::MaskSignal sem = moRes.maskSignal[index];
191 :
192 0 : for (uint32_t i = 0; i < size; i++) {
193 0 : LocalCopy(bufs[i], src[i], len, sem, 1 << i);
194 : }
195 0 : LocalWait(sem, (1 << size) - 1);
196 :
197 0 : if (size > 1) {
198 0 : LocalReduce(bufs, size, dataType, outputDataType, opType, sem, len);
199 0 : LocalWait(sem);
200 : }
201 :
202 0 : LocalCopy(dst, bufs[0], lenForExpansion, sem);
203 0 : LocalWait(sem);
204 0 : }
205 :
206 0 : registeredLoop.insert(loopType);
207 0 : }
208 :
209 0 : void CcuContextReduceScatterMesh1D2Die::ReduceLoopGroup(CcuRep::Memory &outDstOrg, std::vector<CcuRep::Memory> &srcOrg,
210 : GroupOpSize goSize, DataType dataType, DataType outputDataType,
211 : ReduceOp opType)
212 : {
213 0 : const uint32_t size = srcOrg.size();
214 :
215 0 : CcuRep::Memory dst = CreateMemory();
216 0 : dst = outDstOrg;
217 :
218 0 : std::vector<CcuRep::Memory> src;
219 0 : for (uint32_t idx = 0; idx < size; idx++) {
220 0 : src.push_back(CreateMemory());
221 0 : src[idx] = srcOrg[idx];
222 : }
223 :
224 0 : CreateReduceLoop(size, dataType, outputDataType, opType);
225 :
226 0 : std::string loopType = CcuRep::GetReduceTypeStr(dataType, opType);
227 0 : CcuRep::Variable sliceSizeExpansion = CreateVariable();
228 0 : loopType = "local_reduce_" + loopType;
229 0 : uint32_t expansionNum = CcuRep::GetReduceExpansionNum(opType, dataType, outputDataType);
230 0 : if (expansionNum != 1) {
231 0 : CcuRep::Variable tmp = CreateVariable();
232 0 : tmp = CcuRep::GetExpansionParam(expansionNum);
233 0 : dst.token += tmp;
234 0 : }
235 :
236 : // m部分
237 0 : CCU_IF(goSize.loopParam != 0) // goSize1
238 : {
239 0 : CcuRep::Variable loopParam = CreateVariable();
240 0 : loopParam = CcuRep::GetLoopParam(0, moConfig.memSlice * moConfig.loopCount, 0);
241 0 : loopParam += goSize.loopParam;
242 :
243 0 : CcuRep::Variable sliceSize = CreateVariable();
244 0 : sliceSize = moConfig.memSlice;
245 0 : sliceSizeExpansion = moConfig.memSlice * expansionNum;
246 0 : auto lc = Loop(GetLoopBlockTag(loopType, 0))(dst, src, sliceSize, sliceSizeExpansion);
247 :
248 0 : CcuRep::Variable paraCfg = CreateVariable();
249 0 : paraCfg = CcuRep::GetParallelParam(moConfig.loopCount - 1, 0, 1);
250 0 : CcuRep::Variable offsetCfg = CreateVariable();
251 0 : offsetCfg = CcuRep::GetOffsetParam(moConfig.memSlice, moConfig.msInterleave, 1);
252 :
253 0 : LoopGroup({lc}, {loopParam}, paraCfg, offsetCfg);
254 0 : }
255 :
256 0 : CCU_IF(goSize.parallelParam != 0) // goSize2
257 : {
258 : // p部分，加m的偏移
259 0 : for (uint32_t i = 0; i < size; i++) {
260 0 : src[i].addr += goSize.addrOffset;
261 : }
262 0 : for (uint32_t i = 0; i < expansionNum; i++) {
263 0 : dst.addr += goSize.addrOffset;
264 : }
265 :
266 0 : sliceSizeExpansion = 0;
267 0 : for (uint32_t i = 0; i < expansionNum; i++) {
268 0 : sliceSizeExpansion += goSize.residual; // goSize3
269 : }
270 :
271 0 : auto lc0 = Loop(GetLoopBlockTag(loopType, 0))(dst, src, goSize.residual, sliceSizeExpansion);
272 :
273 : // n部分，再加p的偏移
274 0 : for (uint32_t i = 0; i < size; i++) {
275 0 : src[i].addr += goSize.residual;
276 : }
277 :
278 0 : for (uint32_t i = 0; i < expansionNum; i++) {
279 0 : dst.addr += goSize.residual;
280 : }
281 :
282 0 : CcuRep::Variable sliceSize = CreateVariable();
283 0 : sliceSize = moConfig.memSlice;
284 0 : sliceSizeExpansion = moConfig.memSlice * expansionNum;
285 :
286 0 : auto lc1 = Loop(GetLoopBlockTag(loopType, 1))(dst, src, sliceSize, sliceSizeExpansion);
287 :
288 0 : CcuRep::Variable loopCfg0 = CreateVariable();
289 0 : loopCfg0 = CcuRep::GetLoopParam(0, 0, 1);
290 0 : CcuRep::Variable loopCfg1 = CreateVariable();
291 0 : loopCfg1 = CcuRep::GetLoopParam(0, 0, 1);
292 0 : CcuRep::Variable offsetCfg = CreateVariable();
293 0 : offsetCfg = CcuRep::GetOffsetParam(moConfig.memSlice, moConfig.msInterleave, 1);
294 :
295 0 : LoopGroup({lc0, lc1}, {loopCfg0, loopCfg1}, goSize.parallelParam, offsetCfg);
296 0 : }
297 0 : }
298 :
299 0 : void CcuContextReduceScatterMesh1D2Die::Algorithm()
300 : {
301 0 : InitResources();
302 0 : LoadArgs();
303 0 : PreSync();
304 0 : RmtReduce();
305 0 : PostSync(CKE_IDX_0);
306 0 : return;
307 : }
308 :
309 0 : std::vector<uint64_t> CcuContextReduceScatterMesh1D2Die::GeneArgs(const CcuTaskArg &arg)
310 : {
311 0 : const CcuTaskArgReduceScatterMesh1D2Die *taskArg = dynamic_cast<const CcuTaskArgReduceScatterMesh1D2Die *>(&arg);
312 0 : if (taskArg == nullptr) {
313 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMesh1D2Die::taskArg ptr is null"));
314 : }
315 0 : moConfig.loopCount = LOOP_NUM;
316 0 : uint64_t myInput = taskArg->inputAddr_;
317 0 : uint64_t myOutput = taskArg->outputAddr_;
318 0 : uint64_t myToken = taskArg->token_;
319 0 : uint64_t myScratch = taskArg->scratchAddr_;
320 :
321 0 : uint64_t sliceSize = taskArg->sliceSize_;
322 :
323 0 : uint64_t rmtReduceSliceOffset = sliceSize * myRankId_;
324 :
325 0 : u32 dataTypeSize = DataTypeSizeGet(dataType_);
326 :
327 0 : uint64_t localRedcueSize0 = (sliceSize / dataTypeSize) / MISSION_NUM * dataTypeSize;
328 0 : uint64_t localRedcueSize1 = sliceSize - localRedcueSize0;
329 :
330 0 : auto rmtReduceGoSize = CalGoSize(sliceSize);
331 0 : auto localReduceGoSize0 = CalGoSize(localRedcueSize0);
332 0 : auto localReduceGoSize1 = CalGoSize(localRedcueSize1);
333 :
334 0 : HCCL_INFO("[CcuContextReduceScatterMesh1D2Die][GeneArgs] myInput[%llu], myOutput[%llu], myScratch[%llu]"
335 : "rmtReduceSliceOffset[%llu], sliceSize[%llu], localRedcueSize0[%llu], localRedcueSize1[%llu]",
336 : myInput, myOutput, myScratch, rmtReduceSliceOffset, sliceSize, localRedcueSize0, localRedcueSize1);
337 :
338 : std::vector<uint64_t> taskArgs = {myInput,
339 : myOutput,
340 : myToken,
341 : myScratch,
342 : sliceSize,
343 0 : rmtReduceSliceOffset};
344 :
345 0 : for (auto &goSize : {rmtReduceGoSize}) {
346 0 : for (auto &element : goSize) {
347 0 : taskArgs.push_back(element);
348 : }
349 0 : }
350 0 : return taskArgs;
351 0 : }
352 : } // namespace Hccl
|
