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 "reduce_scatter_nb.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 :
16 10 : ReduceScatterNB::ReduceScatterNB(const HcclDispatcher dispatcher)
17 10 : :NBBase(dispatcher)
18 : {
19 10 : }
20 :
21 11 : ReduceScatterNB::~ReduceScatterNB()
22 : {
23 11 : }
24 :
25 1 : HcclResult ReduceScatterNB::Prepare(u64 reduceAttrBitMap, HcomCollOpInfo *opInfo)
26 : {
27 : (void)opInfo;
28 1 : reduceAttr_ = reduceAttrBitMap;
29 1 : return HCCL_SUCCESS;
30 : }
31 :
32 0 : HcclResult ReduceScatterNB::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
33 : {
34 : // 参数校验
35 0 : CHK_RET(SimpleCheck(rank, rankSize, links));
36 0 : HCCL_INFO("ReduceScatterNB run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
37 : inputMem_.ptr(), outputMem_.ptr(), count_);
38 :
39 0 : if (rankSize == 1) {
40 0 : if (inputMem_ != outputMem_) {
41 0 : return HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
42 : }
43 0 : return HCCL_SUCCESS;
44 : }
45 :
46 0 : u32 unitSize = DataUnitSize(dataType_);
47 0 : CHK_PRT_RET(unitSize == 0, HCCL_ERROR("[ReduceScatterRing][RunAsync] rank[%u] unit data size is zero", rank),
48 : HCCL_E_INTERNAL);
49 :
50 0 : std::vector<Slice> outputSlices(slices_);
51 :
52 : // 处理和检查Slices
53 0 : if (slices_.size() == 0) {
54 0 : slices_.resize(rankSize);
55 0 : outputSlices.resize(rankSize);
56 :
57 : // 生成std::vector<Slice> slices_
58 0 : u64 sliceSize = count_ * unitSize;
59 0 : HCCL_DEBUG("[ReduceScatterNB][RunAsync]sliceSize is %llu", sliceSize);
60 :
61 0 : for (u32 i = 0; i < rankSize; i++) {
62 0 : slices_[i].size = sliceSize;
63 0 : slices_[i].offset = (i * sliceSize);
64 :
65 0 : outputSlices[i].size = sliceSize;
66 0 : outputSlices[i].offset = (inputMem_.size() > outputMem_.size()) ? 0 : (i * sliceSize);
67 0 : HCCL_DEBUG("rank[%u], slices[%u].offset=[%llu], slices[%u].size=[%llu] outputSlices[%u].offset=[%llu], \
68 : outputSlices[%u].size=[%llu] count_[%llu] unitSize[%llu]",
69 : rank, i, slices_[i].offset, i, slices_[i].size, i, outputSlices[i].offset, i, outputSlices[i].size,
70 : count_, unitSize);
71 : }
72 : }
73 :
74 0 : CHK_RET(CheckSlices(slices_, rankSize));
75 :
76 : // 创建reducer & sender
77 0 : senderInfo_.reset(new (std::nothrow) Sender(dataType_, reductionOp_, reduceAttr_));
78 0 : CHK_SMART_PTR_NULL(senderInfo_);
79 :
80 0 : reducerInfo_.reset(new (std::nothrow) Reducer(dataType_, reductionOp_, reduceAttr_));
81 0 : CHK_SMART_PTR_NULL(reducerInfo_);
82 :
83 : // 运行reduce-scatter, NB 算法
84 0 : CHK_RET(RunReduceScatterNB(rank, rankSize, links, slices_, outputSlices));
85 :
86 0 : HCCL_INFO("ReduceScatterNB finished: rank[%u] end", rank);
87 0 : return HCCL_SUCCESS;
88 0 : }
89 :
90 0 : HcclResult ReduceScatterNB::SimpleCheck(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
91 : {
92 : // 判断stream, dispatcher是否为空
93 0 : CHK_SMART_PTR_NULL(dispatcher_);
94 0 : CHK_PTR_NULL(stream_.ptr());
95 :
96 : // 检查memory
97 0 : CHK_PRT_RET(!outputMem_ || !inputMem_, HCCL_ERROR("[ReduceScatterNB]rank[%u] inputmem or outputmem is null", rank),
98 : HCCL_E_PTR);
99 :
100 : // 判断links数量是否正确
101 0 : CHK_PRT_RET(links.size() < rankSize,
102 : HCCL_ERROR("[ReduceScatterNB]rank[%u] link size[%llu] is less than "
103 : "rank size[%u]", rank, links.size(), rankSize), HCCL_E_INTERNAL);
104 0 : return HCCL_SUCCESS;
105 : }
106 :
107 0 : HcclResult ReduceScatterNB::CheckSlices(const std::vector<Slice> &checkSlices, const u32 rankSize)
108 : {
109 0 : CHK_PRT_RET(checkSlices.size() % rankSize != 0,
110 : HCCL_ERROR("[ReduceScatterNB]slices.size[%u] should be divided by rankSize[%u]", checkSlices.size(), rankSize),
111 : HCCL_E_INTERNAL);
112 0 : return HCCL_SUCCESS;
113 : }
114 :
115 0 : HcclResult ReduceScatterNB::RunReduceScatterNB(const u32 rank, const u32 rankSize,
116 : const std::vector<LINK> &links,
117 : const std::vector<Slice> &inputSlices,
118 : const std::vector<Slice> &outputSlices)
119 : {
120 0 : bool bRetSize = (inputSlices.size() < rankSize);
121 0 : CHK_PRT_RET(bRetSize,
122 : HCCL_ERROR("[Run][ReduceScatter]rank[%u] inputslice size[%llu] is less than rank size[%u]",
123 : rank, inputSlices.size(), rankSize), HCCL_E_INTERNAL);
124 :
125 0 : bRetSize = (outputSlices.size() < rankSize);
126 0 : CHK_PRT_RET(bRetSize,
127 : HCCL_ERROR("[Run][ReduceScatter]rank[%u] outputslice size[%llu] is less than rank size[%u]",
128 : rank, outputSlices.size(), rankSize), HCCL_E_INTERNAL);
129 :
130 0 : HcclResult ret = HCCL_SUCCESS;
131 :
132 : // 计算通信步数：ceiling(log2(rankSize))
133 0 : u32 nSteps = CalcCeilLog2(rankSize);
134 0 : u32 sliceSize = inputSlices.size() / rankSize;
135 0 : HCCL_DEBUG("ReduceScatter debug-1: rank[%u] rankSize[%u] nSteps[%u] sliceSize[%u]", rank, rankSize, nSteps,
136 : sliceSize);
137 : // 逐步编排任务
138 0 : for (u32 step = 0; step < nSteps; step++) {
139 : // 计算通信对象
140 0 : u32 deltaRank = 1 << step;
141 0 : u32 recvFrom = (rankSize + rank - deltaRank) % rankSize;
142 0 : u32 sendTo = (rank + deltaRank) % rankSize;
143 :
144 : // 数据份数和数据编号增量
145 0 : u32 nSlices = (rankSize - 1 + (1 << step)) / (1 << (step + 1));
146 0 : u32 deltaSliceIndex = 1 << (step + 1);
147 0 : u32 txSliceIdx = (rank + (1 << step)) % rankSize;
148 0 : u32 rxSliceIdx = rank;
149 :
150 0 : LINK linkLeft = links[recvFrom];
151 0 : CHK_SMART_PTR_NULL(linkLeft);
152 :
153 0 : LINK linkRight = links[sendTo];
154 0 : CHK_SMART_PTR_NULL(linkRight);
155 :
156 : // 当前每个数据块发送一次ACK、reduce一次、同步一次
157 0 : HCCL_DEBUG("ReduceScatter debug-2: recvFrom[%u] sendTo[%u] step[%u] nSlices[%u] deltaSliceIndex[%u] "
158 : "rxSliceIdx[%u] txSliceIdx[%u]",
159 : recvFrom, sendTo, step, nSlices, deltaSliceIndex, rxSliceIdx, txSliceIdx);
160 :
161 0 : u32 txCount = 0;
162 0 : u32 txSliceIdxTmp = txSliceIdx;
163 0 : for (u32 i = 0; i < nSlices; i++) {
164 0 : for (u32 j = 0; j < sliceSize; j++) {
165 0 : if (inputSlices[txSliceIdxTmp * sliceSize + j].size > 0) {
166 0 : txCount++;
167 : }
168 : }
169 0 : txSliceIdxTmp = (txSliceIdxTmp + deltaSliceIndex) % rankSize;
170 : }
171 :
172 0 : u32 rxCount = 0;
173 0 : u32 rxSliceIdxTmp = rxSliceIdx;
174 0 : for (u32 i = 0; i < nSlices; i++) {
175 0 : for (u32 j = 0; j < sliceSize; j++){
176 0 : if (inputSlices[rxSliceIdxTmp * sliceSize + j].size > 0) {
177 0 : rxCount++;
178 : }
179 : }
180 0 : rxSliceIdxTmp = (rxSliceIdxTmp + deltaSliceIndex) % rankSize;
181 : }
182 0 : if (rxCount > 0) {
183 0 : CHK_RET(linkLeft->TxAck(stream_));
184 : }
185 0 : if (txCount > 0) {
186 0 : CHK_RET(linkRight->RxAck(stream_));
187 0 : RunSrcReducerNB(step, nSlices, sliceSize, txSliceIdx, deltaSliceIndex, linkRight, rank, rankSize,
188 : inputSlices, outputSlices);
189 : }
190 0 : if (rxCount > 0) {
191 0 : RunDestReducerNB(step, nSteps, sliceSize, nSlices, rxSliceIdx, deltaSliceIndex, linkLeft, rank, rankSize,
192 : inputSlices, outputSlices);
193 0 : ret = linkLeft->RxWaitDone(stream_);
194 0 : CHK_PRT_RET(ret != HCCL_SUCCESS,
195 : HCCL_ERROR("[Run][ReduceScatterNB]rank[%u] step[%u] blocknum[%u] rx wait done failed", rank, step,
196 : nSlices),
197 : ret);
198 0 : ret = linkLeft->PostFinAck(stream_);
199 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][ReduceScatter]PostFinAck failed"), ret);
200 : }
201 0 : if (txCount > 0) {
202 0 : ret = linkRight->TxWaitDone(stream_);
203 0 : CHK_PRT_RET(ret != HCCL_SUCCESS,
204 : HCCL_ERROR("[Run][ReduceScatterNB]rank[%u] step[%u] blocknum[%u] tx wait done failed", rank, step,
205 : nSlices),
206 : ret);
207 0 : ret = linkRight->WaitFinAck(stream_);
208 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][ReduceScatter]WaitFinAck failed"), ret);
209 : }
210 0 : if (linkRight->IsSpInlineReduce() || linkLeft->IsSpInlineReduce()) {
211 : // SDMA场景同步
212 0 : CHK_RET(ExecuteBarrier(linkLeft, linkRight));
213 : }
214 0 : }
215 0 : return HCCL_SUCCESS;
216 : }
217 :
218 0 : HcclResult ReduceScatterNB::RunSrcReducerNB(const u32 step, const u32 nSlices, const u32 sliceSize,
219 : u32 txSliceIdx, const u32 deltaSliceIndex,
220 : const LINK linkRight, const u32 rank,
221 : const u32 rankSize, const std::vector<Slice> &inputSlices,
222 : const std::vector<Slice> &outputSlices)
223 : {
224 0 : HcclResult ret = HCCL_SUCCESS;
225 :
226 0 : std::vector<Slice> txSlices;
227 0 : std::vector<Slice> txSlicestemp;
228 0 : for (u32 i = 0; i < nSlices; i++) {
229 0 : for (u32 j = 0; j < sliceSize; j++) {
230 0 : u32 txIndex = txSliceIdx * sliceSize + j;
231 0 : if (inputSlices[txIndex].size > 0) {
232 0 : txSlices.push_back(inputSlices[txIndex]);
233 0 : txSlicestemp.push_back(outputSlices[txIndex]);
234 : }
235 : }
236 0 : txSliceIdx = (txSliceIdx + deltaSliceIndex) % rankSize;
237 : }
238 :
239 0 : ret = RunSourceReducer(linkRight, txSlices, txSlicestemp);
240 0 : CHK_PRT_RET(ret != HCCL_SUCCESS,
241 : HCCL_ERROR("[Run][ReduceScatterNB]rank[%u] step[%u] blocknum[%u] tx multi blocks failed",
242 : rank, step, nSlices), ret);
243 0 : return HCCL_SUCCESS;
244 0 : }
245 :
246 0 : HcclResult ReduceScatterNB::RunDestReducerNB( const u32 step, const u32 nSteps, const u32 sliceSize,
247 : const u32 nSlices, u32 rxSliceIdx,
248 : const u32 deltaSliceIndex, const LINK linkLeft,
249 : const u32 rank, const u32 rankSize,
250 : const std::vector<Slice> &inputSlices,
251 : const std::vector<Slice> &outputSlices)
252 : {
253 0 : HcclResult ret = HCCL_SUCCESS;
254 :
255 0 : if (step == (nSteps - 1)) {
256 0 : std::vector<ReducerMemoryInfo> rxReduceMems;
257 0 : for (u32 i = 0; i < nSlices; i++) {
258 0 : for (u32 j = 0; j < sliceSize; j++) {
259 0 : u32 rxIndex = rxSliceIdx * sliceSize + j;
260 0 : if (inputSlices[rxIndex].size > 0) {
261 0 : DeviceMem dstMem = outputMem_.range(outputSlices[rxIndex].offset, outputSlices[rxIndex].size);
262 0 : DeviceMem srcMem = inputMem_.range(inputSlices[rxIndex].offset, inputSlices[rxIndex].size);
263 : DeviceMem scratchMem =
264 0 : scratchMem_.range(outputSlices[rxIndex].offset, outputSlices[rxIndex].size);
265 0 : HCCL_DEBUG("final reduce rxSliceIdx[%u] will reduce with inputMem_ offset[%llu] to ouput_mem_ "
266 : "offset[%llu] size[%llu]",
267 : rxIndex, inputSlices[rxIndex].offset, outputSlices[rxIndex].offset,
268 : outputSlices[rxIndex].size);
269 :
270 0 : rxReduceMems.emplace_back(
271 0 : ReducerMemoryInfo{ baseOffset_ + inputSlices[rxIndex].offset, srcMem, dstMem, scratchMem });
272 0 : }
273 : }
274 0 : rxSliceIdx = (rxSliceIdx + deltaSliceIndex) % rankSize;
275 : }
276 :
277 0 : ret = reducerInfo_->run(dispatcher_, linkLeft, rxReduceMems, stream_);
278 0 : } else {
279 0 : std::vector<Slice> rxSlices;
280 0 : std::vector<Slice> rxSlicestemp;
281 0 : for (u32 i = 0; i < nSlices; i++) {
282 0 : for (u32 j = 0; j < sliceSize; j++) {
283 0 : u32 rxIndex = rxSliceIdx * sliceSize + j;
284 0 : if (inputSlices[rxIndex].size > 0) {
285 0 : rxSlices.push_back(inputSlices[rxIndex]);
286 0 : rxSlicestemp.push_back(outputSlices[rxIndex]);
287 : }
288 : }
289 0 : rxSliceIdx = (rxSliceIdx + deltaSliceIndex) % rankSize;
290 : }
291 :
292 0 : ret = RunDestReducer(linkLeft, rxSlices, rxSlicestemp);
293 0 : }
294 :
295 0 : CHK_PRT_RET(ret != HCCL_SUCCESS,
296 : HCCL_ERROR("[Run][ReduceScatterNB]rank[%u] step[%u] blocknum[%u] rx multi blocks failed", rank, step, nSlices),
297 : ret);
298 0 : return HCCL_SUCCESS;
299 : }
300 :
301 0 : HcclResult ReduceScatterNB::RunDestReducer(const LINK &link, const std::vector<Slice> &rxSlices,
302 : const std::vector<Slice> &rxSlicestemp)
303 : {
304 0 : std::vector<ReducerMemoryInfo> rxReduceMems;
305 :
306 0 : for (u64 i = 0; i < rxSlices.size(); i++) {
307 0 : DeviceMem dstMem = inputMem_.range(rxSlices[i].offset, rxSlices[i].size);
308 0 : DeviceMem srcMemTemp = scratchMem_.range(rxSlicestemp[i].offset, rxSlicestemp[i].size);
309 0 : HCCL_DEBUG("rcv offset[%llu], size[%llu] ,then reduce with "
310 : "offset[%llu] size[%llu] ",
311 : rxSlicestemp[i].offset, rxSlicestemp[i].size, rxSlices[i].offset, rxSlices[i].size);
312 0 : rxReduceMems.emplace_back(ReducerMemoryInfo{baseOffset_ + rxSlices[i].offset, dstMem, dstMem, srcMemTemp});
313 0 : }
314 0 : CHK_RET(reducerInfo_->run(dispatcher_, link, rxReduceMems, stream_));
315 0 : return HCCL_SUCCESS;
316 0 : }
317 :
318 0 : HcclResult ReduceScatterNB::RunSourceReducer(const LINK &link, const std::vector<Slice> &txSlices,
319 : const std::vector<Slice> &txSlicestemp)
320 : {
321 0 : std::vector<SenderMemoryInfo> txMems;
322 :
323 0 : for (u64 i = 0; i < txSlices.size(); i++) {
324 0 : DeviceMem srcMem = inputMem_.range(txSlices[i].offset, txSlices[i].size);
325 0 : HCCL_DEBUG(" send inputmem range[%llu], size[%llu] tx dstmem offset[%llu]", txSlices[i].offset,
326 : txSlices[i].size, txSlicestemp[i].offset);
327 0 : txMems.emplace_back(SenderMemoryInfo{baseOffset_ + txSlicestemp[i].offset, srcMem});
328 0 : }
329 0 : CHK_RET(senderInfo_->run(link, txMems, stream_));
330 0 : return HCCL_SUCCESS;
331 0 : }
332 0 : HcclResult ReduceScatterNB::GetNslbAdjInfo(const u32 rank, const u32 rankSize,
333 : const std::vector<LINK> &links, AdjInfo& nslbAdjInfo)
334 : {
335 0 : if (rankSize == 1) {
336 0 : return HCCL_SUCCESS;
337 : }
338 0 : if (links.size() < rankSize) {
339 0 : return HCCL_SUCCESS;
340 : }
341 0 : u32 nSteps = 0;
342 0 : for(u32 temp = rankSize - 1; temp != 0; temp >>= 1, ++nSteps){}
343 :
344 0 : for (u32 step = 0; step < nSteps; step++) {
345 0 : u32 deltaRank = 1 << step;
346 0 : u32 sendTo =(rank + deltaRank) % rankSize;
347 0 : LINK linkRight = links[sendTo];
348 0 : CHK_SMART_PTR_NULL(linkRight);
349 :
350 0 : NslbDpAdjInfo adjInfoStep = {0};
351 0 : adjInfoStep.dstLocalRankId = linkRight->GetRemoteRank();
352 0 : adjInfoStep.phaseId = step + 1;
353 0 : adjInfoStep.rev = 0;
354 0 : nslbAdjInfo.nsAdjInfo.push_back(adjInfoStep);
355 0 : }
356 0 : nslbAdjInfo.dstRankNum = nSteps;
357 0 : return HCCL_SUCCESS;
358 : }
359 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_NB, ReduceScatterNB);
360 : } // ~~ namespace hccl
|
