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