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_ring_direct.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 0 : AllGatherRingDirect::AllGatherRingDirect(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
16 :
17 0 : AllGatherRingDirect::~AllGatherRingDirect() {}
18 :
19 0 : HcclResult AllGatherRingDirect::Prepare(
20 : HcomCollOpInfo* opInfo, u32 userRank, const std::vector<Slice>& userMemOutputSlices, bool isSdma)
21 : {
22 0 : opInfo_ = opInfo;
23 0 : userRank_ = userRank;
24 0 : userMemOutputSlices_ = userMemOutputSlices;
25 0 : isSdma_ = isSdma;
26 0 : return HCCL_SUCCESS;
27 : }
28 :
29 : // allgather的入口函数
30 0 : HcclResult AllGatherRingDirect::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
31 : {
32 : // 基本的检查
33 0 : CHK_RET(CheckParameters(rank, rankSize, links));
34 :
35 0 : if (rankSize == 1) {
36 0 : CHK_RET(OneRankMemcpy());
37 0 : return HCCL_SUCCESS;
38 : }
39 : // 收集邻居信息
40 0 : CHK_RET(GetInitializedNeighborLinks(rank, rankSize, links));
41 :
42 : // 填充slice_
43 0 : CHK_RET(SetSlices(rank, rankSize));
44 :
45 : // 运行all-gather, ring算法
46 0 : CHK_RET(RunAllGather(rank, rankSize));
47 :
48 0 : if (barrierSwitchOn_) {
49 : // 执行barrier，保证数据收发完成
50 0 : CHK_RET(ExecuteBarrier(leftLink_, rightLink_));
51 : }
52 :
53 0 : HCCL_INFO("AllGatherRingDirect finished: rank[%u] end", rank);
54 0 : return HCCL_SUCCESS;
55 : }
56 :
57 0 : HcclResult AllGatherRingDirect::CheckParameters(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
58 : {
59 0 : CHK_PTR_NULL(opInfo_);
60 0 : CHK_RET(CheckConcurrentDirectParameters(rank, rankSize, links));
61 : // 判断userMemInputSlices数量是否正确
62 0 : CHK_PRT_RET(
63 : userMemOutputSlices_.size() % rankSize != 0,
64 : HCCL_ERROR(
65 : "[AllGatherRingDirect] userMemOutputSlices size[%u] can not be divided by rank size[%u]",
66 : userMemOutputSlices_.size(), rankSize),
67 : HCCL_E_PARA);
68 :
69 0 : HCCL_INFO("AllGatherRingDirect finished to CheckParameters");
70 0 : return HCCL_SUCCESS;
71 : }
72 :
73 : // 单卡场景
74 0 : HcclResult AllGatherRingDirect::OneRankMemcpy()
75 : {
76 0 : for (u32 sliceIdx = 0; sliceIdx < slices_.size(); sliceIdx++) {
77 0 : const Slice& srcSlice = slices_[sliceIdx];
78 0 : const Slice& dstSlice = userMemOutputSlices_[sliceIdx];
79 0 : DeviceMem src;
80 0 : DeviceMem dst = DeviceMem::create(static_cast<u8*>(opInfo_->outputAddr) + dstSlice.offset, dstSlice.size);
81 0 : if (opInfo_->inputAddr != nullptr) {
82 : // opInfo_->inputAddr != nullptr指示要从user input获取输入
83 0 : u64 stepOffset = slices_[0].offset;
84 0 : HCCL_DEBUG(
85 : "Memcpy operation: stream[main], rank[%u] starts to copy offset[%llu], size[%llu] at userInput",
86 : userRank_, stepOffset, srcSlice.size);
87 0 : src = DeviceMem::create(static_cast<u8*>(opInfo_->inputAddr) + stepOffset, srcSlice.size);
88 : } else {
89 : // opInfo_->inputAddr == nullptr指示要从CCL buffer获取输入
90 0 : HCCL_DEBUG(
91 : "Memcpy operation: stream[main], rank[%u] starts to copy offset[%llu], size[%llu] at inputMem_",
92 : userRank_, srcSlice.offset, srcSlice.size);
93 0 : src = inputMem_.range(srcSlice.offset, srcSlice.size);
94 : }
95 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
96 0 : }
97 :
98 0 : return HCCL_SUCCESS;
99 : }
100 :
101 : HcclResult
102 0 : AllGatherRingDirect::GetInitializedNeighborLinks(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
103 : {
104 : // 收集左邻居信息
105 0 : leftLink_ = links[(rank + rankSize - 1) % rankSize];
106 0 : CHK_SMART_PTR_NULL(leftLink_);
107 :
108 : // 收集右邻居信息
109 0 : rightLink_ = links[(rank + 1) % rankSize];
110 0 : CHK_SMART_PTR_NULL(rightLink_);
111 :
112 0 : HCCL_INFO("AllGatherRingDirect finished to GetInitializedNeighborLinks");
113 0 : return HCCL_SUCCESS;
114 : }
115 :
116 0 : HcclResult AllGatherRingDirect::SetSlices(const u32 rank, const u32 rankSize)
117 : {
118 0 : inputSlices_ = slices_;
119 0 : if (slices_.size() == 0) {
120 0 : slices_.resize(rankSize);
121 0 : inputSlices_.resize(rankSize);
122 :
123 0 : u64 sliceSize = count_ * DataUnitSize(dataType_);
124 0 : for (u32 i = 0; i < rankSize; i++) {
125 0 : slices_[i].size = sliceSize;
126 0 : slices_[i].offset = sliceSize * i;
127 0 : inputSlices_[i].size = sliceSize;
128 0 : inputSlices_[i].offset = (inputMem_.size() < outputMem_.size()) ? 0 : (sliceSize * i);
129 0 : HCCL_DEBUG(
130 : "rank[%u], slices[%u].offset=%llu, slices[%u].size=[%llu]", rank, i, slices_[i].offset, i,
131 : slices_[i].size);
132 : }
133 : }
134 :
135 0 : if (UNLIKELY(HcclCheckLogLevel(DLOG_DEBUG))) {
136 0 : for (u32 i = 0; i < slices_.size(); i++) {
137 0 : HCCL_DEBUG(
138 : "[AllGatherRingDirect][SetSlices]rank[%u], slices[%u].offset=[%llu], slices[%u].size=[%llu]", rank, i,
139 : slices_[i].offset, i, slices_[i].size);
140 : }
141 : }
142 :
143 0 : HCCL_INFO("AllGatherRingDirect finished to SetSlices");
144 0 : return HCCL_SUCCESS;
145 : }
146 :
147 0 : HcclResult AllGatherRingDirect::RunInitStep(const u32 rank, const u32 rankSize)
148 : {
149 : // 第一步搬到userMemIn_的offset
150 0 : auto firstStepOffset = slices_[0].offset;
151 :
152 : // 第-1步，片内将部分数据从userIn搬到cclIn
153 0 : DeviceMem srcInit;
154 0 : DeviceMem dstInit;
155 0 : u32 initSliceIdx = rank;
156 0 : u32 sliceSize = slices_.size() / rankSize;
157 :
158 0 : for (u32 sliceIdx = 0; sliceIdx < sliceSize; sliceIdx++) {
159 0 : Slice initSlice = slices_[initSliceIdx * sliceSize + sliceIdx];
160 :
161 : // 需要+userMemIn_的offset
162 0 : if (opInfo_->inputAddr != nullptr) {
163 : // AllGather算子调用AllGatherRingDirect场景
164 0 : srcInit = DeviceMem::create(static_cast<u8*>(opInfo_->inputAddr) + firstStepOffset, initSlice.size);
165 : } else {
166 : // AllReduce算子调用AllGatherRingDirect场景
167 0 : srcInit = inputMem_.range(initSlice.offset, initSlice.size);
168 : }
169 :
170 0 : dstInit = outputMem_.range(initSlice.offset, initSlice.size);
171 0 : HCCL_DEBUG(
172 : "Memcpy operation: step[-1] stream[main] src rank[%u] starts to copy(rcv) offset[%llu], "
173 : "size[%llu] on userMemOutput to offset[%llu], size[%llu] on CCL",
174 : userRank_, firstStepOffset, initSlice.size, initSlice.offset, initSlice.size);
175 :
176 : // 若src与dst一样，则不需要搬运
177 0 : if (srcInit != dstInit) {
178 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstInit, srcInit, stream_));
179 : }
180 : }
181 :
182 0 : return HCCL_SUCCESS;
183 0 : }
184 :
185 : // 本端cclout -> 本端userout
186 0 : HcclResult AllGatherRingDirect::RunAllGatherPartOne(const u32 sliceSize, const u32 step, const u32 txSliceIdx)
187 : {
188 0 : std::vector<Slice> txSliceVector;
189 0 : std::vector<Slice> sliceVector;
190 :
191 0 : for (u32 sliceIdx = 0; sliceIdx < sliceSize; sliceIdx++) {
192 0 : txSliceVector.push_back(slices_[txSliceIdx * sliceSize + sliceIdx]);
193 0 : sliceVector.push_back(userMemOutputSlices_[txSliceIdx * sliceSize + sliceIdx]);
194 : }
195 :
196 0 : for (u32 sliceIdx = 0; sliceIdx < sliceSize; sliceIdx++) {
197 0 : DeviceMem src = outputMem_.range(txSliceVector[sliceIdx].offset, txSliceVector[sliceIdx].size);
198 : DeviceMem dst = DeviceMem::create(
199 0 : static_cast<u8*>(opInfo_->outputAddr) + sliceVector[sliceIdx].offset, sliceVector[sliceIdx].size);
200 :
201 0 : HCCL_DEBUG(
202 : "Memcpy operation: step[%u] stream[sub], src rank[%u] starts to send offset[%llu] size[%llu], "
203 : "dst rank starts to rcv offset[%llu] size[%llu] at userMemOutput_",
204 : step, userRank_, sliceVector[sliceIdx].offset, sliceVector[sliceIdx].size, txSliceVector[sliceIdx].offset,
205 : txSliceVector[sliceIdx].size);
206 :
207 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
208 0 : }
209 :
210 0 : return HCCL_SUCCESS;
211 0 : }
212 :
213 : // 对端cclout -> 本端cclout, 如果最后一步则：对端cclout -> 本端userout （DMA消减）
214 0 : HcclResult AllGatherRingDirect::RunAllGatherPartTwo(
215 : const u32 sliceSize, const u32 step, const u32 txSliceIdx, const u32 rxSliceIdx, const u32 rankSize)
216 : {
217 0 : std::vector<Slice> txSliceVector;
218 0 : std::vector<Slice> rxSliceVector;
219 0 : std::vector<Slice> sliceVector;
220 :
221 0 : for (u32 sliceIdx = 0; sliceIdx < sliceSize; sliceIdx++) {
222 0 : txSliceVector.push_back(slices_[txSliceIdx * sliceSize + sliceIdx]);
223 0 : rxSliceVector.push_back(slices_[rxSliceIdx * sliceSize + sliceIdx]);
224 0 : sliceVector.push_back(userMemOutputSlices_[rxSliceIdx * sliceSize + sliceIdx]);
225 : }
226 :
227 0 : CHK_RET(leftLink_->TxAck(stream_));
228 0 : CHK_RET(rightLink_->RxAck(stream_));
229 :
230 0 : std::vector<TxMemoryInfo> txMems;
231 0 : std::vector<RxMemoryInfo> rxMems;
232 :
233 0 : for (u32 sliceIdx = 0; sliceIdx < sliceSize; sliceIdx++) {
234 0 : DeviceMem src = outputMem_.range(txSliceVector[sliceIdx].offset, txSliceVector[sliceIdx].size);
235 0 : HCCL_DEBUG(
236 : "tx srcMem[%p] range[%llu] size[%llu] ", src.ptr(), txSliceVector[sliceIdx].offset,
237 : txSliceVector[sliceIdx].size);
238 0 : txMems.emplace_back(TxMemoryInfo{
239 0 : UserMemType::OUTPUT_MEM, txSliceVector[sliceIdx].offset + baseOffset_, src.ptr(),
240 0 : txSliceVector[sliceIdx].size});
241 :
242 0 : DeviceMem dst;
243 0 : if (isSdma_ && step == rankSize - DMA_REDUCE_TWO_OFFSET) {
244 : // 最后一步实现DMA消减：对端cclout -> 本端userout
245 0 : HCCL_DEBUG(
246 : "DMAReduce(sdma) MemcpyAsync operation: step[%u] stream[main], dst rank[%u] starts to rcv "
247 : "offset[%llu] size[%llu] at userMemOutput_",
248 : step, userRank_, sliceVector[sliceIdx].offset, sliceVector[sliceIdx].size);
249 :
250 0 : dst = DeviceMem::create(
251 0 : static_cast<u8*>(opInfo_->outputAddr) + sliceVector[sliceIdx].offset, sliceVector[sliceIdx].size);
252 : } else {
253 0 : HCCL_DEBUG(
254 : "MemcpyAsync operation: step[%u] stream[main], dst rank[%u] starts to rcv offset[%llu] size[%llu] "
255 : "at outputMem_",
256 : step, userRank_, rxSliceVector[sliceIdx].offset, rxSliceVector[sliceIdx].size);
257 :
258 : // 中间步数无DMA消减
259 0 : dst = outputMem_.range(rxSliceVector[sliceIdx].offset, rxSliceVector[sliceIdx].size);
260 0 : if (!isSdma_ && step == rankSize - DMA_REDUCE_TWO_OFFSET) {
261 : // 最后一步实现DMA消减：对端cclout -> 本端userout
262 0 : HCCL_DEBUG("DMAReduce(rdma) record final addr");
263 :
264 0 : finalSrc_.push_back(outputMem_.range(rxSliceVector[sliceIdx].offset, rxSliceVector[sliceIdx].size));
265 0 : finalDst_.push_back(DeviceMem::create(
266 0 : static_cast<u8*>(opInfo_->outputAddr) + sliceVector[sliceIdx].offset, sliceVector[sliceIdx].size));
267 : }
268 : }
269 :
270 0 : rxMems.emplace_back(RxMemoryInfo{
271 0 : UserMemType::OUTPUT_MEM, rxSliceVector[sliceIdx].offset + baseOffset_, dst.ptr(),
272 0 : rxSliceVector[sliceIdx].size});
273 0 : }
274 :
275 0 : CHK_RET(rightLink_->TxAsync(txMems, stream_));
276 :
277 0 : if (!isSdma_) {
278 0 : CHK_RET(leftLink_->RxAsync(rxMems, stream_));
279 : } else {
280 0 : CHK_RET(leftLink_->RxDataSignal(stream_));
281 :
282 0 : for (auto& mem : rxMems) {
283 0 : CHK_PTR_NULL(mem.dst);
284 0 : void* srcMemPtr = nullptr;
285 0 : CHK_RET(leftLink_->GetRemoteMem(mem.srcMemType, &srcMemPtr));
286 :
287 0 : DeviceMem srcDevMem(static_cast<s8*>(srcMemPtr) + mem.srcOffset, mem.len);
288 0 : DeviceMem dstDevMem(static_cast<s8*>(mem.dst), mem.len);
289 :
290 0 : CHK_RET(HcclD2DMemcpyAsync(
291 : dispatcher_, dstDevMem, srcDevMem, stream_, leftLink_->GetRemoteRank(), leftLink_->GetLinkType()));
292 0 : }
293 : }
294 :
295 0 : return HCCL_SUCCESS;
296 0 : }
297 :
298 0 : HcclResult AllGatherRingDirect::RunAllGather(const u32 rank, const u32 rankSize)
299 : {
300 0 : HCCL_INFO("AllGatherRingDirect starts, the input param rank[%u]", rank);
301 0 : CHK_RET(RunInitStep(rank, rankSize));
302 :
303 0 : finalSrc_.clear();
304 0 : finalDst_.clear();
305 :
306 0 : u32 txSliceIdx = rank;
307 0 : u32 sliceSize = slices_.size() / rankSize;
308 0 : u32 rxSliceIdx = (rank + rankSize - 1) % rankSize;
309 :
310 0 : for (u32 step = 0; step < rankSize - 1; step++) {
311 : // 本端cclout -> 本端userout
312 0 : CHK_RET(RunAllGatherPartOne(sliceSize, step, txSliceIdx));
313 : // 对端cclout -> 本端cclout, 如果最后一步则：对端cclout -> 本端userout （DMA消减）
314 0 : CHK_RET(RunAllGatherPartTwo(sliceSize, step, txSliceIdx, rxSliceIdx, rankSize));
315 : // 更新索引
316 0 : txSliceIdx = (txSliceIdx + rankSize - 1) % rankSize;
317 0 : rxSliceIdx = (rxSliceIdx + rankSize - 1) % rankSize;
318 : }
319 :
320 0 : if (!isSdma_) {
321 0 : for (u32 vecIdx = 0; vecIdx < finalSrc_.size(); vecIdx++) {
322 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, finalDst_[vecIdx], finalSrc_[vecIdx], stream_));
323 : }
324 : }
325 :
326 0 : HCCL_INFO("AllGatherRingDirect finished to RunAllGather");
327 :
328 0 : return HCCL_SUCCESS;
329 : }
330 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_GATHER_RING_DIRECT, AllGatherRingDirect);
331 : } // namespace hccl
|
