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