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 "reduce_scatter_pipeline.h"
12 : #include "alg_template_register.h"
13 :
14 : constexpr u32 STEP_OFFSET_TWO = 2;
15 :
16 :
17 : namespace hccl {
18 1 : ReduceScatterPipeline::ReduceScatterPipeline(const HcclDispatcher dispatcher)
19 1 : : AlgTemplateBase(dispatcher) {}
20 :
21 2 : ReduceScatterPipeline::~ReduceScatterPipeline() {}
22 :
23 0 : HcclResult ReduceScatterPipeline::MainWaitSub(u32 begin)
24 : {
25 0 : u32 subStreamNum = intraRankSize_;
26 0 : for (u32 signalIndex = begin; signalIndex < subStreamNum; signalIndex++) {
27 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, streamNotifyMain_[signalIndex], INVALID_VALUE_STAGE));
28 : }
29 0 : return HCCL_SUCCESS;
30 : }
31 :
32 0 : HcclResult ReduceScatterPipeline::SubRecordMain(u32 begin)
33 : {
34 0 : u32 subStreamNum = intraRankSize_;
35 0 : for (u32 streamIndex = begin; streamIndex < subStreamNum; streamIndex++) {
36 0 : CHK_RET(LocalNotify::Post(subStream_[streamIndex], dispatcher_, streamNotifyMain_[streamIndex], -1));
37 : }
38 0 : return HCCL_SUCCESS;
39 : }
40 :
41 0 : HcclResult ReduceScatterPipeline::MainRecordSub(u32 begin)
42 : {
43 0 : u32 subStreamNum = intraRankSize_;
44 0 : for (u32 signalIndex = begin; signalIndex < subStreamNum; signalIndex++) {
45 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, streamNotifySub_[signalIndex], -1));
46 : }
47 0 : return HCCL_SUCCESS;
48 : }
49 :
50 0 : HcclResult ReduceScatterPipeline::SubWaitMain(u32 begin)
51 : {
52 0 : u32 subStreamNum = intraRankSize_;
53 0 : for (u32 streamIndex = begin; streamIndex < subStreamNum; streamIndex++) {
54 0 : CHK_RET(LocalNotify::Wait(subStream_[streamIndex], dispatcher_, streamNotifySub_[streamIndex],
55 : INVALID_VALUE_STAGE));
56 : }
57 0 : return HCCL_SUCCESS;
58 : }
59 :
60 0 : HcclResult ReduceScatterPipeline::RunIntraServer(u32 step, u64 remoteOffset)
61 : {
62 0 : for (u32 i = 1; i < intraRankSize_; i++) {
63 0 : u32 remIntraRankId = (intraRankId_ + i) % intraRankSize_;
64 0 : CHK_RET(intraLinks_[remIntraRankId]->TxAck(subStream_[i]));
65 0 : CHK_RET(intraLinks_[remIntraRankId]->RxAck(subStream_[i]));
66 0 : void* remoteMemPtr = nullptr;
67 0 : CHK_RET(intraLinks_[remIntraRankId]->GetRemoteMem(UserMemType::INPUT_MEM, &remoteMemPtr));
68 0 : u64 srcOffset = (((interRankId_ + step + 1) % interRankSize_) * intraRankSize_ + remIntraRankId) \
69 0 : * memSliceSize_;
70 0 : u64 offset = (srcOffset + offset_) % HCCL_MIN_SLICE_ALIGN_910B;
71 0 : DeviceMem src = DeviceMem::create(static_cast<u8 *>(usrInMem_) + srcOffset, curSize_);
72 0 : DeviceMem dst = DeviceMem::create(static_cast<u8 *>(remoteMemPtr) + remoteOffset + offset, curSize_);
73 :
74 0 : CHK_RET(HcclReduceAsync(dispatcher_, src.ptr(), count_, dataType_, reductionOp_,
75 : subStream_[i], dst.ptr(), intraLinks_[remIntraRankId]->GetRemoteRank(),
76 : intraLinks_[remIntraRankId]->GetLinkType(), INLINE_REDUCE_BIT));
77 :
78 0 : CHK_RET(intraLinks_[remIntraRankId]->TxDataSignal(subStream_[i]));
79 0 : CHK_RET(intraLinks_[remIntraRankId]->RxDataSignal(subStream_[i]));
80 0 : }
81 0 : return HCCL_SUCCESS;
82 : }
83 :
84 0 : HcclResult ReduceScatterPipeline::RunInterServer(u32 step,
85 : const LINK &prevInterLink,
86 : const LINK &nextInterLink)
87 : {
88 0 : u32 dmaMemSliceNum = dmaMem_.size();
89 0 : u32 rxDMAMemSliceId = (step + 1) % dmaMemSliceNum;
90 0 : u32 txDMAMemSliceId = step % dmaMemSliceNum;
91 0 : u64 sliceMemOffset = memSliceSize_ * (((interRankId_ + 1 + step) % interRankSize_) * intraRankSize_ + intraRankId_);
92 0 : u64 offset = (sliceMemOffset + offset_) % HCCL_MIN_SLICE_ALIGN_910B;
93 0 : u64 rxInterOffset = rxDMAMemSliceId * blockSize_ + offset;
94 0 : void* txLocalAddr = static_cast<u8 *>(dmaMem_[txDMAMemSliceId].ptr()) + offset;
95 0 : DeviceMem srcMem = DeviceMem::create(txLocalAddr, curSize_);
96 0 : CHK_RET(senderInfo_->run(nextInterLink, rxInterOffset, srcMem, subStream_[0]));
97 0 : HCCL_DEBUG("[ReduceScatterPipeline][RunInterServer] local rank[%u] localOffset[%llu]tx with slice[%llu]",
98 : rankId_, rxInterOffset, curSize_);
99 :
100 0 : u64 rxSliceOffset = memSliceSize_ * (((interRankId_ + 2 + step) % interRankSize_) * intraRankSize_ + intraRankId_);
101 0 : u64 rxOffset = (rxSliceOffset + offset_) % HCCL_MIN_SLICE_ALIGN_910B;
102 0 : u64 rxMemOffset = txDMAMemSliceId * blockSize_ + rxOffset;
103 0 : void* rxLocalAddr = static_cast<u8 *>(dmaMem_[rxDMAMemSliceId].ptr()) + rxOffset;
104 0 : DeviceMem rxLocalMem = DeviceMem::create(rxLocalAddr, curSize_);
105 0 : CHK_RET(reducerInfo_->run(dispatcher_, prevInterLink, rxMemOffset, rxLocalMem, rxLocalMem, rxLocalMem,
106 : subStream_[0]));
107 0 : return HCCL_SUCCESS;
108 0 : }
109 :
110 0 : HcclResult ReduceScatterPipeline::CopyToScratchBuffer(u32 step)
111 : {
112 0 : u32 dmaMemSliceNum = dmaMem_.size();
113 0 : u32 dmaMemSliceId = step % dmaMemSliceNum;
114 0 : u64 sliceMemOffset = memSliceSize_ * (((interRankId_ + 1 + step) % interRankSize_) * intraRankSize_ + intraRankId_);
115 0 : u64 offset = (sliceMemOffset + offset_) % HCCL_MIN_SLICE_ALIGN_910B;
116 : // 把一块切片从userIn 做拷贝到CCLBuffer
117 0 : void* srcAddr = static_cast<u8 *>(usrInMem_) + sliceMemOffset;
118 0 : DeviceMem locSrc = DeviceMem::create(srcAddr, curSize_);
119 0 : DeviceMem locDst = DeviceMem::create(static_cast<u8 *>(dmaMem_[dmaMemSliceId].ptr()) + offset, curSize_);
120 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, locDst, locSrc, stream_));
121 0 : return HCCL_SUCCESS;
122 0 : }
123 :
124 0 : HcclResult ReduceScatterPipeline::RunAsync()
125 : {
126 : // inter ring algo
127 0 : u32 prevInterRankId = (interRankId_ + 1) % interRankSize_;
128 0 : u32 nextInterRankId = (interRankId_ - 1 + interRankSize_) % interRankSize_;
129 0 : LINK prevInterLink = interLinks_[prevInterRankId];
130 0 : LINK nextInterLink = interLinks_[nextInterRankId];
131 : // 当前使用3块DMAMem buffer
132 0 : u32 dmaMemSliceNum = dmaMem_.size();
133 0 : HCCL_DEBUG("RunAsync begin.");
134 :
135 0 : for (u32 step = 0; step < interRankSize_; step ++) {
136 0 : u32 begin = 0;
137 0 : if (step == 0) {
138 : // 把第一块切片从userIn 做拷贝到CCLBuffer
139 0 : begin = 1;
140 0 : CHK_RET(CopyToScratchBuffer(step));
141 0 : CHK_RET(MainRecordSub(begin));
142 0 : CHK_RET(SubWaitMain(begin));
143 : }
144 : // server内做SDMA的reduce
145 0 : u64 remoteOffset = (step % dmaMemSliceNum) * blockSize_;
146 0 : HCCL_DEBUG("[RunAsync]remoteOffset is [%llu]", remoteOffset);
147 0 : CHK_RET(RunIntraServer(step, remoteOffset));
148 0 : CHK_RET(SubRecordMain(begin));
149 0 : CHK_RET(MainWaitSub(begin));
150 0 : if (step < interRankSize_ - 1) {
151 : // 把下一块切片从userIn 做拷贝到CCLBuffer
152 0 : CHK_RET(CopyToScratchBuffer(step + 1));
153 : // 全部流同步，确保SDMA执行完成
154 0 : CHK_RET(MainRecordSub(0));
155 0 : CHK_RET(SubWaitMain(0));
156 0 : CHK_RET(prevInterLink->TxAck(subStream_[0]));
157 0 : CHK_RET(nextInterLink->RxAck(subStream_[0]));
158 : // server间做RDMA的reduce，可与下一个step的SDMA并发执行
159 0 : CHK_RET(RunInterServer(step, prevInterLink, nextInterLink));
160 0 : CHK_RET(prevInterLink->PostFinAck(subStream_[0]));
161 0 : CHK_RET(nextInterLink->WaitFinAck(subStream_[0]));
162 : // inter的最后一步需要barrier确保数据发完
163 0 : if (step == interRankSize_ - STEP_OFFSET_TWO) {
164 0 : CHK_RET(ExecuteBarrier(prevInterLink, nextInterLink, subStream_[0]));
165 : }
166 : }
167 : }
168 : // 把对应的切片从CCLBuffer拷贝到userOut
169 0 : DeviceMem locDst = DeviceMem::create(static_cast<u8 *>(usrOutMem_), curSize_);
170 0 : u64 srcOffset = (memSliceSize_ * (interRankId_ * intraRankSize_ + intraRankId_) \
171 0 : + offset_) % HCCL_MIN_SLICE_ALIGN_910B;
172 0 : void* locSrcAddr = static_cast<u8 *>(dmaMem_[(interRankSize_ - 1) % dmaMemSliceNum].ptr()) + srcOffset;
173 0 : DeviceMem locSrc = DeviceMem::create(locSrcAddr, curSize_);
174 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, locDst, locSrc, stream_));
175 0 : HCCL_INFO("[ReduceScatterPipeline][RunAsync]ReduceScatterPipeline finished groupRankId[%u] ", rankId_);
176 0 : return HCCL_SUCCESS;
177 0 : }
178 :
179 : // 适配新CollExecutor接口
180 1 : HcclResult ReduceScatterPipeline::Prepare(HcomCollOpInfo *opInfo, DeviceMem &cclBuffer, const u64 count,
181 : const u64 bufferSize, const u64 offset, const SubCommInfo &level0CommInfo, const SubCommInfo &level1CommInfo,
182 : Stream &mainStream, std::vector<Stream> &subStream, std::vector<std::shared_ptr<LocalNotify>> ¬ifyMain,
183 : std::vector<std::shared_ptr<LocalNotify>> ¬ifySub, u64 reduceAttrBitMap)
184 : {
185 1 : reduceAttr_ = reduceAttrBitMap;
186 1 : opInfo_ = opInfo;
187 :
188 1 : unitSize_ = SIZE_TABLE[opInfo_->dataType];
189 1 : memSliceSize_ = opInfo_->count * unitSize_;
190 1 : usrInMem_ = opInfo_->inputAddr;
191 1 : usrOutMem_ = opInfo_->outputAddr;
192 1 : reductionOp_ = opInfo_->reduceOp;
193 1 : dataType_ = opInfo_->dataType;
194 1 : offset_ = offset;
195 :
196 : // needed resource
197 : // stream: 1 * mainStream + n * subStream
198 : // mem: usrInMem_, usrOutMem, DMAMem
199 : // interNotify, streamNotify
200 :
201 : // stream
202 : // mainStream负责locMemCpy以及subStream同步控制
203 1 : stream_ = mainStream;
204 : // subStream负责:
205 : // streamId[0]: inter执行
206 : // streamId[1:intraRankSize]: intraRankSize-1个intra执行
207 1 : subStream_ = subStream;
208 :
209 : // DMAMem + interNotify from Link
210 1 : intraRankSize_ = level0CommInfo.localRankSize;
211 1 : interRankSize_ = level1CommInfo.localRankSize;
212 1 : intraRankId_ = level0CommInfo.localRank;
213 1 : interRankId_ = level1CommInfo.localRank;
214 1 : rankId_ = intraRankId_ + interRankId_ * intraRankSize_;
215 :
216 : // streamNotify, size: n
217 1 : streamNotifyMain_ = notifyMain;
218 1 : if (streamNotifyMain_.size() < intraRankSize_) {
219 0 : HCCL_ERROR("[ReduceScatterPipeline][Prepare]rank[%u] streamNotifyMain_ size [%u] error, is smaller than," \
220 : "intraRankSize_[%u]", rankId_, streamNotifyMain_.size(), intraRankSize_);
221 0 : return HCCL_E_INTERNAL;
222 : }
223 1 : streamNotifySub_ = notifySub;
224 1 : if (streamNotifySub_.size() < intraRankSize_) {
225 0 : HCCL_ERROR("[ReduceScatterPipeline][Prepare]rank[%u] streamNotifySub_ size [%u] error, is smaller than," \
226 : "intraRankSize_[%u]", rankId_, streamNotifySub_.size(), intraRankSize_);
227 0 : return HCCL_E_INTERNAL;
228 : }
229 : // usrMem
230 :
231 1 : intraLinks_ = level0CommInfo.links;
232 1 : interLinks_ = level1CommInfo.links;
233 :
234 : // 3级流水,使用3块DMAMem
235 1 : cclBuffer_ = cclBuffer;
236 1 : count_ = count;
237 1 : curSize_ = count_ * unitSize_;
238 1 : bufferSize_ = bufferSize;
239 1 : blockSize_ = (bufferSize_ / (HCCL_MIN_SLICE_ALIGN_910B * PIPELINE_DEPTH)) * HCCL_MIN_SLICE_ALIGN_910B;
240 :
241 4 : for (u32 i = 0; i < pipDepth_; i ++) {
242 3 : DeviceMem mem = DeviceMem::create(static_cast<u8 *>(cclBuffer_.ptr()) + blockSize_ * i, blockSize_);
243 3 : dmaMem_.push_back(mem);
244 3 : }
245 :
246 1 : HCCL_INFO("[ReduceScatterPipeline][Prepare]streamNum[%u], streamNotifyMainNum[%u], streamNotifySubNum[%u]",
247 : subStream_.size(), streamNotifyMain_.size(), streamNotifySub_.size());
248 1 : HCCL_INFO("[ReduceScatterPipeline][Prepare]interLinksNum[%u], intraLinksNum[%u]",
249 : interLinks_.size(), intraLinks_.size());
250 1 : senderInfo_.reset(new (std::nothrow) Sender(dataType_, reductionOp_, reduceAttr_));
251 1 : CHK_SMART_PTR_NULL(senderInfo_);
252 1 : reducerInfo_.reset(new (std::nothrow) Reducer(dataType_, reductionOp_, reduceAttr_));
253 1 : CHK_SMART_PTR_NULL(reducerInfo_);
254 1 : return HCCL_SUCCESS;
255 : }
256 :
257 0 : HcclResult ReduceScatterPipeline::GetNslbAdjInfo(const u32 rank, const u32 rankSize,
258 : const std::vector<LINK> &links, AdjInfo& nslbAdjInfo)
259 : {
260 0 : u32 ringNextRank = (rank + 1) % rankSize;
261 0 : LINK nslbNext = links[ringNextRank];
262 0 : HCCL_DEBUG("[ReduceScatterPipeline]GetNslbAdjInfo starts");
263 :
264 : // Pipeline 步长合并 等同于 ring
265 0 : NslbDpAdjInfo adjInfoStep = {0};
266 0 : nslbAdjInfo.dstRankNum = 1;
267 0 : adjInfoStep.dstLocalRankId = nslbNext->GetRemoteRank();
268 0 : adjInfoStep.phaseId = 1;
269 0 : adjInfoStep.rev = 0;
270 0 : nslbAdjInfo.nsAdjInfo.push_back(adjInfoStep);
271 :
272 0 : return HCCL_SUCCESS;
273 0 : }
274 :
275 :
276 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_PIPELINE, ReduceScatterPipeline);
277 : } // namespace hccl
|
