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 "reduce_nhr_oneshot.h"
13 :
14 : namespace hccl {
15 :
16 3 : ReduceNHROneshot::ReduceNHROneshot(const HcclDispatcher dispatcher) : NHRBase(dispatcher) {}
17 :
18 6 : ReduceNHROneshot::~ReduceNHROneshot() {}
19 :
20 3 : HcclResult ReduceNHROneshot::Prepare(u64 reduceAttrBitMap, [[maybe_unused]] HcomCollOpInfo* opInfo)
21 : {
22 3 : reduceAttr_ = reduceAttrBitMap;
23 3 : return HCCL_SUCCESS;
24 : }
25 :
26 3 : HcclResult ReduceNHROneshot::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
27 : {
28 : // 基本的检查
29 3 : CHK_RET(SimpleCheck(rank, rankSize, links));
30 3 : HCCL_INFO(
31 : "[ReduceNHROneshot][RunAsync] run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank,
32 : rankSize, inputMem_.ptr(), outputMem_.ptr(), count_);
33 :
34 3 : u32 unitSize = DataUnitSize(dataType_);
35 3 : CHK_PRT_RET(
36 : unitSize == 0, HCCL_ERROR("[ReduceNHROneshot][RunAsync] rank[%u] unit data size is zero", rank),
37 : HCCL_E_INTERNAL);
38 :
39 : // 如果ranksize为1, 从input->output
40 3 : if (rankSize == 1) {
41 0 : if (inputMem_ != outputMem_) {
42 0 : return HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
43 : }
44 0 : HCCL_DEBUG("[ReduceNHROneshot]RunAsync for rankSize is 1 success");
45 0 : return HCCL_SUCCESS;
46 : }
47 :
48 : // 创建reducer & sender
49 3 : senderInfo_.reset(new (std::nothrow) Sender(dataType_, reductionOp_, reduceAttr_));
50 3 : CHK_SMART_PTR_NULL(senderInfo_);
51 :
52 3 : reducerInfo_.reset(new (std::nothrow) Reducer(dataType_, reductionOp_, reduceAttr_));
53 3 : CHK_SMART_PTR_NULL(reducerInfo_);
54 :
55 : // 运行reduce, NHR 算法
56 3 : CHK_RET(RunReduceNHROneshot(rank, rankSize, links));
57 :
58 3 : HCCL_INFO("[ReduceNHROneshot][RunAsync] finished: rank[%u] end", rank);
59 3 : return HCCL_SUCCESS;
60 : }
61 :
62 3 : HcclResult ReduceNHROneshot::SimpleCheck(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
63 : {
64 : // 判断stream, dispatcher是否为空
65 3 : CHK_SMART_PTR_NULL(dispatcher_);
66 3 : CHK_PTR_NULL(stream_.ptr());
67 :
68 : // 检查memory
69 3 : CHK_PRT_RET(
70 : !outputMem_ || !inputMem_,
71 : HCCL_ERROR("[ReduceNHROneshot][SimpleCheck] rank[%u] inputmem or outputmem is null", rank), HCCL_E_PTR);
72 :
73 : // 判断links数量是否正确
74 3 : CHK_PRT_RET(
75 : links.size() < rankSize,
76 : HCCL_ERROR(
77 : "[ReduceNHROneshot][SimpleCheck] rank[%u] link size[%llu] is "
78 : "less than rank size[%u]",
79 : rank, links.size(), rankSize),
80 : HCCL_E_INTERNAL);
81 3 : return HCCL_SUCCESS;
82 : }
83 :
84 0 : HcclResult ReduceNHROneshot::SdmaRx(
85 : LINK& linkLeft, LINK& linkRight, [[maybe_unused]] InterServerAlgoStep& stepInfo,
86 : [[maybe_unused]] const std::vector<LINK>& links)
87 : {
88 0 : HcclResult ret = HCCL_SUCCESS;
89 0 : u64 totalSize = count_ * SIZE_TABLE[dataType_];
90 0 : DeviceMem srcMem = inputMem_.range(0, totalSize);
91 0 : DeviceMem tempMem = scratchMem_.range(0, totalSize);
92 :
93 0 : if (linkRight != nullptr) {
94 0 : CHK_RET(linkRight->TxAck(stream_));
95 : }
96 :
97 0 : if (linkLeft != nullptr) {
98 0 : CHK_RET(linkLeft->RxAck(stream_));
99 0 : void* remoteMem = nullptr;
100 0 : CHK_RET(linkLeft->GetRemoteMem(UserMemType::INPUT_MEM, &remoteMem));
101 0 : if ((INLINE_REDUCE_BITMASK & reduceAttr_) == 1) { // inlineReduce
102 0 : CHK_RET(HcclReduceAsync(
103 : dispatcher_, static_cast<s8*>(remoteMem) + baseOffset_, tempMem.size() / SIZE_TABLE[dataType_],
104 : dataType_, reductionOp_, stream_, srcMem.ptr(), linkLeft->GetRemoteRank(), linkLeft->GetLinkType(),
105 : INLINE_REDUCE_BIT));
106 : } else { // tbeReduce
107 0 : DeviceMem srcMemLeft(static_cast<s8*>(remoteMem) + baseOffset_, totalSize);
108 0 : CHK_RET(HcclD2DMemcpyAsync(
109 : dispatcher_, tempMem, srcMemLeft, stream_,
110 : linkLeft->GetRemoteRank(), // left的inputMem拷到本端的scratchMem
111 : linkLeft->GetLinkType()));
112 0 : u64 dataCount = srcMem.size() / SIZE_TABLE[dataType_];
113 0 : ret = HcclReduceAsync(
114 0 : dispatcher_, tempMem.ptr(), dataCount, dataType_, reductionOp_, stream_, srcMem.ptr(),
115 : INVALID_VALUE_RANKID, LinkType::LINK_ONCHIP, reduceAttr_);
116 0 : }
117 0 : CHK_RET(linkLeft->TxDataSignal(stream_));
118 : }
119 0 : if (linkRight != nullptr) {
120 0 : CHK_RET(linkRight->RxDataSignal(stream_));
121 : }
122 0 : return ret;
123 0 : }
124 :
125 9 : HcclResult ReduceNHROneshot::RdmaTxRx(
126 : LINK& linkLeft, LINK& linkRight, [[maybe_unused]] InterServerAlgoStep& stepInfo,
127 : [[maybe_unused]] const std::vector<LINK>& links)
128 : {
129 9 : HcclResult ret = HCCL_SUCCESS;
130 9 : u64 totalSize = count_ * SIZE_TABLE[dataType_];
131 9 : DeviceMem srcMem = inputMem_.range(0, totalSize);
132 9 : DeviceMem tempMem = scratchMem_.range(0, totalSize);
133 :
134 9 : if (linkLeft != nullptr) {
135 9 : ret = linkLeft->TxAck(stream_);
136 9 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[ReduceNHROneshot][RunReduceNHROneshot] TxAck failed"), ret);
137 : }
138 :
139 9 : if (linkRight != nullptr) {
140 0 : ret = linkRight->RxAck(stream_);
141 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[ReduceNHROneshot][RunReduceNHROneshot] RxAck failed"), ret);
142 0 : ret = senderInfo_->run(linkRight, baseOffset_, srcMem, stream_);
143 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[ReduceNHROneshot][RunReduceNHROneshot] Tx failed"), ret);
144 : }
145 :
146 9 : if (linkLeft != nullptr) {
147 9 : ret = reducerInfo_->run(dispatcher_, linkLeft, baseOffset_, srcMem, srcMem, tempMem, stream_);
148 9 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[ReduceNHROneshot][RunReduceNHROneshot] Rx failed"), ret);
149 : }
150 :
151 9 : if (barrierSwitchOn_) {
152 9 : CHK_RET(ExecuteBarrier(linkLeft, linkRight));
153 : }
154 9 : return HCCL_SUCCESS;
155 9 : }
156 :
157 3 : HcclResult ReduceNHROneshot::RunReduceNHROneshot(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
158 : {
159 : // 计算通信步数
160 3 : u32 nSteps = GetStepNumInterServer(rankSize);
161 3 : HCCL_DEBUG("[ReduceNHROneshot][RunReduceNHROneshot] rank[%u] rankSize[%u] nSteps[%u]", rank, rankSize, nSteps);
162 :
163 : // 逐步编排任务
164 12 : for (u32 step = 0; step < nSteps; step++) {
165 9 : InterServerAlgoStep stepInfo;
166 9 : GetStepInfo(step, nSteps, rank, rankSize, stepInfo);
167 :
168 9 : u32 sendTo = stepInfo.toRank;
169 9 : u32 recvFrom = stepInfo.fromRank;
170 :
171 : // 当前每个数据块发送一次ACK、reduce一次、同步一次
172 9 : HCCL_DEBUG("[ReduceNHROneshot][RunReduceNHROneshot] recvFrom[%u] sendTo[%u] step[%u]", recvFrom, sendTo, step);
173 :
174 9 : LINK linkLeft;
175 9 : LINK linkRight;
176 9 : if (stepInfo.txSliceIdxs.size() > 0) {
177 0 : linkRight = links[stepInfo.toRank];
178 0 : CHK_SMART_PTR_NULL(linkRight);
179 : }
180 9 : if (stepInfo.rxSliceIdxs.size() > 0) {
181 9 : linkLeft = links[stepInfo.fromRank];
182 9 : CHK_SMART_PTR_NULL(linkLeft);
183 : }
184 :
185 9 : if ((linkRight != nullptr && linkRight->IsSpInlineReduce())
186 9 : || (linkLeft != nullptr && linkLeft->IsSpInlineReduce())) {
187 0 : CHK_RET(SdmaRx(linkLeft, linkRight, stepInfo, links));
188 : } else {
189 9 : CHK_RET(RdmaTxRx(linkLeft, linkRight, stepInfo, links));
190 : }
191 9 : }
192 3 : return HCCL_SUCCESS;
193 : }
194 :
195 : // NHR每步的算法描述原理函数
196 9 : HcclResult ReduceNHROneshot::GetStepInfo(u32 step, u32 nSteps, u32 rank, u32 rankSize, InterServerAlgoStep& stepInfo)
197 : {
198 : (void)nSteps;
199 9 : stepInfo.txSliceIdxs.clear();
200 9 : stepInfo.rxSliceIdxs.clear();
201 9 : stepInfo.nSlices = 1;
202 9 : stepInfo.toRank = rankSize;
203 9 : stepInfo.fromRank = rankSize;
204 9 : stepInfo.step = step;
205 9 : stepInfo.myRank = rank;
206 :
207 9 : u32 nRanks = (rankSize - 1 + (1 << step)) / (1 << (step + 1)); // 本步需要进行收/发的rank数
208 :
209 : // 以0为root，第i步，0+deltaRankPair开始，每隔deltaRankGroup的rank需要发给rank-deltaRankPair
210 9 : u32 deltaRoot = (rank + rankSize - root_) % rankSize;
211 :
212 9 : u32 deltaRankPair = 1 << step;
213 9 : u32 deltaRankGroup = 1 << (step + 1);
214 :
215 9 : if (deltaRoot / deltaRankGroup < nRanks) {
216 9 : if ((deltaRoot + deltaRankPair) % deltaRankGroup == 0) {
217 0 : stepInfo.toRank = (rank + rankSize - deltaRankPair) % rankSize;
218 0 : stepInfo.txSliceIdxs.push_back(0);
219 : }
220 :
221 9 : if (deltaRoot % deltaRankGroup == 0) {
222 9 : stepInfo.fromRank = (rank + deltaRankPair) % rankSize;
223 9 : stepInfo.rxSliceIdxs.push_back(0);
224 : }
225 : }
226 9 : return HCCL_SUCCESS;
227 : }
228 :
229 : HcclResult
230 9 : ReduceNHROneshot::ExecuteBarrier(const std::shared_ptr<Transport>& preLink, const std::shared_ptr<Transport>& aftLink)
231 : {
232 9 : if (preLink != nullptr) {
233 9 : CHK_RET(preLink->TxAck(stream_));
234 : }
235 9 : if (aftLink != nullptr) {
236 0 : CHK_RET(aftLink->RxAck(stream_));
237 0 : CHK_RET(aftLink->TxDataSignal(stream_));
238 : }
239 9 : if (preLink != nullptr) {
240 9 : CHK_RET(preLink->RxDataSignal(stream_));
241 9 : CHK_RET(preLink->PostFinAck(stream_));
242 : }
243 9 : if (aftLink != nullptr) {
244 0 : CHK_RET(aftLink->WaitFinAck(stream_));
245 : }
246 :
247 9 : return HCCL_SUCCESS;
248 : }
249 :
250 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCE_NHR_ONE_SHOT, ReduceNHROneshot);
251 : } // namespace hccl
|
