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.h"
14 :
15 : namespace hccl {
16 : constexpr u32 DEVICE_EIGHT = 8;
17 : constexpr u32 FACTOR_NUM_TWO = 2;
18 0 : ReduceScatterPlantLocalReduce::ReduceScatterPlantLocalReduce(const HcclDispatcher dispatcher)
19 0 : : AlgTemplateBase(dispatcher)
20 0 : {}
21 :
22 0 : ReduceScatterPlantLocalReduce::~ReduceScatterPlantLocalReduce() {}
23 :
24 0 : HcclResult ReduceScatterPlantLocalReduce::Prepare(
25 : void* inputMemPtr, 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 : GroupSlicesInfo& grouSlicesInfo, const HcclReduceOp reductionOp, u32 all2allOffset, const HcclDataType dataType,
28 : bool isNeedSpaceBorrow, bool reverseMemUsage, bool isA3CrossNode)
29 : {
30 0 : inputMemPtr_ = inputMemPtr; // UserInPtr，All2All使用
31 0 : inputMem_ = cclInMem; // 空拷贝 & 存放最后一块数据（Allreduce非整除场景）
32 0 : outputMem_ = outputMem; // 单算子CclOut 图模式Scrach/UserOut，LocalReduce使用
33 0 : stream_ = stream;
34 0 : subStreams_ = subStreams;
35 0 : meshSignalPtr_ = &meshSignal;
36 0 : meshSignalAuxPtr_ = &meshSignalAux;
37 0 : groupSlicesInfo_ = std::move(grouSlicesInfo);
38 0 : reductionOp_ = reductionOp;
39 0 : all2allOffset_ = all2allOffset;
40 0 : dataType_ = dataType;
41 0 : isNeedSpaceBorrow_ = isNeedSpaceBorrow;
42 0 : isA3CrossNode_ = isA3CrossNode;
43 0 : if (reverseMemUsage) {
44 : // 交换两块buffer的用途，in buffer作为输出buffer
45 0 : HCCL_INFO("[%s] reverse memory usage.", __func__);
46 0 : std::swap(scratchMemType_, outputMemType_);
47 0 : std::swap(inputMem_, outputMem_);
48 : }
49 0 : return HCCL_SUCCESS;
50 : }
51 :
52 0 : HcclResult ReduceScatterPlantLocalReduce::MainRecordSub(Stream& mainStream, u32 firstSubStreamIndex, u32 totalTask)
53 : {
54 0 : for (u32 streamIndex = firstSubStreamIndex; streamIndex < totalTask; streamIndex++) {
55 0 : CHK_RET(LocalNotify::Post(mainStream, dispatcher_, (*meshSignalAuxPtr_)[streamIndex], profilerInput_.stage));
56 : }
57 0 : return HCCL_SUCCESS;
58 : }
59 :
60 0 : HcclResult ReduceScatterPlantLocalReduce::SubWaitMain(u32 firstSubStreamIndex, u32 totalTask)
61 : {
62 0 : for (u32 streamIndex = firstSubStreamIndex; streamIndex < totalTask; streamIndex++) {
63 0 : CHK_RET(LocalNotify::Wait(
64 : subStreams_[streamIndex], dispatcher_, (*meshSignalAuxPtr_)[streamIndex], profilerInput_.stage));
65 : }
66 0 : return HCCL_SUCCESS;
67 : }
68 :
69 0 : HcclResult ReduceScatterPlantLocalReduce::MainWaitSub(Stream& mainStream, u32 firstSubStreamIndex, u32 totalTask)
70 : {
71 0 : for (u32 streamIndex = firstSubStreamIndex; streamIndex < totalTask; streamIndex++) {
72 0 : CHK_RET(LocalNotify::Wait(mainStream, dispatcher_, (*meshSignalPtr_)[streamIndex], profilerInput_.stage));
73 : }
74 0 : return HCCL_SUCCESS;
75 : }
76 :
77 0 : HcclResult ReduceScatterPlantLocalReduce::SubRecordMain(u32 firstSubStreamIndex, u32 totalTask)
78 : {
79 0 : for (u32 streamIndex = firstSubStreamIndex; streamIndex < totalTask; streamIndex++) {
80 0 : CHK_RET(LocalNotify::Post(
81 : subStreams_[streamIndex], dispatcher_, (*meshSignalPtr_)[streamIndex], profilerInput_.stage));
82 : }
83 0 : return HCCL_SUCCESS;
84 : }
85 :
86 0 : HcclResult ReduceScatterPlantLocalReduce::MainRecordLocalReduceWait(u32 lRMainStreamIndex)
87 : {
88 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAuxPtr_)[lRMainStreamIndex], profilerInput_.stage));
89 0 : CHK_RET(LocalNotify::Wait(
90 : subStreams_[lRMainStreamIndex], dispatcher_, (*meshSignalAuxPtr_)[lRMainStreamIndex], profilerInput_.stage));
91 0 : return HCCL_SUCCESS;
92 : }
93 :
94 0 : u32 ReduceScatterPlantLocalReduce::CalcOutputIndex(const u32 round)
95 : {
96 0 : return (all2allOffset_ + round + localRank_) % rankSize_;
97 : }
98 :
99 0 : bool ReduceScatterPlantLocalReduce::isLastGroup(const u32 groupId) { return groupId == groupSlicesInfo_.size() - 1; }
100 :
101 0 : bool ReduceScatterPlantLocalReduce::isLastRank(const u32 rankId) const { return rankId == rankSize_ - 1; }
102 :
103 0 : bool ReduceScatterPlantLocalReduce::isLastBlockData(const u32 outputIndex) const
104 : {
105 0 : return outputIndex == rankSize_ - 1;
106 : }
107 :
108 0 : HcclResult ReduceScatterPlantLocalReduce::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
109 : {
110 0 : HCCL_INFO(
111 : "ReduceScatterPlantLocalReduce run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p].", rank, rankSize,
112 : inputMem_.ptr(), outputMem_.ptr());
113 0 : CHK_SMART_PTR_NULL(dispatcher_);
114 0 : CHK_PTR_NULL(stream_.ptr());
115 0 : CHK_PRT_RET(
116 : links.size() < rankSize,
117 : HCCL_ERROR("[%s]rank[%u] linksize[%llu] is less than rankSize[%u]", __func__, rank, links.size(), rankSize),
118 : HCCL_E_INTERNAL);
119 :
120 0 : rankSize_ = rankSize;
121 0 : localRank_ = rank;
122 :
123 : // All2All主流（主流）通知LocalReduce主流开始准备执行,
124 : // All2All需要rankSize条流，其中主流完成LocalCopy&第一个A2A任务，因此主从同步需要rankSize-2个任务。lRMainStreamId_需要-2
125 0 : all2allSubStreamNum_ = isA3CrossNode_ ? std::min(rankSize, DEVICE_EIGHT) - 1 : rankSize - 2;
126 0 : lRMainStreamId_ = all2allSubStreamNum_;
127 :
128 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
129 0 : CHK_RET(MainRecordLocalReduceWait(lRMainStreamId_));
130 : // 额外一次LocalReduce主流通知All2All主流准备好接受信息（通知第一次执行完的All2AllWait）
131 0 : CHK_RET(LocalNotify::Post(
132 : subStreams_[lRMainStreamId_], dispatcher_, (*meshSignalPtr_)[lRMainStreamId_], profilerInput_.stage));
133 :
134 0 : HcclResult ret = HCCL_SUCCESS;
135 0 : for (u32 groupId = 0; groupId < groupSlicesInfo_.size(); groupId++) {
136 0 : const MemBlockInfo& memBlockInfo = groupSlicesInfo_[groupId];
137 0 : if (isA3CrossNode_) {
138 0 : ret = RunGroupAlltoAll(links, groupId, memBlockInfo);
139 : } else {
140 0 : ret = RunAlltoAll(links, groupId, memBlockInfo);
141 : }
142 0 : CHK_PRT_RET(
143 : ret != HCCL_SUCCESS,
144 : HCCL_ERROR(
145 : "[%s]RunAlltoAll or RunGroupAlltoAll failed, localRank[%u], groupId[%u]", __func__, localRank_,
146 : groupId),
147 : ret);
148 :
149 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignalPtr_)[lRMainStreamId_], profilerInput_.stage));
150 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
151 0 : CHK_RET(MainRecordLocalReduceWait(lRMainStreamId_));
152 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, subStreams_[lRMainStreamId_], dispatcher_));
153 :
154 0 : ret = RunLocalReduce(groupId, memBlockInfo);
155 0 : CHK_PRT_RET(
156 : ret != HCCL_SUCCESS,
157 : HCCL_ERROR("[%s]LocalReduce failed, localRank[%u], groupId[%u]", __func__, localRank_, groupId), ret);
158 :
159 : // LocalReduce主流通知All2All主流执行完成，可以下发下一次LocalReduce操作
160 0 : CHK_RET(LocalNotify::Post(
161 : subStreams_[lRMainStreamId_], dispatcher_, (*meshSignalPtr_)[lRMainStreamId_], profilerInput_.stage));
162 : }
163 :
164 : // All2All主流等待最后一次LocalReduce执行完成
165 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignalPtr_)[lRMainStreamId_], profilerInput_.stage));
166 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
167 0 : HCCL_INFO("ReduceScatterPlantLocalReduce finished: localRank[%u] ranksize[%u]", localRank_, rankSize_);
168 0 : return HCCL_SUCCESS;
169 : }
170 :
171 0 : HcclResult ReduceScatterPlantLocalReduce::LocalCopy(u32 groupId, const MemBlockInfo& memBlockInfo)
172 : {
173 0 : u64 sliceSize = memBlockInfo.size[localRank_];
174 0 : if (sliceSize == 0) {
175 0 : return HCCL_SUCCESS;
176 : }
177 :
178 : DeviceMem src
179 0 : = DeviceMem::create(static_cast<u8*>(inputMemPtr_) + memBlockInfo.userInputOffsets[localRank_], sliceSize);
180 :
181 : // 当非最后一组最后一卡且outputIndex是最后一块时，Copy至CclIn/UserIn预留位
182 0 : DeviceMem dst;
183 0 : u32 outputIndex = CalcOutputIndex(localRank_);
184 0 : if (isNeedSpaceBorrow_ && isLastBlockData(outputIndex) && !(isLastRank(localRank_) && isLastGroup(groupId))) {
185 0 : dst = inputMem_.range(memBlockInfo.outputOffsets[localRank_], sliceSize);
186 : } else {
187 0 : dst = outputMem_.range(memBlockInfo.outputOffsets[outputIndex], sliceSize);
188 : }
189 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
190 0 : return HCCL_SUCCESS;
191 0 : }
192 :
193 0 : HcclResult ReduceScatterPlantLocalReduce::RunAlltoAll(
194 : const std::vector<LINK>& links, u32 groupId, const MemBlockInfo& memBlockInfo)
195 : {
196 : // 本卡优先拷贝同号位数据
197 0 : CHK_RET(LocalCopy(groupId, memBlockInfo));
198 : // 主流通知从流可以开始接受数据
199 0 : u32 all2allfirstSubStreamId = 0;
200 0 : CHK_RET(MainRecordSub(stream_, all2allfirstSubStreamId, all2allSubStreamNum_));
201 0 : CHK_RET(SubWaitMain(all2allfirstSubStreamId, all2allSubStreamNum_));
202 :
203 : // 开始数据拷贝
204 0 : u32 streamIndex = 0;
205 0 : for (u32 round = 0; round < rankSize_; round++) {
206 0 : if (round == localRank_) {
207 0 : continue;
208 : }
209 0 : Stream& subStream = (streamIndex == 0) ? stream_ : subStreams_[streamIndex - 1];
210 0 : CHK_SMART_PTR_NULL(links[round]);
211 0 : CHK_RET(links[round]->TxAck(subStream));
212 0 : CHK_RET(links[round]->RxAck(subStream));
213 0 : streamIndex++;
214 : }
215 :
216 0 : CHK_RET(SubRecordMain(all2allfirstSubStreamId, all2allSubStreamNum_));
217 0 : CHK_RET(MainWaitSub(stream_, all2allfirstSubStreamId, all2allSubStreamNum_));
218 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, subStreams_[lRMainStreamId_], dispatcher_));
219 :
220 0 : CHK_RET(MainRecordSub(stream_, all2allfirstSubStreamId, all2allSubStreamNum_));
221 0 : CHK_RET(SubWaitMain(all2allfirstSubStreamId, all2allSubStreamNum_));
222 0 : streamIndex = 0;
223 0 : for (u32 round = 0; round < rankSize_; round++) {
224 0 : if (round == localRank_) {
225 0 : continue;
226 : }
227 0 : Stream& subStream = (streamIndex == 0) ? stream_ : subStreams_[streamIndex - 1];
228 0 : CHK_SMART_PTR_NULL(links[round]);
229 :
230 0 : u64 sliceSize = memBlockInfo.size[round];
231 0 : if (sliceSize != 0) {
232 0 : u64 userMemInOffset = memBlockInfo.userInputOffsets[round];
233 0 : DeviceMem src = DeviceMem::create(static_cast<u8*>(inputMemPtr_) + userMemInOffset, sliceSize);
234 0 : u32 outputIndex = CalcOutputIndex(round);
235 0 : u64 dstOffset = 0;
236 0 : void* remMemPtr = nullptr;
237 0 : if (isNeedSpaceBorrow_ && isLastBlockData(outputIndex) && !(isLastRank(round) && isLastGroup(groupId))) {
238 0 : CHK_RET(links[round]->GetRemoteMem(scratchMemType_, &remMemPtr));
239 0 : dstOffset = memBlockInfo.outputOffsets[round];
240 : } else {
241 0 : CHK_RET(links[round]->GetRemoteMem(outputMemType_, &remMemPtr));
242 0 : dstOffset = memBlockInfo.outputOffsets[outputIndex];
243 : }
244 0 : DeviceMem dst = DeviceMem::create(static_cast<u8*>(remMemPtr) + dstOffset, sliceSize);
245 0 : CHK_RET(HcclD2DMemcpyAsync(
246 : dispatcher_, dst, src, subStream, links[round]->GetRemoteRank(), links[round]->GetLinkType()));
247 0 : }
248 0 : CHK_RET(links[round]->TxDataSignal(subStream));
249 0 : CHK_RET(links[round]->RxDataSignal(subStream));
250 0 : streamIndex++;
251 : }
252 :
253 : // 从流通知主流完成拷贝
254 0 : CHK_RET(SubRecordMain(all2allfirstSubStreamId, all2allSubStreamNum_));
255 0 : CHK_RET(MainWaitSub(stream_, all2allfirstSubStreamId, all2allSubStreamNum_));
256 0 : return HCCL_SUCCESS;
257 : }
258 :
259 0 : HcclResult ReduceScatterPlantLocalReduce::RunGroupAlltoAll(
260 : const std::vector<LINK>& links, u32 groupId, const MemBlockInfo& memBlockInfo)
261 : {
262 0 : constexpr u32 numInGroup = DEVICE_EIGHT;
263 0 : u32 numOfGroups = (rankSize_ + numInGroup - 1) / numInGroup;
264 :
265 : // 本卡优先拷贝同号位数据
266 0 : CHK_RET(LocalCopy(groupId, memBlockInfo));
267 :
268 0 : for (u32 idGroup = 0; idGroup < numOfGroups; ++idGroup) {
269 : // 主流通知从流可以开始接受数据
270 0 : u32 all2allfirstSubStreamId = 0;
271 0 : CHK_RET(MainRecordSub(stream_, all2allfirstSubStreamId, all2allSubStreamNum_));
272 0 : CHK_RET(SubWaitMain(all2allfirstSubStreamId, all2allSubStreamNum_));
273 :
274 : // 开始数据拷贝
275 0 : u32 streamIndex = 0;
276 0 : for (u32 cnt = 0, round = idGroup * numInGroup; round < rankSize_ && cnt < numInGroup; ++round, ++cnt) {
277 0 : if (round == 0) {
278 0 : continue;
279 : }
280 0 : u32 sendRank = (localRank_ + round) % rankSize_;
281 0 : u32 recvRank = (rankSize_ + localRank_ - round) % rankSize_;
282 0 : Stream& subStream = (streamIndex == 0) ? stream_ : subStreams_[streamIndex - 1];
283 0 : CHK_SMART_PTR_NULL(links[sendRank]);
284 0 : CHK_SMART_PTR_NULL(links[recvRank]);
285 0 : CHK_RET(links[recvRank]->TxAck(subStream));
286 0 : CHK_RET(links[sendRank]->RxAck(subStream));
287 0 : streamIndex++;
288 : }
289 :
290 0 : CHK_RET(SubRecordMain(all2allfirstSubStreamId, all2allSubStreamNum_));
291 0 : CHK_RET(MainWaitSub(stream_, all2allfirstSubStreamId, all2allSubStreamNum_));
292 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, subStreams_[lRMainStreamId_], dispatcher_));
293 :
294 0 : CHK_RET(MainRecordSub(stream_, all2allfirstSubStreamId, all2allSubStreamNum_));
295 0 : CHK_RET(SubWaitMain(all2allfirstSubStreamId, all2allSubStreamNum_));
296 0 : streamIndex = 0;
297 0 : for (u32 cnt = 0, round = idGroup * numInGroup; round < rankSize_ && cnt < numInGroup; ++round, ++cnt) {
298 0 : if (round == 0) {
299 0 : continue;
300 : }
301 0 : u32 sendRank = (localRank_ + round) % rankSize_;
302 0 : u32 recvRank = (rankSize_ + localRank_ - round) % rankSize_;
303 0 : Stream& subStream = (streamIndex == 0) ? stream_ : subStreams_[streamIndex - 1];
304 0 : CHK_SMART_PTR_NULL(links[sendRank]);
305 0 : CHK_SMART_PTR_NULL(links[recvRank]);
306 :
307 0 : u64 sliceSize = memBlockInfo.size[sendRank];
308 0 : if (sliceSize != 0) {
309 0 : u64 userMemInOffset = memBlockInfo.userInputOffsets[sendRank];
310 0 : DeviceMem src = DeviceMem::create(static_cast<u8*>(inputMemPtr_) + userMemInOffset, sliceSize);
311 0 : u32 outputIndex = CalcOutputIndex(sendRank);
312 0 : u64 dstOffset = 0;
313 0 : void* remMemPtr = nullptr;
314 0 : if (isNeedSpaceBorrow_ && isLastBlockData(outputIndex)
315 0 : && !(isLastRank(sendRank) && isLastGroup(groupId))) {
316 0 : CHK_RET(links[sendRank]->GetRemoteMem(scratchMemType_, &remMemPtr));
317 0 : dstOffset = memBlockInfo.outputOffsets[sendRank];
318 : } else {
319 0 : CHK_RET(links[sendRank]->GetRemoteMem(outputMemType_, &remMemPtr));
320 0 : dstOffset = memBlockInfo.outputOffsets[outputIndex];
321 : }
322 0 : DeviceMem dst = DeviceMem::create(static_cast<u8*>(remMemPtr) + dstOffset, sliceSize);
323 0 : CHK_RET(HcclD2DMemcpyAsync(
324 : dispatcher_, dst, src, subStream, links[sendRank]->GetRemoteRank(),
325 : links[sendRank]->GetLinkType()));
326 0 : }
327 0 : CHK_RET(links[sendRank]->TxDataSignal(subStream));
328 0 : CHK_RET(links[recvRank]->RxDataSignal(subStream));
329 0 : streamIndex++;
330 : }
331 :
332 : // 从流通知主流完成拷贝
333 0 : CHK_RET(SubRecordMain(all2allfirstSubStreamId, all2allSubStreamNum_));
334 0 : CHK_RET(MainWaitSub(stream_, all2allfirstSubStreamId, all2allSubStreamNum_));
335 : }
336 :
337 0 : return HCCL_SUCCESS;
338 : }
339 :
340 0 : HcclResult ReduceScatterPlantLocalReduce::RunLocalReduce(u32 groupId, const MemBlockInfo& memBlockInfo)
341 : {
342 0 : u32 reduceStep = static_cast<u32>(std::ceil(log2(rankSize_)));
343 0 : u64 srcOffset = memBlockInfo.inputOffsets[localRank_];
344 0 : u64 sliceSize = memBlockInfo.size[localRank_];
345 0 : u32 dataUnitSize = DataUnitSize(dataType_);
346 0 : if (dataUnitSize == 0) {
347 0 : HCCL_ERROR(
348 : "[ReduceScatterPlantLocalReduce][RunLocalReduce]data type[%s] out of range[%d, %d]",
349 : GetDataTypeEnumStr(dataType_).c_str(), HCCL_DATA_TYPE_INT8, static_cast<int>(HCCL_DATA_TYPE_RESERVED) - 1);
350 0 : return HCCL_E_INTERNAL;
351 : }
352 0 : u64 count = sliceSize / dataUnitSize;
353 :
354 0 : for (u32 round = 0; round < reduceStep; round++) {
355 0 : u32 tailIndex = std::min(rankSize_, static_cast<u32>(1 << static_cast<int>(reduceStep - round))) - 1;
356 0 : u32 headIndex = static_cast<u32>(1 << static_cast<int>((reduceStep - round - 1)));
357 0 : u32 reduceSubStreamNum = std::min(tailIndex - headIndex, DEVICE_EIGHT / FACTOR_NUM_TWO - 1);
358 : // LR主流通知从流可以开始接受数据
359 0 : for (u32 offset = 0; offset < reduceSubStreamNum; offset++) {
360 0 : u32 streamId = lRMainStreamId_ + offset + 1;
361 : // 只有reduce任务 > 1时才需要主从流同步: LR主流通知从流, 从流Wait LR主流
362 0 : CHK_RET(LocalNotify::Post(
363 : subStreams_[lRMainStreamId_], dispatcher_, (*meshSignalAuxPtr_)[streamId], profilerInput_.stage));
364 0 : CHK_RET(LocalNotify::Wait(
365 : subStreams_[streamId], dispatcher_, (*meshSignalAuxPtr_)[streamId], profilerInput_.stage));
366 : }
367 :
368 : // LocalReduce操作
369 0 : for (u32 offset = 0; offset <= tailIndex - headIndex; offset++) {
370 0 : u32 inputIndex = CalcOutputIndex(headIndex + offset); // reduce的源数据offset
371 0 : u32 outputIndex = CalcOutputIndex(offset); // reduce的目标offset
372 0 : u32 streamOffset = offset % (reduceSubStreamNum + 1);
373 0 : Stream& subStream = subStreams_[lRMainStreamId_ + streamOffset];
374 0 : if (sliceSize == 0) {
375 0 : continue;
376 : }
377 : void* srcPtr;
378 : void* dstPtr;
379 0 : if (isNeedSpaceBorrow_ && !(isLastRank(localRank_) && isLastGroup(groupId))
380 0 : && isLastBlockData(inputIndex)) {
381 0 : srcPtr = static_cast<u8*>(inputMem_.ptr()) + srcOffset;
382 : } else {
383 0 : srcPtr = static_cast<u8*>(outputMem_.ptr()) + memBlockInfo.outputOffsets[inputIndex];
384 : }
385 :
386 0 : if (isNeedSpaceBorrow_ && !(isLastRank(localRank_) && isLastGroup(groupId))
387 0 : && isLastBlockData(outputIndex)) {
388 0 : dstPtr = static_cast<u8*>(inputMem_.ptr()) + srcOffset;
389 : } else {
390 0 : dstPtr = static_cast<u8*>(outputMem_.ptr()) + memBlockInfo.outputOffsets[outputIndex];
391 : }
392 :
393 0 : CHK_RET(HcclReduceAsync(
394 : dispatcher_, srcPtr, count, dataType_, reductionOp_, subStream, dstPtr, INVALID_VALUE_RANKID,
395 : LinkType::LINK_ONCHIP, INLINE_REDUCE_BIT));
396 : }
397 :
398 : // 从流通知LR主流可以开始下一轮
399 0 : for (u32 offset = 0; offset < reduceSubStreamNum; offset++) {
400 0 : u32 streamId = lRMainStreamId_ + offset + 1;
401 : // 只有reduce任务 > 1时才需要主从流同步: LR主流通知从流, 从流Wait LR主流
402 0 : CHK_RET(LocalNotify::Post(
403 : subStreams_[streamId], dispatcher_, (*meshSignalPtr_)[streamId], profilerInput_.stage));
404 0 : CHK_RET(LocalNotify::Wait(
405 : subStreams_[lRMainStreamId_], dispatcher_, (*meshSignalPtr_)[streamId], profilerInput_.stage));
406 : }
407 :
408 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, subStreams_[lRMainStreamId_], dispatcher_));
409 : }
410 :
411 0 : return HCCL_SUCCESS;
412 : }
413 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_PLANT_LOCAL_REDUCE, ReduceScatterPlantLocalReduce);
414 : } // namespace hccl
|
