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