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 "alg_template_register.h"
12 : #include "all_reduce_mesh_opbase.h"
13 :
14 : namespace hccl {
15 3 : AllReduceMeshDirect::AllReduceMeshDirect(const HcclDispatcher dispatcher) : AlgTemplateBase(dispatcher) {}
16 :
17 6 : AllReduceMeshDirect::~AllReduceMeshDirect() {}
18 :
19 3 : HcclResult AllReduceMeshDirect::Prepare(
20 : u64 reduceAttrBitMap, std::vector<Stream>& meshStreams, std::vector<std::shared_ptr<LocalNotify>>& meshSignal,
21 : std::vector<std::shared_ptr<LocalNotify>>& meshSignalAux, u32 interRank, u32 interRankSize, u32 userRank,
22 : HcomCollOpInfo* opInfo)
23 : {
24 3 : reduceAttr_ = reduceAttrBitMap;
25 3 : localRank_ = interRank;
26 3 : localRankSize_ = interRankSize;
27 3 : userRank_ = userRank;
28 3 : meshStreams_ = meshStreams;
29 3 : meshSignal_ = &meshSignal;
30 3 : meshSignalAux_ = &meshSignalAux;
31 3 : opInfo_ = opInfo;
32 3 : return HCCL_SUCCESS;
33 : }
34 :
35 0 : HcclResult AllReduceMeshDirect::MainRecordSub()
36 : {
37 0 : for (u32 signalIndex = 0; signalIndex < meshSignalAux_->size(); signalIndex++) {
38 0 : CHK_RET(LocalNotify::Post(stream_, dispatcher_, (*meshSignalAux_)[signalIndex], profilerInput_.stage));
39 : }
40 0 : return HCCL_SUCCESS;
41 : }
42 :
43 0 : HcclResult AllReduceMeshDirect::SubWaitMain()
44 : {
45 0 : for (u32 streamIndex = 0; streamIndex < meshSignalAux_->size(); streamIndex++) {
46 0 : CHK_RET(LocalNotify::Wait(
47 : meshStreams_[streamIndex], dispatcher_, (*meshSignalAux_)[streamIndex], profilerInput_.stage));
48 : }
49 0 : return HCCL_SUCCESS;
50 : }
51 :
52 0 : HcclResult AllReduceMeshDirect::MainWaitSub()
53 : {
54 0 : for (u32 signalIndex = 0; signalIndex < meshSignal_->size(); signalIndex++) {
55 0 : CHK_RET(LocalNotify::Wait(stream_, dispatcher_, (*meshSignal_)[signalIndex], profilerInput_.stage));
56 : }
57 0 : return HCCL_SUCCESS;
58 : }
59 :
60 0 : HcclResult AllReduceMeshDirect::SubRecordMain()
61 : {
62 0 : for (u32 streamIndex = 0; streamIndex < meshSignal_->size(); streamIndex++) {
63 0 : CHK_RET(LocalNotify::Post(
64 : meshStreams_[streamIndex], dispatcher_, (*meshSignal_)[streamIndex], profilerInput_.stage));
65 : }
66 0 : return HCCL_SUCCESS;
67 : }
68 :
69 : // 将数据均分，最小单位是128
70 3 : HcclResult AllReduceMeshDirect::PrepareSlice(u64 dataCount, u32 unitSize, u32 sliceNum, std::vector<Slice>& dataSlice)
71 : {
72 3 : Slice temp;
73 3 : u64 totalSize = dataCount * unitSize;
74 3 : dataSlice.clear();
75 3 : dataSlice.reserve(sliceNum);
76 3 : if (sliceNum == 0) {
77 0 : HCCL_ERROR("[Prepare][SliceData]data slice prepare, sliceNum is 0");
78 0 : return HCCL_E_PARA;
79 : }
80 3 : u64 sizePerSlice = (totalSize + sliceNum - 1) / sliceNum; /* 1是为了向上取整 */
81 3 : sizePerSlice = RoundUpWithDivisor(sizePerSlice, HCCL_MIN_SLICE_ALIGN_910B);
82 3 : u64 residueSize = totalSize;
83 3 : u32 i = 0;
84 6 : while (residueSize > 0) {
85 3 : u64 sliceSize = sizePerSlice < residueSize ? sizePerSlice : residueSize;
86 3 : temp.size = sliceSize;
87 3 : temp.offset = totalSize - residueSize;
88 3 : i++;
89 3 : if (sliceSize <= 0) {
90 0 : HCCL_ERROR("[Prepare][SliceData]data_slice_prepare sliceSize[%llu]", sliceSize);
91 0 : return HCCL_E_PARA;
92 : }
93 3 : residueSize -= sliceSize;
94 3 : dataSlice.push_back(temp);
95 : }
96 24 : while (i < sliceNum) {
97 21 : temp.size = 0;
98 21 : temp.offset = totalSize;
99 21 : i++;
100 21 : dataSlice.push_back(temp);
101 : }
102 3 : return HCCL_SUCCESS;
103 : }
104 :
105 0 : HcclResult AllReduceMeshDirect::PrepareAllreduceSliceData()
106 : {
107 0 : return PrepareSlice(count_, DataUnitSize(dataType_), localRankSize_, slices_);
108 : }
109 :
110 : // ringallreduce算法的函数入口
111 0 : HcclResult AllReduceMeshDirect::RunAsync(const u32 rank, const u32 rankSize, const std::vector<LINK>& links)
112 : {
113 0 : HcclResult ret = HCCL_SUCCESS;
114 0 : CHK_SMART_PTR_NULL(dispatcher_);
115 0 : CHK_PTR_NULL(stream_.ptr());
116 0 : HCCL_INFO(
117 : "AllReduceMeshDirect run: rank[%u] ranksize[%u] inputMem[%p] outputMem[%p] count[%llu]", rank, rankSize,
118 : inputMem_.ptr(), outputMem_.ptr(), count_);
119 :
120 0 : if (links.size() < rankSize) {
121 0 : HCCL_ERROR(
122 : "[AllReduceMeshDirect][RunAsync]rank[%u] linksize[%llu] is less than rankSize[%u]", rank, links.size(),
123 : rankSize);
124 0 : return HCCL_E_INTERNAL;
125 : }
126 :
127 : // 如果ranksize为1, inline reduce和普通跨片reduce操作一致，从input->output
128 0 : if (rankSize == 1) {
129 0 : if (opInfo_->inputAddr != opInfo_->outputAddr) {
130 0 : DeviceMem userMemIn = DeviceMem::create(opInfo_->inputAddr, count_ * DataUnitSize(dataType_));
131 0 : DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, count_ * DataUnitSize(dataType_));
132 0 : ret = HcclD2DMemcpyAsync(dispatcher_, userMemOut, userMemIn, stream_);
133 0 : CHK_PRT_RET(
134 : ret != HCCL_SUCCESS, HCCL_ERROR("[AllReduceRing][RunAsync]rank[%u] memcpy async failed", rank), ret);
135 0 : }
136 0 : return ret;
137 : }
138 :
139 0 : ret = PrepareAllreduceSliceData();
140 0 : CHK_PRT_RET(
141 : ret != HCCL_SUCCESS,
142 : HCCL_ERROR(
143 : "[AllReduceRing][RunAsync]rank[%u] count[%llu] failed in PrepareSliceData "
144 : "step",
145 : rank, count_),
146 : ret);
147 :
148 0 : ret = RunReduceScatter(rank, rankSize, links);
149 0 : CHK_PRT_RET(
150 : ret != HCCL_SUCCESS,
151 : HCCL_ERROR(
152 : "[AllReduceRing][RunAsync]rank[%u] count[%llu] failed in reducescater "
153 : "step",
154 : rank, count_),
155 : ret);
156 :
157 0 : ret = RunAllGather(rank, rankSize, links);
158 0 : CHK_PRT_RET(
159 : ret != HCCL_SUCCESS,
160 : HCCL_ERROR(
161 : "[AllReduceRing][RunAsync]rank[%u] count[%llu] failed in AllGather "
162 : "step",
163 : rank, count_),
164 : ret);
165 :
166 0 : HCCL_INFO("AllReduceMeshDirect finished: rank[%u] ranksize[%u].", rank, rankSize);
167 0 : return HCCL_SUCCESS;
168 : }
169 :
170 0 : HcclResult AllReduceMeshDirect::RunReduceScatter(u32 rank, u32 rankSize, const std::vector<LINK>& links)
171 : {
172 0 : HCCL_INFO(
173 : "ReduceScatterMeshAtomicOpbase run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu].", rank,
174 : rankSize, inputMem_.ptr(), outputMem_.ptr(), count_);
175 :
176 : // 数据准备
177 0 : u32 unitSize = DataUnitSize(dataType_);
178 :
179 0 : DeviceMem userMemIn = DeviceMem::create(opInfo_->inputAddr, count_ * unitSize);
180 0 : DeviceMem commMemOut = DeviceMem::create(outputMem_.ptr(), outputMem_.size());
181 :
182 0 : DeviceMem src;
183 0 : DeviceMem dst;
184 :
185 0 : src = DeviceMem::create(static_cast<char*>(opInfo_->inputAddr) + slices_[rank].offset, slices_[rank].size);
186 0 : dst = commMemOut.range(0, slices_[rank].size);
187 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
188 :
189 0 : CHK_RET(MainRecordSub());
190 0 : CHK_RET(SubWaitMain());
191 :
192 0 : for (u32 round = 1; round < rankSize; round++) {
193 0 : u32 dstRank = (round + rank) % rankSize;
194 0 : Stream& subStream = meshStreams_[round - 1];
195 0 : CHK_RET(links[dstRank]->TxAck(subStream));
196 0 : CHK_RET(links[dstRank]->RxAck(subStream));
197 : }
198 :
199 0 : CHK_RET(SubRecordMain());
200 0 : CHK_RET(MainWaitSub());
201 0 : DeviceMem srcTmp = DeviceMem::create(inputMem_.ptr(), 0);
202 0 : DeviceMem dstTmp = DeviceMem::create(outputMem_.ptr(), 0);
203 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstTmp, srcTmp, stream_));
204 0 : CHK_RET(SubWaitMain());
205 0 : CHK_RET(MainRecordSub());
206 :
207 0 : for (u32 round = 1; round < rankSize; round++) {
208 0 : u32 dstRank = (round + rank) % rankSize;
209 0 : Stream& subStream = meshStreams_[round - 1];
210 :
211 0 : void* remMemPtr = nullptr;
212 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
213 :
214 0 : dst = DeviceMem::create(static_cast<char*>(remMemPtr), slices_[dstRank].size);
215 0 : src = userMemIn.range(slices_[dstRank].offset, slices_[dstRank].size);
216 0 : CHK_RET(HcclReduceAsync(
217 : dispatcher_, static_cast<void*>(src.ptr()), slices_[dstRank].size / unitSize, dataType_, reductionOp_,
218 : subStream, static_cast<void*>(dst.ptr()), links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType(),
219 : INLINE_REDUCE_BIT));
220 :
221 0 : CHK_RET(links[dstRank]->TxDataSignal(subStream));
222 0 : CHK_RET(links[dstRank]->RxDataSignal(subStream));
223 : }
224 :
225 0 : CHK_RET(SubRecordMain());
226 0 : CHK_RET(MainWaitSub());
227 0 : return HCCL_SUCCESS;
228 0 : }
229 :
230 0 : HcclResult AllReduceMeshDirect::RunAllGather(u32 rank, u32 rankSize, const std::vector<LINK>& links)
231 : {
232 0 : HCCL_INFO(
233 : "AllGatherMesh run: rank[%u] totalrank[%u] inputMem[%p] outputMem[%p] count[%llu].", rank, rankSize,
234 : inputMem_.ptr(), outputMem_.ptr(), count_);
235 0 : u32 unitSize = DataUnitSize(dataType_);
236 :
237 0 : DeviceMem userMemOut = DeviceMem::create(opInfo_->outputAddr, count_ * unitSize);
238 0 : DeviceMem commMemOut = DeviceMem::create(outputMem_.ptr(), outputMem_.size());
239 :
240 0 : DeviceMem src;
241 0 : DeviceMem dst;
242 :
243 0 : src = commMemOut.range(0, slices_[rank].size);
244 0 : dst = userMemOut.range(slices_[rank].offset, slices_[rank].size);
245 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dst, src, stream_));
246 :
247 0 : CHK_RET(MainRecordSub());
248 0 : CHK_RET(SubWaitMain());
249 :
250 0 : for (u32 round = 1; round < rankSize; round++) {
251 0 : u32 dstRank = BackwardRank(rank, rankSize, round);
252 0 : Stream& subStream = meshStreams_[round - 1];
253 0 : CHK_RET(links[dstRank]->TxAck(subStream));
254 0 : CHK_RET(links[dstRank]->RxAck(subStream));
255 : }
256 :
257 0 : CHK_RET(SubRecordMain());
258 0 : CHK_RET(MainWaitSub());
259 0 : DeviceMem srcTmp = DeviceMem::create(inputMem_.ptr(), 0);
260 0 : DeviceMem dstTmp = DeviceMem::create(outputMem_.ptr(), 0);
261 0 : CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstTmp, srcTmp, stream_));
262 0 : CHK_RET(SubWaitMain());
263 0 : CHK_RET(MainRecordSub());
264 :
265 0 : for (u32 round = 1; round < rankSize; round++) {
266 0 : u32 dstRank = BackwardRank(rank, rankSize, round);
267 0 : Stream& subStream = meshStreams_[round - 1];
268 0 : void* remMemPtr = nullptr;
269 0 : CHK_RET(links[dstRank]->GetRemoteMem(UserMemType::OUTPUT_MEM, &remMemPtr));
270 0 : src = DeviceMem::create(static_cast<char*>(remMemPtr), slices_[dstRank].size);
271 0 : dst = userMemOut.range(slices_[dstRank].offset, slices_[dstRank].size);
272 0 : CHK_RET(HcclD2DMemcpyAsync(
273 : dispatcher_, dst, src, subStream, links[dstRank]->GetRemoteRank(), links[dstRank]->GetLinkType()));
274 0 : CHK_RET(links[dstRank]->TxDataSignal(subStream));
275 0 : CHK_RET(links[dstRank]->RxDataSignal(subStream));
276 : }
277 :
278 0 : CHK_RET(SubRecordMain());
279 0 : CHK_RET(MainWaitSub());
280 :
281 0 : HCCL_INFO("AllGatherMesh finished: rank[%u]", rank);
282 0 : return HCCL_SUCCESS;
283 0 : }
284 : REGISTER_TEMPLATE(TemplateType::TEMPLATE_ALL_REDUCE_MESH_DIRECT, AllReduceMeshDirect);
285 : } // namespace hccl
|
