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_gather_nhr.h"
12 : #include <cmath>
13 : #include "alg_template_register.h"
14 :
15 : namespace hccl {
16 0 : AllGatherNHR::AllGatherNHR(const HcclDispatcher dispatcher) : NHRBase(dispatcher) {}
17 :
18 0 : AllGatherNHR::~AllGatherNHR() {}
19 :
20 0 : HcclResult AllGatherNHR::Prepare(bool needMerge)
21 : {
22 0 : isNeedMerge = needMerge;
23 0 : return HCCL_SUCCESS;
24 : }
25 :
26 : // 服务器间allgather的入口函数
27 0 : HcclResult AllGatherNHR::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
28 : {
29 : // 从AllReduce或者Broadcast调用AllGather需要merge
30 0 : if (isNeedMerge) {
31 : // 获取tree映射，存储到类对象的成员变量中
32 0 : GetRankMapping(rankSize);
33 : }
34 0 : CHK_SMART_PTR_NULL(dispatcher_);
35 0 : CHK_PTR_NULL(stream_.ptr());
36 0 : HCCL_INFO(
37 : "[AllGatherNHR][RunAsync] rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
38 : inputMem_.ptr(), outputMem_.ptr(), count_);
39 :
40 0 : if (rankSize == 1) {
41 0 : if (inputMem_ != outputMem_) {
42 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_));
43 : }
44 0 : return HCCL_SUCCESS;
45 : }
46 :
47 0 : CHK_PRT_RET(
48 : links.size() < rankSize, HCCL_ERROR("[AllGatherNHR][RunAsync] rank[%u] linkSize is less than rankSize", rank),
49 : HCCL_E_INTERNAL);
50 :
51 0 : u32 unitSize = DataUnitSize(dataType_);
52 0 : CHK_PRT_RET(unitSize == 0, HCCL_ERROR("[AllGatherNHR][RunAsync]rank[%u] unitSize is zero", rank), HCCL_E_INTERNAL);
53 :
54 0 : std::vector<Slice> inputSlices(slices_);
55 0 : if (slices_.size() == 0) {
56 0 : slices_.resize(rankSize);
57 0 : inputSlices.resize(rankSize);
58 :
59 0 : u64 sliceSize = count_ * unitSize;
60 0 : for (u32 i = 0; i < rankSize; i++) {
61 0 : slices_[i].size = sliceSize;
62 0 : slices_[i].offset = sliceSize * i;
63 0 : inputSlices[i].size = sliceSize;
64 0 : inputSlices[i].offset = (inputMem_.size() < outputMem_.size()) ? 0 : (sliceSize * i);
65 0 : HCCL_DEBUG(
66 : "[AllGatherNHR][RunAsync] rank[%u], slices[%u].offset=%llu, slices[%u].size=%llu", rank, i,
67 : slices_[i].offset, i, slices_[i].size);
68 : }
69 : }
70 :
71 0 : if (sliceMap_.size() != rankSize) {
72 0 : GetRankMapping(rankSize, true); // 没有初始化过，说明不是由allreduce或者bcast调入，需要保序
73 : }
74 :
75 : // 双buffer下, 先将input拷贝到output的合适位置
76 0 : if (inputMem_ != outputMem_) {
77 0 : u32 targetIdx = sliceMap_[rank];
78 0 : DeviceMem dst = outputMem_.range(slices_[targetIdx].offset, slices_[targetIdx].size);
79 0 : DeviceMem src = inputMem_.range(inputSlices[targetIdx].offset, inputSlices[targetIdx].size);
80 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
81 0 : }
82 :
83 : // 运行all-gather, ring算法
84 0 : CHK_RET(RunAllGather(rank, rankSize, slices_, links));
85 :
86 0 : HCCL_INFO("[AllGatherNHR][RunAsync] AllGatherNHR finished: rank[%u] end", rank);
87 0 : return HCCL_SUCCESS;
88 0 : }
89 :
90 0 : HcclResult AllGatherNHR::SdmaRx(const LINK& linkLeft, const LINK& linkRight, std::vector<Slice>& rxSlices)
91 : {
92 0 : CHK_RET(linkRight->TxAck(stream_)); // 告知right可以从我这读了
93 0 : CHK_RET(linkLeft->RxAck(stream_)); // 等left告知可以从他那读了
94 0 : void* srcMemPtr = nullptr;
95 0 : CHK_RET(linkLeft->GetRemoteMem(UserMemType::OUTPUT_MEM, &srcMemPtr));
96 0 : for (const Slice& rxSlice : rxSlices) {
97 0 : DeviceMem dstMem = outputMem_.range(rxSlice.offset, rxSlice.size);
98 0 : DeviceMem srcMem(static_cast<s8*>(srcMemPtr) + baseOffset_ + rxSlice.offset, rxSlice.size);
99 0 : HCCL_DEBUG("[AllGatherNHR] rx dstMem[%p] range[%llu], size[%llu] ", dstMem.ptr(), rxSlice.offset, rxSlice.size);
100 0 : CHK_RET(HcclD2DMemcpyAsync(
101 : dispatcher_, dstMem, srcMem, stream_, linkLeft->GetRemoteRank(), // Memecpy
102 : linkLeft->GetLinkType()));
103 0 : }
104 0 : CHK_RET(linkLeft->TxDataSignal(stream_)); // 告知left我读完了
105 0 : CHK_RET(linkRight->RxDataSignal(stream_)); // 等right读完
106 :
107 0 : return HCCL_SUCCESS;
108 : }
109 :
110 0 : HcclResult AllGatherNHR::RdmaTxRx(
111 : const LINK& linkLeft, const LINK& linkRight, InterServerAlgoStep& stepInfo, std::vector<Slice>& txSlices,
112 : std::vector<Slice>& rxSlices)
113 : {
114 0 : HcclResult ret = HCCL_SUCCESS;
115 0 : CHK_RET(linkLeft->TxAck(stream_));
116 0 : CHK_RET(linkRight->RxAck(stream_));
117 0 : ret = Tx(linkRight, txSlices);
118 0 : CHK_PRT_RET(
119 : ret != HCCL_SUCCESS,
120 : HCCL_ERROR(
121 : "[AllGatherNHR][RunAllGather] rank[%u] round[%u] "
122 : "tx %u slices failed",
123 : stepInfo.myRank, stepInfo.step, stepInfo.nSlices),
124 : ret);
125 0 : ret = Rx(linkLeft, rxSlices);
126 0 : CHK_PRT_RET(
127 : ret != HCCL_SUCCESS,
128 : HCCL_ERROR(
129 : "[AllGatherNHR][RunAllGather] rank[%u] round[%u] "
130 : "rx %u slices failed",
131 : stepInfo.myRank, stepInfo.step, stepInfo.nSlices),
132 : ret);
133 :
134 0 : ret = linkLeft->PostFinAck(stream_);
135 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[AllGatherNHR][RunAllGather] PostFinAck failed"), ret);
136 0 : ret = linkRight->WaitFinAck(stream_);
137 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[AllGatherNHR][RunAllGather] WaitFinAck failed"), ret);
138 0 : return HCCL_SUCCESS;
139 : }
140 :
141 0 : HcclResult AllGatherNHR::RunAllGather(
142 : u32 rank, u32 rankSize, const std::vector<Slice>& outputSlices, const std::vector<LINK>& links)
143 : {
144 0 : CHK_PRT_RET(
145 : outputSlices.size() < rankSize,
146 : HCCL_ERROR("[AllGatherNHR][RunAllGather] rank[%u] OutputSlice Size is less than rank size", rank),
147 : HCCL_E_INTERNAL);
148 0 : HcclResult ret = HCCL_SUCCESS;
149 :
150 : // 计算通信步数
151 0 : u32 nSteps = GetStepNumInterServer(rankSize);
152 :
153 : // 逐步编排任务
154 0 : for (u32 step = 0; step < nSteps; step++) {
155 0 : InterServerAlgoStep stepInfo;
156 0 : GetStepInfo(step, nSteps, rank, rankSize, stepInfo);
157 :
158 0 : LINK linkLeft = links[stepInfo.fromRank];
159 0 : CHK_SMART_PTR_NULL(linkLeft);
160 :
161 0 : LINK linkRight = links[stepInfo.toRank];
162 0 : CHK_SMART_PTR_NULL(linkRight);
163 :
164 0 : std::vector<Slice> txSlices;
165 0 : std::vector<Slice> rxSlices;
166 :
167 0 : HCCL_DEBUG(
168 : "[AllGatherNHR][RunAllGather] rank[%u] rankSize[%u] recvFrom[%u] sendTo[%u] step[%u] nSteps[%u] "
169 : "nSlices[%u]",
170 : rank, rankSize, stepInfo.fromRank, stepInfo.toRank, step, nSteps, stepInfo.nSlices);
171 :
172 0 : for (u32 i = 0; i < stepInfo.nSlices; i++) {
173 0 : txSlices.push_back(outputSlices[stepInfo.txSliceIdxs[i]]);
174 0 : rxSlices.push_back(outputSlices[stepInfo.rxSliceIdxs[i]]);
175 :
176 0 : HCCL_DEBUG(
177 : "[AllGatherNHR][RunAllGather] i[%u] rxSliceIndex[%u] txSliceIndex[%u] rx data offset[%llu] "
178 : "size[%llu]",
179 : i, stepInfo.rxSliceIdxs[i], stepInfo.txSliceIdxs[i], outputSlices[stepInfo.rxSliceIdxs[i]].offset,
180 : outputSlices[stepInfo.rxSliceIdxs[i]].size);
181 : }
182 :
183 : // 合并连续slices
184 0 : MergeSlices(rxSlices);
185 0 : MergeSlices(txSlices);
186 :
187 0 : if (linkLeft->IsSpInlineReduce() && linkRight->IsSpInlineReduce()) {
188 0 : ret = SdmaRx(linkLeft, linkRight, rxSlices);
189 : } else {
190 0 : ret = RdmaTxRx(linkLeft, linkRight, stepInfo, txSlices, rxSlices);
191 : }
192 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[AllGatherNHR][RunAllGather] RunAllGather failed"), ret);
193 0 : }
194 0 : return HCCL_SUCCESS;
195 : }
196 :
197 0 : HcclResult AllGatherNHR::Tx(const LINK& link, std::vector<Slice>& txSlices)
198 : {
199 0 : std::vector<TxMemoryInfo> txMems;
200 0 : for (const Slice& txSlice : txSlices) {
201 0 : DeviceMem srcMem = outputMem_.range(txSlice.offset, txSlice.size);
202 0 : HCCL_DEBUG(
203 : "[AllGatherNHR][Tx] tx srcMem[%p] offset[%llu] size[%llu]", srcMem.ptr(), txSlice.offset, txSlice.size);
204 0 : txMems.emplace_back(
205 0 : TxMemoryInfo{UserMemType::OUTPUT_MEM, txSlice.offset + baseOffset_, srcMem.ptr(), txSlice.size});
206 0 : }
207 :
208 0 : CHK_RET(link->TxAsync(txMems, stream_));
209 0 : return HCCL_SUCCESS;
210 0 : }
211 :
212 0 : HcclResult AllGatherNHR::Rx(const LINK& link, std::vector<Slice>& rxSlices)
213 : {
214 0 : std::vector<RxMemoryInfo> rxMems;
215 0 : for (const Slice& rxSlice : rxSlices) {
216 0 : DeviceMem dstMem = outputMem_.range(rxSlice.offset, rxSlice.size);
217 0 : HCCL_DEBUG(
218 : "[AllGatherNHR][Rx] rx dstMem[%p] range[%llu], size[%llu] ", dstMem.ptr(), rxSlice.offset, rxSlice.size);
219 0 : rxMems.emplace_back(
220 0 : RxMemoryInfo{UserMemType::OUTPUT_MEM, rxSlice.offset + baseOffset_, dstMem.ptr(), rxSlice.size});
221 0 : }
222 :
223 0 : CHK_RET(link->RxAsync(rxMems, stream_));
224 0 : return HCCL_SUCCESS;
225 0 : }
226 :
227 : // NHR每步的算法描述原理函数
228 0 : HcclResult AllGatherNHR::GetStepInfo(u32 step, u32 nSteps, u32 rank, u32 rankSize, InterServerAlgoStep& stepInfo)
229 : {
230 0 : stepInfo.txSliceIdxs.clear();
231 0 : stepInfo.rxSliceIdxs.clear();
232 0 : u32 sliceSize = slices_.size() / rankSize;
233 0 : stepInfo.step = step;
234 0 : stepInfo.myRank = rank;
235 :
236 : // 计算通信对象
237 0 : u32 deltaRank = 1 << (nSteps - 1 - step);
238 0 : u32 recvFrom = (rank + rankSize - deltaRank) % rankSize;
239 0 : u32 sendTo = (rank + deltaRank) % rankSize;
240 :
241 : // 数据份数和数据编号增量
242 0 : u32 nSlices = (rankSize - 1 + (1 << (nSteps - 1 - step))) / (1 << (nSteps - step));
243 0 : u32 deltaSliceIndex = 1 << (nSteps - step);
244 0 : u32 txSliceIdx = rank;
245 0 : u32 rxSliceIdx = (rank - (1 << (nSteps - 1 - step)) + rankSize) % rankSize;
246 :
247 0 : stepInfo.nSlices = nSlices * sliceSize;
248 0 : stepInfo.toRank = sendTo;
249 0 : stepInfo.fromRank = recvFrom;
250 :
251 0 : for (u32 i = 0; i < nSlices; i++) {
252 0 : for (u32 j = 0; j < sliceSize; j++) {
253 0 : u32 targetTxSliceIdx = sliceMap_[txSliceIdx];
254 0 : stepInfo.txSliceIdxs.push_back(targetTxSliceIdx * sliceSize + j);
255 :
256 0 : u32 targetRxSliceIdx = sliceMap_[rxSliceIdx];
257 0 : stepInfo.rxSliceIdxs.push_back(targetRxSliceIdx * sliceSize + j);
258 :
259 0 : HCCL_DEBUG(
260 : "[AllGatherNHR][GetStepInfo] i[%u] txSliceIdx[%u]->targetTxSliceIdx[%u] rxSliceIdx[%u]->"
261 : "targetRxSliceIdx[%u]",
262 : i, txSliceIdx, targetTxSliceIdx, rxSliceIdx, targetRxSliceIdx);
263 : }
264 0 : txSliceIdx = (txSliceIdx + rankSize - deltaSliceIndex) % rankSize;
265 0 : rxSliceIdx = (rxSliceIdx + rankSize - deltaSliceIndex) % rankSize;
266 : }
267 0 : return HCCL_SUCCESS;
268 : }
269 :
270 : HcclResult
271 0 : AllGatherNHR::GetNslbAdjInfo(const u32 rank, const u32 rankSize, const std::vector<LINK>& links, AdjInfo& nslbAdjInfo)
272 : {
273 0 : if (rankSize == 1) {
274 0 : return HCCL_SUCCESS;
275 : }
276 0 : if (links.size() < rankSize) {
277 0 : return HCCL_SUCCESS;
278 : }
279 0 : u32 nSteps = 0;
280 0 : for (u32 temp = rankSize - 1; temp != 0; temp >>= 1, ++nSteps) {
281 : }
282 :
283 0 : for (u32 step = 0; step < nSteps; step++) {
284 0 : u32 deltaRank = 1 << (nSteps - 1 - step);
285 0 : u32 sendTo = (rank + deltaRank) % rankSize;
286 0 : LINK linkRight = links[sendTo];
287 0 : CHK_SMART_PTR_NULL(linkRight);
288 :
289 0 : NslbDpAdjInfo adjInfoStep = {};
290 0 : adjInfoStep.dstLocalRankId = linkRight->GetRemoteRank();
291 0 : adjInfoStep.phaseId = step + 1;
292 0 : adjInfoStep.rev = 0;
293 0 : nslbAdjInfo.nsAdjInfo.push_back(adjInfoStep);
294 0 : }
295 0 : nslbAdjInfo.dstRankNum = nSteps;
296 0 : return HCCL_SUCCESS;
297 : }
298 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_GATHER_NHR, AllGatherNHR);
299 : } // namespace hccl
|
