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 "aligned_all_gather_double_ring.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 34 : AlignedAllGatherDoubleRing::AlignedAllGatherDoubleRing(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher)
16 : {
17 34 : }
18 :
19 68 : AlignedAllGatherDoubleRing::~AlignedAllGatherDoubleRing()
20 : {
21 68 : }
22 :
23 34 : HcclResult AlignedAllGatherDoubleRing::Prepare(HcomCollOpInfo *opInfo, const u32 userRank,
24 : std::vector<Stream> &subStreams, std::vector<std::shared_ptr<LocalNotify>> &mainSignals,
25 : std::vector<std::shared_ptr<LocalNotify>> &subSignals, const std::vector<std::vector<u32>> &ringsOrders,
26 : const std::vector<std::vector<Slice>> &userMemOutputSlicesOfDoubleRing)
27 : {
28 34 : opInfo_ = opInfo;
29 34 : userRank_ = userRank;
30 34 : subStreams_ = subStreams;
31 34 : mainSignals_ = mainSignals;
32 34 : subSignals_ = subSignals;
33 34 : ringsOrders_ = ringsOrders;
34 34 : userMemOutputSlicesOfDoubleRing_ = userMemOutputSlicesOfDoubleRing;
35 34 : return HCCL_SUCCESS;
36 : }
37 :
38 : // 服务器间allgather的入口函数
39 0 : HcclResult AlignedAllGatherDoubleRing::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
40 : {
41 : // 基本的检查
42 0 : CHK_RET(CheckParameters(rank, rankSize, links));
43 :
44 0 : if (rankSize == 1) {
45 0 : CHK_RET(OneRankMemcpy());
46 0 : return HCCL_SUCCESS;
47 : }
48 : // 收集邻居信息
49 0 : CHK_RET(GetInitializedNeighborLinks(rank, rankSize, links));
50 :
51 : // 填充slice_
52 0 : CHK_RET(SetSlices(rank, rankSize));
53 :
54 0 : HCCL_DEBUG("[AlignedAllGatherDoubleRing]RunAsync begins");
55 :
56 : // 运行all-gather, ring算法
57 0 : CHK_RET(RunAllGather(rank, rankSize));
58 :
59 0 : if (barrierSwitchOn_) {
60 : // 执行barrier，保证数据收发完成
61 0 : CHK_RET(ExecuteBarrier(leftLink_, rightLink_));
62 : }
63 :
64 0 : CHK_RET(LaunchTaskExtend(dispatcher_, stream_, subStreams_));
65 :
66 0 : HCCL_INFO("AlignedAllGatherDoubleRing finished: rank[%u] end", rank);
67 0 : return HCCL_SUCCESS;
68 : }
69 :
70 0 : HcclResult AlignedAllGatherDoubleRing::CheckParameters(const u32 rank, const u32 rankSize,
71 : const std::vector<LINK> &links)
72 : {
73 0 : CHK_PTR_NULL(opInfo_);
74 0 : CHK_RET(CheckConcurrentDirectParameters(rank, rankSize, links));
75 : // 判断subStreams数量是否正确
76 0 : CHK_PRT_RET(subStreams_.size() < 1,
77 : HCCL_ERROR("[AlignedAllGatherDoubleRing] subStreams size[%u] is less than 1", subStreams_.size()),
78 : HCCL_E_PARA);
79 0 : for (auto &s : subStreams_) {
80 0 : CHK_PTR_NULL(s.ptr());
81 : }
82 : // 判断mainSignals数量是否正确
83 0 : CHK_PRT_RET(mainSignals_.size() < 1,
84 : HCCL_ERROR("[AlignedAllGatherDoubleRing] mainSignals size[%u] is less than 1", mainSignals_.size()),
85 : HCCL_E_PARA);
86 : // 判断subSignals数量是否正确
87 0 : CHK_PRT_RET(subSignals_.size() < 1,
88 : HCCL_ERROR("[AlignedAllGatherDoubleRing] subSignals size[%u] is less than 1", subSignals_.size()),
89 : HCCL_E_PARA);
90 : // 判断ringsOrder数量是否正确
91 0 : for (u32 ringIndex = 0; ringIndex < ringsOrders_.size(); ringIndex++) {
92 0 : CHK_PRT_RET(ringsOrders_[ringIndex].size() % rankSize != 0,
93 : HCCL_ERROR("[AlignedAllGatherDoubleRing] ringsOrders[%u] size[%u] can not be divided by rank size[%u]",
94 : ringIndex, ringsOrders_[ringIndex].size(), rankSize), HCCL_E_PARA);
95 : }
96 : // 判断userMemOutputSlices数量是否正确
97 0 : for (u32 ringIndex = 0; ringIndex < userMemOutputSlicesOfDoubleRing_.size(); ringIndex++) {
98 0 : CHK_PRT_RET(userMemOutputSlicesOfDoubleRing_[ringIndex].size() % rankSize != 0,
99 : HCCL_ERROR("[AlignedAllGatherDoubleRing] userMemOutputSlicesOfDoubleRing[%u] size[%u] can not be divided by rank size[%u]",
100 : ringIndex, userMemOutputSlicesOfDoubleRing_[ringIndex].size(), rankSize), HCCL_E_PARA);
101 : }
102 0 : u32 mainSliceSize = multRingsSlices_[ALIGNED_MAIN_RING_INDEX].size() / rankSize;
103 0 : u32 subSliceSize = multRingsSlices_[ALIGNED_SUB_RING_INDEX].size() / rankSize;
104 0 : CHK_PRT_RET(mainSliceSize != subSliceSize,
105 : HCCL_ERROR("[AlignedAllGatherDoubleRing] mainSliceSize[%u] is not equal to subSliceSize[%u].",
106 : mainSliceSize, subSliceSize),
107 : HCCL_E_PARA);
108 0 : HCCL_INFO("AlignedAllGatherDoubleRing finished to CheckParameters");
109 0 : return HCCL_SUCCESS;
110 : }
111 :
112 0 : HcclResult AlignedAllGatherDoubleRing::OneRankMemcpy()
113 : {
114 0 : CHK_RET(MainRecordSub()); // 主流通知从流开始通信
115 0 : CHK_RET(SubWaitMain()); // 从流等待主流通知
116 0 : for (u32 ringIndex = 0; ringIndex < multRingsSlices_.size(); ringIndex++) {
117 0 : for (u32 sliceIdx = 0; sliceIdx < multRingsSlices_[ringIndex].size(); sliceIdx++) {
118 0 : const Slice &srcSlice = multRingsSlices_[ringIndex][sliceIdx];
119 0 : const Slice &dstSlice = userMemOutputSlicesOfDoubleRing_[ringIndex][sliceIdx];
120 0 : DeviceMem src;
121 0 : DeviceMem dst = DeviceMem::create(static_cast<u8 *>(opInfo_->outputAddr) + dstSlice.offset, dstSlice.size);
122 0 : if (opInfo_->inputAddr != nullptr) {
123 : // opInfo_->inputAddr != nullptr指示要从user input获取输入
124 0 : u64 stepOffset = multRingsSlices_[ringIndex][ringsOrders_[ringIndex][0]].offset;
125 0 : HCCL_DEBUG("Memcpy operation: stream[main], rank[%u] starts to copy offset[%llu], size[%llu] at userInput",
126 : userRank_, stepOffset, srcSlice.size);
127 0 : src = DeviceMem::create(static_cast<u8 *>(opInfo_->inputAddr) + stepOffset, srcSlice.size);
128 : } else {
129 : // opInfo_->inputAddr == nullptr指示要从CCL buffer获取输入
130 0 : HCCL_DEBUG("Memcpy operation: stream[main], rank[%u] starts to copy offset[%llu], size[%llu] at inputMem_",
131 : userRank_, srcSlice.offset, srcSlice.size);
132 0 : src = inputMem_.range(srcSlice.offset, srcSlice.size);
133 : }
134 0 : if (ringIndex == 1) {
135 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
136 : } else {
137 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, subStreams_[0]));
138 : }
139 0 : }
140 : }
141 0 : CHK_RET(SubRecordMain()); // 从流通知主流通信完成
142 0 : CHK_RET(MainWaitSub()); // 主流等待从流通知
143 0 : return HCCL_SUCCESS;
144 : }
145 :
146 0 : HcclResult AlignedAllGatherDoubleRing::GetInitializedNeighborLinks(const u32 rank, const u32 rankSize,
147 : const std::vector<LINK> &links)
148 : {
149 : // 收集左邻居信息
150 0 : leftLink_ = links[(rank + rankSize - 1) % rankSize];
151 0 : CHK_SMART_PTR_NULL(leftLink_);
152 :
153 : // 收集右邻居信息
154 0 : rightLink_ = links[(rank + 1) % rankSize];
155 0 : CHK_SMART_PTR_NULL(rightLink_);
156 0 : HCCL_INFO("AlignedAllGatherDoubleRing finished to GetInitializedNeighborLinks");
157 0 : return HCCL_SUCCESS;
158 : }
159 :
160 0 : HcclResult AlignedAllGatherDoubleRing::SetSlices(const u32 rank, const u32 rankSize)
161 : {
162 0 : for (u32 ringIndex = 0; ringIndex < multRingsSlices_.size(); ringIndex++) {
163 0 : if (multRingsSlices_[ringIndex].size() == 0) {
164 0 : multRingsSlices_[ringIndex].resize(rankSize);
165 :
166 0 : u64 sliceSize = count_ * DataUnitSize(dataType_);
167 0 : for (u32 i = 0; i < rankSize; i++) {
168 0 : multRingsSlices_[ringIndex][i].size = sliceSize;
169 0 : multRingsSlices_[ringIndex][i].offset = sliceSize * i;
170 0 : HCCL_DEBUG("multRingsSlices_[%u], rank[%u], slices[%u].offset=%llu, slices[%u].size=[%llu]",
171 : ringIndex, rank, i, multRingsSlices_[ringIndex][i].offset, i,
172 : multRingsSlices_[ringIndex][i].size);
173 : }
174 : }
175 0 : for (u32 i = 0; i < multRingsSlices_[ringIndex].size(); i++) {
176 0 : HCCL_DEBUG(
177 : "[AlignedAllGatherDoubleRing][SetSlices] multRingsSlices_[%u], rank[%u], slices[%u].offset=[%llu], slices[%u].size=[%llu]",
178 : ringIndex, rank, i, multRingsSlices_[ringIndex][i].offset, i, multRingsSlices_[ringIndex][i].size);
179 : }
180 : }
181 0 : HCCL_INFO("AlignedAllGatherDoubleRing finished to SetSlices");
182 0 : return HCCL_SUCCESS;
183 : }
184 :
185 0 : HcclResult AlignedAllGatherDoubleRing::RunInitStep(const u32 rank, const u32 rankSize)
186 : {
187 0 : for (u32 ringIndex = 0; ringIndex < multRingsSlices_.size(); ringIndex++) {
188 : // 第一步搬到userMemIn_的offset, 不同的ring环offset不一样
189 : u64 firstStepOffset;
190 0 : if (ringIndex == 0) {
191 0 : firstStepOffset = multRingsSlices_[ringIndex][ringsOrders_[ringIndex][0]].offset;
192 : } else {
193 0 : const auto &prevRingSlice = multRingsSlices_[ringIndex - 1][ringsOrders_[ringIndex - 1][rank]];
194 0 : const auto &slice = multRingsSlices_[ringIndex][ringsOrders_[ringIndex][rank]];
195 0 : firstStepOffset = slice.offset - prevRingSlice.offset;
196 : }
197 : // 第-1步，片内将部分数据从userIn搬到cclIn
198 0 : DeviceMem srcInit;
199 0 : DeviceMem dstInit;
200 : u32 initSliceIdx;
201 0 : if (ringIndex == 0) {
202 0 : initSliceIdx = rank;
203 : } else {
204 0 : initSliceIdx = (rankSize - rank) % rankSize;
205 : }
206 0 : u32 sliceSize = multRingsSlices_[ringIndex].size() / rankSize;
207 0 : for (u32 sliceIdx = 0; sliceIdx < sliceSize; sliceIdx++) {
208 0 : Slice initSlice = multRingsSlices_[ringIndex][initSliceIdx * sliceSize + sliceIdx];
209 : // 需要+userMemIn_的offset
210 0 : if (opInfo_->inputAddr != nullptr) {
211 : // AllGather算子调用AlignedAllGatherDoubleRing场景
212 0 : HCCL_DEBUG("Memcpy operation: step[-1] stream[main] src rank[%u] starts to copy(rcv) offset[%llu], "
213 : "size[%llu] on userMemOutput to offset[%llu], size[%llu] on CCL",
214 : userRank_, firstStepOffset, initSlice.size, initSlice.offset, initSlice.size);
215 0 : srcInit = DeviceMem::create(static_cast<u8 *>(opInfo_->inputAddr) + firstStepOffset, initSlice.size);
216 : } else {
217 : // AllReduce算子调用AlignedAllGatherDoubleRing场景
218 0 : HCCL_DEBUG("Memcpy operation: step[-1] stream[main] src rank[%u] starts to copy(rcv) offset[%llu], "
219 : "size[%llu] on CCL to offset[%llu], size[%llu] on CCL",
220 : userRank_, initSlice.offset, initSlice.size, initSlice.offset, initSlice.size);
221 0 : srcInit = inputMem_.range(initSlice.offset, initSlice.size);
222 : }
223 0 : dstInit = outputMem_.range(initSlice.offset, initSlice.size);
224 : // 若src与dst一样，则不需要搬运
225 0 : if (srcInit == dstInit) {
226 0 : continue;
227 : }
228 0 : if (ringIndex == 1) {
229 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstInit, srcInit, stream_));
230 : } else {
231 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstInit, srcInit, subStreams_[0]));
232 : }
233 : }
234 0 : }
235 0 : return HCCL_SUCCESS;
236 : }
237 :
238 0 : HcclResult AlignedAllGatherDoubleRing::PrepareRunMainStream(u32 ringIndex, Stream &stream,
239 : LINK &preLink, LINK &nextLink)
240 : {
241 0 : HCCL_DEBUG("AlignedAllGatherDoubleRing PrepareRunMainStream start");
242 0 : if (ringIndex == 1) {
243 0 : stream = stream_;
244 0 : preLink = rightLink_;
245 0 : nextLink = leftLink_;
246 : } else {
247 0 : stream = subStreams_[0];
248 0 : preLink = leftLink_;
249 0 : nextLink = rightLink_;
250 : }
251 0 : HCCL_DEBUG("AlignedAllGatherDoubleRing PrepareRunMainStream end");
252 0 : return HCCL_SUCCESS;
253 : }
254 :
255 0 : HcclResult AlignedAllGatherDoubleRing::RunAllStreams(const u32 step, const u32 rankSize,
256 : const std::vector<TxMemoryInfo> &mainTxMems, std::vector<RxMemoryInfo> &mainRxMems,
257 : const std::vector<TxMemoryInfo> &subTxMems, std::vector<RxMemoryInfo> &subRxMems,
258 : std::vector<DeviceMem> &mainLocalSrcMems, std::vector<DeviceMem> &mainLocalDstMems,
259 : std::vector<DeviceMem> &subLocalSrcMems, std::vector<DeviceMem> &subLocalDstMems)
260 : {
261 : (void)mainTxMems;
262 : (void)subTxMems;
263 0 : Stream mainStream;
264 0 : LINK mainPreLink;
265 0 : LINK mainNextLink;
266 0 : Stream subStream;
267 0 : LINK subPreLink;
268 0 : LINK subNextLink;
269 0 : CHK_RET(PrepareRunMainStream(ALIGNED_MAIN_RING_INDEX, mainStream, mainPreLink, mainNextLink));
270 0 : CHK_RET(PrepareRunMainStream(ALIGNED_SUB_RING_INDEX, subStream, subPreLink, subNextLink));
271 :
272 0 : CHK_RET(mainNextLink->TxAck(mainStream));
273 0 : CHK_RET(subNextLink->TxAck(subStream));
274 :
275 0 : CHK_RET(mainPreLink->RxAck(mainStream));
276 0 : CHK_RET(subPreLink->RxAck(subStream));
277 0 : u32 sliceSize = multRingsSlices_[ALIGNED_MAIN_RING_INDEX].size() / rankSize;
278 0 : for (u32 memIdx = 0; memIdx < sliceSize; memIdx++) {
279 0 : CHK_RET(RxAsyncMemcpy(step, ALIGNED_SUB_RING_INDEX, subRxMems[memIdx], subStream, subPreLink));
280 0 : CHK_RET(LocalMemcpy(ALIGNED_MAIN_RING_INDEX, mainLocalSrcMems[memIdx], mainLocalDstMems[memIdx]));
281 0 : CHK_RET(LocalMemcpy(ALIGNED_SUB_RING_INDEX, subLocalSrcMems[memIdx], subLocalDstMems[memIdx]));
282 0 : CHK_RET(RxAsyncMemcpy(step, ALIGNED_MAIN_RING_INDEX, mainRxMems[memIdx], mainStream, mainPreLink));
283 : }
284 0 : CHK_RET(mainPreLink->TxDataSignal(mainStream));
285 0 : CHK_RET(subPreLink->TxDataSignal(subStream));
286 :
287 0 : CHK_RET(mainNextLink->RxDataSignal(mainStream));
288 0 : CHK_RET(subNextLink->RxDataSignal(subStream));
289 0 : return HCCL_SUCCESS;
290 0 : }
291 :
292 0 : HcclResult AlignedAllGatherDoubleRing::RxAsyncMemcpy(const u32 step, const u32 ringIndex, RxMemoryInfo& mem, Stream &stream, LINK &link)
293 : {
294 : (void)step;
295 : // PreSync
296 0 : if (ringIndex == 1) {
297 0 : CHK_RET(MainWaitSub());
298 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
299 0 : CHK_RET(MainRecordSub());
300 : } else {
301 0 : CHK_RET(LocalNotify::Post(subStreams_[0], dispatcher_, mainSignals_[0], profilerInput_.stage));
302 0 : CHK_RET(LocalNotify::Wait(subStreams_[0], dispatcher_, subSignals_[0], profilerInput_.stage));
303 : }
304 0 : CHK_PTR_NULL(mem.dst);
305 0 : void *srcMemPtr = nullptr;
306 0 : CHK_RET(link->GetRemoteMem(mem.srcMemType, &srcMemPtr));
307 :
308 0 : DeviceMem srcDevMem(static_cast<s8 *>(srcMemPtr) + mem.srcOffset, mem.len);
309 0 : DeviceMem dstDevMem(static_cast<s8 *>(mem.dst), mem.len);
310 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstDevMem, srcDevMem,
311 : stream, link->GetRemoteRank(), link->GetLinkType()));
312 0 : return HCCL_SUCCESS;
313 0 : }
314 :
315 0 : HcclResult AlignedAllGatherDoubleRing::LocalMemcpy(const u32 ringIndex,
316 : DeviceMem &localSrcMem, DeviceMem &localDstMem)
317 : {
318 : // 校验流数
319 0 : if (GetWorkflowMode() != HcclWorkflowMode::HCCL_WORKFLOW_MODE_OPS_KERNEL_INFO_LIB) {
320 0 : CHK_RET(LocalNotify::Post(subStreams_[ringIndex + 1], dispatcher_, mainSignals_[ringIndex + 1], profilerInput_.stage));
321 0 : CHK_RET(LocalNotify::Wait(subStreams_[ringIndex + 1], dispatcher_, subSignals_[ringIndex + 1], profilerInput_.stage));
322 0 : if (localSrcMem != localDstMem) {
323 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, localDstMem, localSrcMem, subStreams_[ringIndex + 1]));
324 : }
325 : }
326 0 : return HCCL_SUCCESS;
327 : }
328 :
329 0 : HcclResult AlignedAllGatherDoubleRing::PrepareDeviceMems(
330 : const u32 step, const u32 ringIndex, const u32 rankSize,
331 : const u32 txSliceIdx, const u32 rxSliceIdx,
332 : std::vector<TxMemoryInfo> &txMems, std::vector<RxMemoryInfo> &rxMems,
333 : std::vector<DeviceMem> &localSrcMems, std::vector<DeviceMem> &localDstMems)
334 : {
335 0 : u32 sliceSize = multRingsSlices_[ringIndex].size() / rankSize;
336 0 : for (u32 sliceIdx = 0; sliceIdx < sliceSize; sliceIdx++) {
337 0 : const Slice &rxSlice = multRingsSlices_[ringIndex][rxSliceIdx * sliceSize + sliceIdx];
338 0 : const Slice &mainSlice = userMemOutputSlicesOfDoubleRing_[ringIndex][rxSliceIdx * sliceSize + sliceIdx];
339 0 : const Slice &txSlice = multRingsSlices_[ringIndex][txSliceIdx * sliceSize + sliceIdx];
340 0 : const Slice &subSlice = userMemOutputSlicesOfDoubleRing_[ringIndex][txSliceIdx * sliceSize + sliceIdx];
341 : // PrepareTxRxMems
342 0 : DeviceMem src = outputMem_.range(txSlice.offset, txSlice.size);
343 0 : HCCL_DEBUG("tx srcMem[%p] range[%llu] size[%llu] ", src.ptr(),
344 : txSlice.offset, txSlice.size);
345 0 : txMems.emplace_back(TxMemoryInfo{UserMemType::OUTPUT_MEM, txSlice.offset + baseOffset_,
346 0 : src.ptr(), txSlice.size});
347 0 : DeviceMem dst;
348 0 : if (step == rankSize - DMA_REDUCE_TWO_OFFSET) {
349 0 : HCCL_DEBUG(
350 : "DMAReduce(sdma) MemcpyAsync operation: step[%u] stream[main], dst rank[%u] starts to rcv "
351 : "offset[%llu] size[%llu] at userMemOutput_",
352 : step, userRank_, mainSlice.offset, mainSlice.size);
353 0 : dst = DeviceMem::create(static_cast<u8 *>(opInfo_->outputAddr) + mainSlice.offset,
354 0 : mainSlice.size);
355 : } else {
356 0 : HCCL_DEBUG(
357 : "MemcpyAsync operation: step[%u] stream[main], dst rank[%u] starts to rcv offset[%llu] size[%llu] "
358 : "at outputMem_",
359 : step, userRank_, rxSlice.offset, rxSlice.size);
360 0 : dst = outputMem_.range(rxSlice.offset, rxSlice.size);
361 : }
362 0 : rxMems.emplace_back(RxMemoryInfo{UserMemType::OUTPUT_MEM, rxSlice.offset + baseOffset_,
363 0 : dst.ptr(), rxSlice.size});
364 : // PrepareLocalCopyDeviceMems
365 : // 从流
366 0 : src = outputMem_.range(txSlice.offset, txSlice.size);
367 0 : dst = DeviceMem::create(static_cast<u8 *>(opInfo_->outputAddr) + subSlice.offset,
368 0 : subSlice.size);
369 0 : HCCL_DEBUG("Memcpy operation: step[%u] stream[sub], src rank[%u] starts to send offset[%llu] size[%llu], "
370 : "dst rank starts to rcv offset[%llu] size[%llu] at userMemOutput_",
371 : step, userRank_, subSlice.offset, subSlice.size, txSlice.offset, txSlice.size);
372 0 : localSrcMems.emplace_back(src);
373 0 : localDstMems.emplace_back(dst);
374 0 : }
375 0 : return HCCL_SUCCESS;
376 : }
377 :
378 0 : HcclResult AlignedAllGatherDoubleRing::RunAllGather(const u32 rank, const u32 rankSize)
379 : {
380 0 : HCCL_INFO("AlignedAllGatherDoubleRing starts, the input param rank[%u]", rank);
381 0 : if (GetWorkflowMode() != HcclWorkflowMode::HCCL_WORKFLOW_MODE_OPS_KERNEL_INFO_LIB) {
382 : // 主环主流通知从环主流开始通信
383 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, subSignals_[0], profilerInput_.stage));
384 : // 从环主流等待主环主流通知
385 0 : CHK_RET(LocalNotify::Wait(subStreams_[0], dispatcher_, subSignals_[0], profilerInput_.stage));
386 0 : CHK_RET(RunInitStep(rank, rankSize));
387 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
388 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, subStreams_[0], dispatcher_));
389 : // 从流通知主流通信完成
390 0 : CHK_RET(LocalNotify::Post(subStreams_[0], dispatcher_, mainSignals_[0], profilerInput_.stage));
391 : // 主流等待从流通知
392 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, mainSignals_[0], profilerInput_.stage));
393 : }
394 : // 主环主流通知从环主流开始通信
395 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
396 0 : CHK_RET(MainRecordSub());
397 : // 从环主流等待主环主流通知
398 0 : CHK_RET(SubWaitMain());
399 0 : u32 txSliceIdxSub = rank;
400 0 : u32 rxSliceIdxSub = (rank + rankSize - 1) % rankSize;
401 0 : u32 txSliceIdxMain = (rankSize - rank) % rankSize;
402 0 : u32 rxSliceIdxMain = (rankSize - rank - 1 + rankSize) % rankSize;
403 : // 空拷贝用于主从流任务并发
404 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
405 0 : CHK_RET(ExecEmptyTasks());
406 0 : for (u32 step = 0; step < rankSize - 1; step++) {
407 0 : std::vector<TxMemoryInfo> txMemsSub;
408 0 : std::vector<RxMemoryInfo> rxMemsSub;
409 0 : std::vector<DeviceMem> localSrcMemsSub;
410 0 : std::vector<DeviceMem> localDstMemsSub;
411 0 : CHK_RET(PrepareDeviceMems(
412 : step, ALIGNED_SUB_RING_INDEX, rankSize,
413 : txSliceIdxSub, rxSliceIdxSub,
414 : txMemsSub, rxMemsSub,
415 : localSrcMemsSub, localDstMemsSub));
416 0 : std::vector<TxMemoryInfo> txMemsMain;
417 0 : std::vector<RxMemoryInfo> rxMemsMain;
418 0 : std::vector<DeviceMem> localSrcMemsMain;
419 0 : std::vector<DeviceMem> localDstMemsMain;
420 0 : CHK_RET(PrepareDeviceMems(
421 : step, ALIGNED_MAIN_RING_INDEX, rankSize,
422 : txSliceIdxMain, rxSliceIdxMain,
423 : txMemsMain, rxMemsMain,
424 : localSrcMemsMain, localDstMemsMain));
425 0 : CHK_RET(RunAllStreams(step, rankSize, txMemsMain, rxMemsMain, txMemsSub, rxMemsSub,
426 : localSrcMemsMain, localDstMemsMain, localSrcMemsSub, localDstMemsSub));
427 :
428 : // 更新索引
429 0 : txSliceIdxSub = (txSliceIdxSub + rankSize - 1) % rankSize;
430 0 : rxSliceIdxSub = (rxSliceIdxSub + rankSize - 1) % rankSize;
431 0 : txSliceIdxMain = (txSliceIdxMain + rankSize - 1) % rankSize;
432 0 : rxSliceIdxMain = (rxSliceIdxMain + rankSize - 1) % rankSize;
433 0 : }
434 : // 从环主流通知主环主流通信完成
435 0 : CHK_RET(SubRecordMain());
436 : // 主环主流等待从环主流通知
437 0 : CHK_RET(MainWaitSub());
438 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
439 0 : HCCL_INFO("AlignedAllGatherDoubleRing finished to RunAllGather");
440 0 : return HCCL_SUCCESS;
441 : }
442 :
443 0 : HcclResult AlignedAllGatherDoubleRing::ExecEmptyTasks()
444 : {
445 0 : for (u32 signalIndex = 0; signalIndex < subStreams_.size(); signalIndex++) {
446 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, subStreams_[signalIndex], dispatcher_));
447 : }
448 0 : return HCCL_SUCCESS;
449 : }
450 :
451 : // 主流通知从流干活
452 0 : HcclResult AlignedAllGatherDoubleRing::MainRecordSub()
453 : {
454 0 : for (u32 signalIndex = 0; signalIndex < subSignals_.size(); signalIndex++) {
455 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, subSignals_[signalIndex],
456 : profilerInput_.stage));
457 : }
458 0 : return HCCL_SUCCESS;
459 : }
460 : // 从流等待主流
461 0 : HcclResult AlignedAllGatherDoubleRing::SubWaitMain()
462 : {
463 0 : for (u32 streamIndex = 0; streamIndex < subSignals_.size(); streamIndex++) {
464 0 : CHK_RET(LocalNotify::Wait(subStreams_[streamIndex], dispatcher_, subSignals_[streamIndex],
465 : profilerInput_.stage));
466 : }
467 0 : return HCCL_SUCCESS;
468 : }
469 : // 主流等待从流
470 0 : HcclResult AlignedAllGatherDoubleRing::MainWaitSub()
471 : {
472 0 : for (u32 signalIndex = 0; signalIndex < mainSignals_.size(); signalIndex++) {
473 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, mainSignals_[signalIndex], profilerInput_.stage));
474 : }
475 0 : return HCCL_SUCCESS;
476 : }
477 : // 从流告诉主流活干完了
478 0 : HcclResult AlignedAllGatherDoubleRing::SubRecordMain()
479 : {
480 0 : for (u32 streamIndex = 0; streamIndex < mainSignals_.size(); streamIndex++) {
481 0 : CHK_RET(LocalNotify::Post(subStreams_[streamIndex], dispatcher_, mainSignals_[streamIndex],
482 : profilerInput_.stage));
483 : }
484 0 : return HCCL_SUCCESS;
485 : }
486 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALIGNED_ALL_GATHER_DOUBLE_RING, AlignedAllGatherDoubleRing);
487 : } // namespace hccl
|
