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 "multi_root_scatter_ring.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : bool DscendSortWithSliceSendEnd(const SliceSendRange& a, const SliceSendRange& b) { return (a.endRank > b.endRank); }
16 :
17 0 : MultiRootScatterRing::MultiRootScatterRing(const HcclDispatcher dispatcher)
18 : : AlgTemplateBase(dispatcher),
19 0 : interRank_(0),
20 0 : interRankSize_(0)
21 0 : {}
22 :
23 0 : MultiRootScatterRing::~MultiRootScatterRing() {}
24 :
25 0 : void MultiRootScatterRing::SlicesDataPrepare(const u32 unitSize, const u64 totalCount, const u32 rankSize) const
26 : {
27 0 : slices_.resize(rankSize);
28 0 : u64 sliceSize = (totalCount / rankSize) * unitSize;
29 0 : for (u32 i = 0; i < rankSize; i++) {
30 0 : slices_[i].offset = i * sliceSize;
31 0 : slices_[i].size = sliceSize;
32 0 : HCCL_DEBUG("rank[%u] default slice[%u]: offset: [%llu] size[%llu]", interRank_, i, i * sliceSize, sliceSize);
33 : }
34 0 : }
35 :
36 : // scatter的入口函数
37 : HcclResult
38 0 : MultiRootScatterRing::RunAsync(const u32 rank, const u32 rankSize, const std::vector<std::shared_ptr<Transport>>& links)
39 : {
40 0 : CHK_SMART_PTR_NULL(dispatcher_);
41 0 : CHK_PTR_NULL(stream_.ptr());
42 0 : if (!outputMem_ || !inputMem_) {
43 0 : HCCL_ERROR("[MultiRootScatterRing][RunAsync]run_async inputmem or outputmem is null");
44 0 : return HCCL_E_PTR;
45 : }
46 :
47 0 : interRank_ = rank;
48 0 : interRankSize_ = rankSize;
49 :
50 0 : HCCL_INFO(
51 : "MultiRootScatterRing run: rank[%u] totalrank[%u] count[%llu] input[%p] output[%p]", interRank_, interRankSize_,
52 : count_, inputMem_.ptr(), outputMem_.ptr());
53 :
54 : // ranksize为1时，只有当input!=output 时候进行拷贝
55 0 : if (interRankSize_ == 1) {
56 0 : if (inputMem_ != outputMem_) {
57 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_));
58 : }
59 0 : return HCCL_SUCCESS;
60 : }
61 :
62 0 : u32 unitSize = DataUnitSize(dataType_);
63 0 : CHK_PRT_RET(
64 : unitSize == 0, HCCL_ERROR("[MultiRootScatterRing][RunAsync]rank[%u] unit data size is zero", rank),
65 : HCCL_E_INTERNAL);
66 :
67 : // 带入vecotr为空，计算每个rank的结果偏移和大小
68 0 : if (slices_.size() == 0) {
69 0 : SlicesDataPrepare(unitSize, count_, interRankSize_);
70 : }
71 :
72 : // 获取link的收、发缓存, 计算chunk_size
73 0 : u32 ringPrevRank = (rank + rankSize - 1) % rankSize;
74 0 : u32 ringNextRank = (rank + 1) % rankSize;
75 :
76 0 : if (links.size() < rankSize) {
77 0 : HCCL_ERROR(
78 : "[MultiRootScatterRing][RunAsync]rank[%u] link size[%llu] is less than rank size", rank, links.size());
79 0 : return HCCL_E_INTERNAL;
80 : }
81 :
82 0 : linkLeft_ = links[ringPrevRank];
83 0 : CHK_SMART_PTR_NULL(linkLeft_);
84 :
85 0 : linkRight_ = links[ringNextRank];
86 0 : CHK_SMART_PTR_NULL(linkRight_);
87 :
88 0 : CHK_RET(MultiRootScatterSlicesPrep(rankSize, nicRankList_.size()));
89 :
90 0 : CHK_RET(RunMultiRootScatterChunk(rank, rankSize, slices_));
91 :
92 0 : if (barrierSwitchOn_) {
93 : // 执行barrier，保证数据收发完成
94 0 : CHK_RET(ExecuteBarrier(linkLeft_, linkRight_));
95 : }
96 :
97 0 : return HCCL_SUCCESS;
98 : }
99 :
100 0 : HcclResult MultiRootScatterRing::RunMultiRootScatterChunk(
101 : const u32 rank, const u32 rankSize, const std::vector<Slice>& outputSlices)
102 : {
103 : HcclResult ret;
104 0 : DeviceMem dstMem;
105 0 : u32 sendSliceLen = rankSliceLists_[rank].size();
106 0 : if (sendSliceLen >= 1) { // 如果slice序列大于等于1则，存在头结点，进行相应slice的发送
107 0 : CHK_RET(HeadScatterChunk(rank, rankSize, outputSlices));
108 0 : for (u32 midRankIdx = 1; midRankIdx < sendSliceLen - 1; midRankIdx++) {
109 0 : ret = MidScatterChunk(rank, rankSize, midRankIdx, outputSlices);
110 0 : CHK_PRT_RET(
111 : ret != HCCL_SUCCESS,
112 : HCCL_ERROR(
113 : "[Run][MultiRootScatterChunk]rank[%u] run mid[%u] ReduceScatter chunk "
114 : "failed",
115 : rank, midRankIdx),
116 : HCCL_E_INTERNAL);
117 : }
118 : }
119 :
120 0 : if (sendSliceLen >= 2) { // 如果slice序列大于等于2则，存在尾结点，进行相应slice的发送
121 0 : CHK_RET(TailScatterChunk(rank, rankSize, sendSliceLen - 1, outputSlices));
122 : }
123 :
124 0 : if (sendSliceLen == 0) { // 如果slice序列长度为0，则接受当前rank会最终保存的slice即可
125 0 : u32 rxSliceIndex = (rank - nicRankList_[0] + HCCL_NIC_MAX_NUM) % HCCL_NIC_MAX_NUM;
126 0 : u64 rxScatterOffset = slices_[rxSliceIndex].offset;
127 0 : u64 rxScatterResult = slices_[rxSliceIndex].size;
128 0 : std::vector<u32> preRankSlices(rankSliceLists_[(rank - 1 + rankSize) % rankSize]);
129 0 : std::vector<u32>::iterator iterSlice = std::find(preRankSlices.begin(), preRankSlices.end(), rxSliceIndex);
130 0 : if (iterSlice != preRankSlices.end()) {
131 0 : CHK_RET(linkLeft_->TxAck(stream_));
132 :
133 0 : dstMem = outputMem_.range(rxScatterOffset, rxScatterResult);
134 0 : ret = linkLeft_->RxAsync(
135 0 : UserMemType::OUTPUT_MEM, baseOffset_ + rxScatterOffset, dstMem.ptr(), rxScatterResult, stream_);
136 0 : CHK_PRT_RET(
137 : ret != HCCL_SUCCESS,
138 : HCCL_ERROR(
139 : "[Run][MultiRootScatterChunk]rank[%u] Left Link rx outputSlices"
140 : "[%u] Failed",
141 : rank, rxSliceIndex),
142 : ret);
143 : }
144 0 : }
145 0 : return HCCL_SUCCESS;
146 0 : }
147 :
148 0 : HcclResult MultiRootScatterRing::HeadScatterChunk(u32 rank, u32 rankSize, const std::vector<Slice>& outputSlices)
149 : {
150 : HcclResult ret;
151 : // 头结点发送及接收slice均为rankSliceLists_的第一个元素
152 0 : u32 rxSliceIndex = rankSliceLists_[rank][0];
153 : // 得到发送及接收slice的偏移和长度
154 0 : u64 scatterOffset = slices_[rxSliceIndex].offset;
155 0 : u64 scatterResult = slices_[rxSliceIndex].size;
156 0 : DeviceMem dstMem = outputMem_.range(scatterOffset, scatterResult);
157 : // 判断当前rank是否需要接收头结点的数据, 得到前一个rank的发送序列，判断当前发送的slice是否在该序列中
158 0 : std::vector<u32> preRankSlices(rankSliceLists_[(rank - 1 + rankSize) % rankSize]);
159 0 : std::vector<u32>::iterator iterSlice = std::find(preRankSlices.begin(), preRankSlices.end(), rxSliceIndex);
160 0 : if (iterSlice != preRankSlices.end()) { // 若发送slice在前一个rank的发送序列中，则需先从前一个rank中接收对应数据
161 0 : CHK_RET(linkLeft_->TxAck(stream_));
162 :
163 0 : ret = linkLeft_->RxAsync(
164 0 : UserMemType::OUTPUT_MEM, baseOffset_ + scatterOffset, dstMem.ptr(), scatterResult, stream_);
165 0 : CHK_PRT_RET(
166 : ret != HCCL_SUCCESS,
167 : HCCL_ERROR(
168 : "[MultiRootScatterRing][HeadScatterChunk]rank[%u] Left Link rx "
169 : "outputSlices[%u] Failed",
170 : rank, rxSliceIndex),
171 : ret);
172 : }
173 :
174 0 : if (rankSliceLists_[rank].size() >= 2) { // 发送序列长度>=2时，需判断发送第一个slice前是否需要接收第二段slice
175 0 : iterSlice = std::find(preRankSlices.begin(), preRankSlices.end(), rankSliceLists_[rank][1]);
176 0 : if (iterSlice != preRankSlices.end()) { // 需要接收第二段slice，此时头结点行为和中间结点一致
177 0 : CHK_RET(MidScatterChunk(rank, rankSize, 0, outputSlices));
178 0 : return HCCL_SUCCESS;
179 : }
180 0 : } else if (
181 0 : rankSliceLists_[rank].size() == 1) { // 发送序列只有一个slice，则头结点同时为尾节点，需要接收最终要保存的数据
182 : u32 rxTailIndex
183 0 : = (rank - nicRankList_[0] + HCCL_NIC_MAX_NUM) % HCCL_NIC_MAX_NUM; // 计算当前rank最终要保存的数据
184 0 : u64 rxTailOffset = slices_[rxTailIndex].offset;
185 0 : u64 rxTailResult = slices_[rxTailIndex].size;
186 : // 判断当前rank是否需要接收最终要保存的数据
187 0 : std::vector<u32>::iterator iterSlice = std::find(preRankSlices.begin(), preRankSlices.end(), rxSliceIndex);
188 0 : if (iterSlice != preRankSlices.end()) { // 接受的数据在前rank的发送序列中
189 0 : CHK_RET(linkLeft_->TxAck(stream_));
190 :
191 0 : CHK_RET(linkRight_->RxAck(stream_));
192 :
193 0 : ret = linkRight_->TxAsync(
194 0 : UserMemType::OUTPUT_MEM, baseOffset_ + scatterOffset, dstMem.ptr(), scatterResult, stream_);
195 0 : CHK_PRT_RET(
196 : ret != HCCL_SUCCESS,
197 : HCCL_ERROR(
198 : "[MultiRootScatterRing][HeadScatterChunk]rank[%u] Right Link tx "
199 : "outputSlices[%u] Failed",
200 : rank, rxSliceIndex),
201 : ret);
202 :
203 0 : dstMem = outputMem_.range(rxTailOffset, rxTailResult);
204 0 : ret = linkLeft_->RxAsync(
205 0 : UserMemType::OUTPUT_MEM, baseOffset_ + rxTailOffset, dstMem.ptr(), rxTailResult, stream_);
206 0 : CHK_PRT_RET(
207 : ret != HCCL_SUCCESS,
208 : HCCL_ERROR(
209 : "[MultiRootScatterRing][HeadScatterChunk]rank[%u] Left Link rx "
210 : "outputSlices[%u] Failed",
211 : rank, rxTailIndex),
212 : ret);
213 0 : return HCCL_SUCCESS;
214 : }
215 : }
216 : // 其他情况直接发送当前头结点slice
217 0 : CHK_RET(linkRight_->RxAck(stream_));
218 :
219 0 : ret = linkRight_->TxAsync(
220 0 : UserMemType::OUTPUT_MEM, baseOffset_ + scatterOffset, dstMem.ptr(), scatterResult, stream_);
221 0 : CHK_PRT_RET(
222 : ret != HCCL_SUCCESS,
223 : HCCL_ERROR(
224 : "[MultiRootScatterRing][HeadScatterChunk]rank[%u] Right Link tx "
225 : "outputSlices[%u] Failed",
226 : rank, rxSliceIndex),
227 : ret);
228 :
229 0 : return HCCL_SUCCESS;
230 0 : }
231 :
232 : HcclResult
233 0 : MultiRootScatterRing::MidScatterChunk(u32 rank, u32 rankSize, u32 sliceIdx, const std::vector<Slice>& outputSlices)
234 : {
235 : (void)outputSlices;
236 : HcclResult ret;
237 0 : DeviceMem dstMem;
238 : // 头结点发送slice为rankSliceLists_的第sliceIdx个元素，接收slice为rankSliceLists_的第sliceIdx+1个元素
239 0 : u32 rxSliceIndex = rankSliceLists_[rank][sliceIdx + 1];
240 0 : u32 txSliceIndex = rankSliceLists_[rank][sliceIdx];
241 0 : u64 rxScatterOffset = slices_[rxSliceIndex].offset;
242 0 : u64 txScatterOffset = slices_[txSliceIndex].offset;
243 0 : u64 rxScatterResult = slices_[rxSliceIndex].size;
244 0 : u64 txScatterResult = slices_[txSliceIndex].size;
245 : // 判断当前rank是否需要接收第sliceIdx+1个元素, 得到前一个rank的发送序列，判断当前发送的slice是否在该序列中
246 0 : std::vector<u32> preRankSlices(rankSliceLists_[(rank - 1 + rankSize) % rankSize]);
247 0 : std::vector<u32>::iterator iterSlice = std::find(preRankSlices.begin(), preRankSlices.end(), rxSliceIndex);
248 0 : if (iterSlice != preRankSlices.end()) { // 若发送slice在前一个rank的发送序列中，则需先从前一个rank中接收对应数据
249 0 : CHK_RET(linkLeft_->TxAck(stream_));
250 :
251 0 : dstMem = outputMem_.range(txScatterOffset, txScatterResult);
252 0 : CHK_RET(linkRight_->RxAck(stream_));
253 :
254 0 : ret = linkRight_->TxAsync(
255 0 : UserMemType::OUTPUT_MEM, baseOffset_ + txScatterOffset, dstMem.ptr(), txScatterResult, stream_);
256 0 : CHK_PRT_RET(
257 : ret != HCCL_SUCCESS,
258 : HCCL_ERROR(
259 : "[MultiRootScatterRing][MidScatterChunk]rank[%u] Right Link tx "
260 : "outputSlices[%u] Failed",
261 : rank, txSliceIndex),
262 : ret);
263 :
264 0 : dstMem = outputMem_.range(rxScatterOffset, rxScatterResult);
265 0 : ret = linkLeft_->RxAsync(
266 0 : UserMemType::OUTPUT_MEM, baseOffset_ + rxScatterOffset, dstMem.ptr(), rxScatterResult, stream_);
267 0 : CHK_PRT_RET(
268 : ret != HCCL_SUCCESS,
269 : HCCL_ERROR(
270 : "[MultiRootScatterRing][MidScatterChunk]rank[%u] Left Link rx "
271 : "outputSlices[%u] Failed",
272 : rank, rxSliceIndex),
273 : ret);
274 : } else { // 其他情况直接发送当前中间结点slice
275 0 : dstMem = outputMem_.range(txScatterOffset, txScatterResult);
276 0 : CHK_RET(linkRight_->RxAck(stream_));
277 :
278 0 : ret = linkRight_->TxAsync(
279 0 : UserMemType::OUTPUT_MEM, baseOffset_ + txScatterOffset, dstMem.ptr(), txScatterResult, stream_);
280 0 : CHK_PRT_RET(
281 : ret != HCCL_SUCCESS,
282 : HCCL_ERROR(
283 : "[MultiRootScatterRing][MidScatterChunk]rank[%u] Right Link tx "
284 : "outputSlices[%u] Failed",
285 : rank, txSliceIndex),
286 : ret);
287 : }
288 0 : return HCCL_SUCCESS;
289 0 : }
290 :
291 : HcclResult
292 0 : MultiRootScatterRing::TailScatterChunk(u32 rank, u32 rankSize, u32 sliceIdx, const std::vector<Slice>& outputSlices)
293 : {
294 : (void)outputSlices;
295 : HcclResult ret;
296 0 : DeviceMem dstMem;
297 : // 尾结点发送slice为rankSliceLists_的第sliceIdx个元素，接收slice为scatter最终会保存的slice位置
298 0 : u32 txSliceIndex = rankSliceLists_[rank][sliceIdx];
299 0 : u64 txScatterOffset = slices_[txSliceIndex].offset;
300 0 : u64 txScatterResult = slices_[txSliceIndex].size;
301 :
302 0 : u32 rxSliceIndex = (rank - nicRankList_[0] + HCCL_NIC_MAX_NUM) % HCCL_NIC_MAX_NUM;
303 0 : u64 rxScatterOffset = slices_[rxSliceIndex].offset;
304 0 : u64 rxScatterResult = slices_[rxSliceIndex].size;
305 :
306 : // 判断当前rank是否需要接收第sliceIdx+1个元素, 得到前一个rank的发送序列，判断当前发送的slice是否在该序列中
307 0 : std::vector<u32> preRankSlices(rankSliceLists_[(rank - 1 + rankSize) % rankSize]);
308 0 : std::vector<u32>::iterator iterSlice = std::find(preRankSlices.begin(), preRankSlices.end(), rxSliceIndex);
309 0 : if (iterSlice != preRankSlices.end()) { // 若接收slice在前一个rank的发送序列中，则需先从前一个rank中接收对应数据
310 0 : CHK_RET(linkLeft_->TxAck(stream_));
311 :
312 0 : dstMem = outputMem_.range(txScatterOffset, txScatterResult);
313 0 : CHK_RET(linkRight_->RxAck(stream_));
314 :
315 0 : ret = linkRight_->TxAsync(
316 0 : UserMemType::OUTPUT_MEM, baseOffset_ + txScatterOffset, dstMem.ptr(), txScatterResult, stream_);
317 0 : CHK_PRT_RET(
318 : ret != HCCL_SUCCESS,
319 : HCCL_ERROR(
320 : "[MultiRootScatterRing][TailScatterChunk]rank[%u] Right Link tx "
321 : "outputSlices[%u] Failed",
322 : rank, txSliceIndex),
323 : ret);
324 :
325 0 : dstMem = outputMem_.range(rxScatterOffset, rxScatterResult);
326 0 : ret = linkLeft_->RxAsync(
327 0 : UserMemType::OUTPUT_MEM, baseOffset_ + rxScatterOffset, dstMem.ptr(), rxScatterResult, stream_);
328 0 : CHK_PRT_RET(
329 : ret != HCCL_SUCCESS,
330 : HCCL_ERROR(
331 : "[MultiRootScatterRing][TailScatterChunk]rank[%u] Left Link rx "
332 : "outputSlices[%u] Failed",
333 : rank, rxSliceIndex),
334 : ret);
335 : } else { // 其他情况直接发送当前尾结点slice
336 0 : dstMem = outputMem_.range(txScatterOffset, txScatterResult);
337 0 : CHK_RET(linkRight_->RxAck(stream_));
338 :
339 0 : ret = linkRight_->TxAsync(
340 0 : UserMemType::OUTPUT_MEM, baseOffset_ + txScatterOffset, dstMem.ptr(), txScatterResult, stream_);
341 0 : CHK_PRT_RET(
342 : ret != HCCL_SUCCESS,
343 : HCCL_ERROR(
344 : "[MultiRootScatterRing][TailScatterChunk]rank[%u] Right Link tx "
345 : "outputSlices[%u] Failed",
346 : rank, txSliceIndex),
347 : ret);
348 : }
349 0 : return HCCL_SUCCESS;
350 0 : }
351 :
352 0 : HcclResult MultiRootScatterRing::MultiRootScatterSlicesPrep(u32 rankSize, u32 nicSize)
353 : {
354 0 : u32 chunkSize = HCCL_NIC_MAX_NUM / nicSize;
355 0 : std::vector<SliceSendRange> sliceSendRangeVec;
356 0 : for (u32 nicIdx = 0; nicIdx < nicSize; nicIdx++) { // 计算每个网口负责的slice发送顺序
357 0 : for (u32 sliceIdx = 0; sliceIdx < chunkSize; sliceIdx++) { // 记录每个网口发送slice的起点和终点
358 0 : SliceSendRange tempSliceSendRange;
359 0 : tempSliceSendRange.sliceIdx = nicIdx * chunkSize + sliceIdx;
360 0 : tempSliceSendRange.startRank = nicRankList_[nicIdx];
361 0 : tempSliceSendRange.endRank = (nicIdx * chunkSize + sliceIdx + nicRankList_[0]) % HCCL_NIC_MAX_NUM;
362 0 : if (tempSliceSendRange.endRank < tempSliceSendRange.startRank) {
363 0 : tempSliceSendRange.endRank = tempSliceSendRange.endRank + HCCL_NIC_MAX_NUM;
364 : }
365 0 : sliceSendRangeVec.push_back(tempSliceSendRange);
366 : }
367 : }
368 :
369 0 : for (u32 rankIdx = 0; rankIdx < rankSize; rankIdx++) { // 计算每个rank发送slice的顺序
370 0 : std::vector<u32> sliceList; // 单个rank上的发送slice编号
371 0 : std::vector<SliceSendRange> rankSliceSendVec;
372 : // 从后往前依次遍历slice, 判断当前rank是否需要发送当前slice
373 0 : std::vector<SliceSendRange>::iterator sliceSendIdx = sliceSendRangeVec.end() - 1;
374 0 : for (; sliceSendIdx >= sliceSendRangeVec.begin(); sliceSendIdx--) {
375 0 : SliceSendRange rankSliceSend;
376 0 : if (rankIdx >= sliceSendIdx->startRank) { // slice终点rank号大于起点rank号
377 0 : if (rankIdx < sliceSendIdx->endRank) {
378 0 : rankSliceSend.sliceIdx = sliceSendIdx->sliceIdx;
379 0 : rankSliceSend.endRank = sliceSendIdx->endRank - rankIdx;
380 0 : rankSliceSendVec.push_back(rankSliceSend);
381 : }
382 : } else { // slice终点rank号小于起点rank号
383 0 : u32 tempRankIdx = rankIdx + HCCL_NIC_MAX_NUM;
384 0 : if (tempRankIdx < sliceSendIdx->endRank) {
385 0 : rankSliceSend.sliceIdx = sliceSendIdx->sliceIdx;
386 0 : rankSliceSend.endRank = sliceSendIdx->endRank - tempRankIdx;
387 0 : rankSliceSendVec.push_back(rankSliceSend);
388 : }
389 : }
390 : }
391 0 : std::sort(rankSliceSendVec.begin(), rankSliceSendVec.end(), DscendSortWithSliceSendEnd);
392 0 : for (u32 sliceIdx = 0; sliceIdx < rankSliceSendVec.size(); sliceIdx++) {
393 0 : sliceList.push_back(rankSliceSendVec[sliceIdx].sliceIdx);
394 : }
395 0 : rankSliceLists_.push_back(sliceList);
396 0 : }
397 :
398 0 : return HCCL_SUCCESS;
399 0 : }
400 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_MULTI_ROOT_SCATTER_RING, MultiRootScatterRing);
401 : } // namespace hccl
|
