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 "alg_template_register.h"
12 : #include "all_reduce_graph_pipeline.h"
13 :
14 : constexpr u32 STEP_OFFSET_TWO = 2;
15 :
16 : namespace hccl {
17 0 : AllReduceGraphPipeline::AllReduceGraphPipeline(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
18 :
19 0 : AllReduceGraphPipeline::~AllReduceGraphPipeline() {}
20 :
21 0 : HcclResult AllReduceGraphPipeline::Prepare(u64 reduceAttrBitMap, [[maybe_unused]] HcomCollOpInfo* opInfo)
22 : {
23 0 : reduceAttr_ = reduceAttrBitMap;
24 0 : return HCCL_SUCCESS;
25 : }
26 :
27 0 : HcclResult AllReduceGraphPipeline::MainWaitSub()
28 : {
29 0 : u32 subStreamNum = intraRankSize_ - 1;
30 0 : for (u32 signalIndex = 0; signalIndex < subStreamNum; signalIndex++) {
31 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, streamNotifyMain_[signalIndex], INVALID_VALUE_STAGE));
32 : }
33 0 : return HCCL_SUCCESS;
34 : }
35 :
36 0 : HcclResult AllReduceGraphPipeline::SubRecordMain()
37 : {
38 0 : u32 subStreamNum = intraRankSize_ - 1;
39 0 : for (u32 streamIndex = 0; streamIndex < subStreamNum; streamIndex++) {
40 0 : CHK_RET(LocalNotify::Post(subStreams_[streamIndex], dispatcher_, streamNotifyMain_[streamIndex], -1));
41 : }
42 0 : return HCCL_SUCCESS;
43 : }
44 :
45 0 : HcclResult AllReduceGraphPipeline::MainRecordSub()
46 : {
47 0 : u32 subStreamNum = intraRankSize_ - 1;
48 0 : for (u32 signalIndex = 0; signalIndex < subStreamNum; signalIndex++) {
49 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, streamNotifySub_[signalIndex], -1));
50 : }
51 0 : return HCCL_SUCCESS;
52 : }
53 :
54 0 : HcclResult AllReduceGraphPipeline::SubWaitMain()
55 : {
56 0 : u32 subStreamNum = intraRankSize_ - 1;
57 0 : for (u32 streamIndex = 0; streamIndex < subStreamNum; streamIndex++) {
58 0 : CHK_RET(LocalNotify::Wait(
59 : subStreams_[streamIndex], dispatcher_, streamNotifySub_[streamIndex], INVALID_VALUE_STAGE));
60 : }
61 0 : return HCCL_SUCCESS;
62 : }
63 :
64 0 : HcclResult AllReduceGraphPipeline::RunReduceScatterIntraServer(u32 step)
65 : {
66 0 : for (u32 i = 1; i < intraRankSize_; i++) {
67 0 : u32 remIntraRankId = (intraRankId_ + i) % intraRankSize_;
68 0 : CHK_RET(intraLinks_[remIntraRankId]->TxAck(subStreams_[i - 1]));
69 0 : CHK_RET(intraLinks_[remIntraRankId]->RxAck(subStreams_[i - 1]));
70 0 : void* remoteMemPtr = nullptr;
71 0 : CHK_RET(intraLinks_[remIntraRankId]->GetRemoteMem(UserMemType::INPUT_MEM, &remoteMemPtr));
72 0 : u32 sliceId = ((interRankId_ + step + 1) % interRankSize_) * intraRankSize_ + remIntraRankId;
73 0 : u64 srcOffset = sliceId * memSliceSize_;
74 0 : u64 dataSize = memSliceSize_;
75 0 : u64 dataCount = sliceCount_;
76 0 : if (sliceId == (interRankSize_ * intraRankSize_ - 1)) {
77 0 : dataSize = lastSliceSize_;
78 0 : dataCount = lastSliceCount_;
79 : }
80 0 : DeviceMem src = DeviceMem::create(static_cast<u8*>(usrInMem_) + srcOffset, dataSize);
81 0 : DeviceMem dst = DeviceMem::create(static_cast<u8*>(remoteMemPtr) + srcOffset, dataSize);
82 :
83 0 : CHK_RET(HcclReduceAsync(
84 : dispatcher_, src.ptr(), dataCount, dataType_, reductionOp_, subStreams_[i - 1], dst.ptr(),
85 : intraLinks_[remIntraRankId]->GetRemoteRank(), intraLinks_[remIntraRankId]->GetLinkType(),
86 : INLINE_REDUCE_BIT));
87 :
88 0 : CHK_RET(intraLinks_[remIntraRankId]->TxDataSignal(subStreams_[i - 1]));
89 0 : CHK_RET(intraLinks_[remIntraRankId]->RxDataSignal(subStreams_[i - 1]));
90 0 : }
91 0 : return HCCL_SUCCESS;
92 : }
93 :
94 0 : HcclResult AllReduceGraphPipeline::RunAllGatherIntraServer(u32 step)
95 : {
96 0 : for (u32 i = 1; i < intraRankSize_; i++) {
97 0 : u32 remIntraRankId = (intraRankId_ + i) % intraRankSize_;
98 0 : CHK_RET(intraLinks_[remIntraRankId]->TxAck(subStreams_[i - 1]));
99 0 : CHK_RET(intraLinks_[remIntraRankId]->RxAck(subStreams_[i - 1]));
100 0 : void* remoteMemPtr = nullptr;
101 0 : CHK_RET(intraLinks_[remIntraRankId]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remoteMemPtr));
102 0 : u32 sliceId = ((interRankId_ + step) % interRankSize_) * intraRankSize_ + remIntraRankId;
103 0 : u64 dstOffset = sliceId * memSliceSize_;
104 0 : u64 dataSize = memSliceSize_;
105 0 : if (sliceId == (interRankSize_ * intraRankSize_ - 1)) {
106 0 : dataSize = lastSliceSize_;
107 : }
108 0 : DeviceMem src = DeviceMem::create(static_cast<u8*>(remoteMemPtr) + dstOffset, dataSize);
109 0 : DeviceMem dst = DeviceMem::create(static_cast<u8*>(usrOutMem_) + dstOffset, dataSize);
110 :
111 0 : CHK_RET(HcclD2DMemcpyAsync(
112 : dispatcher_, dst, src, subStreams_[i - 1], intraLinks_[remIntraRankId]->GetRemoteRank(),
113 : intraLinks_[remIntraRankId]->GetLinkType()));
114 :
115 0 : CHK_RET(intraLinks_[remIntraRankId]->TxDataSignal(subStreams_[i - 1]));
116 0 : CHK_RET(intraLinks_[remIntraRankId]->RxDataSignal(subStreams_[i - 1]));
117 0 : }
118 0 : return HCCL_SUCCESS;
119 : }
120 :
121 : HcclResult
122 0 : AllReduceGraphPipeline::RunReduceScatterInterServer(u32 step, const LINK& prevInterLink, const LINK& nextInterLink)
123 : {
124 0 : u32 txSliceId = ((interRankId_ + 1 + step) % interRankSize_) * intraRankSize_ + intraRankId_;
125 0 : u64 txSliceOffset = memSliceSize_ * txSliceId;
126 0 : u64 txDataSize = memSliceSize_;
127 0 : if (txSliceId == (interRankSize_ * intraRankSize_ - 1)) {
128 0 : txDataSize = lastSliceSize_;
129 : }
130 0 : DeviceMem srcMem = DeviceMem::create(static_cast<u8*>(usrInMem_) + txSliceOffset, txDataSize);
131 0 : CHK_RET(senderInfo_->run(nextInterLink, txSliceOffset, srcMem, stream_, UserMemType::INPUT_MEM));
132 0 : HCCL_DEBUG(
133 : "[AllReduceGraphPipeline][RunReduceScatterInterServer] local rank[%u], localOffset[%llu],"
134 : "tx with slice[%llu]",
135 : rankId_, txSliceOffset, txDataSize);
136 :
137 0 : u32 rxSliceId = ((interRankId_ + 2 + step) % interRankSize_) * intraRankSize_ + intraRankId_;
138 0 : u64 rxSliceOffset = memSliceSize_ * rxSliceId;
139 0 : u64 rxDataSize = memSliceSize_;
140 0 : if (rxSliceId == (interRankSize_ * intraRankSize_ - 1)) {
141 0 : rxDataSize = lastSliceSize_;
142 : }
143 0 : DeviceMem rxLocalMem = DeviceMem::create(static_cast<u8*>(usrInMem_) + rxSliceOffset, rxDataSize);
144 0 : CHK_RET(reducerInfo_->run(dispatcher_, prevInterLink, rxSliceOffset, rxLocalMem, rxLocalMem, rxLocalMem, stream_));
145 0 : return HCCL_SUCCESS;
146 0 : }
147 :
148 : HcclResult
149 0 : AllReduceGraphPipeline::RunAllGatherInterServer(u32 step, const LINK& prevInterLink, const LINK& nextInterLink)
150 : {
151 0 : u32 txSliceId = ((interRankId_ + step) % interRankSize_) * intraRankSize_ + intraRankId_;
152 0 : u64 txSliceOffset = memSliceSize_ * txSliceId;
153 0 : u64 txDataSize = memSliceSize_;
154 0 : if (txSliceId == (interRankSize_ * intraRankSize_ - 1)) {
155 0 : txDataSize = lastSliceSize_;
156 : }
157 0 : CHK_RET(nextInterLink->TxAsync(
158 : UserMemType::OUTPUT_MEM, txSliceOffset, static_cast<u8*>(usrOutMem_) + txSliceOffset, txDataSize, stream_));
159 0 : HCCL_DEBUG(
160 : "[AllReduceGraphPipeline][RunAllGatherInterServer] local rank[%u], localOffset[%llu],"
161 : "tx with slice[%llu]",
162 : rankId_, txSliceOffset, txDataSize);
163 :
164 0 : u32 rxSliceId = ((interRankId_ + step + 1) % interRankSize_) * intraRankSize_ + intraRankId_;
165 0 : u64 rxSliceOffset = memSliceSize_ * rxSliceId;
166 0 : u64 rxDataSize = memSliceSize_;
167 0 : if (rxSliceId == (interRankSize_ * intraRankSize_ - 1)) {
168 0 : rxDataSize = lastSliceSize_;
169 : }
170 0 : CHK_RET(prevInterLink->RxAsync(
171 : UserMemType::OUTPUT_MEM, rxSliceOffset, static_cast<u8*>(usrOutMem_) + rxSliceOffset, rxDataSize, stream_));
172 0 : return HCCL_SUCCESS;
173 : }
174 :
175 0 : HcclResult AllReduceGraphPipeline::RunAsync()
176 : {
177 : // inter ring algo
178 0 : u32 prevInterRankId = (interRankId_ + 1) % interRankSize_;
179 0 : u32 nextInterRankId = (interRankId_ - 1 + interRankSize_) % interRankSize_;
180 0 : LINK prevInterLink = interLinks_[prevInterRankId];
181 0 : LINK nextInterLink = interLinks_[nextInterRankId];
182 : // 在user in执行reducescatter pipeline
183 0 : for (u32 step = 0; step < interRankSize_; step++) {
184 0 : if (step == 0) {
185 0 : CHK_RET(MainRecordSub());
186 0 : CHK_RET(SubWaitMain());
187 : }
188 : // server内做SDMA的reduce
189 0 : CHK_RET(RunReduceScatterIntraServer(step));
190 0 : CHK_RET(SubRecordMain());
191 0 : CHK_RET(MainWaitSub());
192 0 : if (step < interRankSize_ - 1) {
193 0 : CHK_RET(MainRecordSub());
194 0 : CHK_RET(SubWaitMain());
195 0 : CHK_RET(prevInterLink->TxAck(stream_));
196 0 : CHK_RET(nextInterLink->RxAck(stream_));
197 : // server间做RDMA的reduce，可与下一个step的SDMA并发执行
198 0 : CHK_RET(RunReduceScatterInterServer(step, prevInterLink, nextInterLink));
199 : // 确保step[n+2]的SDMA之前step[n]的RDMA已经完成，防止内存踩踏
200 0 : CHK_RET(prevInterLink->PostFinAck(stream_));
201 0 : CHK_RET(nextInterLink->WaitFinAck(stream_));
202 : }
203 : }
204 :
205 : // reducescatter通信结束，将数据从user in拷贝到user out
206 0 : u64 localOffsetByte = memSliceSize_ * rankId_;
207 0 : u64 dataSize = memSliceSize_;
208 0 : if (rankId_ == (interRankSize_ * intraRankSize_ - 1)) {
209 0 : dataSize = lastSliceSize_;
210 : }
211 0 : DeviceMem locSrc = DeviceMem::create(static_cast<u8*>(usrInMem_) + localOffsetByte, dataSize);
212 0 : DeviceMem locDst = DeviceMem::create(static_cast<u8*>(usrOutMem_) + localOffsetByte, dataSize);
213 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, locDst, locSrc, stream_));
214 :
215 : // 在user out执行allgather pipeline
216 0 : for (u32 step = 0; step < interRankSize_; step++) {
217 0 : CHK_RET(MainRecordSub());
218 0 : CHK_RET(SubWaitMain());
219 0 : if (step < interRankSize_ - 1) {
220 0 : CHK_RET(prevInterLink->TxAck(stream_));
221 0 : CHK_RET(nextInterLink->RxAck(stream_));
222 0 : CHK_RET(RunAllGatherInterServer(step, prevInterLink, nextInterLink));
223 0 : CHK_RET(prevInterLink->PostFinAck(stream_));
224 0 : CHK_RET(nextInterLink->WaitFinAck(stream_));
225 : // inter的最后一步需要barrier确保数据发完
226 0 : if (step == interRankSize_ - STEP_OFFSET_TWO) {
227 0 : CHK_RET(ExecuteBarrier(prevInterLink, nextInterLink));
228 : }
229 : }
230 0 : HCCL_DEBUG("[AllReducePipeline][RunAsync]step %u runAllGatherInterServer success", step);
231 0 : CHK_RET(RunAllGatherIntraServer(step));
232 0 : CHK_RET(SubRecordMain());
233 0 : CHK_RET(MainWaitSub());
234 0 : HCCL_INFO("[AllReducePipeline][RunAsync]AllReducePipeline finished groupRankId[%u] ", rankId_);
235 : }
236 0 : return HCCL_SUCCESS;
237 0 : }
238 :
239 0 : HcclResult AllReduceGraphPipeline::Prepare(
240 : const HcomCollOpInfo* opInfo, [[maybe_unused]] DeviceMem& cclBufferA, [[maybe_unused]] DeviceMem& cclBufferB,
241 : const u64 count, const SubCommInfo& level1CommInfo, const SubCommInfo& level0CommInfo, Stream& mainStream,
242 : std::vector<Stream>& subStream, std::vector<std::shared_ptr<LocalNotify>>& notifyMain,
243 : std::vector<std::shared_ptr<LocalNotify>>& notifySub)
244 : {
245 0 : unitSize_ = SIZE_TABLE[opInfo->dataType];
246 0 : sliceCount_ = count / (level0CommInfo.localRankSize * level1CommInfo.localRankSize);
247 0 : memSliceSize_ = sliceCount_ * unitSize_;
248 0 : lastSliceCount_ = count - sliceCount_ * (level0CommInfo.localRankSize * level1CommInfo.localRankSize - 1);
249 0 : lastSliceSize_ = lastSliceCount_ * unitSize_;
250 0 : HCCL_DEBUG("[%s] PrepareSliceDataWithAlignSize for data_slice_prepare", __func__);
251 :
252 0 : usrInMem_ = opInfo->inputAddr;
253 0 : usrOutMem_ = opInfo->outputAddr;
254 0 : reductionOp_ = opInfo->reduceOp;
255 0 : dataType_ = opInfo->dataType;
256 :
257 : // needed resource
258 : // stream: 1 * mainStream + (n -1) * subStream
259 : // interNotify, streamNotify
260 :
261 : // stream
262 : // mainStream负责locMemCpy、inter执行以及subStream同步控制
263 0 : stream_ = mainStream;
264 : // subStream负责:
265 : // streamId[0:intraRankSize-1]: intraRankSize-1个intra执行
266 0 : subStreams_ = subStream;
267 :
268 0 : intraRankSize_ = level0CommInfo.localRankSize;
269 0 : interRankSize_ = level1CommInfo.localRankSize;
270 0 : intraRankId_ = level0CommInfo.localRank;
271 0 : interRankId_ = level1CommInfo.localRank;
272 0 : rankId_ = intraRankId_ + interRankId_ * intraRankSize_;
273 :
274 : // streamNotify, size: n
275 0 : streamNotifyMain_ = notifyMain;
276 0 : if (streamNotifyMain_.size() < intraRankSize_ - 1) {
277 0 : HCCL_ERROR(
278 : "[AllReduceGraphPipeline][Prepare]rank[%u] streamNotifyMain_ size [%u] error, is smaller than,"
279 : "intraRankSize_[%u]",
280 : rankId_, streamNotifyMain_.size(), intraRankSize_);
281 0 : return HCCL_E_INTERNAL;
282 : }
283 0 : streamNotifySub_ = notifySub;
284 0 : if (streamNotifySub_.size() < intraRankSize_ - 1) {
285 0 : HCCL_ERROR(
286 : "[AllReduceGraphPipeline][Prepare]rank[%u] streamNotifySub_ size [%u] error, is smaller than,"
287 : "intraRankSize_[%u]",
288 : rankId_, streamNotifySub_.size(), intraRankSize_);
289 0 : return HCCL_E_INTERNAL;
290 : }
291 :
292 0 : intraLinks_ = level0CommInfo.links;
293 0 : interLinks_ = level1CommInfo.links;
294 :
295 0 : HCCL_INFO(
296 : "[AllReduceGraphPipeline][Prepare]streamNum[%u], streamNotifyMainNum[%u], streamNotifySubNum[%u]",
297 : subStreams_.size(), streamNotifyMain_.size(), streamNotifySub_.size());
298 0 : HCCL_INFO(
299 : "[AllReduceGraphPipeline][Prepare]interLinksNum[%u], intraLinksNum[%u]", interLinks_.size(),
300 : intraLinks_.size());
301 0 : senderInfo_.reset(new (std::nothrow) Sender(dataType_, reductionOp_, reduceAttr_));
302 0 : CHK_SMART_PTR_NULL(senderInfo_);
303 0 : reducerInfo_.reset(new (std::nothrow) Reducer(dataType_, reductionOp_, reduceAttr_));
304 0 : CHK_SMART_PTR_NULL(reducerInfo_);
305 0 : return HCCL_SUCCESS;
306 : }
307 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALLREDUCE_GRAPH_PIPELINE, AllReduceGraphPipeline);
308 : } // namespace hccl
|
