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