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 <cmath>
12 : #include "alg_template_register.h"
13 : #include "reduce_scatter_plant_local_reduce_combine.h"
14 :
15 : namespace hccl {
16 : constexpr u32 DEVICE_EIGHT = 8;
17 : constexpr u32 FACTOR_NUM_TWO = 2;
18 0 : ReduceScatterPlantLocalReduceCombine::ReduceScatterPlantLocalReduceCombine(const HcclDispatcher dispatcher)
19 0 : : AlgTemplateBase(dispatcher)
20 0 : {}
21 :
22 0 : ReduceScatterPlantLocalReduceCombine::~ReduceScatterPlantLocalReduceCombine() {}
23 :
24 0 : HcclResult ReduceScatterPlantLocalReduceCombine::Prepare(
25 : DeviceMem& cclInMem, DeviceMem& outputMem, const Stream& stream, std::vector<Stream>& subStreams,
26 : std::vector<std::shared_ptr<LocalNotify>>& meshSignal, std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux,
27 : MemBlockInfo& memBlockInfo, const HcclReduceOp reductionOp, const HcclDataType dataType, bool isUseCclIn,
28 : bool isLevel0LastRank, bool isNeedSpaceBorrow)
29 : {
30 0 : inputMem_ = cclInMem; // 空拷贝 & 存放最后一块数据（Allreduce非整除场景）
31 0 : outputMem_ = outputMem; // 单算子CclOut 图模式Scrach/UserOut，LocalReduce使用
32 0 : stream_ = stream;
33 0 : subStreams_ = subStreams;
34 0 : meshSignalPtr_ = &meshSignal;
35 0 : meshSignalAuxPtr_ = &meshSignalAux;
36 0 : memBlockInfo_ = std::move(memBlockInfo);
37 0 : reductionOp_ = reductionOp;
38 0 : dataType_ = dataType;
39 0 : isUseCclIn_ = isUseCclIn; // 本卡在level0执行完毕后,需要告知level1数据是否存放被存放在CCLin的标识(rank维度)
40 0 : isLevel0LastRank_ = isLevel0LastRank;
41 0 : isNeedSpaceBorrow_ = isNeedSpaceBorrow; // 是否需要借用CCLIN空间完成LocalReuce\alltoall(算子维度)
42 0 : return HCCL_SUCCESS;
43 : }
44 :
45 : HcclResult
46 0 : ReduceScatterPlantLocalReduceCombine::MainRecordSub(Stream& mainStream, u32 firstSubStreamIndex, u32 totalTask)
47 : {
48 0 : for (u32 streamIndex = firstSubStreamIndex; streamIndex < totalTask; streamIndex++) {
49 0 : CHK_RET(LocalNotify::Post(mainStream, dispatcher_, (*meshSignalAuxPtr_)[streamIndex], profilerInput_.stage));
50 : }
51 0 : return HCCL_SUCCESS;
52 : }
53 :
54 0 : HcclResult ReduceScatterPlantLocalReduceCombine::SubWaitMain(u32 firstSubStreamIndex, u32 totalTask)
55 : {
56 0 : for (u32 streamIndex = firstSubStreamIndex; streamIndex < totalTask; streamIndex++) {
57 0 : CHK_RET(LocalNotify::Wait(
58 : subStreams_[streamIndex], dispatcher_, (*meshSignalAuxPtr_)[streamIndex], profilerInput_.stage));
59 : }
60 0 : return HCCL_SUCCESS;
61 : }
62 :
63 0 : HcclResult ReduceScatterPlantLocalReduceCombine::MainWaitSub(Stream& mainStream, u32 firstSubStreamIndex, u32 totalTask)
64 : {
65 0 : for (u32 streamIndex = firstSubStreamIndex; streamIndex < totalTask; streamIndex++) {
66 0 : CHK_RET(LocalNotify::Wait(mainStream, dispatcher_, (*meshSignalPtr_)[streamIndex], profilerInput_.stage));
67 : }
68 0 : return HCCL_SUCCESS;
69 : }
70 :
71 0 : HcclResult ReduceScatterPlantLocalReduceCombine::SubRecordMain(u32 firstSubStreamIndex, u32 totalTask)
72 : {
73 0 : for (u32 streamIndex = firstSubStreamIndex; streamIndex < totalTask; streamIndex++) {
74 0 : CHK_RET(LocalNotify::Post(
75 : subStreams_[streamIndex], dispatcher_, (*meshSignalPtr_)[streamIndex], profilerInput_.stage));
76 : }
77 0 : return HCCL_SUCCESS;
78 : }
79 :
80 0 : u32 ReduceScatterPlantLocalReduceCombine::CalcOutputIndex(const u32 round) const
81 : {
82 0 : return (round + localRank_) % rankSize_;
83 : }
84 :
85 0 : bool ReduceScatterPlantLocalReduceCombine::isLastRank(const u32 rankId) const { return rankId == rankSize_ - 1; }
86 :
87 0 : bool ReduceScatterPlantLocalReduceCombine::isLastBlockData(const u32 outputIndex)
88 : {
89 0 : return outputIndex == rankSize_ - 1;
90 : }
91 :
92 : HcclResult
93 0 : ReduceScatterPlantLocalReduceCombine::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
94 : {
95 0 : HCCL_INFO(
96 : "ReduceScatterPlantLocalReduceCombine run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p].", rank, rankSize,
97 : inputMem_.ptr(), outputMem_.ptr());
98 0 : CHK_SMART_PTR_NULL(dispatcher_);
99 0 : CHK_PTR_NULL(stream_.ptr());
100 0 : CHK_PRT_RET(
101 : links.size() < rankSize,
102 : HCCL_ERROR(
103 : "[ReduceScatterPlantLocalReduceCombine][RunAsync]rank[%u] "
104 : "linksize[%llu] is less than rankSize[%u]",
105 : rank, links.size(), rankSize),
106 : HCCL_E_INTERNAL);
107 :
108 0 : rankSize_ = rankSize;
109 0 : localRank_ = rank;
110 :
111 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
112 0 : CHK_RET(RunAlltoAll(links));
113 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
114 :
115 : // 执行LocalReduce
116 0 : HcclResult ret = RunLocalReduce();
117 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[%s]localRank[%u] LocalReduce failed", __func__, localRank_), ret);
118 0 : return HCCL_SUCCESS;
119 : }
120 :
121 0 : HcclResult ReduceScatterPlantLocalReduceCombine::LocalCopy()
122 : {
123 0 : u64 sliceSize = memBlockInfo_.size[localRank_];
124 0 : if (sliceSize == 0) {
125 0 : return HCCL_SUCCESS;
126 : }
127 :
128 0 : DeviceMem src;
129 0 : if (isNeedSpaceBorrow_ && isUseCclIn_) {
130 0 : src = inputMem_.range(memBlockInfo_.userInputOffsets[localRank_], sliceSize);
131 : } else {
132 0 : src = outputMem_.range(memBlockInfo_.inputOffsets[localRank_], sliceSize);
133 : }
134 :
135 0 : DeviceMem dst;
136 0 : u32 outputIndex = CalcOutputIndex(localRank_);
137 0 : if (isNeedSpaceBorrow_ && isLevel0LastRank_ && isLastBlockData(outputIndex) && !isLastRank(localRank_)) {
138 0 : dst = inputMem_.range(memBlockInfo_.userInputOffsets[localRank_], sliceSize);
139 : } else {
140 0 : dst = outputMem_.range(memBlockInfo_.outputOffsets[outputIndex], sliceSize);
141 : }
142 :
143 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
144 0 : return HCCL_SUCCESS;
145 0 : }
146 :
147 : HcclResult
148 0 : ReduceScatterPlantLocalReduceCombine::RunAlltoAllRDMA(u32 round, u64 sliceSize, const std::vector<LINK>& links)
149 : {
150 0 : u64 srcOffset = memBlockInfo_.inputOffsets[round];
151 0 : void* srcPtr = static_cast<u8*>(outputMem_.ptr()) + srcOffset;
152 0 : if (isNeedSpaceBorrow_ && isUseCclIn_) {
153 0 : srcOffset = sliceSize == 0 ? 0 : memBlockInfo_.userInputOffsets[round];
154 0 : srcPtr = static_cast<u8*>(inputMem_.ptr()) + srcOffset;
155 : }
156 :
157 0 : u32 outputIndex = CalcOutputIndex(round);
158 0 : u64 dstOffset = memBlockInfo_.outputOffsets[outputIndex];
159 0 : if (isNeedSpaceBorrow_ && isLevel0LastRank_ && !isLastRank(round) && isLastBlockData(outputIndex)) {
160 : // 只有level0最后一组的最后一块数据需要放到对方的input上（且非全局最后一张卡）
161 0 : dstOffset = sliceSize == 0 ? 0 : memBlockInfo_.userInputOffsets[round];
162 0 : CHK_RET(links[round]->TxAsync(UserMemType::INPUT_MEM, dstOffset, srcPtr, sliceSize, stream_));
163 : } else {
164 0 : CHK_RET(links[round]->TxAsync(UserMemType::OUTPUT_MEM, dstOffset, srcPtr, sliceSize, stream_));
165 : }
166 :
167 0 : u32 localOutputIndex = CalcOutputIndex(localRank_);
168 0 : u64 localDstOffset = memBlockInfo_.outputOffsets[localRank_];
169 0 : void* dstPtr = static_cast<u8*>(outputMem_.ptr()) + localDstOffset;
170 0 : if (isNeedSpaceBorrow_ && isLevel0LastRank_ && isLastBlockData(localOutputIndex) && !isLastRank(localRank_)) {
171 0 : localDstOffset = memBlockInfo_.userInputOffsets[round];
172 0 : dstPtr = static_cast<u8*>(inputMem_.ptr()) + localDstOffset;
173 : }
174 :
175 0 : u64 remoteSrcOffset = memBlockInfo_.inputOffsets[round];
176 0 : CHK_RET(links[round]->RxAsync(UserMemType::OUTPUT_MEM, remoteSrcOffset, dstPtr, sliceSize, stream_));
177 :
178 0 : CHK_RET(links[round]->PostFinAck(stream_));
179 0 : CHK_RET(links[round]->WaitFinAck(stream_));
180 0 : return HCCL_SUCCESS;
181 : }
182 :
183 : HcclResult
184 0 : ReduceScatterPlantLocalReduceCombine::RunAlltoAllSDMA(u32 round, u64 sliceSize, const std::vector<LINK>& links)
185 : {
186 0 : if (sliceSize != 0) {
187 0 : DeviceMem src;
188 0 : if (isNeedSpaceBorrow_ && isUseCclIn_) {
189 0 : src = inputMem_.range(memBlockInfo_.userInputOffsets[round], sliceSize);
190 : } else {
191 0 : src = outputMem_.range(memBlockInfo_.inputOffsets[round], sliceSize);
192 : }
193 :
194 0 : u32 outputIndex = CalcOutputIndex(round);
195 0 : u64 dstOffset = memBlockInfo_.outputOffsets[outputIndex];
196 0 : void* remMemPtr = nullptr;
197 0 : if (isNeedSpaceBorrow_ && isLevel0LastRank_ && !isLastRank(round) && isLastBlockData(outputIndex)) {
198 0 : CHK_RET(links[round]->GetRemoteMem(UserMemType::INPUT_MEM, &remMemPtr));
199 0 : dstOffset = memBlockInfo_.userInputOffsets[round];
200 : } else {
201 0 : CHK_RET(links[round]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
202 : }
203 0 : DeviceMem dst = DeviceMem::create(static_cast<u8*>(remMemPtr) + dstOffset, sliceSize);
204 :
205 0 : CHK_RET(HcclD2DMemcpyAsync(
206 : dispatcher_, dst, src, stream_, links[round]->GetRemoteRank(), links[round]->GetLinkType()));
207 0 : }
208 :
209 0 : CHK_RET(links[round]->TxDataSignal(stream_));
210 0 : CHK_RET(links[round]->RxDataSignal(stream_));
211 0 : return HCCL_SUCCESS;
212 : }
213 :
214 0 : HcclResult ReduceScatterPlantLocalReduceCombine::RunAlltoAll(const std::vector<LINK>& links)
215 : {
216 0 : CHK_RET(LocalCopy());
217 0 : for (u32 round = 0; round < rankSize_; round++) {
218 0 : if (round == localRank_) {
219 0 : continue;
220 : }
221 0 : CHK_SMART_PTR_NULL(links[round]);
222 0 : CHK_RET(links[round]->TxAck(stream_));
223 0 : CHK_RET(links[round]->RxAck(stream_));
224 :
225 0 : u64 sliceSize = memBlockInfo_.size[round];
226 0 : if (links[round]->GetLinkType() == LinkType::LINK_ROCE) {
227 0 : CHK_RET(RunAlltoAllRDMA(round, sliceSize, links));
228 : } else {
229 0 : CHK_RET(RunAlltoAllSDMA(round, sliceSize, links));
230 : }
231 : }
232 0 : return HCCL_SUCCESS;
233 : }
234 :
235 0 : HcclResult ReduceScatterPlantLocalReduceCombine::RunLocalReduce()
236 : {
237 0 : u32 reduceStep = static_cast<u32>(std::ceil(log2(rankSize_)));
238 0 : u64 sliceSize = memBlockInfo_.size[localRank_];
239 0 : u32 dataUnitSize = DataUnitSize(dataType_);
240 0 : if (dataUnitSize == 0) {
241 0 : HCCL_ERROR(
242 : "[ReduceScatterPlantLocalReduceCombine][RunLocalReduce]data type[%s] out of range[%d, %d]",
243 : GetDataTypeEnumStr(dataType_).c_str(), HCCL_DATA_TYPE_INT8, static_cast<int>(HCCL_DATA_TYPE_RESERVED) - 1);
244 0 : return HCCL_E_INTERNAL;
245 : }
246 0 : u64 count = sliceSize / dataUnitSize;
247 :
248 0 : for (u32 round = 0; round < reduceStep; round++) {
249 0 : u32 tailIndex = std::min(rankSize_, static_cast<u32>(1 << static_cast<int>(reduceStep - round))) - 1;
250 0 : u32 headIndex = static_cast<u32>(1 << static_cast<int>((reduceStep - round - 1)));
251 0 : u32 reduceSubStreamNum = std::min(tailIndex - headIndex, DEVICE_EIGHT / FACTOR_NUM_TWO - 1);
252 :
253 : // LR主流通知从流可以开始接受数据
254 0 : for (u32 offset = 0; offset < reduceSubStreamNum; offset++) {
255 0 : u32 streamId = offset;
256 : // 只有reduce任务 > 1时才需要主从流同步: LR主流通知从流, 从流Wait LR主流
257 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAuxPtr_)[streamId], profilerInput_.stage));
258 0 : CHK_RET(LocalNotify::Wait(
259 : subStreams_[streamId], dispatcher_, (*meshSignalAuxPtr_)[streamId], profilerInput_.stage));
260 : }
261 :
262 : // LocalReduce操作
263 0 : for (u32 offset = 0; offset <= tailIndex - headIndex; offset++) {
264 0 : u32 inputIndex = CalcOutputIndex(headIndex + offset); // reduce的源数据offset
265 0 : u32 outputIndex = CalcOutputIndex(offset); // reduce的目标offset
266 :
267 0 : u32 streamOffset = offset % (reduceSubStreamNum + 1);
268 0 : Stream& subStream = streamOffset == 0 ? stream_ : subStreams_[streamOffset - 1];
269 :
270 0 : if (sliceSize == 0) {
271 0 : continue;
272 : }
273 :
274 : void* srcPtr;
275 : void* dstPtr;
276 0 : if (isNeedSpaceBorrow_ && isLevel0LastRank_ && !isLastRank(localRank_) && isLastBlockData(inputIndex)) {
277 0 : srcPtr = static_cast<u8*>(inputMem_.ptr()) + memBlockInfo_.userInputOffsets[localRank_];
278 : } else {
279 0 : srcPtr = static_cast<u8*>(outputMem_.ptr()) + memBlockInfo_.outputOffsets[inputIndex];
280 : }
281 :
282 0 : if (isNeedSpaceBorrow_ && isLevel0LastRank_ && !isLastRank(localRank_) && isLastBlockData(outputIndex)) {
283 0 : dstPtr = static_cast<u8*>(inputMem_.ptr()) + memBlockInfo_.userInputOffsets[localRank_];
284 : } else {
285 0 : dstPtr = static_cast<u8*>(outputMem_.ptr()) + memBlockInfo_.outputOffsets[outputIndex];
286 : }
287 :
288 0 : CHK_RET(HcclReduceAsync(
289 : dispatcher_, srcPtr, count, dataType_, reductionOp_, subStream, dstPtr, INVALID_VALUE_RANKID,
290 : LinkType::LINK_ONCHIP, INLINE_REDUCE_BIT));
291 : }
292 :
293 : // 从流通知LR主流可以开始下一轮
294 0 : for (u32 offset = 0; offset < reduceSubStreamNum; offset++) {
295 0 : u32 streamId = offset;
296 : // 只有reduce任务 > 1时才需要主从流同步: LR主流通知从流, 从流Wait LR主流
297 0 : CHK_RET(LocalNotify::Post(
298 : subStreams_[streamId], dispatcher_, (*meshSignalPtr_)[streamId], profilerInput_.stage));
299 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignalPtr_)[streamId], profilerInput_.stage));
300 : }
301 :
302 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
303 : }
304 :
305 0 : return HCCL_SUCCESS;
306 : }
307 : REGISTER_TEMPLATE(
308 : TemplateType::TEMPLATE_REDUCESCATTER_PLANT_LOCAL_REDUCE_COMBINE, ReduceScatterPlantLocalReduceCombine);
309 : } // namespace hccl
|
