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 = dynamic_cast<const CcuCtxArgReduceScatterMeshMem2Mem1D*>(&arg);
30 0 : if (ctxArg == nullptr) {
31 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMeshMem2Mem1D::ctxArg ptr is null"));
32 : }
33 0 : rankId_ = ctxArg->rankId_;
34 0 : rankSize_ = ctxArg->dimSize_[0];
35 0 : dataType_ = ctxArg->op_.dataType;
36 0 : outputDataType_ = ctxArg->op_.outputDataType;
37 0 : if (outputDataType_ == DataType::INVALID) {
38 0 : outputDataType_ = dataType_;
39 0 : HCCL_DEBUG(
40 : "[CcuContextReduceScatterMeshMem2Mem1D] outputDataType is [INVALID], set outputDataType to[%s]",
41 : outputDataType_.Describe().c_str());
42 : }
43 0 : reduceOp_ = ctxArg->op_.reduceOp;
44 0 : HCCL_INFO(
45 : "[CcuContextReduceScatterMeshMem2Mem1D] Init, CtxArgs are rankId[%u], rankSize_[%u], dataType[%s], "
46 : "outputDataType[%s], reduceOp[%s]",
47 : rankId_, rankSize_, dataType_.Describe().c_str(), outputDataType_.Describe().c_str(),
48 : reduceOp_.Describe().c_str());
49 0 : }
50 :
51 0 : void CcuContextReduceScatterMeshMem2Mem1D::InitResource()
52 : {
53 0 : uint16_t transportIdx = 0;
54 0 : if (transports.size() == 0) {
55 0 : THROW<NullPtrException>(StringFormat("CcuContextReduceScatterMeshMem2Mem1D transports is empty"));
56 : }
57 :
58 : // 按照rank号从小到大遍历transports，遇到本rank就填充本地资源，否则依次取远端资源，要求给框架返回的Link同样是按顺序排列的
59 0 : for (uint64_t peerId = 0; peerId < rankSize_; peerId++) {
60 0 : if (peerId == rankId_) {
61 0 : input_.push_back(CreateVariable());
62 0 : scratch_.push_back(CreateVariable());
63 0 : token_.push_back(CreateVariable());
64 : } else {
65 0 : HCCL_DEBUG(
66 : "[CcuContextReduceScatterMeshMem2Mem1D] MyRank[%u], PeerId[%u], TransportId[%u]", rankId_, peerId,
67 : transportIdx);
68 0 : CHK_PRT_THROW(
69 : transports[transportIdx] == nullptr,
70 : HCCL_ERROR(
71 : "[CcuContextReduceScatterMeshMem2Mem1D][InitResource] transports[%u] is nullptr", 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(
153 0 : *transports[transportId], remoteMem_[rankIdx], localMem_[rankIdx], normalSliceSize_, localSignal_,
154 0 : 1 << rankIdx);
155 0 : transportId++;
156 : }
157 : }
158 : // 等读完所有对端
159 0 : LocalWait(localSignal_, (1 << rankSize_) - 1);
160 :
161 0 : ReduceLoopGroup(outDst, localMem_[rankId_], remoteMem_, normalGoSize_, dataType_, outputDataType_, reduceOp_);
162 0 : }
163 :
164 0 : void CcuContextReduceScatterMeshMem2Mem1D::DoRepeatReduceScatter()
165 : {
166 0 : CcuRep::Variable scratchOffset = CreateVariable();
167 0 : scratchOffset = 0;
168 :
169 0 : for (uint32_t rankIdx = 0; rankIdx < rankSize_; rankIdx++) {
170 0 : localMem_[rankIdx].addr = input_[rankIdx];
171 0 : localMem_[rankIdx].addr += currentRankSliceInputOffset_;
172 0 : localMem_[rankIdx].token = token_[rankIdx];
173 :
174 0 : remoteMem_[rankIdx].addr = scratch_[rankId_];
175 0 : remoteMem_[rankIdx].addr += scratchOffset;
176 0 : scratchOffset += normalSliceSize_;
177 0 : remoteMem_[rankIdx].token = token_[rankId_];
178 : }
179 :
180 0 : CcuRep::Variable repeatNumAdd = CreateVariable();
181 0 : repeatNumAdd = 1;
182 0 : flag_ = 0;
183 0 : CCU_WHILE(repeatNum_ != UINT64_MAX)
184 : {
185 0 : repeatNum_ += repeatNumAdd;
186 0 : CCU_IF(flag_ == 1)
187 : {
188 : // 非第一轮执行时，src和dst已经初始化，需要添加偏移量
189 0 : for (auto& s : localMem_) {
190 0 : s.addr += inputRepeatStride_;
191 : }
192 0 : output_ += outputRepeatStride_;
193 0 : }
194 0 : CCU_IF(normalSliceSize_ != 0) { DoReduceScatter(); }
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(
205 : uint32_t size, DataType dataType, DataType outputDataType, 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 : std::vector<CcuRep::CcuBuffer> bufs
231 0 : = {moRes.ccuBuffer.begin() + index * moConfig.msInterleave,
232 0 : moRes.ccuBuffer.begin() + index * moConfig.msInterleave + usedBufNum};
233 0 : CcuRep::MaskSignal sem = moRes.maskSignal[index];
234 :
235 0 : for (uint32_t i = 0; i < size; i++) {
236 0 : if (i == rankId_) {
237 0 : LocalCopy(bufs[i], src, len, sem, 1 << i);
238 : } else {
239 0 : LocalCopy(bufs[i], scratch[i], len, sem, 1 << i);
240 : }
241 : }
242 0 : LocalWait(sem, (1 << size) - 1);
243 :
244 0 : if (size > 1) {
245 0 : LocalReduce(bufs, size, dataType, outputDataType, opType, sem, len);
246 0 : LocalWait(sem);
247 : }
248 :
249 0 : LocalCopy(dst, bufs[0], lenForExpansion, sem);
250 0 : LocalWait(sem);
251 0 : }
252 :
253 0 : registeredLoop.insert(loopType);
254 0 : }
255 :
256 0 : void CcuContextReduceScatterMeshMem2Mem1D::ReduceLoopGroup(
257 : CcuRep::Memory outDstOrg, CcuRep::Memory srcOrg, std::vector<CcuRep::Memory>& scratchOrg, GroupOpSize goSize,
258 : DataType dataType, DataType outputDataType, ReduceOp opType)
259 : {
260 0 : const uint32_t size = scratchOrg.size();
261 :
262 0 : CcuRep::Memory dst = CreateMemory();
263 0 : dst = outDstOrg;
264 :
265 0 : CcuRep::Memory src = CreateMemory();
266 0 : src = srcOrg;
267 :
268 0 : std::vector<CcuRep::Memory> scratch;
269 0 : for (uint32_t idx = 0; idx < size; idx++) {
270 0 : scratch.push_back(CreateMemory());
271 0 : scratch[idx] = scratchOrg[idx];
272 : }
273 :
274 0 : CreateReduceLoop(size, dataType, outputDataType, opType);
275 :
276 0 : std::string loopType = CcuRep::GetReduceTypeStr(dataType, opType);
277 0 : uint32_t expansionNum = CcuRep::GetReduceExpansionNum(opType, dataType, outputDataType);
278 0 : CcuRep::Variable sliceSizeExpansion = CreateVariable();
279 :
280 0 : if (expansionNum != 1) {
281 0 : CcuRep::Variable tmp = CreateVariable();
282 0 : tmp = CcuRep::GetExpansionParam(expansionNum);
283 0 : dst.token += tmp;
284 0 : }
285 :
286 : // m部分
287 0 : CCU_IF(goSize.loopParam != 0) // goSize1
288 : {
289 0 : CcuRep::Variable loopParam = CreateVariable();
290 0 : loopParam = CcuRep::GetLoopParam(0, moConfig.memSlice * moConfig.loopCount, 0);
291 0 : loopParam += goSize.loopParam;
292 :
293 0 : CcuRep::Variable sliceSize = CreateVariable();
294 0 : sliceSize = moConfig.memSlice;
295 0 : sliceSizeExpansion = moConfig.memSlice * expansionNum;
296 :
297 0 : auto lc = Loop(GetLoopBlockTag(loopType, 0))(dst, src, scratch, sliceSize, sliceSizeExpansion);
298 :
299 0 : CcuRep::Variable paraCfg = CreateVariable();
300 0 : paraCfg = CcuRep::GetParallelParam(moConfig.loopCount - 1, 0, 1);
301 0 : CcuRep::Variable offsetCfg = CreateVariable();
302 0 : offsetCfg = CcuRep::GetOffsetParam(moConfig.memSlice, moConfig.msInterleave, 1);
303 :
304 0 : LoopGroup({lc}, {loopParam}, paraCfg, offsetCfg);
305 0 : }
306 :
307 0 : CCU_IF(goSize.parallelParam != 0) // goSize2
308 : {
309 : // p部分，加m的偏移
310 0 : for (uint32_t i = 0; i < size; i++) {
311 0 : scratch[i].addr += goSize.addrOffset;
312 : }
313 0 : src.addr += goSize.addrOffset; // goSize0
314 0 : for (uint32_t i = 0; i < expansionNum; i++) {
315 0 : dst.addr += goSize.addrOffset;
316 : }
317 :
318 0 : sliceSizeExpansion = 0;
319 0 : for (uint32_t i = 0; i < expansionNum; i++) {
320 0 : sliceSizeExpansion += goSize.residual; // goSize3
321 : }
322 :
323 0 : auto lc0 = Loop(GetLoopBlockTag(loopType, 0))(dst, src, scratch, goSize.residual, sliceSizeExpansion);
324 :
325 : // n部分，再加p的偏移
326 0 : for (uint32_t i = 0; i < size; i++) {
327 0 : scratch[i].addr += goSize.residual;
328 : }
329 0 : src.addr += goSize.residual;
330 0 : for (uint32_t i = 0; i < expansionNum; i++) {
331 0 : dst.addr += goSize.residual;
332 : }
333 :
334 0 : CcuRep::Variable sliceSize = CreateVariable();
335 0 : sliceSize = moConfig.memSlice;
336 0 : sliceSizeExpansion = moConfig.memSlice * expansionNum;
337 :
338 0 : auto lc1 = Loop(GetLoopBlockTag(loopType, 1))(dst, src, scratch, sliceSize, sliceSizeExpansion);
339 :
340 0 : CcuRep::Variable loopCfg0 = CreateVariable();
341 0 : loopCfg0 = CcuRep::GetLoopParam(0, 0, 1);
342 0 : CcuRep::Variable loopCfg1 = CreateVariable();
343 0 : loopCfg1 = CcuRep::GetLoopParam(0, 0, 1);
344 0 : CcuRep::Variable offsetCfg = CreateVariable();
345 0 : offsetCfg = CcuRep::GetOffsetParam(moConfig.memSlice, moConfig.msInterleave, 1);
346 :
347 0 : LoopGroup({lc0, lc1}, {loopCfg0, loopCfg1}, goSize.parallelParam, offsetCfg);
348 0 : }
349 0 : }
350 :
351 0 : void CcuContextReduceScatterMeshMem2Mem1D::Algorithm()
352 : {
353 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem1D] ReduceScatterMesh1DMem2Mem run");
354 :
355 0 : InitResource();
356 :
357 0 : LoadArgs();
358 :
359 0 : PreSync();
360 :
361 0 : DoRepeatReduceScatter();
362 :
363 0 : PostSync();
364 :
365 0 : HCCL_INFO("[CcuContextReduceScatterMeshMem2Mem1D] ReduceScatterMesh1DMem2Mem end");
366 0 : return;
367 : }
368 :
369 0 : std::vector<uint64_t> CcuContextReduceScatterMeshMem2Mem1D::GeneArgs(const CcuTaskArg& arg)
370 : {
371 0 : const CcuTaskArgReduceScatterMeshMem2Mem1D* taskArg
372 0 : = dynamic_cast<const CcuTaskArgReduceScatterMeshMem2Mem1D*>(&arg);
373 0 : uint64_t inputAddr = taskArg->inputAddr_;
374 0 : uint64_t outputAddr = taskArg->outputAddr_;
375 0 : uint64_t tokenInfo = taskArg->token_;
376 0 : uint64_t scratchAddr = taskArg->scratchAddr_;
377 0 : uint64_t currentRankSliceInputOffset = taskArg->inputSliceStride_ * rankId_;
378 0 : uint64_t inputRepeatStride = taskArg->inputRepeatStride_;
379 0 : uint64_t outputRepeatStride = taskArg->outputRepeatStride_;
380 0 : uint64_t normalSliceSize = taskArg->normalSliceSize_;
381 0 : uint64_t repeatNum = taskArg->repeatNum_;
382 0 : auto normalGoSize = CalGoSize(normalSliceSize);
383 :
384 : std::vector<uint64_t> taskArgs
385 : = {inputAddr, outputAddr, tokenInfo, scratchAddr, currentRankSliceInputOffset,
386 0 : inputRepeatStride, outputRepeatStride, normalSliceSize, repeatNum};
387 :
388 0 : HCCL_INFO(
389 : "[CcuContextReduceScatterMeshMem2Mem1D] TaskArgs: inputAddr[%llu], outputAddr[%llu], "
390 : "scratchAddr[%llu], currentRankSliceInputOffset[%llu], inputRepeatStride[%llu],"
391 : "outputRepeatStride[%llu], normalSliceSize[%llu], repeatNum[%llu]",
392 : inputAddr, outputAddr, scratchAddr, currentRankSliceInputOffset, inputRepeatStride, outputRepeatStride,
393 : normalSliceSize, repeatNum);
394 :
395 0 : taskArgs.insert(taskArgs.cend(), normalGoSize.cbegin(), normalGoSize.cend());
396 :
397 0 : return taskArgs;
398 0 : }
399 :
400 : } // namespace Hccl
|
