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_mesh1d_mem2mem.h"
12 : #include "ccu_instruction_reduce_scatter_mesh1d_mem2mem.h"
13 : #include "ccu_assist.h"
14 :
15 : namespace Hccl {
16 :
17 : constexpr int INPUT_XN_ID = 0;
18 : constexpr int SCRATCH_XN_ID = 1;
19 : constexpr int TOKEN_XN_ID = 2;
20 : constexpr int CKE_IDX_0 = 0;
21 : constexpr int CKE_IDX_1 = 1;
22 : constexpr int CKE_IDX_2 = 2;
23 : constexpr int CKE_IDX_3 = 3;
24 :
25 0 : CcuContextReduceScatterMeshMem2Mem1D::CcuContextReduceScatterMeshMem2Mem1D(
26 0 : const CcuCtxArg &arg, const std::vector<CcuTransport *> &transports, const CcuTransportGroup &group)
27 0 : : CcuContextAlgBase(arg, transports, group)
28 : {
29 0 : const CcuCtxArgReduceScatterMeshMem2Mem1D *ctxArg
30 0 : = dynamic_cast<const CcuCtxArgReduceScatterMeshMem2Mem1D *>(&arg);
31 0 : if (ctxArg == nullptr) {
32 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMeshMem2Mem1D::ctxArg ptr is null"));
33 : }
34 0 : rankId_ = ctxArg->rankId_;
35 0 : rankSize_ = ctxArg->dimSize_[0];
36 0 : dataType_ = ctxArg->op_.dataType;
37 0 : outputDataType_ = ctxArg->op_.outputDataType;
38 0 : if (outputDataType_ == DataType::INVALID) {
39 0 : outputDataType_ = dataType_;
40 0 : HCCL_DEBUG(
41 : "[CcuContextReduceScatterMeshMem2Mem1D] outputDataType is [INVALID], set outputDataType to[%s]",
42 : outputDataType_.Describe().c_str());
43 : }
44 0 : reduceOp_ = ctxArg->op_.reduceOp;
45 0 : HCCL_INFO(
46 : "[CcuContextReduceScatterMeshMem2Mem1D] Init, CtxArgs are rankId[%u], rankSize_[%u], dataType[%s], "
47 : "outputDataType[%s], reduceOp[%s]",
48 : rankId_, rankSize_, dataType_.Describe().c_str(), outputDataType_.Describe().c_str(),
49 : reduceOp_.Describe().c_str());
50 0 : }
51 :
52 0 : void CcuContextReduceScatterMeshMem2Mem1D::InitResource()
53 : {
54 0 : uint16_t transportIdx = 0;
55 0 : if (transports.size() == 0) {
56 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMeshMem2Mem1D transports is empty"));
57 : }
58 :
59 : // 按照rank号从小到大遍历transports，遇到本rank就填充本地资源，否则依次取远端资源，要求给框架返回的Link同样是按顺序排列的
60 0 : for (uint64_t peerId = 0; peerId < rankSize_; peerId++) {
61 0 : if (peerId == rankId_) {
62 0 : input_.push_back(CreateVariable());
63 0 : scratch_.push_back(CreateVariable());
64 0 : token_.push_back(CreateVariable());
65 : } else {
66 0 : HCCL_DEBUG("[CcuContextReduceScatterMeshMem2Mem1D] MyRank[%u], PeerId[%u], TransportId[%u]",
67 : rankId_, peerId, transportIdx);
68 0 : CHK_PRT_THROW(
69 : transports[transportIdx] == nullptr,
70 : HCCL_ERROR("[CcuContextReduceScatterMeshMem2Mem1D][InitResource] transports[%u] is nullptr",
71 : transportIdx),
72 : NullPtrException, "transport is null");
73 :
74 0 : input_.push_back(
75 0 : CreateVariable((*transports[transportIdx]), INPUT_XN_ID)); // 获取transport中id=0的Var来传递output
76 0 : scratch_.push_back(CreateVariable((*transports[transportIdx]), SCRATCH_XN_ID));
77 0 : token_.push_back(CreateVariable((*transports[transportIdx]), TOKEN_XN_ID));
78 0 : transportIdx++;
79 : }
80 : }
81 0 : output_ = CreateVariable();
82 0 : currentRankSliceInputOffset_ = CreateVariable();
83 0 : normalSliceSize_ = CreateVariable();
84 0 : inputRepeatStride_ = CreateVariable();
85 0 : outputRepeatStride_ = CreateVariable();
86 0 : repeatNum_ = CreateVariable();
87 0 : flag_ = CreateVariable();
88 :
89 0 : normalGoSize_ = CreateGroupOpSize();
90 :
91 0 : selfBit_ = 1 << rankId_; // 仅rankid位为1，其他位为0，代表本端准备好了
92 0 : allBit_ = ((1 << rankSize_) - 1) & (~(1 << rankId_)); // 仅rankid位为0，其他位为1，代表远端准备好了
93 0 : localMem_.reserve(rankSize_);
94 0 : remoteMem_.reserve(rankSize_);
95 0 : for (uint64_t rankIdx = 0; rankIdx < rankSize_; rankIdx++) {
96 0 : remoteMem_.push_back(CreateMemory());
97 0 : localMem_.push_back(CreateMemory());
98 : }
99 :
100 0 : localSignal_ = CreateMaskSignal();
101 0 : return;
102 : }
103 :
104 0 : void CcuContextReduceScatterMeshMem2Mem1D::LoadArgs()
105 : {
106 0 : Load(input_[rankId_]);
107 0 : Load(output_);
108 0 : Load(token_[rankId_]);
109 0 : Load(scratch_[rankId_]);
110 0 : Load(currentRankSliceInputOffset_);
111 0 : Load(inputRepeatStride_);
112 0 : Load(outputRepeatStride_);
113 0 : Load(normalSliceSize_);
114 0 : Load(repeatNum_);
115 0 : Load(normalGoSize_);
116 0 : return;
117 : }
118 :
119 0 : void CcuContextReduceScatterMeshMem2Mem1D::PreSync()
120 : {
121 0 : for (auto &t : transports) {
122 0 : WriteVariableWithSignal(*t, input_[rankId_], INPUT_XN_ID, CKE_IDX_1, selfBit_); // index = 1，传递input信息
123 0 : WriteVariableWithSignal(*t, scratch_[rankId_], SCRATCH_XN_ID, CKE_IDX_2, selfBit_);
124 0 : WriteVariableWithSignal(*t, token_[rankId_], TOKEN_XN_ID, CKE_IDX_3, selfBit_);
125 : }
126 0 : GroupWait(*transportGroup, CKE_IDX_1, allBit_);
127 0 : GroupWait(*transportGroup, CKE_IDX_2, allBit_);
128 0 : GroupWait(*transportGroup, CKE_IDX_3, allBit_);
129 0 : return;
130 : }
131 :
132 0 : void CcuContextReduceScatterMeshMem2Mem1D::PostSync()
133 : {
134 0 : for (auto &t : transports) {
135 0 : RemotePost(*t, CKE_IDX_0, selfBit_);
136 : }
137 0 : GroupWait(*transportGroup, CKE_IDX_0, allBit_);
138 0 : }
139 :
140 0 : void CcuContextReduceScatterMeshMem2Mem1D::DoReduceScatter()
141 : {
142 0 : u32 transportId = 0;
143 :
144 0 : CcuRep::Memory outDst = CreateMemory();
145 0 : outDst.addr = output_;
146 0 : outDst.token = token_[rankId_];
147 :
148 0 : for (uint32_t rankIdx = 0; rankIdx < rankSize_; rankIdx++) {
149 0 : if (rankIdx == rankId_) {
150 0 : LocalPost(localSignal_, 1 << rankIdx);
151 : } else {
152 0 : Read(*transports[transportId], remoteMem_[rankIdx], localMem_[rankIdx], normalSliceSize_,
153 0 : localSignal_, 1 << rankIdx);
154 0 : transportId++;
155 : }
156 : }
157 : // 等读完所有对端
158 0 : LocalWait(localSignal_, (1 << rankSize_) - 1);
159 :
160 0 : ReduceLoopGroup(outDst, localMem_[rankId_], remoteMem_, normalGoSize_, dataType_, outputDataType_, reduceOp_);
161 0 : }
162 :
163 0 : void CcuContextReduceScatterMeshMem2Mem1D::DoRepeatReduceScatter()
164 : {
165 0 : CcuRep::Variable scratchOffset = CreateVariable();
166 0 : scratchOffset = 0;
167 :
168 0 : for (uint32_t rankIdx = 0; rankIdx < rankSize_; rankIdx++) {
169 0 : localMem_[rankIdx].addr = input_[rankIdx];
170 0 : localMem_[rankIdx].addr += currentRankSliceInputOffset_;
171 0 : localMem_[rankIdx].token = token_[rankIdx];
172 :
173 0 : remoteMem_[rankIdx].addr = scratch_[rankId_];
174 0 : remoteMem_[rankIdx].addr += scratchOffset;
175 0 : scratchOffset += normalSliceSize_;
176 0 : remoteMem_[rankIdx].token = token_[rankId_];
177 : }
178 :
179 0 : CcuRep::Variable repeatNumAdd = CreateVariable();
180 0 : repeatNumAdd = 1;
181 0 : flag_ = 0;
182 0 : CCU_WHILE(repeatNum_ != UINT64_MAX) {
183 0 : repeatNum_ += repeatNumAdd;
184 0 : CCU_IF(flag_ == 1) {
185 : // 非第一轮执行时，src和dst已经初始化，需要添加偏移量
186 0 : for (auto &s : localMem_) {
187 0 : s.addr += inputRepeatStride_;
188 : }
189 0 : output_ += outputRepeatStride_;
190 0 : }
191 0 : CCU_IF(normalSliceSize_ != 0)
192 : {
193 0 : DoReduceScatter();
194 0 : }
195 0 : flag_ = 1;
196 0 : }
197 0 : }
198 :
199 0 : std::string CcuContextReduceScatterMeshMem2Mem1D::GetLoopBlockTag(std::string loopType, int32_t index)
200 : {
201 0 : return loopType + LOOP_BLOCK_TAG + std::to_string(index);
202 : }
203 :
204 0 : void CcuContextReduceScatterMeshMem2Mem1D::CreateReduceLoop(uint32_t size, DataType dataType, DataType outputDataType,
205 : ReduceOp opType)
206 : {
207 0 : constexpr uint32_t LOOP_NUM = 16;
208 0 : AllocGoResource(LOOP_NUM);
209 :
210 0 : std::string loopType = CcuRep::GetReduceTypeStr(dataType, opType);
211 0 : if (registeredLoop.find(loopType) != registeredLoop.end()) {
212 0 : return;
213 : }
214 :
215 0 : uint32_t expansionNum = CcuRep::GetReduceExpansionNum(opType, dataType, outputDataType);
216 0 : uint32_t usedBufNum = size > expansionNum ? size : expansionNum;
217 :
218 0 : for (int32_t index = 0; index < 2; index++) { // 需要实例化2个Loop
219 0 : CcuRep::Memory dst = CreateMemory();
220 0 : CcuRep::Memory src = CreateMemory();
221 0 : std::vector<CcuRep::Memory> scratch;
222 0 : for (uint32_t i = 0; i < size; i++) {
223 0 : scratch.emplace_back(CreateMemory());
224 : }
225 0 : CcuRep::Variable len = CreateVariable();
226 0 : CcuRep::Variable lenForExpansion = CreateVariable();
227 0 : CcuRep::LoopBlock lb(this, GetLoopBlockTag(loopType, index));
228 0 : lb(dst, src, scratch, len, lenForExpansion);
229 :
230 0 : std::vector<CcuRep::CcuBuffer> bufs = {moRes.ccuBuffer.begin() + index * moConfig.msInterleave,
231 0 : moRes.ccuBuffer.begin() + index * moConfig.msInterleave + usedBufNum};
232 0 : CcuRep::MaskSignal sem = moRes.maskSignal[index];
233 :
234 0 : for (uint32_t i = 0; i < size; i++) {
235 0 : if (i == rankId_) {
236 0 : LocalCopy(bufs[i], src, len, sem, 1 << i);
237 : } else {
238 0 : LocalCopy(bufs[i], scratch[i], len, sem, 1 << i);
239 : }
240 : }
241 0 : LocalWait(sem, (1 << size) - 1);
242 :
243 0 : if (size > 1) {
244 0 : LocalReduce(bufs, size, dataType, outputDataType, opType, sem, len);
245 0 : LocalWait(sem);
246 : }
247 :
248 0 : LocalCopy(dst, bufs[0], lenForExpansion, sem);
249 0 : LocalWait(sem);
250 0 : }
251 :
252 0 : registeredLoop.insert(loopType);
253 0 : }
254 :
255 0 : void CcuContextReduceScatterMeshMem2Mem1D::ReduceLoopGroup(CcuRep::Memory outDstOrg, CcuRep::Memory srcOrg,
256 : std::vector<CcuRep::Memory> &scratchOrg, GroupOpSize goSize, DataType dataType, DataType outputDataType,
257 : ReduceOp opType)
258 : {
259 0 : const uint32_t size = scratchOrg.size();
260 :
261 0 : CcuRep::Memory dst = CreateMemory();
262 0 : dst = outDstOrg;
263 :
264 0 : CcuRep::Memory src = CreateMemory();
265 0 : src = srcOrg;
266 :
267 0 : std::vector<CcuRep::Memory> scratch;
268 0 : for (uint32_t idx = 0; idx < size; idx++) {
269 0 : scratch.push_back(CreateMemory());
270 0 : scratch[idx] = scratchOrg[idx];
271 : }
272 :
273 0 : CreateReduceLoop(size, dataType, outputDataType, opType);
274 :
275 0 : std::string loopType = CcuRep::GetReduceTypeStr(dataType, opType);
276 0 : uint32_t expansionNum = CcuRep::GetReduceExpansionNum(opType, dataType, outputDataType);
277 0 : CcuRep::Variable sliceSizeExpansion = CreateVariable();
278 :
279 0 : if (expansionNum != 1) {
280 0 : CcuRep::Variable tmp = CreateVariable();
281 0 : tmp = CcuRep::GetExpansionParam(expansionNum);
282 0 : dst.token += tmp;
283 0 : }
284 :
285 : // m部分
286 0 : CCU_IF(goSize.loopParam != 0) // goSize1
287 : {
288 0 : CcuRep::Variable loopParam = CreateVariable();
289 0 : loopParam = CcuRep::GetLoopParam(0, moConfig.memSlice * moConfig.loopCount, 0);
290 0 : loopParam += goSize.loopParam;
291 :
292 0 : CcuRep::Variable sliceSize = CreateVariable();
293 0 : sliceSize = moConfig.memSlice;
294 0 : sliceSizeExpansion = moConfig.memSlice * expansionNum;
295 :
296 0 : auto lc = Loop(GetLoopBlockTag(loopType, 0))(dst, src, scratch, sliceSize, sliceSizeExpansion);
297 :
298 0 : CcuRep::Variable paraCfg = CreateVariable();
299 0 : paraCfg = CcuRep::GetParallelParam(moConfig.loopCount - 1, 0, 1);
300 0 : CcuRep::Variable offsetCfg = CreateVariable();
301 0 : offsetCfg = CcuRep::GetOffsetParam(moConfig.memSlice, moConfig.msInterleave, 1);
302 :
303 0 : LoopGroup({lc}, {loopParam}, paraCfg, offsetCfg);
304 0 : }
305 :
306 0 : CCU_IF(goSize.parallelParam != 0) // goSize2
307 : {
308 : // p部分，加m的偏移
309 0 : for (uint32_t i = 0; i < size; i++) {
310 0 : scratch[i].addr += goSize.addrOffset;
311 : }
312 0 : src.addr += goSize.addrOffset; // goSize0
313 0 : for (uint32_t i = 0; i < expansionNum; i++) {
314 0 : dst.addr += goSize.addrOffset;
315 : }
316 :
317 0 : sliceSizeExpansion = 0;
318 0 : for (uint32_t i = 0; i < expansionNum; i++) {
319 0 : sliceSizeExpansion += goSize.residual; // goSize3
320 : }
321 :
322 0 : auto lc0 = Loop(GetLoopBlockTag(loopType, 0))(dst, src, scratch, goSize.residual, sliceSizeExpansion);
323 :
324 : // n部分，再加p的偏移
325 0 : for (uint32_t i = 0; i < size; i++) {
326 0 : scratch[i].addr += goSize.residual;
327 : }
328 0 : src.addr += goSize.residual;
329 0 : for (uint32_t i = 0; i < expansionNum; i++) {
330 0 : dst.addr += goSize.residual;
331 : }
332 :
333 0 : CcuRep::Variable sliceSize = CreateVariable();
334 0 : sliceSize = moConfig.memSlice;
335 0 : sliceSizeExpansion = moConfig.memSlice * expansionNum;
336 :
337 0 : auto lc1 = Loop(GetLoopBlockTag(loopType, 1))(dst, src, scratch, sliceSize, sliceSizeExpansion);
338 :
339 0 : CcuRep::Variable loopCfg0 = CreateVariable();
340 0 : loopCfg0 = CcuRep::GetLoopParam(0, 0, 1);
341 0 : CcuRep::Variable loopCfg1 = CreateVariable();
342 0 : loopCfg1 = CcuRep::GetLoopParam(0, 0, 1);
343 0 : CcuRep::Variable offsetCfg = CreateVariable();
344 0 : offsetCfg = CcuRep::GetOffsetParam(moConfig.memSlice, moConfig.msInterleave, 1);
345 :
346 0 : LoopGroup({lc0, lc1}, {loopCfg0, loopCfg1}, goSize.parallelParam, offsetCfg);
347 0 : }
348 0 : }
349 :
350 0 : void CcuContextReduceScatterMeshMem2Mem1D::Algorithm()
351 : {
352 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem1D] ReduceScatterMesh1DMem2Mem run");
353 :
354 0 : InitResource();
355 :
356 0 : LoadArgs();
357 :
358 0 : PreSync();
359 :
360 0 : DoRepeatReduceScatter();
361 :
362 0 : PostSync();
363 :
364 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem1D] ReduceScatterMesh1DMem2Mem end");
365 0 : return;
366 : }
367 :
368 0 : std::vector<uint64_t> CcuContextReduceScatterMeshMem2Mem1D::GeneArgs(const CcuTaskArg &arg)
369 : {
370 0 : const CcuTaskArgReduceScatterMeshMem2Mem1D *taskArg
371 0 : = dynamic_cast<const CcuTaskArgReduceScatterMeshMem2Mem1D *>(&arg);
372 0 : uint64_t inputAddr = taskArg->inputAddr_;
373 0 : uint64_t outputAddr = taskArg->outputAddr_;
374 0 : uint64_t tokenInfo = taskArg->token_;
375 0 : uint64_t scratchAddr = taskArg->scratchAddr_;
376 0 : uint64_t currentRankSliceInputOffset = taskArg->inputSliceStride_ * rankId_;
377 0 : uint64_t inputRepeatStride = taskArg->inputRepeatStride_;
378 0 : uint64_t outputRepeatStride = taskArg->outputRepeatStride_;
379 0 : uint64_t normalSliceSize = taskArg->normalSliceSize_;
380 0 : uint64_t repeatNum = taskArg->repeatNum_;
381 0 : auto normalGoSize = CalGoSize(normalSliceSize);
382 :
383 : std::vector<uint64_t> taskArgs = {
384 : inputAddr, outputAddr, tokenInfo,
385 : scratchAddr, currentRankSliceInputOffset,
386 : inputRepeatStride, outputRepeatStride, normalSliceSize,
387 : repeatNum
388 0 : };
389 :
390 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem1D] TaskArgs: inputAddr[%llu], outputAddr[%llu], "
391 : "scratchAddr[%llu], currentRankSliceInputOffset[%llu], inputRepeatStride[%llu],"
392 : "outputRepeatStride[%llu], normalSliceSize[%llu], repeatNum[%llu]",
393 : inputAddr, outputAddr, scratchAddr, currentRankSliceInputOffset,
394 : inputRepeatStride, outputRepeatStride, normalSliceSize, repeatNum);
395 :
396 0 : taskArgs.insert(taskArgs.cend(), normalGoSize.cbegin(), normalGoSize.cend());
397 :
398 0 : return taskArgs;
399 0 : }
400 :
401 : } // namespace Hccl
|
