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_nhr.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : AllReduceNHR::AllReduceNHR(const HcclDispatcher dispatcher) : NHRBase(dispatcher) {}
16 :
17 0 : AllReduceNHR::~AllReduceNHR() {}
18 :
19 0 : HcclResult AllReduceNHR::Prepare(u64 reduceAttrBitMap, [[maybe_unused]] HcomCollOpInfo* opInfo)
20 : {
21 0 : reduceAttr_ = reduceAttrBitMap;
22 0 : return HCCL_SUCCESS;
23 : }
24 :
25 0 : HcclResult AllReduceNHR::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
26 : {
27 : // 基本的检查
28 0 : CHK_RET(SimpleCheck(rank, rankSize, links));
29 0 : HCCL_INFO(
30 : "[AllReduceNHR][RunAsync] run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
31 : inputMem_.ptr(), outputMem_.ptr(), count_);
32 :
33 0 : HcclResult ret = HCCL_SUCCESS;
34 : // 如果ranksize为1, inline reduce和普通跨片reduce操作一致，从input->output
35 0 : if (rankSize == 1) {
36 0 : if (inputMem_ != outputMem_) {
37 0 : ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
38 0 : CHK_PRT_RET(
39 : ret != HCCL_SUCCESS, HCCL_ERROR("[AllReduceNHR][RunAsync] rank[%u] memcpy async failed", rank), ret);
40 : }
41 :
42 0 : return ret;
43 : }
44 :
45 : // reducescatter + allgather
46 0 : ret = PrepareRunAsync(rank, rankSize, links);
47 0 : CHK_PRT_RET(
48 : ret != HCCL_SUCCESS,
49 : HCCL_ERROR(
50 : "[AllReduceNHR][RunAsync] rank[%u] count[%llu] "
51 : "failed in PrepareRunAsync step",
52 : rank, count_),
53 : ret);
54 :
55 : // 先执行reducescater
56 0 : ret = RunReduceScatter(rank, rankSize, links);
57 0 : CHK_PRT_RET(
58 : ret != HCCL_SUCCESS,
59 : HCCL_ERROR(
60 : "[AllReduceNHR][RunAsync] rank[%u] count[%llu] failed in reducescater "
61 : "step",
62 : rank, count_),
63 : ret);
64 :
65 : // 再执行allgather
66 0 : ret = RunAllGather(rank, rankSize, links);
67 0 : CHK_PRT_RET(
68 : ret != HCCL_SUCCESS,
69 : HCCL_ERROR(
70 : "[AllReduceNHR][RunAsync] rank[%u] count[%llu] failed in AllGather "
71 : "step",
72 : rank, count_),
73 : ret);
74 :
75 0 : HCCL_INFO("[AllReduceNHR][RunAsync] finished: rank[%u] ranksize[%u]", rank, rankSize);
76 0 : return HCCL_SUCCESS;
77 : }
78 :
79 0 : HcclResult AllReduceNHR::SimpleCheck(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
80 : {
81 : // 判断stream, dispatcher是否为空
82 0 : CHK_SMART_PTR_NULL(dispatcher_);
83 0 : CHK_PTR_NULL(stream_.ptr());
84 :
85 : // 检查memory
86 0 : CHK_PRT_RET(
87 : !outputMem_ || !inputMem_,
88 : HCCL_ERROR("[AllReduceNHR][SimpleCheck] rank[%u] inputmem or outputmem is null", rank), HCCL_E_PTR);
89 :
90 : // 判断links数量是否正确
91 0 : CHK_PRT_RET(
92 : links.size() < rankSize,
93 : HCCL_ERROR(
94 : "[AllReduceNHR][SimpleCheck] rank[%u] link size[%llu] is less than "
95 : "rank size[%u]",
96 : rank, links.size(), rankSize),
97 : HCCL_E_INTERNAL);
98 0 : return HCCL_SUCCESS;
99 : }
100 :
101 0 : HcclResult AllReduceNHR::PrepareRunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
102 : {
103 : (void)links;
104 : // 计算reducescatter阶段每个rank结果上的offset和size
105 0 : if (slices_.size() == 0) {
106 0 : slices_.resize(rankSize);
107 0 : u64 totalSize = count_ * SIZE_TABLE[dataType_];
108 0 : u64 sliceSizeCalculated = (totalSize + (rankSize - 1)) / rankSize;
109 0 : u64 sliceSizeAligned = AlgTemplateBase::RoundUpWithDivisor(sliceSizeCalculated, HCCL_MIN_SLICE_ALIGN);
110 :
111 0 : u64 residueSize = totalSize;
112 :
113 0 : HCCL_DEBUG(
114 : "[AllReduceNHR][PrepareRunAsync]residueSize is %llu, sliceSizeAligned is %llu", residueSize,
115 : sliceSizeAligned);
116 0 : for (u32 i = 0; i < rankSize; i++) {
117 0 : slices_[i].size = (residueSize > sliceSizeAligned) ? sliceSizeAligned : residueSize;
118 0 : slices_[i].offset = totalSize - residueSize;
119 0 : residueSize -= slices_[i].size;
120 : }
121 :
122 0 : if (HcclCheckLogLevel(HCCL_LOG_DEBUG)) {
123 0 : for (size_t j = 0; j < slices_.size(); j++) {
124 0 : HCCL_DEBUG(
125 : "[AllReduceNHR][PrepareRunAsync] rank[%u] slice[%u]: offset[%llu] size[%llu]", rank, j,
126 : slices_[j].offset, slices_[j].size);
127 : }
128 : }
129 : }
130 0 : return HCCL_SUCCESS;
131 : }
132 :
133 0 : HcclResult AllReduceNHR::RunReduceScatter(u32 rank, u32 rankSize, const std::vector<LINK>& links)
134 : {
135 : std::unique_ptr<AlgTemplateBase> tempAlg
136 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_REDUCESCATTER_NHR, dispatcher_);
137 0 : CHK_SMART_PTR_NULL(tempAlg);
138 0 : CHK_RET(tempAlg->Prepare(reduceAttr_, true));
139 0 : HCCL_INFO(
140 : "[AllReduceNHR][RunReduceScatter] rank[%u] tempAlg ReduceScatterNHR inputMem[%p] outputMem[%p] "
141 : "mem_size[%llu] count[%llu] planeID:[%d]",
142 : rank, inputMem_.ptr(), outputMem_.ptr(), outputMem_.size(), count_, profilerInput_.planeID);
143 :
144 0 : if (!barrierSwitchOn_) {
145 0 : tempAlg->CloseBarrier();
146 : }
147 :
148 0 : CHK_RET(tempAlg->Prepare(
149 : inputMem_, inputMem_, outputMem_, count_, dataType_, stream_, reductionOp_, root_, slices_, baseOffset_));
150 :
151 0 : CHK_RET(tempAlg->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
152 :
153 0 : return tempAlg->RunAsync(rank, rankSize, links);
154 0 : }
155 :
156 0 : HcclResult AllReduceNHR::RunAllGather(u32 rank, u32 rankSize, const std::vector<LINK>& links)
157 : {
158 : std::unique_ptr<AlgTemplateBase> tempAlg
159 0 : = AlgTemplateRegistry::Instance().GetAlgTemplate(TemplateType::TEMPLATE_ALL_GATHER_NHR, dispatcher_);
160 0 : CHK_SMART_PTR_NULL(tempAlg);
161 0 : CHK_RET(tempAlg->Prepare(true));
162 0 : HCCL_INFO(
163 : "[AllReduceNHR][RunAllGather] rank[%u] tempAlg AllGatherNHR inputMem[%p] outputMem[%p] mem_size[%llu] "
164 : "count[%llu] planeID:[%d]",
165 : rank, inputMem_.ptr(), outputMem_.ptr(), outputMem_.size(), count_, profilerInput_.planeID);
166 :
167 0 : CHK_RET(tempAlg->Prepare(
168 : inputMem_, outputMem_, outputMem_, count_, dataType_, stream_, reductionOp_, root_, slices_, baseOffset_));
169 :
170 0 : CHK_RET(tempAlg->RegisterProfiler(profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
171 :
172 0 : return tempAlg->RunAsync(rank, rankSize, links);
173 0 : }
174 :
175 : HcclResult
176 0 : AllReduceNHR::GetNslbAdjInfo(const u32 rank, const u32 rankSize, const std::vector<LINK>& links, AdjInfo& nslbAdjInfo)
177 : {
178 0 : if (rankSize == 1) {
179 0 : return HCCL_SUCCESS;
180 : }
181 0 : if (links.size() < rankSize) {
182 0 : return HCCL_SUCCESS;
183 : }
184 0 : u32 nSteps = 0;
185 0 : for (u32 temp = rankSize - 1; temp != 0; temp >>= 1, ++nSteps) {
186 : }
187 :
188 : // 先执行ReduceScatter的NHR流程
189 0 : for (u32 step = 0; step < nSteps; step++) {
190 0 : u32 deltaRank = 1 << step;
191 0 : u32 sendTo = (rank + rankSize - deltaRank) % rankSize;
192 : ;
193 0 : LINK linkRight = links[sendTo];
194 0 : CHK_SMART_PTR_NULL(linkRight);
195 :
196 0 : NslbDpAdjInfo adjInfoStep = {};
197 0 : adjInfoStep.dstLocalRankId = linkRight->GetRemoteRank();
198 0 : adjInfoStep.phaseId = step + 1;
199 0 : adjInfoStep.rev = 0;
200 0 : nslbAdjInfo.nsAdjInfo.push_back(adjInfoStep);
201 0 : }
202 0 : u32 begin = nSteps;
203 : // 后续执行AllGather的NB流程
204 0 : for (u32 step = 0; step < nSteps; step++) {
205 0 : u32 deltaRank = 1 << (nSteps - 1 - step);
206 0 : u32 sendTo = (rank + deltaRank) % rankSize;
207 0 : LINK linkRight = links[sendTo];
208 0 : CHK_SMART_PTR_NULL(linkRight);
209 0 : NslbDpAdjInfo allGatherInfoStep = {};
210 0 : allGatherInfoStep.dstLocalRankId = linkRight->GetRemoteRank();
211 0 : allGatherInfoStep.phaseId = step + begin + 1;
212 0 : allGatherInfoStep.rev = 0;
213 0 : nslbAdjInfo.nsAdjInfo.push_back(allGatherInfoStep);
214 0 : }
215 0 : nslbAdjInfo.dstRankNum = nslbAdjInfo.nsAdjInfo.size();
216 0 : return HCCL_SUCCESS;
217 : }
218 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_REDUCE_NHR, AllReduceNHR);
219 : } // namespace hccl
|
