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