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 "scatter_nb.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : ScatterNB::ScatterNB(const HcclDispatcher dispatcher) : NBBase(dispatcher), interRank_(0), interRankSize_(0) {}
16 :
17 0 : ScatterNB::~ScatterNB() {}
18 :
19 0 : HcclResult ScatterNB::RunScatterNB(const std::vector<std::shared_ptr<Transport>>& links)
20 : {
21 0 : HcclResult ret = HCCL_SUCCESS;
22 : // 需要判断input不等于outputmem，scatter 输入只有一个input时不用拷贝
23 0 : if (inputMem_ != outputMem_) {
24 0 : ret = HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_);
25 0 : CHK_PRT_RET(
26 : ret != HCCL_SUCCESS,
27 : HCCL_ERROR(
28 : "[Run][ScatterOnRootRank]root rank[%u] memcpy async from input[%p] "
29 : "failed to output[%p]",
30 : interRank_, inputMem_.ptr(), outputMem_.ptr()),
31 : ret);
32 : }
33 :
34 : // 计算通信步数：ceiling(log2(rankSize))
35 0 : u32 nSteps = CalcCeilLog2(interRankSize_);
36 :
37 : // 逐步编排任务
38 0 : u32 deltaRoot = (interRank_ + interRankSize_ - root_) % interRankSize_;
39 0 : for (u32 step = 0; step < nSteps; step++) {
40 0 : if (deltaRoot < u32(1 << step)) {
41 0 : if (step != nSteps - 1 || deltaRoot < (interRankSize_ - (1 << step))) {
42 0 : RunScatterTx(step, links);
43 : }
44 0 : } else if (deltaRoot < u32(1 << (step + 1)) && interRank_ != root_) {
45 0 : RunScatterRx(step, links);
46 : }
47 : }
48 0 : return HCCL_SUCCESS;
49 : }
50 :
51 0 : HcclResult ScatterNB::RunScatterTx(const u32 step, const std::vector<std::shared_ptr<Transport>>& links)
52 : {
53 0 : HcclResult ret = HCCL_SUCCESS;
54 :
55 : // 计算通信对象
56 0 : u32 deltaRank = 1 << step;
57 0 : u32 sendTo = (interRank_ + deltaRank) % interRankSize_;
58 : // 数据份数和数据编号增量
59 0 : u32 nSlices = (interRankSize_ - 1 + (1 << step)) / (1 << (step + 1));
60 0 : u32 deltaSliceIndex = 1 << (step + 1);
61 0 : u32 sliceIdx = (interRank_ + (1 << step)) % interRankSize_;
62 0 : LINK linkRight = links[sendTo];
63 0 : CHK_SMART_PTR_NULL(linkRight);
64 :
65 0 : std::vector<Slice> txSlices;
66 0 : for (u32 i = 0; i < nSlices; i++) {
67 0 : if (slicesFlag_[sliceIdx] == false) {
68 0 : continue;
69 : }
70 0 : txSlices.push_back(slices_[sliceIdx]);
71 0 : sliceIdx = (sliceIdx + deltaSliceIndex) % interRankSize_;
72 : }
73 :
74 0 : CHK_RET(linkRight->RxAck(stream_));
75 0 : ret = Tx(linkRight, txSlices);
76 0 : CHK_PRT_RET(
77 : ret != HCCL_SUCCESS,
78 : HCCL_ERROR("[Run][Scatter]rank[%u] step[%u] RightLink tx slices count [%u] Failed", interRank_, step, nSlices),
79 : ret);
80 0 : ret = linkRight->TxWaitDone(stream_);
81 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][Scatter]TxWaitDone failed"), ret);
82 :
83 : // 为了避免在大数据量场景下触发网卡轮询机制，这里添加一组Data Notify，确保对端数据接收完成才进行下一次通信任务
84 0 : CHK_RET(linkRight->RxDataSignal(stream_));
85 0 : return HCCL_SUCCESS;
86 0 : }
87 :
88 0 : HcclResult ScatterNB::RunScatterRx(const u32 step, const std::vector<std::shared_ptr<Transport>>& links)
89 : {
90 0 : HcclResult ret = HCCL_SUCCESS;
91 :
92 : // 计算通信对象
93 0 : u32 deltaRank = 1 << step;
94 0 : u32 recvFrom = (interRank_ + interRankSize_ - deltaRank) % interRankSize_;
95 : // 数据份数和数据编号增量
96 0 : u32 nSlices = (interRankSize_ - 1 + (1 << step)) / (1 << (step + 1));
97 0 : u32 deltaSliceIndex = 1 << (step + 1);
98 0 : u32 sliceIdx = interRank_;
99 0 : LINK linkLeft = links[recvFrom];
100 0 : CHK_SMART_PTR_NULL(linkLeft);
101 :
102 0 : std::vector<Slice> rxSlices;
103 0 : for (u32 i = 0; i < nSlices; i++) {
104 0 : rxSlices.push_back(slices_[sliceIdx]);
105 0 : sliceIdx = (sliceIdx + deltaSliceIndex) % interRankSize_;
106 0 : slicesFlag_[sliceIdx] = true;
107 : }
108 :
109 0 : CHK_RET(linkLeft->TxAck(stream_));
110 0 : ret = Rx(linkLeft, rxSlices);
111 0 : CHK_PRT_RET(
112 : ret != HCCL_SUCCESS,
113 : HCCL_ERROR(
114 : "[Run][Scatter]rank[%u] step[%u] Right Link rx slices count [%u] "
115 : "Failed",
116 : interRank_, step, nSlices),
117 : ret);
118 :
119 0 : ret = linkLeft->RxWaitDone(stream_);
120 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Run][Scatter]RxWaitDone failed"), ret);
121 :
122 : // 为了避免在大数据量场景下触发网卡轮询机制，这里添加一组Data Notify，确保对端数据接收完成才进行下一次通信任务
123 0 : CHK_RET(linkLeft->TxDataSignal(stream_));
124 0 : return HCCL_SUCCESS;
125 0 : }
126 :
127 0 : void ScatterNB::PrepareSlicesData(const u32 unitSize, const u64 totalCount, const u32 rankSize) const
128 : {
129 0 : slices_.resize(rankSize);
130 0 : u64 sliceSize = (totalCount / rankSize) * unitSize;
131 :
132 0 : for (u32 i = 0; i < rankSize; i++) {
133 0 : slices_[i].offset = i * sliceSize;
134 0 : slices_[i].size = sliceSize;
135 : }
136 0 : }
137 :
138 : // scatter的入口函数
139 0 : HcclResult ScatterNB::RunAsync(const u32 rank, const u32 rankSize, const std::vector<std::shared_ptr<Transport>>& links)
140 : {
141 0 : CHK_SMART_PTR_NULL(dispatcher_);
142 0 : CHK_PTR_NULL(stream_.ptr());
143 0 : if (!outputMem_ || !inputMem_) {
144 0 : HCCL_ERROR("[ScatterNB][RunAsync]run_async inputmem or outputmem is null");
145 0 : return HCCL_E_PTR;
146 : }
147 :
148 0 : interRank_ = rank;
149 0 : interRankSize_ = rankSize;
150 0 : if (interRank_ == root_) {
151 0 : slicesFlag_.resize(interRankSize_, true);
152 : } else {
153 0 : slicesFlag_.resize(interRankSize_, false);
154 : }
155 :
156 : // ranksize为1时，只有当input!=output 时候进行拷贝
157 0 : if (interRankSize_ == 1) {
158 0 : if (inputMem_ != outputMem_) {
159 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, outputMem_, inputMem_, stream_));
160 : }
161 0 : return HCCL_SUCCESS;
162 : }
163 :
164 0 : u32 unitSize = DataUnitSize(dataType_);
165 0 : CHK_PRT_RET(
166 : unitSize == 0, HCCL_ERROR("[ScatterNB][RunAsync]rank[%u] unit data size is zero", rank), HCCL_E_INTERNAL);
167 :
168 : // 带入vecotr为空，计算每个rank的结果偏移和大小
169 0 : if (slices_.size() == 0) {
170 0 : PrepareSlicesData(unitSize, count_, interRankSize_);
171 : }
172 :
173 0 : CHK_PRT_RET(
174 : links.size() < rankSize,
175 : HCCL_ERROR("[ScatterNB][RunAsync]rank[%u] link size[%llu] is less than rank size", rank, links.size()),
176 : HCCL_E_INTERNAL);
177 :
178 0 : CHK_RET(RunScatterNB(links));
179 :
180 0 : if (barrierSwitchOn_) {
181 : // 执行barrier，保证数据收发完成
182 0 : CHK_RET(ExecuteBarrier(
183 : links[(interRank_ + interRankSize_ - 1) % interRankSize_], links[(interRank_ + 1) % interRankSize_]));
184 : }
185 0 : return HCCL_SUCCESS;
186 : }
187 :
188 0 : HcclResult ScatterNB::Tx(const LINK& link, const std::vector<Slice>& txSlices)
189 : {
190 0 : std::vector<TxMemoryInfo> txMems;
191 0 : for (const Slice& txSlice : txSlices) {
192 0 : DeviceMem srcMem = outputMem_.range(txSlice.offset, txSlice.size);
193 0 : HCCL_DEBUG("tx srcMem[%p] range[%llu] size[%llu] ", srcMem.ptr(), txSlice.offset, txSlice.size);
194 0 : txMems.emplace_back(
195 0 : TxMemoryInfo{UserMemType::OUTPUT_MEM, txSlice.offset + baseOffset_, srcMem.ptr(), txSlice.size});
196 0 : }
197 :
198 0 : CHK_RET(link->TxAsync(txMems, stream_));
199 0 : return HCCL_SUCCESS;
200 0 : }
201 :
202 0 : HcclResult ScatterNB::Rx(const LINK& link, const std::vector<Slice>& rxSlices)
203 : {
204 0 : std::vector<RxMemoryInfo> rxMems;
205 0 : for (const Slice& rxSlice : rxSlices) {
206 0 : DeviceMem dstMem = outputMem_.range(rxSlice.offset, rxSlice.size);
207 0 : HCCL_DEBUG("rx dstMem[%p] range[%llu], size[%llu] ", dstMem.ptr(), rxSlice.offset, rxSlice.size);
208 0 : rxMems.emplace_back(
209 0 : RxMemoryInfo{UserMemType::OUTPUT_MEM, rxSlice.offset + baseOffset_, dstMem.ptr(), rxSlice.size});
210 0 : }
211 :
212 0 : CHK_RET(link->RxAsync(rxMems, stream_));
213 0 : return HCCL_SUCCESS;
214 0 : }
215 :
216 : HcclResult
217 0 : ScatterNB::GetNslbAdjInfo(const u32 rank, const u32 rankSize, const std::vector<LINK>& links, AdjInfo& nslbAdjInfo)
218 : {
219 0 : if (rankSize == 1) {
220 0 : return HCCL_SUCCESS;
221 : }
222 0 : if (links.size() < rankSize) {
223 0 : return HCCL_SUCCESS;
224 : }
225 0 : HCCL_DEBUG("[ScatterNB]GetNslbAdjInfo start");
226 0 : u32 nSteps = 0;
227 0 : for (u32 temp = rankSize - 1; temp != 0; temp >>= 1, ++nSteps) {
228 : }
229 :
230 0 : u32 deltaRoot = (rank + rankSize - root_) % rankSize;
231 0 : for (u32 step = 0; step < nSteps; step++) {
232 0 : if (deltaRoot >= u32(1 << step)) {
233 0 : continue;
234 : }
235 0 : HCCL_DEBUG("[ScatterNB]now step is %u start", step);
236 0 : if (step != nSteps - 1 || deltaRoot < (rankSize - (1 << step))) {
237 0 : u32 deltaRank = 1 << step;
238 0 : u32 sendTo = (rank + deltaRank) % rankSize;
239 0 : HCCL_DEBUG("[ScatterNB]sendTo is %u", sendTo);
240 0 : LINK linkRight = links[sendTo];
241 0 : CHK_SMART_PTR_NULL(linkRight);
242 :
243 0 : NslbDpAdjInfo adjInfoStep = {};
244 0 : adjInfoStep.dstLocalRankId = linkRight->GetRemoteRank();
245 0 : adjInfoStep.phaseId = step + 1;
246 0 : adjInfoStep.rev = 0;
247 0 : nslbAdjInfo.nsAdjInfo.push_back(adjInfoStep);
248 0 : }
249 : }
250 0 : nslbAdjInfo.dstRankNum = nslbAdjInfo.nsAdjInfo.size();
251 0 : return HCCL_SUCCESS;
252 : }
253 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_SCATTER_NB, ScatterNB);
254 : } // namespace hccl
|
