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_graph_pipeline.h"
12 : #include "alg_template_register.h"
13 :
14 : constexpr u32 STEP_OFFSET_TWO = 2;
15 :
16 : namespace hccl {
17 0 : ReduceScatterGraphPipeline::ReduceScatterGraphPipeline(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher)
18 0 : {}
19 :
20 0 : ReduceScatterGraphPipeline::~ReduceScatterGraphPipeline()
21 0 : {}
22 :
23 0 : HcclResult ReduceScatterGraphPipeline::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 ReduceScatterGraphPipeline::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 ReduceScatterGraphPipeline::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 ReduceScatterGraphPipeline::SubWaitMain(u32 begin)
51 : {
52 0 : u32 subStreamNum = intraRankSize_;
53 0 : for (u32 streamIndex = begin; streamIndex < subStreamNum; streamIndex++) {
54 0 : CHK_RET(LocalNotify::Wait(
55 : subStream_[streamIndex], dispatcher_, streamNotifySub_[streamIndex], INVALID_VALUE_STAGE));
56 : }
57 0 : return HCCL_SUCCESS;
58 : }
59 :
60 0 : HcclResult ReduceScatterGraphPipeline::RunIntraServer(u64 blockIdx)
61 : {
62 0 : u64 blockOff = blockIdx * intraRankSize_;
63 0 : u64 memOffset = (blockOff + intraRankId_) * memSliceSize_;
64 0 : for (u32 i = 1; i < intraRankSize_; i++) {
65 0 : u32 remIntraRankId = (intraRankId_ + i) % intraRankSize_;
66 0 : CHK_RET(intraLinks_[remIntraRankId]->TxAck(subStream_[i]));
67 0 : CHK_RET(intraLinks_[remIntraRankId]->RxAck(subStream_[i]));
68 0 : void *remoteMemPtr = nullptr;
69 0 : CHK_RET(intraLinks_[remIntraRankId]->GetRemoteMem(UserMemType::INPUT_MEM, &remoteMemPtr));
70 0 : DeviceMem dst = DeviceMem::create(static_cast<u8 *>(usrInMem_) + memOffset, memSliceSize_);
71 0 : DeviceMem src = DeviceMem::create(static_cast<u8 *>(remoteMemPtr) + memOffset, memSliceSize_);
72 :
73 0 : CHK_RET(HcclReduceAsync(dispatcher_,
74 : src.ptr(),
75 : count_,
76 : dataType_,
77 : reductionOp_,
78 : subStream_[i],
79 : dst.ptr(),
80 : intraLinks_[remIntraRankId]->GetRemoteRank(),
81 : intraLinks_[remIntraRankId]->GetLinkType(),
82 : INLINE_REDUCE_BIT));
83 :
84 0 : CHK_RET(intraLinks_[remIntraRankId]->TxDataSignal(subStream_[i]));
85 0 : CHK_RET(intraLinks_[remIntraRankId]->RxDataSignal(subStream_[i]));
86 0 : }
87 0 : return HCCL_SUCCESS;
88 : }
89 :
90 0 : HcclResult ReduceScatterGraphPipeline::RunInterServer(
91 : u64 blockIdx, const LINK &prevInterLink, const LINK &nextInterLink)
92 : {
93 0 : u64 blockOff = blockIdx * intraRankSize_;
94 0 : u64 memOffset = (blockOff + intraRankId_) * memSliceSize_;
95 0 : u64 preBlockOff = ((blockIdx + 1) % interRankSize_) * intraRankSize_;
96 0 : u64 preMemOffset = (preBlockOff + intraRankId_) * memSliceSize_;
97 :
98 0 : DeviceMem srcMem = DeviceMem::create(static_cast<u8 *>(usrInMem_) + memOffset, memSliceSize_);
99 0 : CHK_RET(senderInfo_->run(nextInterLink, memOffset, srcMem, subStream_[0], UserMemType::INPUT_MEM));
100 0 : HCCL_DEBUG("[ReduceScatterGraphPipeline][RunInterServer] local rank[%u] localOffset[%llu]tx with slice[%llu]",
101 : rankId_,
102 : memOffset,
103 : memSliceSize_);
104 :
105 0 : DeviceMem rxLocalMem = DeviceMem::create(static_cast<u8 *>(usrInMem_) + preMemOffset, memSliceSize_);
106 0 : CHK_RET(
107 : reducerInfo_->run(dispatcher_, prevInterLink, preMemOffset, rxLocalMem, rxLocalMem, rxLocalMem, subStream_[0]));
108 0 : return HCCL_SUCCESS;
109 0 : }
110 :
111 0 : HcclResult ReduceScatterGraphPipeline::RunAsync()
112 : {
113 : // inter ring algo
114 0 : u32 prevInterRankId = (interRankId_ + 1) % interRankSize_;
115 0 : u32 nextInterRankId = (interRankId_ - 1 + interRankSize_) % interRankSize_;
116 0 : LINK prevInterLink = interLinks_[prevInterRankId];
117 0 : LINK nextInterLink = interLinks_[nextInterRankId];
118 :
119 0 : for (u32 step = 0; step < interRankSize_; step++) {
120 0 : u32 begin = 0;
121 0 : if (step == 0) {
122 0 : begin = 1;
123 0 : CHK_RET(MainRecordSub(begin));
124 0 : CHK_RET(SubWaitMain(begin));
125 : }
126 : // server内做SDMA的reduce
127 0 : u64 blockIdx = ((interRankId_ + step + 1) % interRankSize_);
128 0 : CHK_RET(RunIntraServer(blockIdx));
129 0 : CHK_RET(SubRecordMain(begin));
130 0 : CHK_RET(MainWaitSub(begin));
131 0 : if (step < interRankSize_ - 1) {
132 : // 全部流同步,确保SDMA执行完成
133 0 : CHK_RET(MainRecordSub(0));
134 0 : CHK_RET(SubWaitMain(0));
135 0 : CHK_RET(prevInterLink->TxAck(subStream_[0]));
136 0 : CHK_RET(nextInterLink->RxAck(subStream_[0]));
137 : // server间做RDMA的reduce,可与下一个step的SDMA并发执行
138 0 : CHK_RET(RunInterServer(blockIdx, prevInterLink, nextInterLink));
139 0 : CHK_RET(prevInterLink->PostFinAck(subStream_[0]));
140 0 : CHK_RET(nextInterLink->WaitFinAck(subStream_[0]));
141 : // inter的最后一步需要barrier确保数据发完
142 0 : if (step == interRankSize_ - STEP_OFFSET_TWO) {
143 0 : CHK_RET(ExecuteBarrier(prevInterLink, nextInterLink, subStream_[0]));
144 : }
145 : }
146 : }
147 : // 把对应的切片从usrIn拷贝到userOut
148 0 : DeviceMem locSrc = DeviceMem::create(static_cast<u8 *>(usrInMem_) + rankId_ * memSliceSize_, memSliceSize_);
149 0 : DeviceMem locDst = DeviceMem::create(static_cast<u8 *>(usrOutMem_), memSliceSize_);
150 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, locDst, locSrc, stream_));
151 0 : HCCL_INFO("[ReduceScatterGraphPipeline][RunAsync]ReduceScatterGraphPipeline finished groupRankId[%u] ", rankId_);
152 0 : return HCCL_SUCCESS;
153 0 : }
154 :
155 : // 适配新CollExecutor接口
156 0 : HcclResult ReduceScatterGraphPipeline::Prepare(HcomCollOpInfo *opInfo, DeviceMem &cclBuffer, const u64 count,
157 : const u64 bufferSize, const u64 offset, const SubCommInfo &level0CommInfo, const SubCommInfo &level1CommInfo,
158 : Stream &mainStream, std::vector<Stream> &subStream, std::vector<std::shared_ptr<LocalNotify>> ¬ifyMain,
159 : std::vector<std::shared_ptr<LocalNotify>> ¬ifySub, u64 reduceAttrBitMap)
160 : {
161 0 : reduceAttr_ = reduceAttrBitMap;
162 0 : opInfo_ = opInfo;
163 :
164 0 : unitSize_ = SIZE_TABLE[opInfo_->dataType];
165 0 : count_ = opInfo_->count;
166 0 : memSliceSize_ = opInfo_->count * unitSize_;
167 0 : usrInMem_ = opInfo_->inputAddr;
168 0 : usrOutMem_ = opInfo_->outputAddr;
169 0 : reductionOp_ = opInfo_->reduceOp;
170 0 : dataType_ = opInfo_->dataType;
171 0 : offset_ = offset;
172 :
173 : // needed resource
174 : // stream: 1 * mainStream + n * subStream
175 : // mem: usrInMem_, usrOutMem
176 : // interNotify, streamNotify
177 :
178 : // stream
179 : // mainStream负责locMemCPY以及subStream同步控制
180 0 : stream_ = mainStream;
181 : // subStream负责:
182 : // streamId[0]: inter执行
183 : // streamId[1:intraRankSize]: intraRankSize-1个intra执行
184 0 : subStream_ = subStream;
185 :
186 : // DMAMem + interNotify from Link
187 0 : intraRankSize_ = level0CommInfo.localRankSize;
188 0 : interRankSize_ = level1CommInfo.localRankSize;
189 0 : intraRankId_ = level0CommInfo.localRank;
190 0 : interRankId_ = level1CommInfo.localRank;
191 0 : rankId_ = intraRankId_ + interRankId_ * intraRankSize_;
192 :
193 : // streamNotify, size: n
194 0 : streamNotifyMain_ = notifyMain;
195 0 : if (streamNotifyMain_.size() < intraRankSize_) {
196 0 : HCCL_ERROR("[ReduceScatterGraphPipeline][Prepare]rank[%u] streamNotifyMain_ size [%u] error, is smaller than,"
197 : "intraRankSize_[%u]",
198 : rankId_,
199 : streamNotifyMain_.size(),
200 : intraRankSize_);
201 0 : return HCCL_E_INTERNAL;
202 : }
203 0 : streamNotifySub_ = notifySub;
204 0 : if (streamNotifySub_.size() < intraRankSize_) {
205 0 : HCCL_ERROR("[ReduceScatterGraphPipeline][Prepare]rank[%u] streamNotifySub_ size [%u] error, is smaller than,"
206 : "intraRankSize_[%u]",
207 : rankId_,
208 : streamNotifySub_.size(),
209 : intraRankSize_);
210 0 : return HCCL_E_INTERNAL;
211 : }
212 :
213 0 : intraLinks_ = level0CommInfo.links;
214 0 : interLinks_ = level1CommInfo.links;
215 :
216 0 : HCCL_INFO("[ReduceScatterGraphPipeline][Prepare]streamNum[%u], streamNotifyMainNum[%u], streamNotifySubNum[%u]",
217 : subStream_.size(),
218 : streamNotifyMain_.size(),
219 : streamNotifySub_.size());
220 0 : HCCL_INFO("[ReduceScatterGraphPipeline][Prepare]interLinksNum[%u], intraLinksNum[%u]",
221 : interLinks_.size(),
222 : intraLinks_.size());
223 0 : senderInfo_.reset(new (std::nothrow) Sender(dataType_, reductionOp_, reduceAttr_));
224 0 : CHK_SMART_PTR_NULL(senderInfo_);
225 0 : reducerInfo_.reset(new (std::nothrow) Reducer(dataType_, reductionOp_, reduceAttr_));
226 0 : CHK_SMART_PTR_NULL(reducerInfo_);
227 0 : return HCCL_SUCCESS;
228 : }
229 :
230 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_REDUCESCATTER_GRAPH_PIPELINE, ReduceScatterGraphPipeline);
231 : } // namespace hccl
|