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 "all_gather_ring_concurrent_direct.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : AllGatherRingConcurrentDirect::AllGatherRingConcurrentDirect(
16 0 : const HcclDispatcher dispatcher)
17 0 : : AlgTemplateBase(dispatcher)
18 : {
19 0 : }
20 :
21 0 : AllGatherRingConcurrentDirect::~AllGatherRingConcurrentDirect()
22 : {
23 0 : }
24 :
25 0 : HcclResult AllGatherRingConcurrentDirect::Prepare(HcomCollOpInfo *opInfo, const u32 userRank,
26 : std::vector<Stream> &subStreams, const std::vector<std::shared_ptr<LocalNotify>> &mainSignals,
27 : const std::vector<std::shared_ptr<LocalNotify>> &subSignals,
28 : const std::vector<u32> &ringsOrder, const std::vector<Slice> &userMemSlices, bool isSdma)
29 : {
30 0 : opInfo_ = opInfo;
31 0 : userRank_ = userRank;
32 0 : subStreams_ = subStreams;
33 0 : mainSignals_ = mainSignals;
34 0 : subSignals_ = subSignals;
35 0 : ringsOrder_ = ringsOrder;
36 0 : userMemOutputSlices_ = userMemSlices;
37 0 : isSdma_ = isSdma;
38 0 : return HCCL_SUCCESS;
39 : }
40 :
41 : // 服务器间allgather的入口函数
42 0 : HcclResult AllGatherRingConcurrentDirect::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
43 : {
44 : // 基本的检查
45 0 : CHK_RET(CheckParameters(rank, rankSize, links));
46 :
47 0 : if (rankSize == 1) {
48 0 : CHK_RET(OneRankMemcpy());
49 0 : return HCCL_SUCCESS;
50 : }
51 : // 收集邻居信息
52 0 : CHK_RET(GetInitializedNeighborLinks(rank, rankSize, links));
53 :
54 : // 填充slice_
55 0 : CHK_RET(SetSlices(rank, rankSize));
56 :
57 : // 运行all-gather, ring算法
58 0 : CHK_RET(RunAllGather(rank, rankSize));
59 :
60 0 : if (barrierSwitchOn_) {
61 : // 执行barrier，保证数据收发完成
62 0 : CHK_RET(ExecuteBarrier(leftLink_, rightLink_));
63 : }
64 :
65 0 : CHK_RET(LaunchTaskExtend(dispatcher_, stream_, subStreams_));
66 :
67 0 : HCCL_INFO("AllGatherRingConcurrentDirect finished: rank[%u] end", rank);
68 0 : return HCCL_SUCCESS;
69 : }
70 :
71 0 : HcclResult AllGatherRingConcurrentDirect::CheckParameters(const u32 rank, const u32 rankSize,
72 : const std::vector<LINK> &links)
73 : {
74 0 : CHK_PTR_NULL(opInfo_);
75 0 : CHK_RET(CheckConcurrentDirectParameters(rank, rankSize, links));
76 : // 判断subStreams数量是否正确
77 0 : CHK_PRT_RET(subStreams_.size() < 1,
78 : HCCL_ERROR("[AllGatherRingConcurrentDirect] subStreams size[%u] is less than 1", subStreams_.size()),
79 : HCCL_E_PARA);
80 0 : for (auto &s : subStreams_) {
81 0 : CHK_PTR_NULL(s.ptr());
82 : }
83 : // 判断mainSignals数量是否正确
84 0 : CHK_PRT_RET(mainSignals_.size() < 1,
85 : HCCL_ERROR("[AllGatherRingConcurrentDirect] mainSignals size[%u] is less than 1", mainSignals_.size()),
86 : HCCL_E_PARA);
87 : // 判断subSignals数量是否正确
88 0 : CHK_PRT_RET(subSignals_.size() < 1,
89 : HCCL_ERROR("[AllGatherRingConcurrentDirect] subSignals size[%u] is less than 1", subSignals_.size()),
90 : HCCL_E_PARA);
91 : // 判断ringsOrder数量是否正确
92 0 : CHK_PRT_RET(ringsOrder_.size() % rankSize != 0,
93 : HCCL_ERROR("[AllGatherRingConcurrentDirect] ringsOrder size[%u] can not be divided by rank size[%u]",
94 : ringsOrder_.size(), rankSize), HCCL_E_PARA);
95 : // 判断userMemInputSlices数量是否正确
96 0 : CHK_PRT_RET(userMemOutputSlices_.size() % rankSize != 0,
97 : HCCL_ERROR("[AllGatherRingConcurrentDirect] userMemOutputSlices size[%u] can not be divided by rank size[%u]",
98 : userMemOutputSlices_.size(), rankSize), HCCL_E_PARA);
99 0 : HCCL_INFO("AllGatherRingConcurrentDirect finished to CheckParameters");
100 0 : return HCCL_SUCCESS;
101 : }
102 :
103 0 : HcclResult AllGatherRingConcurrentDirect::OneRankMemcpy()
104 : {
105 0 : for (u32 sliceIdx = 0; sliceIdx < slices_.size(); sliceIdx++) {
106 0 : const Slice &srcSlice = slices_[sliceIdx];
107 0 : const Slice &dstSlice = userMemOutputSlices_[sliceIdx];
108 0 : DeviceMem src;
109 0 : DeviceMem dst = DeviceMem::create(static_cast<u8 *>(opInfo_->outputAddr) + dstSlice.offset, dstSlice.size);
110 0 : if (opInfo_->inputAddr != nullptr) {
111 : // opInfo_->inputAddr != nullptr指示要从user input获取输入
112 0 : u64 stepOffset = slices_[ringsOrder_[0]].offset;
113 0 : HCCL_DEBUG("Memcpy operation: stream[main], rank[%u] starts to copy offset[%llu], size[%llu] at userInput",
114 : userRank_, stepOffset, srcSlice.size);
115 0 : src = DeviceMem::create(static_cast<u8 *>(opInfo_->inputAddr) + stepOffset, srcSlice.size);
116 : } else {
117 : // opInfo_->inputAddr == nullptr指示要从CCL buffer获取输入
118 0 : HCCL_DEBUG("Memcpy operation: stream[main], rank[%u] starts to copy offset[%llu], size[%llu] at inputMem_",
119 : userRank_, srcSlice.offset, srcSlice.size);
120 0 : src = inputMem_.range(srcSlice.offset, srcSlice.size);
121 : }
122 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
123 0 : HCCL_DEBUG("[AllGatherRingConcurrentDirect][OneRankMemcpy]sliceIdx[%u] for Memcpy success", sliceIdx);
124 0 : }
125 :
126 0 : return HCCL_SUCCESS;
127 : }
128 :
129 0 : HcclResult AllGatherRingConcurrentDirect::GetInitializedNeighborLinks(const u32 rank, const u32 rankSize,
130 : const std::vector<LINK> &links)
131 : {
132 : // 收集左邻居信息
133 0 : leftLink_ = links[(rank + rankSize - 1) % rankSize];
134 0 : CHK_SMART_PTR_NULL(leftLink_);
135 :
136 : // 收集右邻居信息
137 0 : rightLink_ = links[(rank + 1) % rankSize];
138 0 : CHK_SMART_PTR_NULL(rightLink_);
139 0 : HCCL_INFO("AllGatherRingConcurrentDirect finished to GetInitializedNeighborLinks");
140 0 : return HCCL_SUCCESS;
141 : }
142 :
143 0 : HcclResult AllGatherRingConcurrentDirect::SetSlices(const u32 rank, const u32 rankSize)
144 : {
145 0 : inputSlices_ = slices_;
146 0 : if (slices_.size() == 0) {
147 0 : slices_.resize(rankSize);
148 0 : inputSlices_.resize(rankSize);
149 :
150 0 : u64 sliceSize = count_ * DataUnitSize(dataType_);
151 0 : for (u32 i = 0; i < rankSize; i++) {
152 0 : slices_[i].size = sliceSize;
153 0 : slices_[i].offset = sliceSize * i;
154 0 : inputSlices_[i].size = sliceSize;
155 0 : inputSlices_[i].offset = (inputMem_.size() < outputMem_.size()) ? 0 : (sliceSize * i);
156 0 : HCCL_DEBUG("rank[%u], slices[%u].offset=%llu, slices[%u].size=[%llu]", rank, i, slices_[i].offset, i,
157 : slices_[i].size);
158 : }
159 : }
160 0 : if (UNLIKELY(HcclCheckLogLevel(DLOG_DEBUG))) {
161 0 : for (u32 i = 0; i < slices_.size(); i++) {
162 0 : HCCL_DEBUG(
163 : "[AllGatherRingConcurrentDirect][SetSlices]rank[%u], slices[%u].offset=[%llu], slices[%u].size=[%llu]",
164 : rank, i, slices_[i].offset, i, slices_[i].size);
165 : }
166 : }
167 0 : HCCL_INFO("AllGatherRingConcurrentDirect finished to SetSlices");
168 0 : return HCCL_SUCCESS;
169 : }
170 :
171 0 : HcclResult AllGatherRingConcurrentDirect::RunInitStep(const u32 rank, const u32 rankSize)
172 : {
173 : // 第一步搬到userMemIn_的offset, 不同的ring环offset不一样
174 0 : auto firstStepOffset = slices_[ringsOrder_[0]].offset;
175 : // 第-1步，片内将部分数据从userIn搬到cclIn
176 0 : DeviceMem srcInit;
177 0 : DeviceMem dstInit;
178 0 : u32 initSliceIdx = rank;
179 0 : u32 sliceSize = slices_.size() / rankSize;
180 0 : for (u32 sliceIdx = 0; sliceIdx < sliceSize; sliceIdx++) {
181 0 : Slice initSlice = slices_[initSliceIdx * sliceSize + sliceIdx];
182 : // 需要+userMemIn_的offset
183 0 : if (opInfo_->inputAddr != nullptr) {
184 : // AllGather算子调用AllGatherRingConcurrentDirect场景
185 0 : srcInit = DeviceMem::create(static_cast<u8 *>(opInfo_->inputAddr) + firstStepOffset, initSlice.size);
186 : } else {
187 : // AllReduce算子调用AllGatherRingConcurrentDirect场景
188 0 : srcInit = inputMem_.range(initSlice.offset, initSlice.size);
189 : }
190 0 : dstInit = outputMem_.range(initSlice.offset, initSlice.size);
191 0 : HCCL_DEBUG("Memcpy operation: step[-1] stream[main] src rank[%u] starts to copy(rcv) offset[%llu], "
192 : "size[%llu] on userMemOutput to offset[%llu], size[%llu] on CCL",
193 : userRank_, firstStepOffset, initSlice.size, initSlice.offset, initSlice.size);
194 : // 若src与dst一样，则不需要搬运
195 0 : if (srcInit != dstInit) {
196 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstInit, srcInit, stream_));
197 : }
198 : }
199 0 : return HCCL_SUCCESS;
200 0 : }
201 :
202 0 : HcclResult AllGatherRingConcurrentDirect::RunAllGather(const u32 rank, const u32 rankSize)
203 : {
204 0 : HCCL_INFO("AllGatherRingConcurrentDirect starts, the input param rank[%u]", rank);
205 : // 空拷贝用于后续操作附着
206 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
207 :
208 0 : CHK_RET(RunInitStep(rank, rankSize));
209 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
210 0 : CHK_RET(MainRecordSub()); // 主流通知从流开始通信
211 0 : CHK_RET(SubWaitMain()); // 从流等待主流通知
212 : // 空拷贝用于主从流任务并发
213 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
214 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, subStreams_[0], dispatcher_));
215 0 : u32 txSliceIdx = rank;
216 0 : u32 sliceSize = slices_.size() / rankSize;
217 0 : u32 rxSliceIdx = (rank + rankSize - 1) % rankSize;
218 :
219 0 : std::vector<DeviceMem> finalSrc;
220 0 : std::vector<DeviceMem> finalDst;
221 0 : for (u32 step = 0; step < rankSize - 1; step++) {
222 0 : std::vector<Slice> rxSliceVector;
223 0 : std::vector<Slice> mainSliceVector;
224 0 : std::vector<Slice> txSliceVector;
225 0 : std::vector<Slice> subSliceVector;
226 0 : for (u32 sliceIdx = 0; sliceIdx < sliceSize; sliceIdx++) {
227 0 : rxSliceVector.push_back(slices_[rxSliceIdx * sliceSize + sliceIdx]);
228 0 : mainSliceVector.push_back(userMemOutputSlices_[rxSliceIdx * sliceSize + sliceIdx]);
229 0 : txSliceVector.push_back(slices_[txSliceIdx * sliceSize + sliceIdx]);
230 0 : subSliceVector.push_back(userMemOutputSlices_[txSliceIdx * sliceSize + sliceIdx]);
231 : }
232 : // 从流
233 0 : if (!isSdma_) {
234 0 : CHK_RET(LocalNotify::Post(subStreams_[0], dispatcher_, mainSignals_[0], profilerInput_.stage));
235 0 : CHK_RET(LocalNotify::Wait(subStreams_[0], dispatcher_, subSignals_[0], profilerInput_.stage));
236 : }
237 0 : for (u32 sliceIdx = 0; sliceIdx < sliceSize; sliceIdx++) {
238 0 : if (isSdma_) {
239 0 : CHK_RET(LocalNotify::Post(subStreams_[0], dispatcher_, mainSignals_[0], profilerInput_.stage));
240 0 : CHK_RET(LocalNotify::Wait(subStreams_[0], dispatcher_, subSignals_[0], profilerInput_.stage));
241 : }
242 0 : DeviceMem src = outputMem_.range(txSliceVector[sliceIdx].offset, txSliceVector[sliceIdx].size);
243 0 : DeviceMem dst = DeviceMem::create(static_cast<u8 *>(opInfo_->outputAddr) + subSliceVector[sliceIdx].offset,
244 0 : subSliceVector[sliceIdx].size);
245 0 : HCCL_DEBUG("Memcpy operation: step[%u] stream[sub], src rank[%u] starts to send offset[%llu] size[%llu], "
246 : "dst rank starts to rcv offset[%llu] size[%llu] at userMemOutput_",
247 : step, userRank_, subSliceVector[sliceIdx].offset, subSliceVector[sliceIdx].size,
248 : txSliceVector[sliceIdx].offset, txSliceVector[sliceIdx].size);
249 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, subStreams_[0]));
250 0 : }
251 :
252 : // 主流
253 : // Ack
254 0 : CHK_RET(leftLink_->TxAck(stream_));
255 0 : CHK_RET(rightLink_->RxAck(stream_));
256 :
257 0 : std::vector<TxMemoryInfo> txMems;
258 0 : std::vector<RxMemoryInfo> rxMems;
259 0 : for (u32 sliceIdx = 0; sliceIdx < sliceSize; sliceIdx++) {
260 0 : DeviceMem src = outputMem_.range(txSliceVector[sliceIdx].offset, txSliceVector[sliceIdx].size);
261 0 : HCCL_DEBUG("tx srcMem[%p] range[%llu] size[%llu] ", src.ptr(),
262 : txSliceVector[sliceIdx].offset, txSliceVector[sliceIdx].size);
263 0 : txMems.emplace_back(TxMemoryInfo{UserMemType::OUTPUT_MEM, txSliceVector[sliceIdx].offset + baseOffset_,
264 0 : src.ptr(), txSliceVector[sliceIdx].size});
265 0 : DeviceMem dst;
266 0 : if (isSdma_ && step == rankSize - DMA_REDUCE_TWO_OFFSET) {
267 0 : HCCL_DEBUG(
268 : "DMAReduce(sdma) MemcpyAsync operation: step[%u] stream[main], dst rank[%u] starts to rcv "
269 : "offset[%llu] size[%llu] at userMemOutput_",
270 : step, userRank_, mainSliceVector[sliceIdx].offset, mainSliceVector[sliceIdx].size);
271 0 : dst = DeviceMem::create(static_cast<u8 *>(opInfo_->outputAddr) + mainSliceVector[sliceIdx].offset,
272 0 : mainSliceVector[sliceIdx].size);
273 : } else {
274 0 : HCCL_DEBUG(
275 : "MemcpyAsync operation: step[%u] stream[main], dst rank[%u] starts to rcv offset[%llu] size[%llu] "
276 : "at outputMem_",
277 : step, userRank_, rxSliceVector[sliceIdx].offset, rxSliceVector[sliceIdx].size);
278 0 : dst = outputMem_.range(rxSliceVector[sliceIdx].offset, rxSliceVector[sliceIdx].size);
279 0 : if (!isSdma_ && step == rankSize - DMA_REDUCE_TWO_OFFSET) {
280 0 : HCCL_DEBUG("DMAReduce(rdma) record final addr");
281 0 : finalSrc.push_back(outputMem_.range(rxSliceVector[sliceIdx].offset, rxSliceVector[sliceIdx].size));
282 0 : finalDst.push_back(DeviceMem::create(static_cast<u8 *>(opInfo_->outputAddr) +
283 0 : mainSliceVector[sliceIdx].offset, mainSliceVector[sliceIdx].size));
284 : }
285 : }
286 0 : rxMems.emplace_back(RxMemoryInfo{UserMemType::OUTPUT_MEM, rxSliceVector[sliceIdx].offset + baseOffset_,
287 0 : dst.ptr(), rxSliceVector[sliceIdx].size});
288 0 : }
289 0 : CHK_RET(rightLink_->TxAsync(txMems, stream_));
290 0 : if (!isSdma_) {
291 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, mainSignals_[0], profilerInput_.stage));
292 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
293 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, subSignals_[0], profilerInput_.stage));
294 0 : CHK_RET(leftLink_->RxAsync(rxMems, stream_));
295 : } else {
296 0 : CHK_RET(leftLink_->RxDataSignal(stream_));
297 0 : for (auto& mem : rxMems) {
298 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, mainSignals_[0], profilerInput_.stage));
299 0 : CHK_RET(AlgTemplateBase::ExecEmptyTask(inputMem_, outputMem_, stream_, dispatcher_));
300 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, subSignals_[0], profilerInput_.stage));
301 0 : CHK_PTR_NULL(mem.dst);
302 0 : void *srcMemPtr = nullptr;
303 0 : CHK_RET(leftLink_->GetRemoteMem(mem.srcMemType, &srcMemPtr));
304 :
305 0 : DeviceMem srcDevMem(static_cast<s8 *>(srcMemPtr) + mem.srcOffset, mem.len);
306 0 : DeviceMem dstDevMem(static_cast<s8 *>(mem.dst), mem.len);
307 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstDevMem, srcDevMem,
308 : stream_, leftLink_->GetRemoteRank(), leftLink_->GetLinkType()));
309 0 : }
310 : }
311 :
312 : // 更新索引
313 0 : txSliceIdx = (txSliceIdx + rankSize - 1) % rankSize;
314 0 : rxSliceIdx = (rxSliceIdx + rankSize - 1) % rankSize;
315 0 : }
316 0 : CHK_RET(SubRecordMain()); // 从流通知主流通信完成
317 0 : CHK_RET(MainWaitSub()); // 主流等待从流通知
318 0 : if (!isSdma_) {
319 0 : for (u32 vecIdx = 0; vecIdx < finalSrc.size(); vecIdx++) {
320 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, finalDst[vecIdx], finalSrc[vecIdx], stream_));
321 : }
322 : }
323 0 : HCCL_INFO("AllGatherRingConcurrentDirect finished to RunAllGather");
324 0 : return HCCL_SUCCESS;
325 0 : }
326 :
327 : // 主流通知从流干活
328 0 : HcclResult AllGatherRingConcurrentDirect::MainRecordSub()
329 : {
330 0 : for (u32 signalIndex = 0; signalIndex < subSignals_.size(); signalIndex++) {
331 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, subSignals_[signalIndex],
332 : profilerInput_.stage));
333 : }
334 0 : return HCCL_SUCCESS;
335 : }
336 : // 从流等待主流
337 0 : HcclResult AllGatherRingConcurrentDirect::SubWaitMain()
338 : {
339 0 : for (u32 streamIndex = 0; streamIndex < subSignals_.size(); streamIndex++) {
340 0 : CHK_RET(LocalNotify::Wait(subStreams_[streamIndex], dispatcher_, subSignals_[streamIndex],
341 : profilerInput_.stage));
342 : }
343 0 : return HCCL_SUCCESS;
344 : }
345 : // 主流等待从流
346 0 : HcclResult AllGatherRingConcurrentDirect::MainWaitSub()
347 : {
348 0 : for (u32 signalIndex = 0; signalIndex < mainSignals_.size(); signalIndex++) {
349 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, mainSignals_[signalIndex], profilerInput_.stage));
350 : }
351 0 : return HCCL_SUCCESS;
352 : }
353 : // 从流告诉主流活干完了
354 0 : HcclResult AllGatherRingConcurrentDirect::SubRecordMain()
355 : {
356 0 : for (u32 streamIndex = 0; streamIndex < mainSignals_.size(); streamIndex++) {
357 0 : CHK_RET(LocalNotify::Post(subStreams_[streamIndex], dispatcher_, mainSignals_[streamIndex],
358 : profilerInput_.stage));
359 : }
360 0 : return HCCL_SUCCESS;
361 : }
362 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_GATHER_RING_CONCURRENT_DIRECT, AllGatherRingConcurrentDirect);
363 : } // namespace hccl
|
