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