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 "allltoall_pipeline_mesh_pairwise_ccl_enough.h"
12 : #include "alg_template_register.h"
13 :
14 : namespace hccl {
15 :
16 : static const u32 INTRA_STREAM_INFO_SENDLEN_INDEX = 0; // intraStreamInfo 中 sendLen 的下标
17 : static const u32 INTRA_STREAM_INFO_RECVLEN_INDEX = 1; // intraStreamInfo 中 recvLen 的下标
18 : static const u32 INTRA_STREAM_INFO_RECV_REMOTE_OFFSET_INDEX = 2; // intraStreamInfo 中 recvRemoteOffset 的下标
19 : static const u32 INTRA_STREAM_INFO_RECV_LOCAL_OFFSET_INDEX = 3; // intraStreamInfo 中 recvLocalOffset 的下标
20 :
21 0 : AlltoallPipelineMeshPairwiseCCLEnough::AlltoallPipelineMeshPairwiseCCLEnough(const HcclDispatcher dispatcher)
22 0 : : AlltoallPipelineBase(dispatcher)
23 0 : {}
24 :
25 0 : AlltoallPipelineMeshPairwiseCCLEnough::~AlltoallPipelineMeshPairwiseCCLEnough() {}
26 :
27 0 : u32 AlltoallPipelineMeshPairwiseCCLEnough::CalcInterNumSteps() { return interRankSize_ - 1; }
28 :
29 : // 适配新CollExecutor接口
30 0 : HcclResult AlltoallPipelineMeshPairwiseCCLEnough::Prepare(
31 : u32 userRank, A2aPipelineMemory A2aPipelineMemory, const SubCommInfo& level0CommInfo,
32 : const SubCommInfo& level1CommInfo, Stream& mainStream, std::vector<Stream>& subStream,
33 : std::vector<std::shared_ptr<LocalNotify>>& notifyMain, std::vector<std::shared_ptr<LocalNotify>>& notifySub,
34 : std::vector<SendRecvInfo>& allMeshAggregationSendRecvInfo, HcclWorkflowMode workMode)
35 : {
36 0 : AlltoallPipelineBase::Prepare(
37 : userRank, A2aPipelineMemory, level0CommInfo, level1CommInfo, mainStream, subStream, notifyMain, notifySub,
38 : allMeshAggregationSendRecvInfo, workMode);
39 0 : GetIntraScratchOffset();
40 0 : CHK_RET(DeviceMemMapping());
41 0 : return HCCL_SUCCESS;
42 : }
43 :
44 : // 统一计算每步 mesh 内收发时从各卡 scratch 读取的 offset 和 length
45 0 : HcclResult AlltoallPipelineMeshPairwiseCCLEnough::GetIntraScratchOffset()
46 : {
47 0 : for (u32 i = 0; i < intraRankSize_; i++) {
48 0 : intraScratchOffsetMap_[i] = std::vector<u64>();
49 0 : intraScratchLengMap_[i] = std::vector<u64>();
50 0 : u64 startOffset = 0;
51 0 : for (u32 remoteRank = i; remoteRank < groupRankSize_; remoteRank += intraRankSize_) {
52 0 : if (remoteRank == userRank_) {
53 0 : localScratchOffset_ = startOffset;
54 : }
55 0 : const std::vector<u64>& remoteSendOffset = (*allMeshAggregationSendRecvInfo_)[remoteRank].sendOffset;
56 0 : const std::vector<u64>& remoteSendLength = (*allMeshAggregationSendRecvInfo_)[remoteRank].sendLength;
57 0 : intraScratchOffsetMap_[i].push_back(
58 0 : startOffset + (remoteSendOffset[userRank_] - remoteSendOffset[meshRankStart_]));
59 : startOffset
60 0 : += (remoteSendOffset[meshRankEnd_] + remoteSendLength[meshRankEnd_] - remoteSendOffset[meshRankStart_]);
61 0 : intraScratchLengMap_[i].push_back(remoteSendLength[userRank_]);
62 : }
63 : }
64 0 : return HCCL_SUCCESS;
65 : }
66 :
67 0 : HcclResult AlltoallPipelineMeshPairwiseCCLEnough::DeviceMemMapping()
68 : {
69 0 : if (workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
70 0 : interTransportSend_ = cclIn_;
71 0 : interTransportRecv_ = cclOut_;
72 0 : intraTransportSend_ = cclOut_;
73 : } else {
74 0 : interTransportSend_ = inputMem_;
75 0 : interTransportRecv_ = scratchMem_;
76 0 : intraTransportSend_ = scratchMem_;
77 : }
78 0 : for (u32 intraRank = 0; intraRank < intraRankSize_; intraRank++) {
79 0 : if (intraRank == intraRankId_) {
80 0 : continue;
81 : }
82 0 : LINK& intraNeighboorTransport = intraLinks_[intraRank];
83 0 : void* remDMAMemPtr = nullptr;
84 0 : CHK_RET(intraNeighboorTransport->GetRemoteMem(UserMemType::INPUT_MEM, &remDMAMemPtr));
85 : DeviceMem remoteAlltoallScratch = DeviceMem::create(
86 : static_cast<u8*>(remDMAMemPtr),
87 0 : workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE ? cclIn_.size() : scratchMem_.size());
88 0 : intraNeighBoorMemory_[intraRank] = {remoteAlltoallScratch};
89 0 : }
90 0 : return HCCL_SUCCESS;
91 0 : }
92 :
93 0 : HcclResult AlltoallPipelineMeshPairwiseCCLEnough::PrepareInterData(u32 step)
94 : {
95 : // 准备 mesh 间发送信息
96 0 : nextInterSendData_.clear();
97 0 : u32 interSendRankStart = ((interRankId_ + 1 + step) % interRankSize_) * intraRankSize_;
98 0 : u32 interSendRankEnd = interSendRankStart + intraRankSize_ - 1;
99 0 : u64 startMemOffset = localSendRecvInfo_.sendOffset[interSendRankStart];
100 0 : u64 meshSendLength = localSendRecvInfo_.sendOffset[interSendRankEnd]
101 0 : + localSendRecvInfo_.sendLength[interSendRankEnd] - startMemOffset;
102 0 : u64 sendDestOffset = 0;
103 0 : for (u32 relatedRank = intraRankId_; relatedRank < userRank_; relatedRank += intraRankSize_) {
104 0 : const SendRecvInfo& info = (*allMeshAggregationSendRecvInfo_)[relatedRank];
105 : sendDestOffset
106 0 : += (info.sendOffset[interSendRankEnd] + info.sendLength[interSendRankEnd]
107 0 : - info.sendOffset[interSendRankStart]);
108 : }
109 0 : DeviceMem srcMem = inputMem_.range(startMemOffset, meshSendLength);
110 0 : DeviceMem dstMem = interTransportSend_.range(startMemOffset, meshSendLength);
111 0 : HCCL_DEBUG(
112 : "[AlltoallPipelineMeshPairwiseCCLEnough][PrepareInterSendData] userRank %u, interRank %u, "
113 : "intraRank %u move from userInput offset %llu length %llu to interTransportSend, send to remote offset "
114 : "%llu",
115 : userRank_, interRankId_, intraRankId_, startMemOffset, meshSendLength, sendDestOffset);
116 :
117 0 : HCCL_DEBUG("user size %u, inter size %u, intra size %u", groupRankSize_, interRankSize_, intraRankSize_);
118 0 : if (workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
119 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstMem, srcMem, mainStream_));
120 : }
121 0 : nextInterSendData_.emplace_back(
122 0 : TxMemoryInfo{UserMemType::OUTPUT_MEM, sendDestOffset, dstMem.ptr(), meshSendLength});
123 :
124 : // 准备 mesh 间接收信息
125 0 : nextInterRecvData_.clear();
126 0 : u32 recvBlockStart = (((interRankId_ + interRankSize_ - 1u - step) % interRankSize_) * intraRankSize_);
127 0 : u64 recvLocalOffset = 0;
128 0 : for (u32 relatedRank = intraRankId_; relatedRank < recvBlockStart; relatedRank += intraRankSize_) {
129 0 : const SendRecvInfo& info = (*allMeshAggregationSendRecvInfo_)[relatedRank];
130 : recvLocalOffset
131 0 : += (info.sendOffset[meshRankEnd_] + info.sendLength[meshRankEnd_] - info.sendOffset[meshRankStart_]);
132 : }
133 0 : const SendRecvInfo& recvRankInfo = (*allMeshAggregationSendRecvInfo_)[recvBlockStart + intraRankId_];
134 0 : u64 recvRemoteOffset = recvRankInfo.sendOffset[meshRankStart_];
135 : u64 recvLength
136 0 : = (recvRankInfo.sendOffset[meshRankEnd_] + recvRankInfo.sendLength[meshRankEnd_]
137 0 : - recvRankInfo.sendOffset[meshRankStart_]);
138 0 : nextInterRecvData_.emplace_back(RxMemoryInfo{
139 0 : UserMemType::INPUT_MEM, recvRemoteOffset, interTransportRecv_.range(recvLocalOffset, recvLength).ptr(),
140 : recvLength});
141 0 : return HCCL_SUCCESS;
142 0 : }
143 :
144 : // 将原先在userInput，且需要发到本mesh内其他卡的数据搬到CCL
145 0 : HcclResult AlltoallPipelineMeshPairwiseCCLEnough::PrepareIntraData()
146 : {
147 0 : u64 startMemOffset = localSendRecvInfo_.sendOffset[meshRankStart_];
148 : u64 meshSendLength
149 0 : = localSendRecvInfo_.sendOffset[meshRankEnd_] + localSendRecvInfo_.sendLength[meshRankEnd_] - startMemOffset;
150 0 : u64 intraSendOffset = 0;
151 0 : for (u32 relatedRank = intraRankId_; relatedRank < meshRankStart_; relatedRank += intraRankSize_) {
152 0 : const SendRecvInfo& info = (*allMeshAggregationSendRecvInfo_)[relatedRank];
153 : intraSendOffset
154 0 : += (info.sendOffset[meshRankEnd_] + info.sendLength[meshRankEnd_] - info.sendOffset[meshRankStart_]);
155 : }
156 0 : DeviceMem srcMem = inputMem_.range(startMemOffset, meshSendLength);
157 0 : DeviceMem dstMem = intraTransportSend_.range(intraSendOffset, meshSendLength);
158 0 : if (workMode_ == HcclWorkflowMode::HCCL_WORKFLOW_MODE_OP_BASE) {
159 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstMem, srcMem, mainStream_));
160 : }
161 0 : HCCL_DEBUG(
162 : "[AlltoallPipelineMeshPairwiseCCLEnough][PrepareIntraData] userRank %u, interRank %u, intraRank %u "
163 : "copy from userInput offset %llu length %llu to intraTransportSend offset %llu",
164 : userRank_, interRankId_, intraRankId_, startMemOffset, meshSendLength, intraSendOffset);
165 0 : return HCCL_SUCCESS;
166 0 : }
167 :
168 0 : void AlltoallPipelineMeshPairwiseCCLEnough::UpdateIntraStreamInfo(u32 step)
169 : {
170 0 : intraStreamInfo_.clear();
171 0 : u32 localMeshIndex = (interRankId_ + interRankSize_ - step) % interRankSize_;
172 0 : u32 firstDataBlockIndex = (meshRankStart_ + groupRankSize_ - step * intraRankSize_) % groupRankSize_;
173 : const std::vector<u64>& sendLengths
174 0 : = (*allMeshAggregationSendRecvInfo_)[firstDataBlockIndex + intraRankId_].sendLength;
175 0 : const std::vector<u64>& recvLengths = localSendRecvInfo_.recvLength;
176 0 : const std::vector<u64>& recvOffsets = localSendRecvInfo_.recvOffset;
177 0 : HCCL_DEBUG(
178 : "[AlltoallPipelineMeshPairwiseCCLEnough][UpdateIntraStreamInfo] userRank %u, "
179 : "interRank %u, intraRank %u, step %u",
180 : userRank_, interRankId_, intraRankId_, step);
181 0 : for (u32 intraRank = 0; intraRank < intraRankSize_; intraRank++) {
182 0 : u64 sendLen = sendLengths[meshRankStart_ + intraRank];
183 0 : u64 recvLen = recvLengths[firstDataBlockIndex + intraRank];
184 0 : u64 recvRemoteOffset = intraScratchOffsetMap_[intraRank][localMeshIndex];
185 0 : u64 recvLocalOffset = recvOffsets[firstDataBlockIndex + intraRank];
186 0 : if (intraRank != intraRankId_) {
187 0 : intraStreamInfo_[intraRank] = {sendLen, recvLen, recvRemoteOffset, recvLocalOffset};
188 0 : HCCL_DEBUG(
189 : "[AlltoallPipelineMeshPairwiseCCLEnough][UpdateIntraStreamInfo] userRank %u, interRank %u, "
190 : "intraRank %u, sdma stream %u need send %llu and read length %llu from remote offset %llu "
191 : "to local offset %llu",
192 : userRank_, interRankId_, intraRankId_, intraRank, sendLen, recvLen, recvRemoteOffset, recvLocalOffset);
193 : }
194 : }
195 0 : }
196 :
197 0 : HcclResult AlltoallPipelineMeshPairwiseCCLEnough::SendRecvDataIntraMesh()
198 : {
199 0 : HCCL_DEBUG(
200 : "[AlltoallPipelineMeshPairwiseCCLEnough][SendRecvDataIntraMesh] userRank %u, "
201 : "interRank %u, intraRank %u, sdma stream %s wait main stream",
202 : userRank_, interRankId_, intraRankId_, GetStreamIndexString().c_str());
203 0 : for (auto& intraInfo : intraStreamInfo_) {
204 0 : u32 streamIndex = intraInfo.first;
205 0 : u64 recvLen = intraInfo.second[INTRA_STREAM_INFO_RECVLEN_INDEX];
206 0 : Stream& currStream = subStream_[streamIndex];
207 0 : LINK& readTransport = intraLinks_[streamIndex];
208 0 : CHK_RET(readTransport->TxAck(currStream));
209 0 : CHK_RET(readTransport->RxAck(currStream));
210 0 : u64 recvRemoteOffset = intraInfo.second[INTRA_STREAM_INFO_RECV_REMOTE_OFFSET_INDEX];
211 0 : u64 recvLocalOffset = intraInfo.second[INTRA_STREAM_INFO_RECV_LOCAL_OFFSET_INDEX];
212 0 : DeviceMem src = intraNeighBoorMemory_[streamIndex][0].range(recvRemoteOffset, recvLen);
213 0 : DeviceMem dst = outputMem_.range(recvLocalOffset, recvLen);
214 0 : CHK_RET(HcclD2DMemcpyAsync(
215 : dispatcher_, dst, src, currStream, readTransport->GetRemoteRank(), readTransport->GetLinkType()));
216 0 : CHK_RET(readTransport->TxDataSignal(currStream));
217 0 : HCCL_DEBUG(
218 : "[AlltoallPipelineMeshPairwiseCCLEnough][SendRecvDataIntraMesh] userRank %u, interRank %u, "
219 : "intraRank %u, sdma stream %llu read data from remote offset %llu len %llu to local %llu",
220 : userRank_, interRankId_, intraRankId_, streamIndex, recvRemoteOffset, recvLen, recvLocalOffset);
221 0 : CHK_RET(readTransport->RxDataSignal(currStream));
222 0 : }
223 0 : HCCL_DEBUG(
224 : "[AlltoallPipelineMeshPairwiseCCLEnough][SendRecvDataIntraMesh] userRank %u, interRank %u, "
225 : "intraRank %u, sdma stream %s notify main stream",
226 : userRank_, interRankId_, intraRankId_, GetStreamIndexString().c_str());
227 0 : return HCCL_SUCCESS;
228 : }
229 :
230 0 : HcclResult AlltoallPipelineMeshPairwiseCCLEnough::SendRecvDataInterMesh(u32 step)
231 : {
232 0 : Stream& interStream = subStream_[intraRankId_];
233 0 : LINK& interRecvTransport = interLinks_[(interRankId_ + interRankSize_ - 1 - step) % interRankSize_];
234 0 : LINK& interSendTransport = interLinks_[(interRankId_ + 1 + step) % interRankSize_];
235 0 : CHK_RET(interRecvTransport->TxAck(interStream));
236 0 : CHK_RET(interSendTransport->RxAck(interStream));
237 0 : CHK_RET(interSendTransport->TxAsync(nextInterSendData_, interStream));
238 0 : CHK_RET(interRecvTransport->RxAsync(nextInterRecvData_, interStream));
239 0 : CHK_RET(interRecvTransport->PostFinAck(interStream));
240 0 : CHK_RET(interSendTransport->WaitFinAck(interStream));
241 0 : CHK_RET(ExecuteBarrier(interRecvTransport, interSendTransport, interStream));
242 0 : return HCCL_SUCCESS;
243 : }
244 :
245 0 : HcclResult AlltoallPipelineMeshPairwiseCCLEnough::LocalCopyDataRecvFromInter(u32 interRankDistance)
246 : {
247 : u64 scratchOffset
248 0 : = intraScratchOffsetMap_[intraRankId_][(interRankId_ + interRankSize_ - interRankDistance) % interRankSize_];
249 : u64 recvLen = localSendRecvInfo_
250 0 : .recvLength[(userRank_ + groupRankSize_ - interRankDistance * intraRankSize_) % groupRankSize_];
251 : u64 userOutOffset
252 : = localSendRecvInfo_
253 0 : .recvOffset[(userRank_ + groupRankSize_ - interRankDistance * intraRankSize_) % groupRankSize_];
254 0 : DeviceMem src = interTransportRecv_.range(scratchOffset, recvLen);
255 0 : DeviceMem dst = outputMem_.range(userOutOffset, recvLen);
256 0 : HCCL_DEBUG(
257 : "[AlltoallPipelineMeshPairwiseCCLEnough][LocalCopyDataRecvFromInter]local move from "
258 : "interTransportRecv_ offset %llu length %llu to outputMem_ %llu",
259 : scratchOffset, recvLen, userOutOffset);
260 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, mainStream_));
261 0 : return HCCL_SUCCESS;
262 0 : }
263 :
264 0 : HcclResult AlltoallPipelineMeshPairwiseCCLEnough::PreProcess()
265 : {
266 : // server 间收发时间较长，先搬 server 间收发所需数据然后马上让server间开始收发
267 0 : CHK_RET(PrepareInterData(0));
268 0 : CHK_RET(NotifyInterStreamStart());
269 : // 之后将 server 内所需数据准备好之后唤醒server内从流收发
270 0 : UpdateIntraStreamInfo(0);
271 0 : CHK_RET(PrepareIntraData());
272 0 : CHK_RET(NotifyIntraStreamStart());
273 0 : CHK_RET(SendRecvDataIntraMesh());
274 0 : return HCCL_SUCCESS;
275 : }
276 :
277 0 : HcclResult AlltoallPipelineMeshPairwiseCCLEnough::PipelineSend(u32 step, bool isLastStep)
278 : {
279 0 : CHK_RET(SendRecvDataInterMesh(step));
280 0 : CHK_RET(PrepareInterData(step + 1u));
281 0 : CHK_RET(WaitInterStreamFinish());
282 0 : CHK_RET(WaitIntraStreamFinish());
283 0 : CHK_RET(ExecEmptyTask(inputMem_, outputMem_, mainStream_, dispatcher_));
284 0 : if (!isLastStep) {
285 0 : CHK_RET(NotifyInterStreamStart());
286 : }
287 0 : UpdateIntraStreamInfo(step + 1u);
288 0 : CHK_RET(NotifyIntraStreamStart());
289 0 : CHK_RET(SendRecvDataIntraMesh());
290 0 : CHK_RET(LocalCopyDataRecvFromInter(step + 1u));
291 0 : return HCCL_SUCCESS;
292 : }
293 :
294 0 : HcclResult AlltoallPipelineMeshPairwiseCCLEnough::PostProcess()
295 : {
296 : // 最后的收尾工作
297 0 : CHK_RET(LocalCopyDataRecvFromInter(0));
298 0 : CHK_RET(WaitIntraStreamFinish());
299 0 : return HCCL_SUCCESS;
300 : }
301 : REGISTER_TEMPLATE(
302 : TemplateType::TEMPLATE_ALL_2_ALL_PIPELINE_MESH_PAIRWISE_CCL_ENOUGH, AlltoallPipelineMeshPairwiseCCLEnough);
303 : } // namespace hccl
|
