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