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_nb.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 9 : AllGatherNB::AllGatherNB(const HcclDispatcher dispatcher) : NBBase(dispatcher)
16 : {
17 9 : }
18 :
19 9 : AllGatherNB::~AllGatherNB()
20 : {
21 9 : }
22 :
23 0 : HcclResult AllGatherNB::Tx(const LINK &link, const std::vector<Slice> &txSlices)
24 : {
25 0 : std::vector<TxMemoryInfo> txMems;
26 0 : for (const Slice &txSlice : txSlices) {
27 0 : DeviceMem srcMem = outputMem_.range(txSlice.offset, txSlice.size);
28 0 : HCCL_DEBUG("tx srcMem[%p] range[%llu] size[%llu] ", srcMem.ptr(), txSlice.offset, txSlice.size);
29 0 : txMems.emplace_back(TxMemoryInfo{UserMemType::OUTPUT_MEM, txSlice.offset + baseOffset_,
30 0 : srcMem.ptr(), txSlice.size});
31 0 : }
32 :
33 0 : CHK_RET(link->TxAsync(txMems, stream_));
34 0 : return HCCL_SUCCESS;
35 0 : }
36 :
37 0 : HcclResult AllGatherNB::Rx(const LINK &link, const std::vector<Slice> &rxSlices)
38 : {
39 0 : std::vector<RxMemoryInfo> rxMems;
40 0 : for (const Slice &rxSlice : rxSlices) {
41 0 : DeviceMem dstMem = outputMem_.range(rxSlice.offset, rxSlice.size);
42 0 : HCCL_DEBUG("rx dstMem[%p] range[%llu], size[%llu] ", dstMem.ptr(), rxSlice.offset, rxSlice.size);
43 0 : rxMems.emplace_back(RxMemoryInfo{UserMemType::OUTPUT_MEM, rxSlice.offset + baseOffset_,
44 0 : dstMem.ptr(), rxSlice.size});
45 0 : }
46 :
47 0 : CHK_RET(link->RxAsync(rxMems, stream_));
48 0 : return HCCL_SUCCESS;
49 0 : }
50 :
51 : // 服务器间allgather的入口函数
52 0 : HcclResult AllGatherNB::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK> &links)
53 : {
54 0 : CHK_SMART_PTR_NULL(dispatcher_);
55 0 : CHK_PTR_NULL(stream_.ptr());
56 0 : HCCL_INFO("[AllGatherNB] run_async rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", \
57 : rank, rankSize, inputMem_.ptr(), outputMem_.ptr(), count_);
58 :
59 0 : if (rankSize == 1) {
60 0 : if (inputMem_ != outputMem_) {
61 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_));
62 : }
63 0 : return HCCL_SUCCESS;
64 : }
65 :
66 0 : CHK_PRT_RET(links.size() < rankSize,
67 : HCCL_ERROR("[AllGatherNB][RunAsync]rank[%u] linkSize is less than rankSize", rank), HCCL_E_INTERNAL);
68 :
69 0 : u32 unitSize = DataUnitSize(dataType_);
70 0 : CHK_PRT_RET(unitSize == 0, HCCL_ERROR("[AllGatherNB][RunAsync]unitSize is zero"), HCCL_E_INTERNAL);
71 :
72 0 : std::vector<Slice> inputSlices(slices_);
73 0 : if (slices_.size() == 0) {
74 0 : slices_.resize(rankSize);
75 0 : inputSlices.resize(rankSize);
76 :
77 0 : u64 sliceSize = count_ * unitSize;
78 0 : for (u32 i = 0; i < rankSize; i++) {
79 0 : slices_[i].size = sliceSize;
80 0 : slices_[i].offset = sliceSize * i;
81 0 : inputSlices[i].size = sliceSize;
82 0 : inputSlices[i].offset = (inputMem_.size() < outputMem_.size()) ? 0 : (sliceSize * i);
83 0 : HCCL_DEBUG("rank[%u], slices[%u].offset=%llu, slices[%u].size=%llu", rank, i, slices_[i].offset,
84 : i, slices_[i].size);
85 : }
86 : }
87 :
88 : // 双buffer下, 先将input拷贝到output的合适位置
89 0 : if (inputMem_ != outputMem_) {
90 0 : DeviceMem dst = outputMem_.range(slices_[rank].offset, slices_[rank].size);
91 0 : DeviceMem src = inputMem_.range(inputSlices[rank].offset, inputSlices[rank].size);
92 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
93 0 : }
94 :
95 : // 运行all-gather, ring算法
96 0 : CHK_RET(RunAllGather(rank, rankSize, slices_, links));
97 :
98 0 : HCCL_INFO("[AllGatherNB] finished: rank[%u] end", rank);
99 0 : return HCCL_SUCCESS;
100 0 : }
101 :
102 0 : HcclResult AllGatherNB::RunAllGather(u32 rank, u32 rankSize, const std::vector<Slice> &outputSlices,
103 : const std::vector<LINK> &links)
104 : {
105 0 : CHK_PRT_RET(outputSlices.size() < rankSize,
106 : HCCL_ERROR("[Run][AllGather]rank[%u] OutputSlice Size is less than rank size", rank), HCCL_E_INTERNAL);
107 0 : HcclResult ret = HCCL_SUCCESS;
108 :
109 0 : u32 nSteps = CalcCeilLog2(rankSize);
110 0 : u32 nRealSliceSize = outputSlices.size() / rankSize;
111 0 : HCCL_DEBUG("[AllGatherNB][RunAllGather]Starts, outputSlices.size[%u], rankSize[%u], nRealSliceSize[%u]",
112 : outputSlices.size(), rankSize, nRealSliceSize);
113 :
114 : // 逐步编排任务
115 0 : for (u32 step = 0; step < nSteps; step++) {
116 : // 计算通信对象
117 0 : u32 deltaRank = 1 << step;
118 0 : u32 recvFrom = (rankSize + rank - deltaRank) % rankSize;
119 0 : u32 sendTo = (rank + deltaRank) % rankSize;
120 :
121 : // 数据份数和数据编号增量
122 : // 节点6个，nSteps = 3，step = 2
123 0 : u32 nSlices = 1<<step;
124 0 : if (step == (nSteps - 1) && rankSize != u32(1<<nSteps)) {
125 0 : nSlices = rankSize - (1<<step);
126 : }
127 :
128 0 : LINK linkLeft = links[recvFrom];
129 0 : CHK_SMART_PTR_NULL(linkLeft);
130 :
131 0 : LINK linkRight = links[sendTo];
132 0 : CHK_SMART_PTR_NULL(linkRight);
133 :
134 0 : std::vector<Slice> txSlices;
135 0 : std::vector<Slice> rxSlices;
136 0 : for (u32 i = 0; i < nSlices; i++) {
137 0 : u32 rxSliceIndex = (rank - (1 << step) + rankSize - i) % rankSize;
138 0 : u32 txSliceIndex = (rank + rankSize - i) % rankSize;
139 0 : for (u32 j = 0; j < nRealSliceSize; j++) {
140 0 : u32 rxIndex = rxSliceIndex * nRealSliceSize + j;
141 0 : u32 txIndex = txSliceIndex * nRealSliceSize + j;
142 0 : if (outputSlices[txIndex].size > 0) {
143 0 : txSlices.push_back(outputSlices[txIndex]);
144 : }
145 0 : if (outputSlices[rxIndex].size > 0) {
146 0 : rxSlices.push_back(outputSlices[rxIndex]);
147 : }
148 0 : HCCL_DEBUG("rank[%u] round[%u] slice[%u] realSlice[%u] rx data outputSlice[%u] offset[%llu] size[%llu]",
149 : rank, step, i, j, rxIndex, outputSlices[rxIndex].offset, outputSlices[rxIndex].size);
150 : }
151 : }
152 0 : if (rxSlices.size() > 0) {
153 0 : CHK_RET(linkLeft->TxAck(stream_));
154 : }
155 0 : if (txSlices.size() > 0) {
156 0 : CHK_RET(linkRight->RxAck(stream_));
157 0 : ret = Tx(linkRight, txSlices);
158 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][AllGather]rank[%u] round[%u] tx %u slices failed",
159 : rank, step, nSlices), ret);
160 : }
161 0 : if (rxSlices.size() > 0) {
162 0 : ret = Rx(linkLeft, rxSlices);
163 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][AllGather]rank[%u] round[%u] rx %u slices failed",
164 : rank, step, nSlices), ret);
165 :
166 0 : ret = linkLeft->RxWaitDone(stream_);
167 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][AllGather]RxWaitDone failed"), ret);
168 0 : ret = linkLeft->PostFinAck(stream_);
169 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][AllGather]PostFinAck failed"), ret);
170 : }
171 0 : if (txSlices.size() > 0) {
172 0 : ret = linkRight->TxWaitDone(stream_);
173 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][AllGather]TxWaitDone failed"), ret);
174 0 : ret = linkRight->WaitFinAck(stream_);
175 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][AllGather]WaitFinAck failed"), ret);
176 : }
177 :
178 0 : if (linkRight->IsSpInlineReduce() || linkLeft->IsSpInlineReduce()) {
179 : // SDMA场景同步
180 0 : CHK_RET(ExecuteBarrier(linkLeft, linkRight));
181 : }
182 0 : }
183 0 : return HCCL_SUCCESS;
184 : }
185 :
186 0 : HcclResult AllGatherNB::GetNslbAdjInfo(const u32 rank, const u32 rankSize,
187 : const std::vector<LINK> &links, AdjInfo& nslbAdjInfo)
188 : {
189 0 : if (rankSize == 1) {
190 0 : return HCCL_SUCCESS;
191 : }
192 0 : if (links.size() < rankSize) {
193 0 : return HCCL_SUCCESS;
194 : }
195 0 : u32 nSteps = 0;
196 0 : for(u32 temp = rankSize - 1; temp != 0; temp >>= 1, ++nSteps){}
197 :
198 0 : for (u32 step = 0; step < nSteps; step++) {
199 0 : u32 deltaRank = 1 << step;
200 0 : u32 sendTo =(rank + deltaRank) % rankSize;
201 0 : LINK linkRight = links[sendTo];
202 0 : CHK_SMART_PTR_NULL(linkRight);
203 :
204 0 : NslbDpAdjInfo adjInfoStep = {0};
205 0 : adjInfoStep.dstLocalRankId = linkRight->GetRemoteRank();
206 0 : adjInfoStep.phaseId = step + 1;
207 0 : adjInfoStep.rev = 0;
208 0 : nslbAdjInfo.nsAdjInfo.push_back(adjInfoStep);
209 0 : }
210 0 : nslbAdjInfo.dstRankNum = nSteps;
211 0 : return HCCL_SUCCESS;
212 : }
213 :
214 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_GATHER_NB, AllGatherNB);
215 : } // namespace hccl
|
