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, 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 : HcclResult
85 0 : ReduceNHROneshot::SdmaRx(LINK& linkLeft, LINK& linkRight, InterServerAlgoStep& stepInfo, const std::vector<LINK>& links)
86 : {
87 0 : HcclResult ret = HCCL_SUCCESS;
88 0 : u64 totalSize = count_ * SIZE_TABLE[dataType_];
89 0 : DeviceMem srcMem = inputMem_.range(0, totalSize);
90 0 : DeviceMem tempMem = scratchMem_.range(0, totalSize);
91 :
92 0 : if (linkRight != nullptr) {
93 0 : CHK_RET(linkRight->TxAck(stream_));
94 : }
95 :
96 0 : if (linkLeft != nullptr) {
97 0 : CHK_RET(linkLeft->RxAck(stream_));
98 0 : void* remoteMem = nullptr;
99 0 : CHK_RET(linkLeft->GetRemoteMem(UserMemType::INPUT_MEM, &remoteMem));
100 0 : if ((INLINE_REDUCE_BITMASK & reduceAttr_) == 1) { // inlineReduce
101 0 : CHK_RET(HcclReduceAsync(
102 : dispatcher_, static_cast<s8*>(remoteMem) + baseOffset_, tempMem.size() / SIZE_TABLE[dataType_],
103 : dataType_, reductionOp_, stream_, srcMem.ptr(), linkLeft->GetRemoteRank(), linkLeft->GetLinkType(),
104 : INLINE_REDUCE_BIT));
105 : } else { // tbeReduce
106 0 : DeviceMem srcMemLeft(static_cast<s8*>(remoteMem) + baseOffset_, totalSize);
107 0 : CHK_RET(HcclD2DMemcpyAsync(
108 : dispatcher_, tempMem, srcMemLeft, stream_,
109 : linkLeft->GetRemoteRank(), // left的inputMem拷到本端的scratchMem
110 : linkLeft->GetLinkType()));
111 0 : u64 dataCount = srcMem.size() / SIZE_TABLE[dataType_];
112 0 : ret = HcclReduceAsync(
113 0 : dispatcher_, tempMem.ptr(), dataCount, dataType_, reductionOp_, stream_, srcMem.ptr(),
114 : INVALID_VALUE_RANKID, LinkType::LINK_ONCHIP, reduceAttr_);
115 0 : }
116 0 : CHK_RET(linkLeft->TxDataSignal(stream_));
117 : }
118 0 : if (linkRight != nullptr) {
119 0 : CHK_RET(linkRight->RxDataSignal(stream_));
120 : }
121 0 : return ret;
122 0 : }
123 :
124 9 : HcclResult ReduceNHROneshot::RdmaTxRx(
125 : LINK& linkLeft, LINK& linkRight, InterServerAlgoStep& stepInfo, const std::vector<LINK>& links)
126 : {
127 9 : HcclResult ret = HCCL_SUCCESS;
128 9 : u64 totalSize = count_ * SIZE_TABLE[dataType_];
129 9 : DeviceMem srcMem = inputMem_.range(0, totalSize);
130 9 : DeviceMem tempMem = scratchMem_.range(0, totalSize);
131 :
132 9 : if (linkLeft != nullptr) {
133 9 : ret = linkLeft->TxAck(stream_);
134 9 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[ReduceNHROneshot][RunReduceNHROneshot] TxAck failed"), ret);
135 : }
136 :
137 9 : if (linkRight != nullptr) {
138 0 : ret = linkRight->RxAck(stream_);
139 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[ReduceNHROneshot][RunReduceNHROneshot] RxAck failed"), ret);
140 0 : ret = senderInfo_->run(linkRight, baseOffset_, srcMem, stream_);
141 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[ReduceNHROneshot][RunReduceNHROneshot] Tx failed"), ret);
142 : }
143 :
144 9 : if (linkLeft != nullptr) {
145 9 : ret = reducerInfo_->run(dispatcher_, linkLeft, baseOffset_, srcMem, srcMem, tempMem, stream_);
146 9 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[ReduceNHROneshot][RunReduceNHROneshot] Rx failed"), ret);
147 : }
148 :
149 9 : if (barrierSwitchOn_) {
150 9 : CHK_RET(ExecuteBarrier(linkLeft, linkRight));
151 : }
152 9 : return HCCL_SUCCESS;
153 9 : }
154 :
155 3 : HcclResult ReduceNHROneshot::RunReduceNHROneshot(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
156 : {
157 : // 计算通信步数
158 3 : u32 nSteps = GetStepNumInterServer(rankSize);
159 3 : HCCL_DEBUG("[ReduceNHROneshot][RunReduceNHROneshot] rank[%u] rankSize[%u] nSteps[%u]", rank, rankSize, nSteps);
160 :
161 : // 逐步编排任务
162 12 : for (u32 step = 0; step < nSteps; step++) {
163 9 : InterServerAlgoStep stepInfo;
164 9 : GetStepInfo(step, nSteps, rank, rankSize, stepInfo);
165 :
166 9 : u32 sendTo = stepInfo.toRank;
167 9 : u32 recvFrom = stepInfo.fromRank;
168 :
169 : // 当前每个数据块发送一次ACK、reduce一次、同步一次
170 9 : HCCL_DEBUG("[ReduceNHROneshot][RunReduceNHROneshot] recvFrom[%u] sendTo[%u] step[%u]", recvFrom, sendTo, step);
171 :
172 9 : LINK linkLeft;
173 9 : LINK linkRight;
174 9 : if (stepInfo.txSliceIdxs.size() > 0) {
175 0 : linkRight = links[stepInfo.toRank];
176 0 : CHK_SMART_PTR_NULL(linkRight);
177 : }
178 9 : if (stepInfo.rxSliceIdxs.size() > 0) {
179 9 : linkLeft = links[stepInfo.fromRank];
180 9 : CHK_SMART_PTR_NULL(linkLeft);
181 : }
182 :
183 9 : if ((linkRight != nullptr && linkRight->IsSpInlineReduce())
184 9 : || (linkLeft != nullptr && linkLeft->IsSpInlineReduce())) {
185 0 : CHK_RET(SdmaRx(linkLeft, linkRight, stepInfo, links));
186 : } else {
187 9 : CHK_RET(RdmaTxRx(linkLeft, linkRight, stepInfo, links));
188 : }
189 9 : }
190 3 : return HCCL_SUCCESS;
191 : }
192 :
193 : // NHR每步的算法描述原理函数
194 9 : HcclResult ReduceNHROneshot::GetStepInfo(u32 step, u32 nSteps, u32 rank, u32 rankSize, InterServerAlgoStep& stepInfo)
195 : {
196 : (void)nSteps;
197 9 : stepInfo.txSliceIdxs.clear();
198 9 : stepInfo.rxSliceIdxs.clear();
199 9 : stepInfo.nSlices = 1;
200 9 : stepInfo.toRank = rankSize;
201 9 : stepInfo.fromRank = rankSize;
202 9 : stepInfo.step = step;
203 9 : stepInfo.myRank = rank;
204 :
205 9 : u32 nRanks = (rankSize - 1 + (1 << step)) / (1 << (step + 1)); // 本步需要进行收/发的rank数
206 :
207 : // 以0为root，第i步，0+deltaRankPair开始，每隔deltaRankGroup的rank需要发给rank-deltaRankPair
208 9 : u32 deltaRoot = (rank + rankSize - root_) % rankSize;
209 :
210 9 : u32 deltaRankPair = 1 << step;
211 9 : u32 deltaRankGroup = 1 << (step + 1);
212 :
213 9 : if (deltaRoot / deltaRankGroup < nRanks) {
214 9 : if ((deltaRoot + deltaRankPair) % deltaRankGroup == 0) {
215 0 : stepInfo.toRank = (rank + rankSize - deltaRankPair) % rankSize;
216 0 : stepInfo.txSliceIdxs.push_back(0);
217 : }
218 :
219 9 : if (deltaRoot % deltaRankGroup == 0) {
220 9 : stepInfo.fromRank = (rank + deltaRankPair) % rankSize;
221 9 : stepInfo.rxSliceIdxs.push_back(0);
222 : }
223 : }
224 9 : return HCCL_SUCCESS;
225 : }
226 :
227 : HcclResult
228 9 : ReduceNHROneshot::ExecuteBarrier(const std::shared_ptr<Transport>& preLink, const std::shared_ptr<Transport>& aftLink)
229 : {
230 9 : if (preLink != nullptr) {
231 9 : CHK_RET(preLink->TxAck(stream_));
232 : }
233 9 : if (aftLink != nullptr) {
234 0 : CHK_RET(aftLink->RxAck(stream_));
235 0 : CHK_RET(aftLink->TxDataSignal(stream_));
236 : }
237 9 : if (preLink != nullptr) {
238 9 : CHK_RET(preLink->RxDataSignal(stream_));
239 9 : CHK_RET(preLink->PostFinAck(stream_));
240 : }
241 9 : if (aftLink != nullptr) {
242 0 : CHK_RET(aftLink->WaitFinAck(stream_));
243 : }
244 :
245 9 : return HCCL_SUCCESS;
246 : }
247 :
248 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCE_NHR_ONE_SHOT, ReduceNHROneshot);
249 : } // namespace hccl
|
