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