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_nhr_v1.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 :
16 0 : ReduceScatterNHRV1::ReduceScatterNHRV1(const HcclDispatcher dispatcher) : NHRV1Base(dispatcher) {}
17 :
18 0 : ReduceScatterNHRV1::~ReduceScatterNHRV1() {}
19 :
20 0 : HcclResult ReduceScatterNHRV1::Prepare(u64 reduceAttrBitMap, HcomCollOpInfo* opInfo)
21 : {
22 : (void)opInfo;
23 0 : reduceAttr_ = reduceAttrBitMap;
24 0 : return HCCL_SUCCESS;
25 : }
26 :
27 0 : HcclResult ReduceScatterNHRV1::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
28 : {
29 : // 基本的检查
30 0 : CHK_RET(SimpleCheck(rank, rankSize, links));
31 0 : HCCL_INFO(
32 : "ReduceScatterNHRV1 run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
33 : inputMem_.ptr(), outputMem_.ptr(), count_);
34 :
35 : // 判断rank_size == 1
36 0 : if (rankSize == 1) {
37 0 : if (inputMem_ != outputMem_) {
38 0 : return HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
39 : }
40 0 : return HCCL_SUCCESS;
41 : }
42 :
43 : // 处理和检查Slices
44 0 : if (slices_.size() == 0) {
45 0 : CHK_RET(SetDefaultSlices(rank, rankSize));
46 : }
47 0 : CHK_RET(CheckSlices(rankSize));
48 :
49 : // 获取通信关系
50 0 : RingInfo info = GetRingInfo(rankSize);
51 :
52 : // 垂直方向做Ring
53 0 : CHK_RET(RunReduceScatterOnVertical(rank, links, info));
54 :
55 : // 水平方向做Ring
56 0 : CHK_RET(RunReduceScatterOnHorizontal(rank, links, info));
57 :
58 : // 额外的搬运（从(x,sqrt-1)搬运到(x,sqrt)）
59 : /* 一个可能的优化点：
60 : 以8节点为例： 0 1 2
61 : 3 4 5
62 : 6 7
63 : 当前的做法是：{0,1}、{3,4}、{6、7}做水平Ring，0/3/6/7拿到各自的那份结果，1拿到1和2的结果，4拿到4和5的结果，
64 : 最后1把2的结果发给2，4把5的结果发给5
65 : 一种可能的更优做法是：以ReduceOp=Sum为例，首先把2和5的数据都置为0，
66 : 然后直接{0,1,2}、{3,4,5}、{6,7}做水平Ring，避免不等分Ring和额外的拷贝步骤，但需要调用TBE-asign
67 : */
68 0 : CHK_RET(RunLastCopyStep(rank, links, info));
69 :
70 : // 搬运数据到OutputMem
71 0 : CHK_RET(RunCopyDataToOutputMem(rank));
72 :
73 0 : HCCL_INFO("ReduceScatterNHRV1 finished: rank[%u] end", rank);
74 0 : return HCCL_SUCCESS;
75 0 : }
76 :
77 0 : HcclResult ReduceScatterNHRV1::SimpleCheck(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
78 : {
79 : // 判断stream, dispatcher是否为空
80 0 : CHK_SMART_PTR_NULL(dispatcher_);
81 0 : CHK_PTR_NULL(stream_.ptr());
82 :
83 : // 检查memory
84 0 : CHK_PRT_RET(
85 : !outputMem_ || !inputMem_, HCCL_ERROR("[ReduceScatterNHRV1]rank[%u] inputmem or outputmem is null", rank),
86 : HCCL_E_PTR);
87 :
88 : // 判断links数量是否正确
89 0 : CHK_PRT_RET(
90 : links.size() < rankSize,
91 : HCCL_ERROR(
92 : "[ReduceScatterNHRV1]rank[%u] link size[%llu] is less than "
93 : "rank size[%u]",
94 : rank, links.size(), rankSize),
95 : HCCL_E_INTERNAL);
96 0 : return HCCL_SUCCESS;
97 : }
98 :
99 0 : HcclResult ReduceScatterNHRV1::SetDefaultSlices(const u32 rank, const u32 rankSize)
100 : {
101 0 : u32 unitSize = DataUnitSize(dataType_);
102 0 : CHK_PRT_RET(
103 : unitSize == 0, HCCL_ERROR("[ReduceScatterNHRV1]rank[%u] unit data size is zero", rank), HCCL_E_INTERNAL);
104 :
105 0 : slices_.resize(rankSize);
106 0 : u64 sliceSize = count_ * unitSize;
107 0 : for (u32 i = 0; i < rankSize; i++) {
108 0 : slices_[i].size = sliceSize;
109 0 : slices_[i].offset = (i * sliceSize);
110 0 : HCCL_DEBUG(
111 : "rank[%u], slices[%u].offset=[%llu], slices[%u].size=[%llu] ", rank, i, slices_[i].offset, i,
112 : slices_[i].size);
113 : }
114 0 : return HCCL_SUCCESS;
115 : }
116 :
117 0 : HcclResult ReduceScatterNHRV1::CheckSlices(const u32 rankSize)
118 : {
119 0 : CHK_PRT_RET(
120 : slices_.size() != rankSize,
121 : HCCL_ERROR("[ReduceScatterNHRV1]slices.size[%u] should be equal to rankSize[%u]", slices_.size(), rankSize),
122 : HCCL_E_INTERNAL);
123 :
124 0 : for (u32 idx = 1; idx < slices_.size(); idx++) {
125 0 : CHK_PRT_RET(
126 : slices_[idx - 1].offset + slices_[idx - 1].size != slices_[idx].offset,
127 : HCCL_ERROR(
128 : "[ReduceScatterNHRV1]only support continuous slices, but get "
129 : "slices[%u].offset[%u], slices[%u].size[%u], slices[%u].offset[%u]",
130 : idx - 1, slices_[idx - 1].offset, idx - 1, slices_[idx - 1].size, idx, slices_[idx].offset),
131 : HCCL_E_INTERNAL);
132 : }
133 0 : return HCCL_SUCCESS;
134 : }
135 :
136 0 : HcclResult ReduceScatterNHRV1::RunReduceScatterBrokenRing(
137 : const u32 rank, const std::vector<LINK>& links, const std::vector<Slice>& slices)
138 : {
139 : std::unique_ptr<AlgTemplateBase> tempAlg
140 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_REDUCESCATTER_RING, dispatcher_);
141 0 : CHK_SMART_PTR_NULL(tempAlg);
142 0 : CHK_RET(tempAlg->Prepare(reduceAttr_));
143 :
144 0 : if (!barrierSwitchOn_) {
145 0 : tempAlg->CloseBarrier();
146 : }
147 :
148 0 : CHK_RET(
149 : tempAlg->Prepare(inputMem_, inputMem_, scratchMem_, count_, dataType_, stream_, reductionOp_, root_, slices));
150 :
151 0 : CHK_RET(tempAlg->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
152 :
153 0 : return tempAlg->RunAsync(rank, links.size(), links);
154 0 : }
155 :
156 : HcclResult
157 0 : ReduceScatterNHRV1::RunReduceScatterOnVertical(const u32 rank, const std::vector<LINK>& links, const RingInfo& info)
158 : {
159 0 : u32 hIndex = info.GetHIndex(rank); // 查找自己位于第几列
160 :
161 : // 构造新的links和slices
162 0 : std::vector<LINK> subLinks;
163 0 : std::vector<Slice> subSlices;
164 0 : u32 sliceIndexOffset = 0; // slice数量的累计偏移
165 0 : u32 hIndexForRing = (hIndex < info.GetRowSize()) ? hIndex : info.GetVIndex(rank); // 属于第几个垂直Ring
166 0 : u32 vSizeForRing = info.GetVSizeByHIndex(hIndexForRing); // 所属垂直Ring的大小
167 0 : for (u32 vIdx = 0; vIdx < vSizeForRing; vIdx++) { // 处理垂直方向上的Rank
168 : // 增加link
169 0 : u32 rankInRing = info.GetRank(vIdx, hIndexForRing);
170 0 : CHK_PRT_RET(
171 : rankInRing >= links.size(),
172 : HCCL_ERROR(
173 : "[ReduceScatterNHRV1][Vertical] rank[%u] out of range, "
174 : "rankInRing=%u, links.size=%u",
175 : rank, rankInRing, links.size()),
176 : HCCL_E_INTERNAL);
177 0 : HCCL_DEBUG("[ReduceScatterNHRV1][Vertical] rank[%u] links[%u]=%u", rank, vIdx, rankInRing);
178 0 : subLinks.push_back(links[rankInRing]);
179 :
180 : // 寻找要合并的slice
181 0 : u32 nSlices = info.GetHSizeByVIndex(vIdx);
182 0 : Slice& headSlice = slices_[sliceIndexOffset];
183 0 : Slice& tailSlice = slices_[sliceIndexOffset + nSlices - 1];
184 :
185 : // 增加slice
186 0 : Slice slice;
187 0 : slice.offset = headSlice.offset;
188 0 : slice.size = tailSlice.offset + tailSlice.size - headSlice.offset;
189 0 : HCCL_DEBUG(
190 : "[ReduceScatterNHRV1][Vertical] rank[%u] subSlices[%u].offset=%llu, subSlices[%u].size=%llu", rank, vIdx,
191 : slice.offset, vIdx, slice.size);
192 0 : subSlices.push_back(slice);
193 :
194 : // 更新偏移
195 0 : sliceIndexOffset += nSlices;
196 : }
197 :
198 : // -- 可能还涉及跳跃的一个链接，比如8节点
199 : // ---- 0 1 2
200 : // ---- 3 4 5
201 : // ---- 6 7
202 : // -- 两个垂直Ring分别是{0,3,6,2}和{1,4,7,5}，而不是{0,3,6}和{1,4,7}
203 0 : if (info.GetHSizeByVIndex(hIndexForRing) > info.GetRowSize()) {
204 : // 添加link
205 0 : u32 rankInRing = info.GetRank(hIndexForRing, info.GetRowSize());
206 0 : CHK_PRT_RET(
207 : rankInRing >= links.size(),
208 : HCCL_ERROR(
209 : "[ReduceScatterNHRV1][Vertical] rank[%u] out of range, "
210 : "rankInRing=%u, links.size=%u",
211 : rank, rankInRing, links.size()),
212 : HCCL_E_INTERNAL);
213 0 : HCCL_DEBUG("[ReduceScatterNHRV1][Vertical] rank[%u] links[%u]=%u", rank, subLinks.size(), rankInRing);
214 0 : subLinks.push_back(links[rankInRing]);
215 :
216 : // 添加slice
217 0 : Slice slice;
218 0 : slice.offset = 0;
219 0 : slice.size = 0;
220 0 : HCCL_DEBUG(
221 : "[ReduceScatterNHRV1][Vertical] rank[%u] subSlices[%u].offset=%llu, subSlices[%u].size=%llu", rank,
222 : subLinks.size(), slice.offset, subLinks.size(), slice.size);
223 0 : subSlices.push_back(slice);
224 : }
225 :
226 : // 长度不足2，直接跳过
227 0 : if (subLinks.size() < 2) {
228 0 : return HCCL_SUCCESS;
229 : }
230 :
231 : // 计算在垂直Ring中的rank号
232 0 : u32 subRank = (hIndex == hIndexForRing) ? info.GetVIndex(rank) : vSizeForRing;
233 0 : HCCL_DEBUG("[ReduceScatterNHRV1][Vertical] rank[%u] subRank=%u", rank, subRank);
234 :
235 : // 执行Broken Ring ReduceScatter
236 0 : return RunReduceScatterBrokenRing(subRank, subLinks, subSlices);
237 0 : }
238 :
239 : HcclResult
240 0 : ReduceScatterNHRV1::RunReduceScatterOnHorizontal(const u32 rank, const std::vector<LINK>& links, const RingInfo& info)
241 : {
242 0 : u32 hIndex = info.GetHIndex(rank);
243 0 : if (hIndex >= info.GetRowSize()) {
244 0 : return HCCL_SUCCESS;
245 : }
246 :
247 : // 构造新的links和slices
248 0 : u32 vIndex = info.GetVIndex(rank);
249 0 : std::vector<LINK> subLinks;
250 0 : std::vector<Slice> subSlices;
251 0 : for (u32 hIdx = 0; hIdx < info.GetRowSize(); hIdx++) {
252 : // 增加link
253 0 : u32 rankInRing = info.GetRank(vIndex, hIdx);
254 0 : CHK_PRT_RET(
255 : rankInRing >= links.size(),
256 : HCCL_ERROR(
257 : "[ReduceScatterNHRV1][Horizontal] rank[%u] out of range, "
258 : "rankInRing=%u, links.size=%u",
259 : rank, rankInRing, links.size()),
260 : HCCL_E_INTERNAL);
261 0 : HCCL_DEBUG("[ReduceScatterNHRV1][Horizontal] rank[%u] links[%u]=%u", rank, hIdx, rankInRing);
262 0 : subLinks.push_back(links[rankInRing]);
263 :
264 : // 增肌slice（CheckSlices()已经约束slices_里的Slice都是连续的）
265 : // -- 比如8节点在做完水平Ring后
266 : // ---- 0 1 2
267 : // ---- 3 4 5
268 : // ---- 6 7
269 : // -- 0/3/6/7节点只拿到自己那份ReduceScatter结果，而1拿到1和2的两份数据、4拿到4和5的两份数据
270 0 : u64 sliceSize = slices_[rankInRing].size;
271 0 : if (hIdx == info.GetRowSize() - 1 && info.GetHSizeByVIndex(vIndex) > info.GetRowSize()) {
272 0 : sliceSize += slices_[rankInRing + 1].size;
273 : }
274 :
275 0 : Slice slice;
276 0 : slice.offset = slices_[rankInRing].offset;
277 0 : slice.size = sliceSize;
278 0 : HCCL_DEBUG(
279 : "[ReduceScatterNHRV1][Horizontal] rank[%u] subSlices[%u].offset=%llu, subSlices[%u].size=%llu", rank, hIdx,
280 : slice.offset, hIdx, slice.size);
281 0 : subSlices.push_back(slice);
282 : }
283 :
284 : // 长度不足2，直接跳过
285 0 : if (subLinks.size() < 2) {
286 0 : return HCCL_SUCCESS;
287 : }
288 :
289 : // 计算在水平Ring中的rank号
290 0 : u32 subRank = hIndex;
291 0 : HCCL_DEBUG("[ReduceScatterNHRV1][Horizontal] rank[%u] subRank=%u", rank, subRank);
292 :
293 : // 执行Broken Ring ReduceScatter
294 0 : return RunReduceScatterBrokenRing(subRank, subLinks, subSlices);
295 0 : }
296 :
297 0 : HcclResult ReduceScatterNHRV1::RunLastCopyStep(const u32 rank, const std::vector<LINK>& links, const RingInfo& info)
298 : {
299 : HcclResult ret;
300 :
301 0 : u32 hIndex = info.GetHIndex(rank); // 查找自己位于第几列
302 0 : u32 vIndex = info.GetVIndex(rank); // 查找自己位于第几行
303 0 : if (hIndex >= info.GetRowSize() - 1 && info.GetHSizeByVIndex(vIndex) > info.GetRowSize()) {
304 0 : u32 peerRank = (hIndex == info.GetRowSize() - 1) ? rank + 1 : rank - 1;
305 :
306 : // 检查指针
307 0 : CHK_SMART_PTR_NULL(links[peerRank]);
308 :
309 : // TxAck
310 0 : ret = links[peerRank]->TxAck(stream_);
311 0 : CHK_PRT_RET(
312 : ret != HCCL_SUCCESS,
313 : HCCL_ERROR("[ReduceScatterNHRV1][RunLastCopyStep]rank[%u] tx ack from peerank[%u] failed", rank, peerRank),
314 : ret);
315 :
316 : // RxAck
317 0 : ret = links[peerRank]->RxAck(stream_);
318 0 : CHK_PRT_RET(
319 : ret != HCCL_SUCCESS,
320 : HCCL_ERROR("[ReduceScatterNHRV1][RunLastCopyStep]rank[%u] rx ack from peerank[%u] failed", rank, peerRank),
321 : ret);
322 :
323 0 : if (hIndex == info.GetRowSize() - 1) { // 发数据
324 0 : Slice& txSlice = slices_[peerRank];
325 0 : DeviceMem srcMem = inputMem_.range(txSlice.offset, txSlice.size);
326 0 : HCCL_DEBUG("tx srcMem[%p] range[%llu] size[%llu] ", srcMem.ptr(), txSlice.offset, txSlice.size);
327 0 : CHK_RET(ExecuteTxSync(
328 : links[peerRank], UserMemType::INPUT_MEM, txSlice.offset + baseOffset_, srcMem.ptr(), srcMem.size(),
329 : stream_));
330 :
331 0 : ret = links[peerRank]->TxWaitDone(stream_);
332 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[ReduceScatterNHRV1][RunLastCopyStep]TxWaitDone failed"), ret);
333 0 : } else { // 收数据
334 0 : Slice& rxSlice = slices_[rank];
335 0 : DeviceMem dstMem = inputMem_.range(rxSlice.offset, rxSlice.size);
336 0 : HCCL_DEBUG("rx dstMem[%p] range[%llu], size[%llu] ", dstMem.ptr(), rxSlice.offset, rxSlice.size);
337 0 : CHK_RET(ExecuteRxSync(
338 : links[peerRank], UserMemType::INPUT_MEM, rxSlice.offset + baseOffset_, dstMem.ptr(), dstMem.size(),
339 : stream_));
340 :
341 0 : ret = links[peerRank]->RxWaitDone(stream_);
342 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[ReduceScatterNHRV1][RunLastCopyStep]RxWaitDone failed"), ret);
343 0 : }
344 :
345 : // 如果不Barrier，SDMA结果在数据量超过CCL Buffer之后结果不正确
346 0 : CHK_RET(ExecuteBarrier(links[peerRank], stream_));
347 : }
348 0 : return HCCL_SUCCESS;
349 : }
350 :
351 0 : HcclResult ReduceScatterNHRV1::RunCopyDataToOutputMem(const u32 rank)
352 : {
353 0 : if (inputMem_ != outputMem_) {
354 0 : Slice& srcSlice = slices_[rank];
355 0 : DeviceMem dst = outputMem_.range(0, srcSlice.size);
356 0 : DeviceMem src = inputMem_.range(srcSlice.offset, srcSlice.size);
357 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
358 0 : }
359 0 : return HCCL_SUCCESS;
360 : }
361 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_NHR_V1, ReduceScatterNHRV1);
362 : } // namespace hccl
|
