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