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 "reduce_scatter_unified_march.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 : static const u32 NEIGHBORS_NUM_TWO = 2; // 2: 邻居数量
16 : static const u32 NEIGHBORS_NUM_ONE = 1; // 1: 邻居数量
17 : static const u32 DIVISOR_NUM_TWO = 2;
18 :
19 1 : ReduceScatterUnifiedMarch::ReduceScatterUnifiedMarch(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
20 :
21 2 : ReduceScatterUnifiedMarch::~ReduceScatterUnifiedMarch() {}
22 :
23 1 : HcclResult ReduceScatterUnifiedMarch::Prepare(
24 : Stream& mainStream, SubCommInfo& level0CommInfo, DeviceMem& userInput, DeviceMem& userOutput, DeviceMem& usrInMem,
25 : DeviceMem& scratchMem, u64 totalCount, std::vector<Stream>& subStreams,
26 : const std::vector<std::shared_ptr<LocalNotify>>& meshSignalMainToSub,
27 : const std::vector<std::shared_ptr<LocalNotify>>& meshSignalSubToMain, const HcclDataType dataType,
28 : const HcclReduceOp reductionOp, const std::vector<std::vector<Slice>>& multRingsUserMemSlice, u64 reduceAttrBitMap)
29 : {
30 1 : reduceAttr_ = reduceAttrBitMap;
31 1 : mainStream_ = mainStream;
32 1 : intraRank_ = level0CommInfo.localRank;
33 1 : intraRankSize_ = level0CommInfo.localRankSize;
34 1 : CHK_PRT_RET(
35 : intraRankSize_ == 0 || (intraRankSize_ % DIVISOR_NUM_TWO != 0),
36 : HCCL_ERROR("[ReduceScatterUnifiedMarch][Prepare]intraRankSize_ is zero or not divisible by 2"), HCCL_E_PARA);
37 1 : links_ = level0CommInfo.links;
38 :
39 1 : userInput_ = userInput;
40 1 : userOutput_ = userOutput;
41 1 : usrInMem_ = usrInMem;
42 1 : scratchMem_ = scratchMem;
43 1 : HCCL_INFO(
44 : "userInput_[%p] size[%llu], userOutput_[%p] size[%llu], usrInMem_[%p] size[%llu], scratchMem_[%p] size[%llu]",
45 : userInput_.ptr(), userInput_.size(), userOutput_.ptr(), userOutput_.size(), usrInMem_.ptr(), usrInMem_.size(),
46 : scratchMem_.ptr(), scratchMem_.size());
47 :
48 1 : subStreams_ = subStreams;
49 1 : meshSignalMainToSub_ = meshSignalMainToSub;
50 1 : meshSignalSubToMain_ = meshSignalSubToMain;
51 1 : CHK_PRT_RET(
52 : subStreams_.size() < NEIGHBORS_NUM_TWO || meshSignalMainToSub_.size() < NEIGHBORS_NUM_TWO
53 : || meshSignalSubToMain_.size() < NEIGHBORS_NUM_TWO,
54 : HCCL_ERROR(
55 : "[AllGatherUnifiedMarch] subStreams_ size[%u] or meshSignalMainToSub_ size[%u] or "
56 : "meshSignalSubToMain_ size[%u] is less than 2",
57 : subStreams_.size(), meshSignalMainToSub_.size(), meshSignalSubToMain_.size()),
58 : HCCL_E_PARA);
59 :
60 1 : totalCount_ = totalCount;
61 1 : dataType_ = dataType;
62 1 : reductionOp_ = reductionOp;
63 1 : blockDataByte_ = totalCount_ * SIZE_TABLE[dataType_];
64 1 : multRingsUserMemSlice_ = multRingsUserMemSlice;
65 1 : CHK_PRT_RET(
66 : multRingsUserMemSlice_[0].size() % intraRankSize_ != 0,
67 : HCCL_ERROR(
68 : "[ReduceScatterUnifiedMarch] multRingsUserMemSlice_[0] size[%u] can not be divided by rank size[%u]",
69 : multRingsUserMemSlice_[0].size(), intraRankSize_),
70 : HCCL_E_PARA);
71 :
72 1 : return HCCL_SUCCESS;
73 : }
74 :
75 0 : std::string ReduceScatterUnifiedMarch::GetStreamIndexString()
76 : {
77 0 : std::string res = "";
78 0 : for (u32 streamIndex = 0; streamIndex < subStreams_.size(); streamIndex++) {
79 0 : res += std::to_string(streamIndex) + ", ";
80 : }
81 0 : return res;
82 0 : }
83 :
84 : // 主流通知所有从流
85 0 : HcclResult ReduceScatterUnifiedMarch::NotifySubStreamStart(u32 streamSize)
86 : {
87 0 : CHK_PRT_RET(
88 : streamSize > subStreams_.size() || streamSize > meshSignalSubToMain_.size(),
89 : HCCL_ERROR(
90 : "[ReduceScatterUnifiedMarch][NotifySubStreamStart] streamSize[%u] is out of range"
91 : "subStreams_ size[%zu] or meshSignalSubToMain_ size[%zu]",
92 : streamSize, subStreams_.size(), meshSignalSubToMain_.size()),
93 : HCCL_E_PARA);
94 0 : for (u32 streamIndex = 0; streamIndex < streamSize; streamIndex++) {
95 0 : CHK_RET(LocalNotify::Post(mainStream_, dispatcher_, meshSignalSubToMain_[streamIndex], INVALID_VALUE_STAGE));
96 0 : CHK_RET(LocalNotify::Wait(
97 : subStreams_[streamIndex], dispatcher_, meshSignalSubToMain_[streamIndex], INVALID_VALUE_STAGE));
98 : }
99 0 : HCCL_DEBUG(
100 : "[ReduceScatterUnifiedMarch][NotifySubStreamStart] intraRank [%u] main stream notify substream [%s]",
101 : intraRank_, GetStreamIndexString().c_str());
102 0 : return HCCL_SUCCESS;
103 : }
104 :
105 0 : HcclResult ReduceScatterUnifiedMarch::WaitSubStreamFinish(u32 streamSize)
106 : {
107 0 : CHK_PRT_RET(
108 : streamSize > subStreams_.size() || streamSize > meshSignalMainToSub_.size(),
109 : HCCL_ERROR(
110 : "[ReduceScatterUnifiedMarch][WaitSubStreamFinish] streamSize[%u] is out of range"
111 : "subStreams_ size[%zu] or meshSignalMainToSub_ size[%zu]",
112 : streamSize, subStreams_.size(), meshSignalMainToSub_.size()),
113 : HCCL_E_PARA);
114 0 : for (u32 streamIndex = 0; streamIndex < streamSize; streamIndex++) {
115 0 : CHK_RET(LocalNotify::Post(
116 : subStreams_[streamIndex], dispatcher_, meshSignalMainToSub_[streamIndex], INVALID_VALUE_STAGE));
117 0 : CHK_RET(LocalNotify::Wait(mainStream_, dispatcher_, meshSignalMainToSub_[streamIndex], INVALID_VALUE_STAGE));
118 : }
119 0 : HCCL_DEBUG(
120 : "[ReduceScatterUnifiedMarch][WaitSubStreamFinish] intraRank [%u] main stream wait substream [%s]", intraRank_,
121 : GetStreamIndexString().c_str());
122 0 : return HCCL_SUCCESS;
123 : }
124 :
125 0 : HcclResult ReduceScatterUnifiedMarch::NotifyNeighborsStart(LINK& prevIntraLink, LINK& nextIntralLink, u32 neighbors)
126 : {
127 : // 图模式保持使用Post/Wait接口
128 0 : if (GetWorkflowMode() != HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
129 : // notify是否越界由平台侧保证
130 0 : for (u32 neighborRankId = 0; neighborRankId < neighbors; neighborRankId++) {
131 0 : if (neighborRankId == 0) {
132 0 : CHK_RET(nextIntralLink->Post(notifyIdx_, subStreams_[neighborRankId])); // AckRecord
133 0 : CHK_RET(prevIntraLink->Wait(notifyIdx_, subStreams_[neighborRankId])); // AckWait
134 0 : } else if (neighborRankId == 1) {
135 0 : CHK_RET(prevIntraLink->Post(notifyIdx_, subStreams_[neighborRankId])); // AckRecord
136 0 : CHK_RET(nextIntralLink->Wait(notifyIdx_, subStreams_[neighborRankId])); // AckWait
137 : }
138 : }
139 0 : HCCL_DEBUG(
140 : "[ReduceScatterUnifiedMarch][NotifyNeighborsStart] intraRank[%u] switch on [%u]neigbhbors done", intraRank_,
141 : neighbors);
142 0 : return HCCL_SUCCESS;
143 : }
144 :
145 : // 一条流负责一个环
146 0 : for (u32 neighborRankId = 0; neighborRankId < neighbors; neighborRankId++) {
147 : // 交替使用Ack和DataSignal两种notify
148 0 : const u32 NOTIFY_IDX_TWO = 2;
149 0 : if (neighborRankId == 0) {
150 0 : if (notifyIdx_ % NOTIFY_IDX_TWO == 0) {
151 0 : CHK_RET(nextIntralLink->TxAck(subStreams_[neighborRankId])); // AckRecord
152 0 : CHK_RET(prevIntraLink->RxAck(subStreams_[neighborRankId])); // AckWait
153 : } else {
154 0 : CHK_RET(nextIntralLink->TxDataSignal(subStreams_[neighborRankId])); // DataRecord
155 0 : CHK_RET(prevIntraLink->RxDataSignal(subStreams_[neighborRankId])); // DataWait
156 : }
157 0 : } else if (neighborRankId == 1) {
158 0 : if (notifyIdx_ % NOTIFY_IDX_TWO == 0) {
159 0 : CHK_RET(prevIntraLink->TxAck(subStreams_[neighborRankId])); // AckRecord
160 0 : CHK_RET(nextIntralLink->RxAck(subStreams_[neighborRankId])); // AckWait
161 : } else {
162 0 : CHK_RET(prevIntraLink->TxDataSignal(subStreams_[neighborRankId])); // DataRecord
163 0 : CHK_RET(nextIntralLink->RxDataSignal(subStreams_[neighborRankId])); // DataWait
164 : }
165 : }
166 : }
167 0 : HCCL_DEBUG(
168 : "[ReduceScatterUnifiedMarch][NotifyNeighborsStart] intraRank[%u] switch on [%u]neigbhbors done", intraRank_,
169 : neighbors);
170 0 : return HCCL_SUCCESS;
171 : }
172 :
173 0 : HcclResult ReduceScatterUnifiedMarch::NotifyNeighborsEnd(LINK& prevIntraLink, LINK& nextIntralLink, u32 neighbors)
174 : {
175 0 : for (u32 neighborRankId = 0; neighborRankId < neighbors; neighborRankId++) {
176 0 : if (neighborRankId == 0) {
177 0 : CHK_RET(prevIntraLink->TxDataSignal(subStreams_[neighborRankId])); // DataRecord
178 0 : CHK_RET(nextIntralLink->RxDataSignal(subStreams_[neighborRankId]));
179 0 : } else if (neighborRankId == 1) {
180 0 : CHK_RET(nextIntralLink->TxDataSignal(subStreams_[neighborRankId]));
181 0 : CHK_RET(prevIntraLink->RxDataSignal(subStreams_[neighborRankId]));
182 : }
183 : }
184 0 : HCCL_DEBUG(
185 : "[ReduceScatterUnifiedMarch][NotifyNeighborsEnd] intraRank[%u] notifys [%u]neigbhbors reduce done", intraRank_,
186 : neighbors);
187 0 : return HCCL_SUCCESS;
188 : }
189 :
190 0 : HcclResult ReduceScatterUnifiedMarch::DoSerialReduce(
191 : void* remDMAMemPtr, void* dstAddr, u64 memSize, u64 dataCount, Stream& tmpStream, LINK& tmpLink,
192 : u64 remoteOffsetByte)
193 : {
194 0 : for (u32 sliceIdx = 0; sliceIdx < (multRingsUserMemSlice_[0].size() / intraRankSize_); sliceIdx++) {
195 : DeviceMem srcMem = DeviceMem::create(
196 0 : static_cast<u8*>(remDMAMemPtr) + remoteOffsetByte + multRingsUserMemSlice_[0][sliceIdx].offset, memSize);
197 : DeviceMem dstMem
198 0 : = DeviceMem::create(static_cast<u8*>(dstAddr) + multRingsUserMemSlice_[0][sliceIdx].offset, memSize);
199 :
200 0 : if ((INLINE_REDUCE_BITMASK & reduceAttr_) == 1) { // inlineReduce
201 0 : struct hccl::Transport::Buffer remoteBuf;
202 0 : remoteBuf.addr = srcMem.ptr();
203 0 : remoteBuf.size = srcMem.size();
204 0 : struct hccl::Transport::Buffer localBuf;
205 0 : localBuf.addr = dstMem.ptr();
206 0 : localBuf.size = dstMem.size();
207 0 : HCCL_DEBUG(
208 : "intralRank[%u] slice[%u] offset[%llu] do inlinereduce with remoteBuf[addr[%p], size[%llu]] and "
209 : "localBuf[addr[%p], size[%llu]]",
210 : intraRank_, sliceIdx, multRingsUserMemSlice_[0][sliceIdx].offset, remoteBuf.addr, remoteBuf.size,
211 : localBuf.addr, localBuf.size);
212 0 : CHK_RET(tmpLink->ReadReduceSync(localBuf, remoteBuf, dataType_, reductionOp_, tmpStream));
213 : } else { // TBE_reduce
214 : // left的inputMem拷到本端的scratchMem
215 0 : DeviceMem tempMem = scratchMem_.range(remoteOffsetByte, srcMem.size());
216 0 : struct hccl::Transport::Buffer remoteBuf;
217 0 : remoteBuf.addr = srcMem.ptr();
218 0 : remoteBuf.size = srcMem.size();
219 0 : struct hccl::Transport::Buffer localBuf;
220 0 : localBuf.addr = tempMem.ptr();
221 0 : localBuf.size = tempMem.size();
222 0 : HCCL_DEBUG(
223 : "intralRank[%u] slice[%u] offset[%llu] do SDMA read with remoteBuf[addr[%p], size[%llu]] and "
224 : "localBuf[addr[%p], size[%llu]]",
225 : intraRank_, sliceIdx, multRingsUserMemSlice_[0][sliceIdx].offset, remoteBuf.addr, remoteBuf.size,
226 : localBuf.addr, localBuf.size);
227 0 : CHK_RET(tmpLink->ReadSync(localBuf, remoteBuf, tmpStream));
228 0 : CHK_RET(HcclReduceAsync(
229 : dispatcher_, tempMem.ptr(), dataCount, dataType_, reductionOp_, tmpStream, dstMem.ptr(),
230 : INVALID_VALUE_RANKID, LinkType::LINK_ONCHIP, reduceAttr_));
231 0 : }
232 0 : }
233 0 : return HCCL_SUCCESS;
234 : }
235 :
236 0 : HcclResult ReduceScatterUnifiedMarch::RunSingleSliceRead(u32 ringPrevRank, u32 ringNextRank, u32 step, u32 totalStep)
237 : {
238 0 : LINK prevIntraLink = links_[ringPrevRank];
239 0 : CHK_SMART_PTR_NULL(prevIntraLink);
240 0 : LINK nextIntralLink = links_[ringNextRank];
241 0 : CHK_SMART_PTR_NULL(nextIntralLink);
242 0 : u32 neighbors = (ringPrevRank == ringNextRank) ? NEIGHBORS_NUM_ONE : NEIGHBORS_NUM_TWO;
243 0 : CHK_RET(NotifyNeighborsStart(prevIntraLink, nextIntralLink, neighbors));
244 :
245 : // 拉齐 从流record主流、主流record从流 保证从流同时开始做SDMA
246 0 : CHK_RET(WaitSubStreamFinish(neighbors));
247 0 : CHK_RET(NotifySubStreamStart(neighbors));
248 :
249 : // 从前向rank读取数据
250 0 : void* preRemDMAMemPtr = nullptr;
251 0 : CHK_RET(prevIntraLink->GetRemoteMem(UserMemType::INPUT_MEM, &preRemDMAMemPtr));
252 0 : u32 preDataIndex = (intraRank_ + intraRankSize_ - step - totalStep) % intraRankSize_;
253 0 : u64 preOffsetByte = preDataIndex * blockDataByte_;
254 0 : void* preDstAddr = static_cast<u8*>(userInput_.ptr()) + preOffsetByte;
255 :
256 0 : CHK_RET(DoSerialReduce(
257 : preRemDMAMemPtr, preDstAddr, blockDataByte_, totalCount_, subStreams_[0], prevIntraLink, preOffsetByte));
258 0 : HCCL_INFO(
259 : "[ReduceScatterUnifiedMarch][RunSingleSliceRead] intralRank [%u] reduce with ringPrevRank [%u] done",
260 : intraRank_, ringPrevRank);
261 :
262 : // 从后向rank读取数据
263 0 : if (neighbors > NEIGHBORS_NUM_ONE) {
264 0 : void* nextRemDMAMemPtr = nullptr;
265 0 : CHK_RET(nextIntralLink->GetRemoteMem(UserMemType::INPUT_MEM, &nextRemDMAMemPtr));
266 0 : u32 nextDataIndex = (intraRank_ + totalStep + step) % intraRankSize_;
267 0 : u64 nextOffsetByte = nextDataIndex * blockDataByte_;
268 0 : void* nextDstAddr = static_cast<u8*>(userInput_.ptr()) + nextOffsetByte;
269 :
270 0 : CHK_RET(DoSerialReduce(
271 : nextRemDMAMemPtr, nextDstAddr, blockDataByte_, totalCount_, subStreams_[1], nextIntralLink,
272 : nextOffsetByte));
273 0 : HCCL_INFO(
274 : "[ReduceScatterUnifiedMarch][RunSingleSliceRead] intralRank [%u]"
275 : "reduce with ringNextRank [%u] done",
276 : intraRank_, ringNextRank);
277 : }
278 :
279 : /* 2卡 场景，在最后一步的notifyDone */
280 0 : if (step == 0) {
281 0 : CHK_RET(NotifyNeighborsEnd(prevIntraLink, nextIntralLink, neighbors));
282 : }
283 0 : notifyIdx_++;
284 :
285 0 : return HCCL_SUCCESS;
286 0 : }
287 :
288 0 : HcclResult ReduceScatterUnifiedMarch::RunHalfSliceRead(u32 ringPrevRank, u32 ringNextRank, u32 step, u32 totalStep)
289 : {
290 0 : LINK prevIntraLink = links_[ringPrevRank];
291 0 : CHK_SMART_PTR_NULL(prevIntraLink);
292 0 : LINK nextIntralLink = links_[ringNextRank];
293 0 : CHK_SMART_PTR_NULL(nextIntralLink);
294 0 : CHK_RET(NotifyNeighborsStart(prevIntraLink, nextIntralLink, NEIGHBORS_NUM_TWO));
295 :
296 : // 拉齐 从流record主流、主流record从流 保证从流同时开始做SDMA
297 0 : CHK_RET(WaitSubStreamFinish(NEIGHBORS_NUM_TWO));
298 0 : CHK_RET(NotifySubStreamStart(NEIGHBORS_NUM_TWO));
299 :
300 : // 从前向rank读取数据
301 0 : void* preRemDMAMemPtr = nullptr;
302 0 : CHK_RET(prevIntraLink->GetRemoteMem(UserMemType::INPUT_MEM, &preRemDMAMemPtr));
303 0 : u32 temIdx = (step == 0) ? totalStep : 0;
304 0 : u32 preDataIndex = (intraRank_ + intraRankSize_ - temIdx) % intraRankSize_;
305 : // 考虑总数据量不能被整除的情况
306 0 : u32 partOneCount
307 0 : = (step != totalStep) ? (totalCount_ / DIVISOR_NUM_TWO) : (totalCount_ - totalCount_ / DIVISOR_NUM_TWO);
308 0 : u64 partOneSize = partOneCount * SIZE_TABLE[dataType_];
309 0 : u64 preOffsetByte = (step != totalStep) ?
310 0 : (preDataIndex * blockDataByte_) :
311 0 : (preDataIndex * blockDataByte_ + totalCount_ / DIVISOR_NUM_TWO * SIZE_TABLE[dataType_]);
312 0 : void* preDstAddr = static_cast<u8*>(userInput_.ptr()) + preOffsetByte;
313 :
314 0 : CHK_RET(DoSerialReduce(
315 : preRemDMAMemPtr, preDstAddr, partOneSize, partOneCount, subStreams_[0], prevIntraLink, preOffsetByte));
316 0 : HCCL_INFO(
317 : "[ReduceScatterUnifiedMarch][RunHalfSliceRead] intralRank [%u] reduce with ringPrevRank [%u] done", intraRank_,
318 : ringPrevRank);
319 :
320 : // 从后向rank读取数据
321 0 : void* nextRemDMAMemPtr = nullptr;
322 0 : CHK_RET(nextIntralLink->GetRemoteMem(UserMemType::INPUT_MEM, &nextRemDMAMemPtr));
323 0 : temIdx = (step == 0) ? totalStep : 0;
324 0 : u32 nextDataIndex = (intraRank_ + temIdx) % intraRankSize_;
325 0 : u32 partTwoCount = totalCount_ - partOneCount;
326 0 : u64 partTwoSize = partTwoCount * SIZE_TABLE[dataType_];
327 0 : u64 nextOffsetByte
328 0 : = (step != totalStep) ? (nextDataIndex * blockDataByte_ + partOneSize) : nextDataIndex * blockDataByte_;
329 0 : void* nextDstAddr = static_cast<u8*>(userInput_.ptr()) + nextOffsetByte;
330 :
331 0 : CHK_RET(DoSerialReduce(
332 : nextRemDMAMemPtr, nextDstAddr, partTwoSize, partTwoCount, subStreams_[1], nextIntralLink, nextOffsetByte));
333 0 : HCCL_INFO(
334 : "[ReduceScatterUnifiedMarch][RunHalfSliceRead] intralRank [%u] reduce with ringNextRank [%u] done", intraRank_,
335 : ringNextRank);
336 :
337 : /* 4卡及以上的 场景，在最后一步的notifyDone */
338 : // 单算子使用Ack/Datasignal接口，必须保证两者交替使用
339 0 : if (step == totalStep) {
340 0 : const u32 NOTIFY_IDX_TWO = 2;
341 0 : if (notifyIdx_ % NOTIFY_IDX_TWO != 0 && GetWorkflowMode() == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
342 0 : notifyIdx_++;
343 0 : CHK_RET(WaitSubStreamFinish(NEIGHBORS_NUM_TWO));
344 0 : CHK_RET(NotifySubStreamStart(NEIGHBORS_NUM_TWO));
345 0 : CHK_RET(NotifyNeighborsStart(prevIntraLink, nextIntralLink, NEIGHBORS_NUM_TWO));
346 : }
347 0 : CHK_RET(NotifyNeighborsEnd(prevIntraLink, nextIntralLink, NEIGHBORS_NUM_TWO));
348 : }
349 0 : notifyIdx_++;
350 :
351 0 : return HCCL_SUCCESS;
352 0 : }
353 :
354 0 : HcclResult ReduceScatterUnifiedMarch::RunAsync()
355 : {
356 0 : HcclOpMetaInfoDef opMeta = HcclOpMetaInfo::GetOneForReduceScatter();
357 0 : CHK_RET(InitTask(dispatcher_, mainStream_, opMeta.isEnableCache, opMeta.GetCacheKey()));
358 :
359 : // 获取link的收、发
360 0 : u32 ringPrevRank = (intraRank_ + intraRankSize_ - 1) % intraRankSize_;
361 0 : u32 ringNextRank = (intraRank_ + 1) % intraRankSize_;
362 :
363 0 : u32 neighbors = (ringPrevRank == ringNextRank) ? NEIGHBORS_NUM_ONE : NEIGHBORS_NUM_TWO;
364 0 : CHK_RET(NotifySubStreamStart(neighbors));
365 :
366 : // 计算所需的总步骤
367 0 : u32 totalStep = intraRankSize_ / DIVISOR_NUM_TWO + 1;
368 0 : u32 step = 0;
369 0 : if (totalStep == DIVISOR_NUM_TWO) {
370 0 : CHK_RET(RunSingleSliceRead(ringPrevRank, ringNextRank, step, totalStep));
371 : } else {
372 : // 进行第1步收发
373 0 : CHK_RET(RunHalfSliceRead(ringPrevRank, ringNextRank, step, totalStep));
374 :
375 : // 进行第k步收发
376 0 : step++;
377 0 : for (; step < totalStep - DIVISOR_NUM_TWO; step++) {
378 0 : CHK_RET(RunSingleSliceRead(ringPrevRank, ringNextRank, step, totalStep));
379 : }
380 :
381 : // 进行第totalStep - 1步收发
382 0 : CHK_RET(RunHalfSliceRead(ringPrevRank, ringNextRank, step, totalStep));
383 :
384 : // 进行第totalStep步收发
385 0 : CHK_RET(RunHalfSliceRead(ringPrevRank, ringNextRank, totalStep, totalStep));
386 : }
387 :
388 0 : CHK_RET(WaitSubStreamFinish(neighbors));
389 0 : CHK_RET(LaunchTaskExtend(dispatcher_, mainStream_, subStreams_));
390 :
391 0 : HCCL_INFO("[ReduceScatterUnifiedMarch][RunAsync] finished.");
392 0 : return HCCL_SUCCESS;
393 : }
394 :
395 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_UNIFIED_MARCH, ReduceScatterUnifiedMarch);
396 : } // namespace hccl
|
