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