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_reduce_nb.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 5 : AllReduceNB::AllReduceNB(const HcclDispatcher dispatcher) : NBBase(dispatcher) {}
16 :
17 5 : AllReduceNB::~AllReduceNB() {}
18 :
19 5 : HcclResult AllReduceNB::Prepare(u64 reduceAttrBitMap, [[maybe_unused]] HcomCollOpInfo* opInfo)
20 : {
21 5 : reduceAttr_ = reduceAttrBitMap;
22 5 : return HCCL_SUCCESS;
23 : }
24 :
25 : // nb allreduce算法的函数入口
26 0 : HcclResult AllReduceNB::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
27 : {
28 0 : HcclResult ret = HCCL_SUCCESS;
29 0 : ret = PrepareRunAsync(rank, rankSize, links);
30 :
31 0 : CHK_PRT_RET(
32 : ret != HCCL_SUCCESS,
33 : HCCL_ERROR("[AllReduceNB][RunAsync]rank[%u] count[%llu] failed in PrepareRunAsync step", rank, count_), ret);
34 :
35 0 : CHK_PRT_RET(rankSize == 1, HCCL_INFO("[AllReduceNB][RunAsync] rankSize[%u], do nothing.", rankSize), HCCL_SUCCESS);
36 :
37 0 : CHK_PRT_RET(count_ == 0, HCCL_INFO("[AllReduceNB][RunAsync] count_[%u], do nothing.", count_), HCCL_SUCCESS);
38 :
39 : // 先执行reducescater
40 0 : ret = RunReduceScatter(rank, rankSize, links);
41 0 : CHK_PRT_RET(
42 : ret != HCCL_SUCCESS,
43 : HCCL_ERROR(
44 : "[AllReduceNB][RunAsync]rank[%u] count[%llu] failed in reducescater "
45 : "step",
46 : rank, count_),
47 : ret);
48 :
49 : // 再执行allgather
50 0 : ret = RunAllGather(rank, rankSize, links);
51 0 : CHK_PRT_RET(
52 : ret != HCCL_SUCCESS,
53 : HCCL_ERROR(
54 : "[AllReduceNB][RunAsync]rank[%u] count[%llu] failed in AllGather "
55 : "step",
56 : rank, count_),
57 : ret);
58 :
59 0 : HCCL_INFO("AllReduceNB finished: rank[%u] ranksize[%u]", rank, rankSize);
60 0 : return HCCL_SUCCESS;
61 : }
62 :
63 : HcclResult
64 0 : AllReduceNB::RunAsyncStaged(const u32 rank, const u32 rankSize, const std::vector<LINK>& links, RunStage stage)
65 : {
66 0 : CHK_PRT_RET(
67 : rankSize == 1 && stage != RunStage::RUN_PREPARE,
68 : HCCL_INFO("[AllReduceNB][RunAsyncStaged] rankSize[%u], stage[%d], do nothing.", rankSize, stage), HCCL_SUCCESS);
69 :
70 0 : HcclResult ret = HCCL_SUCCESS;
71 0 : switch (stage) {
72 0 : case RunStage::RUN_PREPARE:
73 0 : ret = PrepareRunAsync(rank, rankSize, links);
74 0 : CHK_PRT_RET(
75 : ret != HCCL_SUCCESS,
76 : HCCL_ERROR(
77 : "[AllReduceNB][RunAsyncStaged]rank[%u] count[%llu] failed in PrepareRunAsync step", rank, count_),
78 : ret);
79 0 : break;
80 0 : case RunStage::RUN_REDUCE_SCATTER:
81 : // 先执行reducescater
82 0 : ret = RunReduceScatter(rank, rankSize, links);
83 0 : CHK_PRT_RET(
84 : ret != HCCL_SUCCESS,
85 : HCCL_ERROR(
86 : "[AllReduceNB][RunAsyncStaged]rank[%u] count[%llu] "
87 : "failed in reducescater step",
88 : rank, count_),
89 : ret);
90 0 : break;
91 0 : case RunStage::RUN_ALLGATHER:
92 : // 再执行AllGather
93 0 : ret = RunAllGather(rank, rankSize, links);
94 0 : CHK_PRT_RET(
95 : ret != HCCL_SUCCESS,
96 : HCCL_ERROR(
97 : "[AllReduceNB][RunAsyncStaged]rank[%u] count[%llu] "
98 : "failed in AllGather step",
99 : rank, count_),
100 : ret);
101 0 : break;
102 0 : default:
103 0 : HCCL_ERROR("[AllReduceNB][RunAsyncStaged]stage[%d]is not support", stage);
104 0 : return HCCL_E_NOT_SUPPORT;
105 : }
106 0 : HCCL_INFO("AllReduceNB RunAsyncStaged stage[%d] finished: rank[%u] ranksize[%u]", stage, rank, rankSize);
107 0 : return HCCL_SUCCESS;
108 : }
109 :
110 0 : HcclResult AllReduceNB::PrepareRunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
111 : {
112 0 : HcclResult ret = HCCL_SUCCESS;
113 0 : CHK_SMART_PTR_NULL(dispatcher_);
114 0 : CHK_PTR_NULL(stream_.ptr());
115 0 : if (!outputMem_ || !inputMem_) {
116 0 : HCCL_ERROR("[AllReduceNB][RunAsync]rank[%u] run_async inputmem or outputmem is null", rank);
117 0 : return HCCL_E_PTR;
118 : }
119 0 : HCCL_INFO(
120 : "AllReduceNB run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
121 : inputMem_.ptr(), outputMem_.ptr(), count_);
122 :
123 0 : if (links.size() < rankSize) {
124 0 : HCCL_ERROR(
125 : "[AllReduceNB][RunAsync]rank[%u] linksize[%llu] is less than rankSize[%u]", rank, links.size(), rankSize);
126 0 : return HCCL_E_INTERNAL;
127 : }
128 :
129 : // 如果ranksize为1, inline reduce和普通跨片reduce操作一致，从input->output
130 0 : if (rankSize == 1) {
131 0 : if (inputMem_ != outputMem_) {
132 0 : ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
133 0 : CHK_PRT_RET(
134 : ret != HCCL_SUCCESS, HCCL_ERROR("[AllReduceNB][RunAsync]rank[%u] memcpy async failed", rank), ret);
135 : }
136 :
137 0 : return ret;
138 : }
139 : // 计算reducescatter 阶段每个rank结果上的offset和size
140 0 : if (slices_.size() == 0) {
141 0 : slices_.resize(rankSize);
142 0 : const u64 totalSize = count_ * SIZE_TABLE[dataType_];
143 0 : const u64 sliceSizeAligned = GetSliceSizeOfNB(totalSize, rankSize);
144 0 : u64 residueSize = totalSize;
145 :
146 0 : for (u32 i = 0; i < rankSize; i++) {
147 0 : slices_[i].size = (residueSize > sliceSizeAligned) ? sliceSizeAligned : residueSize;
148 0 : slices_[i].offset = totalSize - residueSize;
149 0 : residueSize -= slices_[i].size;
150 : }
151 :
152 0 : if (HcclCheckLogLevel(HCCL_LOG_DEBUG)) {
153 0 : for (size_t j = 0; j < slices_.size(); j++) {
154 0 : HCCL_DEBUG("rank[%u] slice[%u]: offset[%llu] size[%llu]", rank, j, slices_[j].offset, slices_[j].size);
155 : }
156 : }
157 : }
158 0 : HCCL_INFO("AllReduceNB PrepareRunAsync finished: rank[%u] ranksize[%u]", rank, rankSize);
159 0 : return HCCL_SUCCESS;
160 : }
161 :
162 0 : HcclResult AllReduceNB::RunReduceScatter(u32 rank, u32 rankSize, const std::vector<LINK>& links)
163 : {
164 : // 调用ReduceScatterNB算法
165 : std::unique_ptr<AlgTemplateBase> tempAlg
166 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_REDUCESCATTER_NB, dispatcher_);
167 0 : CHK_SMART_PTR_NULL(tempAlg);
168 0 : CHK_RET(tempAlg->Prepare(reduceAttr_));
169 0 : HCCL_INFO(
170 : "rank[%u] tempAlg ReduceScatterNB inputMem[%p] outputMem[%p] mem_size[%llu] "
171 : "count[%llu] planeID:[%d]",
172 : rank, inputMem_.ptr(), outputMem_.ptr(), outputMem_.size(), count_, profilerInput_.planeID);
173 0 : tempAlg->CloseBarrier();
174 0 : CHK_RET(tempAlg->Prepare(
175 : inputMem_, inputMem_, outputMem_, count_, dataType_, stream_, reductionOp_, root_, slices_, baseOffset_));
176 :
177 0 : CHK_RET(tempAlg->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
178 :
179 0 : return tempAlg->RunAsync(rank, rankSize, links);
180 0 : }
181 :
182 0 : HcclResult AllReduceNB::RunAllGather(u32 rank, u32 rankSize, const std::vector<LINK>& links)
183 : {
184 : std::unique_ptr<AlgTemplateBase> tempAlg
185 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_ALL_GATHER_NB, dispatcher_);
186 0 : CHK_SMART_PTR_NULL(tempAlg);
187 0 : HCCL_INFO(
188 : "rank[%u] tempAlg AllGatherNB inputMem[%p] outputMem[%p] mem_size[%llu] "
189 : "count[%llu] planeID:[%d]",
190 : rank, inputMem_.ptr(), outputMem_.ptr(), outputMem_.size(), count_, profilerInput_.planeID);
191 : // 判断是否关闭allgather的barrier
192 0 : tempAlg->CloseBarrier();
193 :
194 : // 调用allgatherring的算法执行
195 0 : CHK_RET(tempAlg->Prepare(
196 : inputMem_, outputMem_, outputMem_, count_, dataType_, stream_, reductionOp_, root_, slices_, baseOffset_));
197 :
198 0 : CHK_RET(tempAlg->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
199 :
200 0 : return tempAlg->RunAsync(rank, rankSize, links);
201 0 : }
202 :
203 0 : u64 GetSliceSizeOfNB(const u64 dataSize, const u32 rankSize)
204 : {
205 0 : const u64 sliceSizeCalculated = (dataSize + (rankSize - 1)) / rankSize;
206 0 : u64 sliceSizeAligned = 0;
207 :
208 : // 优化小包性能，小于128k不切片
209 0 : if (sliceSizeCalculated > NB_ALLREDUCE_SMALL_SIZE) {
210 0 : sliceSizeAligned = AlgTemplateBase::RoundUpWithDivisor(sliceSizeCalculated, HCCL_MIN_SLICE_ALIGN);
211 : } else {
212 0 : sliceSizeAligned = AlgTemplateBase::RoundUpWithDivisor(sliceSizeCalculated, NB_ALLREDUCE_SMALL_SIZE);
213 : }
214 0 : HCCL_INFO(
215 : "dataSize[%llu], rankSize[%u], sliceSizeCalculated[%llu], sliceSizeAligned[%llu]", dataSize, rankSize,
216 : sliceSizeCalculated, sliceSizeAligned);
217 :
218 0 : return sliceSizeAligned;
219 : }
220 :
221 : HcclResult
222 0 : AllReduceNB::GetNslbAdjInfo(const u32 rank, const u32 rankSize, const std::vector<LINK>& links, AdjInfo& nslbAdjInfo)
223 : {
224 0 : if (rankSize == 1) {
225 0 : return HCCL_SUCCESS;
226 : }
227 0 : if (links.size() < rankSize) {
228 0 : return HCCL_SUCCESS;
229 : }
230 0 : u32 nSteps = 0;
231 0 : for (u32 temp = rankSize - 1; temp != 0; temp >>= 1, ++nSteps) {
232 : }
233 :
234 : // 先执行ReduceScatter的NB流程
235 0 : for (u32 step = 0; step < nSteps; step++) {
236 0 : u32 deltaRank = 1 << step;
237 0 : u32 sendTo = (rank + deltaRank) % rankSize;
238 0 : LINK linkRight = links[sendTo];
239 0 : CHK_SMART_PTR_NULL(linkRight);
240 0 : NslbDpAdjInfo adjInfoStep = {};
241 0 : adjInfoStep.dstLocalRankId = linkRight->GetRemoteRank();
242 0 : adjInfoStep.phaseId = step + 1;
243 0 : adjInfoStep.rev = 0;
244 0 : nslbAdjInfo.nsAdjInfo.push_back(adjInfoStep);
245 0 : }
246 0 : u32 begin = nSteps;
247 : // 后续执行AllGather的NB流程
248 0 : for (u32 step = 0; step < nSteps; step++) {
249 0 : u32 deltaRank = 1 << step;
250 0 : u32 sendTo = (rank + deltaRank) % rankSize;
251 0 : LINK linkRight = links[sendTo];
252 0 : CHK_SMART_PTR_NULL(linkRight);
253 0 : NslbDpAdjInfo allGatherInfoStep = {};
254 0 : allGatherInfoStep.dstLocalRankId = linkRight->GetRemoteRank();
255 0 : allGatherInfoStep.phaseId = step + begin + 1;
256 0 : allGatherInfoStep.rev = 0;
257 0 : nslbAdjInfo.nsAdjInfo.push_back(allGatherInfoStep);
258 0 : }
259 0 : nslbAdjInfo.dstRankNum = nslbAdjInfo.nsAdjInfo.size();
260 0 : return HCCL_SUCCESS;
261 : }
262 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_REDUCE_NB, AllReduceNB);
263 : } // namespace hccl
|
