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