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