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_graph_pipeline.h"
12 : #include "alg_template_register.h"
13 :
14 : constexpr u32 STEP_OFFSET_TWO = 2;
15 :
16 : namespace hccl {
17 0 : AllGatherGraphPipeline::AllGatherGraphPipeline(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher)
18 0 : {}
19 :
20 0 : AllGatherGraphPipeline::~AllGatherGraphPipeline()
21 0 : {}
22 :
23 0 : HcclResult AllGatherGraphPipeline::Prepare(HcomCollOpInfo *opInfo, u32 userRank, u64 &count, DeviceMem &inputMem,
24 : DeviceMem &outputMem, SubCommInfo &level0CommInfo, SubCommInfo &level1CommInfo, Stream &mainStream,
25 : std::vector<Stream> &subStream, std::vector<std::shared_ptr<LocalNotify>> ¬ifyMain,
26 : std::vector<std::shared_ptr<LocalNotify>> ¬ifySub)
27 : {
28 0 : opInfo_ = opInfo;
29 0 : memSliceCount_ = count;
30 0 : userRank_ = userRank;
31 :
32 0 : u32 unitSize = SIZE_TABLE[opInfo->dataType];
33 0 : u64 memSliceSize = memSliceCount_ * unitSize;
34 :
35 0 : usrInMemAddr_ = opInfo_->inputAddr;
36 0 : usrOutMemAddr_ = opInfo_->outputAddr;
37 :
38 : // needed resource
39 : // stream: 1 * mainStream + n * subStream
40 : // mem: usrInMem, usrOutMem
41 : // intorNotify, streamNotify
42 :
43 : // stream
44 : // mainStream负责locMemCpy和subStream同步控制
45 0 : stream_ = mainStream;
46 : // subStream负责:
47 : // streamId[0]: inter执行
48 : // streamId[1:intraRankSize]: intraRankSize-1个intra执行
49 0 : subStream_ = subStream;
50 :
51 0 : HCCL_DEBUG("[AllGatherGraphPipeline]prepare for userRank is %u, memSliceCount is %llu", userRank_, memSliceCount_);
52 0 : intraRankSize_ = level0CommInfo.localRankSize;
53 0 : interRankSize_ = level1CommInfo.localRankSize;
54 0 : intraRankId_ = level0CommInfo.localRank;
55 0 : interRankId_ = level1CommInfo.localRank;
56 0 : intraLinks_ = level0CommInfo.links;
57 0 : interLinks_ = level1CommInfo.links;
58 :
59 : // streamNotify, size: n
60 0 : streamNotifyMain_ = notifyMain;
61 0 : if (streamNotifyMain_.size() < intraRankSize_) {
62 0 : HCCL_ERROR("[AllGatherGraphPipeline][Prepare]rank[%u] streamNotifyMain_ size[%u] error, is smaller than,"
63 : "intraRankSize_[%u]",
64 : userRank_,
65 : streamNotifyMain_.size(),
66 : intraRankSize_);
67 0 : return HCCL_E_INTERNAL;
68 : }
69 0 : streamNotifySub_ = notifySub;
70 0 : if (streamNotifySub_.size() < intraRankSize_) {
71 0 : HCCL_ERROR("[AllGatherGraphPipeline][Prepare]rank[%u] streamNotifySub_ size[%u] error, is smaller than,"
72 : "intraRankSize_[%u]",
73 : userRank_,
74 : streamNotifySub_.size(),
75 : intraRankSize_);
76 0 : return HCCL_E_INTERNAL;
77 : }
78 :
79 0 : DeviceMem dmaMem0 = DeviceMem::create(inputMem.ptr(), memSliceSize);
80 0 : DeviceMem dmaMem1 = DeviceMem::create(outputMem.ptr(), memSliceSize);
81 :
82 0 : dmaMem_.push_back(dmaMem0);
83 0 : dmaMem_.push_back(dmaMem1);
84 :
85 0 : HCCL_INFO("[AllGatherGraphPipeline][Prepare]streamNum[%zu], streamNotifyMainNum[%zu], streamNotifySubNum[%zu]",
86 : subStream_.size(),
87 : streamNotifyMain_.size(),
88 : streamNotifySub_.size());
89 0 : HCCL_INFO("[AllGatherGraphPipeline][Prepare]interLinksNum[%zu], intraLinksNum[%zu]",
90 : interLinks_.size(),
91 : intraLinks_.size());
92 0 : return HCCL_SUCCESS;
93 0 : }
94 :
95 0 : HcclResult AllGatherGraphPipeline::MainWaitSub()
96 : {
97 0 : u32 subStreamNum = intraRankSize_;
98 0 : for (u32 signalIndex = 0; signalIndex < subStreamNum; signalIndex++) {
99 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, streamNotifyMain_[signalIndex], INVALID_VALUE_STAGE));
100 : }
101 0 : return HCCL_SUCCESS;
102 : }
103 :
104 0 : HcclResult AllGatherGraphPipeline::SubRecordMain()
105 : {
106 0 : u32 subStreamNum = intraRankSize_;
107 0 : for (u32 streamIndex = 0; streamIndex < subStreamNum; streamIndex++) {
108 0 : CHK_RET(LocalNotify::Post(
109 : subStream_[streamIndex], dispatcher_, streamNotifyMain_[streamIndex], INVALID_VALUE_STAGE));
110 : }
111 0 : return HCCL_SUCCESS;
112 : }
113 :
114 0 : HcclResult AllGatherGraphPipeline::MainRecordSub()
115 : {
116 0 : u32 subStreamNum = intraRankSize_;
117 0 : for (u32 signalIndex = 0; signalIndex < subStreamNum; signalIndex++) {
118 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, streamNotifySub_[signalIndex], INVALID_VALUE_STAGE));
119 : }
120 0 : return HCCL_SUCCESS;
121 : }
122 :
123 0 : HcclResult AllGatherGraphPipeline::SubWaitMain()
124 : {
125 0 : u32 subStreamNum = intraRankSize_;
126 0 : for (u32 streamIndex = 0; streamIndex < subStreamNum; streamIndex++) {
127 0 : CHK_RET(LocalNotify::Wait(
128 : subStream_[streamIndex], dispatcher_, streamNotifySub_[streamIndex], INVALID_VALUE_STAGE));
129 : }
130 0 : return HCCL_SUCCESS;
131 : }
132 :
133 0 : HcclResult AllGatherGraphPipeline::RunAsync()
134 : {
135 0 : HCCL_INFO("[AllGatherGraphPipeline][RunAsync]AllGatherRingMesh starts groupRankId[%u]", userRank_);
136 : // inter ring algo
137 0 : u32 prevInterRankId = (interRankId_ + interRankSize_ - 1) % interRankSize_;
138 0 : u32 nextInterRankId = (interRankId_ + 1) % interRankSize_;
139 0 : LINK prevInterLink = interLinks_[prevInterRankId];
140 0 : CHK_SMART_PTR_NULL(prevInterLink);
141 0 : LINK nextInterLink = interLinks_[nextInterRankId];
142 0 : CHK_SMART_PTR_NULL(nextInterLink);
143 :
144 : // intra fullmesh algo
145 : // intra 使用全部连接,不再映射
146 :
147 0 : u32 unitSize = SIZE_TABLE[opInfo_->dataType];
148 0 : u64 memSliceSize = memSliceCount_ * unitSize;
149 :
150 : // step 0 前置操作:所有卡本地数据从userIn-->userOut
151 0 : DeviceMem locSrc = DeviceMem::create(usrInMemAddr_, memSliceSize);
152 0 : u64 localOffsetByte = memSliceCount_ * userRank_ * unitSize;
153 0 : DeviceMem locDst = DeviceMem::create(static_cast<u8 *>(dmaMem_[1].ptr()) + localOffsetByte, memSliceSize);
154 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, locDst, locSrc, stream_));
155 :
156 0 : for (u32 step = 0; step < interRankSize_; step++) {
157 : // 主从流同步
158 0 : CHK_RET(MainRecordSub());
159 0 : CHK_RET(SubWaitMain());
160 :
161 0 : u64 serverRankOffset = intraRankId_ + (interRankId_ + interRankSize_ - step) % interRankSize_ * intraRankSize_;
162 0 : u64 serverOffsetByte = memSliceCount_ * serverRankOffset * unitSize;
163 0 : u64 readRemoteOffset = intraRankId_ + (prevInterRankId + interRankSize_ - step) % interRankSize_ *
164 0 : intraRankSize_; // server间前通信rank偏移
165 0 : u64 readRemoteOffsetByte = memSliceCount_ * readRemoteOffset * unitSize;
166 :
167 0 : if (step < interRankSize_ - 1) {
168 0 : CHK_RET(prevInterLink->TxAck(subStream_[0])); // AckRecord
169 0 : CHK_RET(nextInterLink->RxAck(subStream_[0])); // AckWait
170 :
171 0 : CHK_RET(nextInterLink->TxAsync(UserMemType::OUTPUT_MEM,
172 : serverOffsetByte,
173 : static_cast<u8 *>(dmaMem_[1].ptr()) + serverOffsetByte,
174 : memSliceSize,
175 : subStream_[0]));
176 0 : HCCL_DEBUG("[AllGatherGraphPipeline][RunAsync] local rank[%u] localOffset[%llu]tx with remoteRank[%u], "
177 : "remoteOffset[%llu] with slice[%llu]",
178 : userRank_,
179 : serverOffsetByte,
180 : nextInterRankId,
181 : serverOffsetByte,
182 : memSliceSize);
183 :
184 0 : CHK_RET(prevInterLink->RxAsync(UserMemType::OUTPUT_MEM,
185 : readRemoteOffsetByte,
186 : static_cast<u8 *>(dmaMem_[1].ptr()) + readRemoteOffsetByte,
187 : memSliceSize,
188 : subStream_[0])); // wait
189 0 : HCCL_DEBUG(
190 : "[AllGatherGraphPipeline][RunAsync] read local rank[%u] localOffset[%llu]tx with remoteRank[%u], "
191 : "remoteOffset[%llu] with slice[%llu]",
192 : userRank_,
193 : readRemoteOffsetByte,
194 : prevInterRankId,
195 : readRemoteOffsetByte,
196 : memSliceSize);
197 :
198 0 : CHK_RET(prevInterLink->PostFinAck(subStream_[0]));
199 0 : CHK_RET(nextInterLink->WaitFinAck(subStream_[0]));
200 :
201 : // inter的最后一步需要barrier确保数据发完
202 0 : if (step == interRankSize_ - STEP_OFFSET_TWO) {
203 0 : CHK_RET(ExecuteBarrier(prevInterLink, nextInterLink, subStream_[0]));
204 : }
205 : }
206 :
207 0 : HCCL_DEBUG("[AllGatherGraphPipeline][RunAsync]now step is %u, intraRankSize is %u", step, intraRankSize_);
208 0 : for (u32 i = 1; i < intraRankSize_; i++) {
209 0 : u32 remIntraRankId = (intraRankId_ + i) % intraRankSize_;
210 0 : CHK_RET(intraLinks_[remIntraRankId]->TxAck(subStream_[i])); // ackrecord
211 0 : CHK_RET(intraLinks_[remIntraRankId]->RxAck(subStream_[i])); // ackwait
212 :
213 0 : void *remDMAMemPtr = nullptr;
214 0 : CHK_RET(intraLinks_[remIntraRankId]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remDMAMemPtr));
215 0 : u64 remoteOffset =
216 0 : (interRankId_ - step + interRankSize_) % interRankSize_ * intraRankSize_ + remIntraRankId;
217 0 : u64 remoteOffsetByte = memSliceCount_ * remoteOffset * unitSize;
218 0 : void *dstAddr = static_cast<u8 *>(usrOutMemAddr_) + remoteOffsetByte;
219 :
220 0 : DeviceMem src = DeviceMem::create(static_cast<u8 *>(remDMAMemPtr) + remoteOffsetByte, memSliceSize);
221 0 : DeviceMem dst = DeviceMem::create(dstAddr, memSliceSize);
222 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, subStream_[i]));
223 :
224 0 : CHK_RET(intraLinks_[remIntraRankId]->TxDataSignal(subStream_[i])); // data record
225 0 : CHK_RET(intraLinks_[remIntraRankId]->RxDataSignal(subStream_[i])); // data wait
226 0 : }
227 :
228 0 : CHK_RET(SubRecordMain());
229 0 : CHK_RET(MainWaitSub());
230 : }
231 :
232 0 : HCCL_INFO("[AllGatherGraphPipeline][RunAsync]AllGatherRingMesh finished groupRankId[%u]", userRank_);
233 0 : return HCCL_SUCCESS;
234 0 : }
235 :
236 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_GATHER_GRAPH_PIPELINE, AllGatherGraphPipeline);
237 : } // namespace hccl
|