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_pipeline.h"
12 : #include "alg_template_register.h"
13 :
14 : constexpr u32 STEP_OFFSET_TWO = 2;
15 :
16 : namespace hccl {
17 0 : AllGatherPipeline::AllGatherPipeline(const HcclDispatcher dispatcher): AlgTemplateBase(dispatcher) {}
18 :
19 0 : AllGatherPipeline::~AllGatherPipeline() {}
20 :
21 0 : HcclResult AllGatherPipeline::Prepare(HcomCollOpInfo *opInfo, u32 userRank, u64 &count, DeviceMem &cclBufferPartOne,
22 : DeviceMem &cclBufferPartTwo, SubCommInfo &level0CommInfo, SubCommInfo &level1CommInfo,
23 : Stream &mainStream, std::vector<Stream> &subStream,
24 : std::vector<std::shared_ptr<LocalNotify>> ¬ifyMain, std::vector<std::shared_ptr<LocalNotify>> ¬ifySub)
25 : {
26 0 : opInfo_ = opInfo;
27 0 : memSliceCount_ = count;
28 0 : userRank_ = userRank;
29 :
30 0 : u32 unitSize = SIZE_TABLE[opInfo_->dataType];
31 0 : u64 memSliceSize = memSliceCount_ * unitSize;
32 :
33 0 : usrInMemAddr_ = opInfo_->inputAddr;
34 0 : usrOutMemAddr_ = opInfo_->outputAddr;
35 :
36 : // needed resource
37 : // stream: 1 * mainStream + n * subStream
38 : // mem: usrInMem, usrOutMem, DMAMem
39 : // interNotify, streamNotify
40 :
41 : // stream
42 : // mainStream负责locMemCpy以及subStream同步控制
43 0 : stream_ = mainStream;
44 : // subStream负责:
45 : // streamId[0]: inter执行
46 : // streamId[1:intraRankSize]: intraRankSize-1个intra执行
47 0 : subStream_ = subStream;
48 :
49 : // DMAMem + interNotify from Link
50 0 : intraRankSize_ = level0CommInfo.localRankSize;
51 0 : interRankSize_ = level1CommInfo.localRankSize;
52 0 : intraRankId_ = level0CommInfo.localRank;
53 0 : interRankId_ = level1CommInfo.localRank;
54 0 : intraLinks_ = level0CommInfo.links;
55 0 : interLinks_ = level1CommInfo.links;
56 :
57 : // streamNotify, size: n
58 0 : streamNotifyMain_ = notifyMain;
59 0 : if (streamNotifyMain_.size() < intraRankSize_) {
60 0 : HCCL_ERROR("[AllGatherPipeline][Prepare]rank[%u] streamNotifyMain_ size[%u] error, is smaller than," \
61 : "intraRankSize_[%u]", userRank_, streamNotifyMain_.size(), intraRankSize_);
62 0 : return HCCL_E_INTERNAL;
63 : }
64 0 : streamNotifySub_ = notifySub;
65 0 : if (streamNotifySub_.size() < intraRankSize_) {
66 0 : HCCL_ERROR("[AllGatherPipeline][Prepare]rank[%u] streamNotifySub_ size[%u] error, is smaller than, " \
67 : "intraRankSize_[%u]", userRank_, streamNotifySub_.size(), intraRankSize_);
68 0 : return HCCL_E_INTERNAL;
69 : }
70 :
71 : // 128byte align offset
72 0 : DeviceMem dmaMem0 = DeviceMem::create(cclBufferPartOne.ptr(), memSliceSize);
73 0 : DeviceMem dmaMem1 = DeviceMem::create(cclBufferPartTwo.ptr(), memSliceSize);
74 :
75 0 : dmaMem_.push_back(dmaMem0);
76 0 : dmaMem_.push_back(dmaMem1);
77 :
78 0 : HCCL_INFO("[AllGatherPipeline][Prepare]streamNum[%zu], streamNotifyMainNum[%zu], streamNotifySubNum[%zu].",
79 : subStream_.size(), streamNotifyMain_.size(), streamNotifySub_.size());
80 0 : HCCL_INFO("[AllGatherPipeline][Prepare]interLinksNum[%zu], intraLinksNum[%zu].",
81 : interLinks_.size(), intraLinks_.size());
82 0 : return HCCL_SUCCESS;
83 0 : }
84 :
85 0 : HcclResult AllGatherPipeline::MainWaitSub()
86 : {
87 0 : u32 subStreamNum = intraRankSize_;
88 0 : for (u32 signalIndex = 0; signalIndex < subStreamNum; signalIndex++) {
89 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, streamNotifyMain_[signalIndex], INVALID_VALUE_STAGE));
90 : }
91 0 : return HCCL_SUCCESS;
92 : }
93 :
94 0 : HcclResult AllGatherPipeline::SubRecordMain()
95 : {
96 0 : u32 subStreamNum = intraRankSize_;
97 0 : for (u32 streamIndex = 0; streamIndex < subStreamNum; streamIndex++) {
98 0 : CHK_RET(LocalNotify::Post(subStream_[streamIndex], dispatcher_, streamNotifyMain_[streamIndex],
99 : INVALID_VALUE_STAGE));
100 : }
101 0 : return HCCL_SUCCESS;
102 : }
103 :
104 0 : HcclResult AllGatherPipeline::MainRecordSub()
105 : {
106 0 : u32 subStreamNum = intraRankSize_;
107 0 : for (u32 signalIndex = 0; signalIndex < subStreamNum; signalIndex++) {
108 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, streamNotifySub_[signalIndex],
109 : INVALID_VALUE_STAGE));
110 : }
111 0 : return HCCL_SUCCESS;
112 : }
113 :
114 0 : HcclResult AllGatherPipeline::SubWaitMain()
115 : {
116 0 : u32 subStreamNum = intraRankSize_;
117 0 : for (u32 streamIndex = 0; streamIndex < subStreamNum; streamIndex++) {
118 0 : CHK_RET(LocalNotify::Wait(subStream_[streamIndex], dispatcher_, streamNotifySub_[streamIndex],
119 : INVALID_VALUE_STAGE));
120 : }
121 0 : return HCCL_SUCCESS;
122 : }
123 :
124 0 : HcclResult AllGatherPipeline::RunAsync()
125 : {
126 0 : HCCL_INFO("[AllGatherPipeline][RunAsync]AllGatherRingMesh starts groupRankId[%u]. ", userRank_);
127 : // inter ring algo
128 0 : u32 prevInterRankId = (interRankId_ - 1 + interRankSize_) % interRankSize_;
129 0 : u32 nextInterRankId = (interRankId_ + 1) % interRankSize_;
130 0 : LINK prevInterLink = interLinks_[prevInterRankId];
131 0 : LINK nextInterLink = interLinks_[nextInterRankId];
132 :
133 : // intra fullmesh algo
134 : // intra使用全部连接,不再映射
135 :
136 0 : u32 unitSize = SIZE_TABLE[opInfo_->dataType];
137 0 : u64 memSliceSize = memSliceCount_ * unitSize;
138 0 : u64 memSliceOffset = opInfo_->count * unitSize;
139 :
140 : // 仅使用两块DMAMem,为了方便切换使用
141 0 : u32 dmaMemSliceId = 0;
142 0 : u32 dmaMemSliceNum = dmaMem_.size();
143 :
144 : // step 0前置操作 : 所有卡本地数据从userIn-->DMAIn
145 0 : DeviceMem locSrc = DeviceMem::create(usrInMemAddr_, memSliceSize);
146 0 : u64 localOffset = (opInfo_->count * userRank_ * unitSize) % HCCL_MIN_SLICE_ALIGN_910B;
147 0 : DeviceMem locDMAInMem = DeviceMem::create(static_cast<u8 *>(dmaMem_[dmaMemSliceId].ptr()) + localOffset,
148 0 : memSliceSize);
149 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, locDMAInMem, locSrc, stream_));
150 :
151 0 : for (u32 step = 0; step < interRankSize_; step++) {
152 : // 主从流同步
153 0 : CHK_RET(MainRecordSub());
154 0 : CHK_RET(SubWaitMain());
155 :
156 : // 数据搬运及后同步
157 0 : u32 srcDMAMemSliceId = dmaMemSliceId;
158 0 : dmaMemSliceId = (dmaMemSliceId + 1) % dmaMemSliceNum;
159 0 : u32 dstDMAMemSliceId = dmaMemSliceId;
160 :
161 0 : u64 serverRankOffset = intraRankId_ + (interRankId_ + interRankSize_ - step) % interRankSize_ * intraRankSize_;
162 0 : u64 serverOffsetByte = (opInfo_->count * serverRankOffset * unitSize) % HCCL_MIN_SLICE_ALIGN_910B;
163 0 : u64 readRemoteOffset = intraRankId_ + (prevInterRankId + interRankSize_ - step) % interRankSize_ *
164 0 : intraRankSize_; // sever间前通信rank偏移
165 0 : u64 readRemoteOffsetByte = (opInfo_->count * readRemoteOffset * unitSize) % HCCL_MIN_SLICE_ALIGN_910B;
166 0 : if (step < interRankSize_ - 1) {
167 0 : CHK_RET(prevInterLink->TxAck(subStream_[0])); // AckRecord
168 0 : CHK_RET(nextInterLink->RxAck(subStream_[0])); // AckWait
169 : // RdmaSend + Record 或 PCIE::Record
170 0 : CHK_RET(nextInterLink->TxAsync((dstDMAMemSliceId == 1? UserMemType::OUTPUT_MEM: UserMemType::INPUT_MEM),
171 : serverOffsetByte, static_cast<u8 *>(dmaMem_[srcDMAMemSliceId].ptr()) + serverOffsetByte,
172 : memSliceSize, subStream_[0]));
173 0 : HCCL_DEBUG("[AllGatherPipeline][RunAsync] local rank[%u] localOffset[%llu]tx with remoteRank[%u]," \
174 : "remoteOffset[%llu] with slice[%llu].", userRank_, serverOffsetByte, nextInterRankId,
175 : serverOffsetByte, memSliceSize);
176 : // 对于RDM RxAsync,内存属性入参无效 RDMA::Wait
177 : // 对于PCIE,需设置内存属性 PCIE::Read + Record
178 0 : CHK_RET(prevInterLink->RxAsync((srcDMAMemSliceId == 0? UserMemType::INPUT_MEM: UserMemType::OUTPUT_MEM),
179 : readRemoteOffsetByte, static_cast<u8 *>(dmaMem_[dstDMAMemSliceId].ptr()) + readRemoteOffsetByte,
180 : memSliceSize, subStream_[0])); // wait
181 0 : HCCL_DEBUG("[AllGatherPipeline][RunAsync]read local rank[%u] localOffset[%llu]tx with remoteRank[%u]," \
182 : "remoteOffset[%llu] with slice[%llu].", userRank_, readRemoteOffsetByte, readRemoteOffset,
183 : readRemoteOffsetByte, memSliceSize);
184 0 : CHK_RET(prevInterLink->PostFinAck(subStream_[0]));
185 0 : CHK_RET(nextInterLink->WaitFinAck(subStream_[0]));
186 : // inter的最后一步需要barrier确保数据发完
187 0 : if (step == interRankSize_ - STEP_OFFSET_TWO) {
188 0 : CHK_RET(ExecuteBarrier(prevInterLink, nextInterLink, subStream_[0]));
189 : }
190 : }
191 :
192 0 : for (u32 i = 1; i < intraRankSize_; i++) {
193 0 : u32 remIntraRankId = (intraRankId_ + i) % intraRankSize_;
194 0 : CHK_RET(intraLinks_[remIntraRankId]->TxAck(subStream_[i])); // ackrecord
195 0 : CHK_RET(intraLinks_[remIntraRankId]->RxAck(subStream_[i]));
196 0 : void* remDMAMemPtr = nullptr;
197 0 : CHK_RET(intraLinks_[remIntraRankId]->GetRemoteMem(srcDMAMemSliceId == 1?
198 : UserMemType::OUTPUT_MEM: UserMemType::INPUT_MEM, &remDMAMemPtr));
199 0 : void* dstAddr = static_cast<u8 *>(usrOutMemAddr_) + ((interRankId_ - step + interRankSize_) %
200 0 : interRankSize_ * intraRankSize_ + remIntraRankId) * memSliceOffset;
201 :
202 0 : u64 remoteOffsetByte = (opInfo_->count * (remIntraRankId + serverRankOffset - intraRankId_) * unitSize) %
203 : HCCL_MIN_SLICE_ALIGN_910B;
204 0 : DeviceMem src = DeviceMem::create(static_cast<u8 *>(remDMAMemPtr) + remoteOffsetByte, memSliceSize);
205 0 : DeviceMem dst = DeviceMem::create(dstAddr, memSliceSize);
206 0 : HCCL_DEBUG("[AllGatherPipeline][RunAsync]remoteOffsetByte is %llu", remoteOffsetByte);
207 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, subStream_[i],
208 : intraLinks_[remIntraRankId]->GetRemoteRank(), intraLinks_[remIntraRankId]->GetLinkType()));
209 0 : CHK_RET(intraLinks_[remIntraRankId]->TxDataSignal(subStream_[i])); // data record
210 0 : CHK_RET(intraLinks_[remIntraRankId]->RxDataSignal(subStream_[i])); // data wait
211 0 : }
212 :
213 0 : CHK_RET(SubRecordMain());
214 0 : CHK_RET(MainWaitSub());
215 :
216 0 : void* dstAddr = static_cast<u8 *>(usrOutMemAddr_) + ((interRankId_ - step + interRankSize_) %
217 0 : interRankSize_ * intraRankSize_ + intraRankId_) * memSliceOffset;
218 0 : DeviceMem locDst = DeviceMem::create(dstAddr, memSliceSize);
219 0 : DeviceMem srcMem = DeviceMem::create(static_cast<u8 *>(dmaMem_[srcDMAMemSliceId].ptr()) + serverOffsetByte,
220 0 : memSliceSize);
221 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, locDst, srcMem, stream_));
222 0 : }
223 :
224 0 : HCCL_INFO("[AllGatherPipeline][RunAsync]AllGatherRingMesh finished groupRankId[%u] ", userRank_);
225 0 : return HCCL_SUCCESS;
226 0 : }
227 :
228 0 : HcclResult AllGatherPipeline::GetNslbAdjInfo(const u32 rank, const u32 rankSize,
229 : const std::vector<LINK> &links, AdjInfo& nslbAdjInfo)
230 : {
231 0 : HCCL_DEBUG("[AllGatherPipeline]GetNslbAdjInfo start");
232 0 : u32 ringNextRank = (rank + 1) % rankSize;
233 0 : LINK nslbNext = links[ringNextRank];
234 0 : CHK_SMART_PTR_NULL(nslbNext);
235 :
236 : // Pipeline 步长合并 等同于 ring
237 0 : NslbDpAdjInfo adjInfoStep = {0};
238 0 : nslbAdjInfo.dstRankNum = 1;
239 0 : adjInfoStep.dstLocalRankId = nslbNext->GetRemoteRank();
240 0 : adjInfoStep.phaseId = 1;
241 0 : adjInfoStep.rev = 0;
242 0 : nslbAdjInfo.nsAdjInfo.push_back(adjInfoStep);
243 :
244 0 : return HCCL_SUCCESS;
245 0 : }
246 :
247 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_GATHER_PIPELINE, AllGatherPipeline);
248 : } // namespace hccl
|