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_data_trans_wrapper.h"
12 : #include "ins_temp_reduce_mesh_2D.h"
13 :
14 : namespace Hccl {
15 :
16 0 : InsTempReduceMesh2D::InsTempReduceMesh2D(const RankId virtualRank, const u32 tempRankSize,
17 : const std::vector <std::vector<RankId>> &tempVTopo,
18 0 : const std::map <RankId, u32> &tempVirtRankMap)
19 0 : : InsAlgTemplateBase(virtualRank, tempRankSize, tempVTopo, tempVirtRankMap)
20 : {
21 0 : }
22 :
23 0 : InsTempReduceMesh2D::~InsTempReduceMesh2D()
24 : {
25 0 : }
26 :
27 0 : HcclResult InsTempReduceMesh2D::CalcRes(AlgTempResReq &tempResReq)
28 : {
29 0 : HCCL_INFO("[InsTempReduceMesh2D] Calculate communication resources start");
30 :
31 0 : CHK_PRT_RET(tempVTopo_.size() != AXIS_NUM,
32 : HCCL_ERROR("[InsTempReduceMesh2D] The dimension of topo is invalid, expect [%u], now is [%u]",
33 : AXIS_NUM, tempVTopo_.size()), HcclResult::HCCL_E_INTERNAL);
34 :
35 0 : axisRankSize_[AXIS_X] = tempVTopo_.at(AXIS_X).size();
36 0 : axisRankSize_[AXIS_Y] = tempVTopo_.at(AXIS_Y).size();
37 :
38 0 : CHK_PRT_RET(axisRankSize_[AXIS_X] == 0 || axisRankSize_[AXIS_Y] == 0,
39 : HCCL_ERROR("[InsTempReduceMesh2D] The rankSize of dimension is invalid, xRankSize is [%u], yRankSize is [%u]",
40 : axisRankSize_[AXIS_X], axisRankSize_[AXIS_Y]), HcclResult::HCCL_E_INTERNAL);
41 :
42 0 : tempResReq.queNum = axisRankSize_[AXIS_X] + axisRankSize_[AXIS_Y];
43 0 : tempResReq.streamNum = tempResReq.queNum;
44 :
45 0 : tempResReq.queNotifys = CreateNotifiesRequest(axisRankSize_[AXIS_X], axisRankSize_[AXIS_Y]);
46 :
47 0 : CHK_RET(CalcResLinksConcurrMesh(myRank_, tempRankSize_, tempVTopo_, linkNumBtwPeers_, tempResReq));
48 :
49 0 : HCCL_INFO("[InsTempReduceMesh2D] Calculate communication resources finished, queNum[%u], streamNum[%u], "
50 : "queNotifyNum[%u] linkNum[%u]", tempResReq.queNum, tempResReq.streamNum, tempResReq.queNotifys.size(),
51 : tempResReq.links.size());
52 :
53 0 : return HcclResult::HCCL_SUCCESS;
54 : }
55 :
56 0 : std::vector<std::tuple<QId, QId, u32>> InsTempReduceMesh2D::CreateNotifiesRequest(u32 xQueueNum, u32 yQueueNum) const
57 : {
58 0 : std::vector<std::tuple<QId, QId, u32>> notifyRequests;
59 0 : if (xQueueNum == 0) {
60 0 : HCCL_INFO("[InsTempReduceMesh2D] xQueueNum is zero, return empty notifyRequests");
61 0 : return notifyRequests;
62 : }
63 0 : if (yQueueNum == 0) {
64 0 : HCCL_INFO("[InsTempReduceMesh2D] yQueueNum is zero, return empty notifyRequests");
65 0 : return notifyRequests;
66 : };
67 0 : u32 queueNum = xQueueNum + yQueueNum;
68 :
69 0 : u32 slaveNum = queueNum - 1;
70 0 : if (slaveNum < 1) {
71 0 : HCCL_INFO("[InsTempReduceMesh2D] slaveNum is zero, return empty notifyRequests");
72 0 : return notifyRequests;
73 : }
74 :
75 0 : u32 ctrlNotfiyReqNum = 2; // X轴向的控制流(主流)和Y轴向的控制流之间的Notify
76 0 : u32 xNotifyReqNum = (xQueueNum - 1) * 2; // X轴向的控制流和业务流之间的Notify
77 0 : u32 yNotifyReqNum = (yQueueNum - 1) * 2; // Y轴向的控制流和业务流之间的Notify
78 0 : u32 totalNotifyReqNum = ctrlNotfiyReqNum + xNotifyReqNum + yNotifyReqNum;
79 0 : notifyRequests.reserve(totalNotifyReqNum);
80 :
81 0 : QId xCtrlId = 0;
82 0 : QId yCtrlId = xCtrlId + xNotifyReqNum;
83 :
84 0 : notifyRequests.emplace_back(std::make_tuple(xCtrlId, yCtrlId, 0));
85 0 : notifyRequests.emplace_back(std::make_tuple(yCtrlId, xCtrlId, 0));
86 :
87 0 : for (QId xId = xCtrlId + 1; xId < xNotifyReqNum; ++xId) {
88 0 : notifyRequests.emplace_back(std::make_tuple(xCtrlId, xId, 0));
89 0 : notifyRequests.emplace_back(std::make_tuple(xId, xCtrlId, 0));
90 : }
91 :
92 0 : for (QId yId = yCtrlId + 1; yId < xNotifyReqNum + yNotifyReqNum; ++yId) {
93 0 : notifyRequests.emplace_back(std::make_tuple(yCtrlId, yId, 0));
94 0 : notifyRequests.emplace_back(std::make_tuple(yId, yCtrlId, 0));
95 : }
96 :
97 0 : HCCL_DEBUG("[InsTempReduceMesh2D] Create notifies request: "
98 : "totalNotifyReqNum[%u], ctrlNotfiyReqNum[%u], xNotifyReqNum[%u], yNotifyReqNum[%u]",
99 : totalNotifyReqNum, ctrlNotfiyReqNum, xNotifyReqNum, yNotifyReqNum);
100 :
101 0 : return notifyRequests;
102 0 : }
103 :
104 0 : HcclResult InsTempReduceMesh2D::CalcResLinksConcurrMesh(const RankId myRank, const u32 tempRankSize,
105 : const std::vector<std::vector<RankId>> &tempVTopo, const u32 linkNumBtwPeers, AlgTempResReq &tempResReq) const
106 : {
107 : (void)tempRankSize;
108 : u32 myAlgRank;
109 0 : for (u32 dim = 0; dim < tempVTopo.size(); dim++) {
110 0 : CHK_RET(GetAlgRank(myRank, tempVTopo[dim], myAlgRank));
111 0 : for (u32 queIdx = 0; queIdx < tempVTopo[dim].size() - 1; queIdx++) {
112 0 : RankId neighborRank = tempVTopo[dim][(myAlgRank + 1 + queIdx) % (tempVTopo[dim].size())];
113 0 : tempResReq.links[neighborRank] = linkNumBtwPeers;
114 : }
115 : }
116 0 : return HcclResult::HCCL_SUCCESS;
117 : }
118 :
119 0 : u32 InsTempReduceMesh2D::CalcScratchMultiple(BufferType inBuffType, BufferType outBuffType)
120 : {
121 : (void)inBuffType;
122 : (void)outBuffType;
123 :
124 : // 数据会在2个维度间切换通信,选择最大的维度切分scratch方便数据处理
125 0 : u32 scratchMultiple = max(tempVTopo_.at(AXIS_X).size(), tempVTopo_.at(AXIS_Y).size());
126 0 : HCCL_INFO("[InsTempReduceMesh2D] Scratch multiple is [%u]", scratchMultiple);
127 0 : return scratchMultiple;
128 : }
129 :
130 0 : HcclResult InsTempReduceMesh2D::GenExtIns(const TempFuncs &tempFuncs, const TemplateDataParams &templateDataParams,
131 : const ResLinks &tempLinks, std::vector<InsQuePtr> &tempInsQues)
132 : {
133 : (void)tempFuncs;
134 0 : HCCL_INFO("[InsTempReduceMesh2D] GenExtIns start rank[%d]", myRank_);
135 :
136 0 : CHK_RET(CalcParams(templateDataParams));
137 :
138 : // 单卡场景可以直接Input拷贝到Output,单独判断
139 0 : if (tempRankSize_ == 1) {
140 0 : CHK_RET(LocalCopyFromInputToOutput(templateDataParams, tempInsQues));
141 0 : return HcclResult::HCCL_SUCCESS;
142 : }
143 :
144 : // 将队列分为2组,一组负责X轴向通信,一组负责Y轴向通信,两组中的第一条流兼任控制流
145 0 : std::vector<InsQuePtr> ctrlTempInsQues;
146 0 : std::vector<InsQuePtr> xTempInsQues;
147 0 : std::vector<InsQuePtr> yTempInsQues;
148 0 : CHK_RET(SplitInsQues(tempInsQues, ctrlTempInsQues, xTempInsQues, yTempInsQues));
149 :
150 0 : CHK_RET(PreSyncInterQueues(ctrlTempInsQues)); // XY轴并行启动
151 :
152 0 : if (u32(myRank_) == root_) {
153 : // 数据片A第一步通信
154 0 : CHK_RET(GatherFromInput(SLICE_A, AXIS_X, tempLinks, xTempInsQues));
155 0 : CHK_RET(ReduceToScratch(SLICE_A, AXIS_X, xTempInsQues));
156 : // 数据片B第一步通信
157 0 : CHK_RET(GatherFromInput(SLICE_B, AXIS_Y, tempLinks, yTempInsQues));
158 0 : CHK_RET(ReduceToScratch(SLICE_B, AXIS_Y, yTempInsQues));
159 : // X轴和Y轴控制流同步,然后交换处理数据
160 0 : CHK_RET(PreSyncInterQueues(ctrlTempInsQues));
161 0 : CHK_RET(PostSyncInterQueues(ctrlTempInsQues));
162 : // 数据片A第二步通信
163 0 : CHK_RET(GatherFromScratch(SLICE_A, AXIS_Y, tempLinks, yTempInsQues));
164 0 : CHK_RET(ReduceToOutput(SLICE_A, AXIS_Y, yTempInsQues));
165 : // 数据片B第二步通信
166 0 : CHK_RET(GatherFromScratch(SLICE_B, AXIS_X, tempLinks, xTempInsQues));
167 0 : CHK_RET(ReduceToOutput(SLICE_B, AXIS_X, xTempInsQues));
168 0 : } else if (axisRank_[AXIS_X] == axisRoot_[AXIS_X]) {
169 : // 数据片A通信
170 0 : CHK_RET(GatherFromInput(SLICE_A, AXIS_X, tempLinks, xTempInsQues));
171 0 : CHK_RET(ReduceToScratch(SLICE_A, AXIS_X, xTempInsQues));
172 0 : CHK_RET(SendFromScratch(SLICE_A, AXIS_Y, tempLinks, xTempInsQues));
173 : // 数据片B通信
174 0 : CHK_RET(SendFromInput(SLICE_B, AXIS_Y, tempLinks, yTempInsQues));
175 0 : } else if (axisRank_[AXIS_Y] == axisRoot_[AXIS_Y]) {
176 : // 数据片A通信
177 0 : CHK_RET(SendFromInput(SLICE_A, AXIS_X, tempLinks, xTempInsQues));
178 : // 数据片B通信
179 0 : CHK_RET(GatherFromInput(SLICE_B, AXIS_Y, tempLinks, yTempInsQues));
180 0 : CHK_RET(ReduceToScratch(SLICE_B, AXIS_Y, yTempInsQues));
181 0 : CHK_RET(SendFromScratch(SLICE_B, AXIS_X, tempLinks, yTempInsQues));
182 : } else {
183 : // 数据片A通信
184 0 : CHK_RET(SendFromInput(SLICE_A, AXIS_X, tempLinks, xTempInsQues));
185 : // 数据片B通信
186 0 : CHK_RET(SendFromInput(SLICE_B, AXIS_Y, tempLinks, yTempInsQues));
187 : }
188 :
189 0 : CHK_RET(PostSyncInterQueues(ctrlTempInsQues)); // 返回主流
190 :
191 0 : HCCL_INFO("[InsTempReduceMesh2D] GenExtIns finished rank[%d]", myRank_);
192 :
193 0 : return HcclResult::HCCL_SUCCESS;
194 0 : }
195 :
196 0 : HcclResult InsTempReduceMesh2D::CalcParams(const TemplateDataParams &templateDataParams)
197 : {
198 0 : axisRankSize_[AXIS_X] = tempVTopo_.at(AXIS_X).size();
199 0 : axisRankSize_[AXIS_Y] = tempVTopo_.at(AXIS_Y).size();
200 0 : axisRank_[AXIS_X] = u32(myRank_) % axisRankSize_[AXIS_X];
201 0 : axisRank_[AXIS_Y] = u32(myRank_) / axisRankSize_[AXIS_X];
202 0 : axisRoot_[AXIS_X] = root_ % axisRankSize_[AXIS_X];
203 0 : axisRoot_[AXIS_Y] = root_ / axisRankSize_[AXIS_X];
204 :
205 0 : u32 dataTypeSize = DataTypeSizeGet(dataType_);
206 : // 用count均分,防止数据截断;并且保证在奇数情况下SLICE_A的切分比SLICE_B大
207 0 : sliceSize_[SLICE_A] = (templateDataParams.sliceSize / dataTypeSize + 1) / SLICE_NUM * dataTypeSize;
208 0 : sliceSize_[SLICE_B] = templateDataParams.sliceSize - sliceSize_[SLICE_A];
209 :
210 0 : sliceInputBaseOffset_[SLICE_A] = templateDataParams.buffInfo.inBuffBaseOff;
211 0 : sliceInputBaseOffset_[SLICE_B] = sliceInputBaseOffset_[SLICE_A] + sliceSize_[SLICE_A];
212 :
213 0 : sliceOutputBaseOffset_[SLICE_A] = templateDataParams.buffInfo.outBuffBaseOff;
214 0 : sliceOutputBaseOffset_[SLICE_B] = sliceOutputBaseOffset_[SLICE_A] + sliceSize_[SLICE_A];
215 :
216 : // Scratch切分时,上下两部分都按照最大的轴向RankSize来切分,从而保证数据换轴通信时有足够的暂存Buffer来做确定性计算
217 0 : u32 maxAxisRankSize = max(axisRankSize_[AXIS_X], axisRankSize_[AXIS_Y]);
218 0 : sliceScratchBaseOffset_[SLICE_A] = templateDataParams.buffInfo.scratchBuffBaseOff;
219 0 : sliceScratchBaseOffset_[SLICE_B] = sliceScratchBaseOffset_[SLICE_A] + sliceSize_[SLICE_A] * maxAxisRankSize;
220 :
221 0 : return HcclResult::HCCL_SUCCESS;
222 : }
223 :
224 0 : HcclResult InsTempReduceMesh2D::SplitInsQues(std::vector<InsQuePtr> &tempInsQues,
225 : std::vector<InsQuePtr> &ctrlTempInsQues, std::vector<InsQuePtr> &xTempInsQues, std::vector<InsQuePtr> &yTempInsQues)
226 : {
227 0 : u32 expectQueNum = axisRankSize_[AXIS_X] + axisRankSize_[AXIS_Y];
228 0 : CHK_PRT_RET(tempInsQues.size() != expectQueNum,
229 : HCCL_ERROR("[InsTempReduceMesh2D] The count of queues is invalid, expect [%u], now is [%u]",
230 : expectQueNum, tempInsQues.size()), HcclResult::HCCL_E_INTERNAL);
231 :
232 0 : ctrlTempInsQues.emplace_back(tempInsQues.at(0));
233 0 : ctrlTempInsQues.emplace_back(tempInsQues.at(axisRankSize_[AXIS_X]));
234 0 : xTempInsQues = std::vector<InsQuePtr>(tempInsQues.begin(), tempInsQues.begin() + axisRankSize_[AXIS_X]);
235 0 : yTempInsQues = std::vector<InsQuePtr>(tempInsQues.begin() + axisRankSize_[AXIS_X], tempInsQues.end());
236 :
237 0 : HCCL_INFO("[InsTempReduceMesh2D] splitInsQues success, ctrlTempInsQuesNum[%u], xTempInsQuesNum[%u], "
238 : "yTempInsQuesNum[%u]", ctrlTempInsQues.size(), xTempInsQues.size(), yTempInsQues.size());
239 :
240 0 : return HcclResult::HCCL_SUCCESS;
241 : }
242 :
243 0 : HcclResult InsTempReduceMesh2D::LocalCopyFromInputToOutput(const TemplateDataParams &templateDataParams,
244 : std::vector<InsQuePtr> &tempInsQues) const
245 : {
246 0 : DataSlice srcLocalSlice(BufferType::INPUT, 0, templateDataParams.sliceSize);
247 0 : DataSlice dstLocalSlice(BufferType::OUTPUT, 0, templateDataParams.sliceSize);
248 0 : CHK_PRT_RET(LocalCopy(tempInsQues[0], srcLocalSlice, dstLocalSlice),
249 : HCCL_ERROR("[InsTempReduceMesh2D] LocalCopy data failed"),
250 : HcclResult::HCCL_E_INTERNAL);
251 0 : return HcclResult::HCCL_SUCCESS;
252 : }
253 :
254 0 : HcclResult InsTempReduceMesh2D::GatherFromInput(const u32 slice, const u32 axis,
255 : const ResLinks &tempLinks, std::vector<InsQuePtr> &axisTempInsQues)
256 : {
257 0 : HCCL_DEBUG("[InsTempReduceMesh2D] Gather from input start.");
258 :
259 0 : CHK_PRT_RET(axisTempInsQues.empty(),
260 : HCCL_ERROR("[InsTempReduceMesh2D][GatherFromInput] axisTempInsQues is empty."), HcclResult::HCCL_E_INTERNAL);
261 0 : CHK_PTR_NULL(axisTempInsQues[0]);
262 0 : u64 sliceSize = sliceSize_[slice];
263 0 : u64 sliceScratchBaseOffset = sliceScratchBaseOffset_[slice];
264 :
265 0 : DataSlice srcDataSlice(BufferType::INPUT, sliceInputBaseOffset_[slice], sliceSize);
266 :
267 0 : if (axisTempInsQues.size() > 1) {
268 0 : CHK_RET(PreSyncInterQueues(axisTempInsQues));
269 : }
270 :
271 : // 主队列本地拷贝,从Input拷贝到Scratch
272 0 : DataSlice dstLocalSlice(BufferType::SCRATCH, sliceScratchBaseOffset + axisRoot_[axis] * sliceSize, sliceSize);
273 0 : CHK_PRT_RET(LocalCopy(axisTempInsQues[0], srcDataSlice, dstLocalSlice),
274 : HCCL_ERROR("[InsTempReduceMesh2D] LocalCopy data failed"),
275 : HcclResult::HCCL_E_INTERNAL);
276 :
277 : // 从队列负责接收来自其它rank的数据
278 0 : u32 queIdx = 1;
279 0 : for (u32 axisRank = 0; axisRank < tempVTopo_.at(axis).size(); ++axisRank) {
280 0 : RankId rmtRank = tempVTopo_.at(axis).at(axisRank);
281 0 : if (rmtRank == myRank_) {
282 0 : continue;
283 : }
284 :
285 0 : const LinkData &recvLink = tempLinks.at(rmtRank).at(0);
286 : // 按照发送rank的序号来计算接收数据存放的偏移
287 0 : DataSlice dstDataSlice(BufferType::SCRATCH, sliceScratchBaseOffset + axisRank * sliceSize, sliceSize);
288 0 : SlicesList recvSlicesList({srcDataSlice}, {dstDataSlice});
289 0 : DataInfo recvInfo(recvLink, recvSlicesList);
290 0 : CHK_PRT_THROW(queIdx >= axisTempInsQues.size(),
291 : HCCL_ERROR("[InsTempReduceMesh2D] queIdx[%u] is bigger than axisTempInsQues size[%zu].", queIdx,
292 : axisTempInsQues.size()),
293 : InvalidParamsException, "queIdx is invalid");
294 0 : CHK_PRT_RET(Recv(recvInfo, axisTempInsQues[queIdx], 0, true, DmaMode::PUT),
295 : HCCL_ERROR("[InsTempReduceMesh2D] Recv data failed"),
296 : HcclResult::HCCL_E_INTERNAL);
297 :
298 0 : queIdx++;
299 0 : }
300 :
301 0 : if (axisTempInsQues.size() > 1) {
302 0 : CHK_RET(PostSyncInterQueues(axisTempInsQues));
303 : }
304 :
305 0 : return HcclResult::HCCL_SUCCESS;
306 : }
307 :
308 0 : HcclResult InsTempReduceMesh2D::GatherFromScratch(const u32 slice, const u32 axis,
309 : const ResLinks &tempLinks, std::vector<InsQuePtr> &axisTempInsQues)
310 : {
311 0 : HCCL_DEBUG("[InsTempReduceMesh2D] Gather from scratch start");
312 :
313 0 : u64 sliceSize = sliceSize_[slice];
314 0 : u64 sliceScratchBaseOffset = sliceScratchBaseOffset_[slice];
315 :
316 0 : DataSlice srcDataSlice(BufferType::SCRATCH, sliceScratchBaseOffset + axisRoot_[axis] * sliceSize, sliceSize);
317 :
318 0 : if (axisTempInsQues.size() > 1) {
319 0 : CHK_RET(PreSyncInterQueues(axisTempInsQues));
320 : }
321 :
322 : // 主队列本地拷贝,从Scratch拷贝到Output
323 0 : DataSlice dstLocalSlice(BufferType::OUTPUT, sliceOutputBaseOffset_[slice], sliceSize);
324 0 : CHK_PRT_RET(LocalCopy(axisTempInsQues[0], srcDataSlice, dstLocalSlice),
325 : HCCL_ERROR("[InsTempReduceMesh2D] LocalCopy data failed"),
326 : HcclResult::HCCL_E_INTERNAL);
327 :
328 : // 从队列负责接收来自其它rank的数据
329 0 : u32 queIdx = 1;
330 0 : for (u32 axisRank = 0; axisRank < tempVTopo_.at(axis).size(); ++axisRank) {
331 0 : RankId rmtRank = tempVTopo_.at(axis).at(axisRank);
332 0 : if (rmtRank == myRank_) {
333 0 : continue;
334 : }
335 :
336 0 : const LinkData &recvLink = tempLinks.at(rmtRank).at(0);
337 : // 按照发送rank的序号来计算接收数据存放的偏移
338 0 : DataSlice dstDataSlice(BufferType::SCRATCH, sliceScratchBaseOffset + axisRank * sliceSize, sliceSize);
339 0 : SlicesList recvSlicesList({srcDataSlice}, {dstDataSlice});
340 0 : DataInfo recvInfo(recvLink, recvSlicesList);
341 :
342 0 : CHK_PRT_RET(Recv(recvInfo, axisTempInsQues[queIdx], 0, true, DmaMode::PUT),
343 : HCCL_ERROR("[InsTempReduceMesh2D] Recv data failed"),
344 : HcclResult::HCCL_E_INTERNAL);
345 :
346 0 : queIdx++;
347 0 : }
348 :
349 0 : if (axisTempInsQues.size() > 1) {
350 0 : CHK_RET(PostSyncInterQueues(axisTempInsQues));
351 : }
352 :
353 0 : return HcclResult::HCCL_SUCCESS;
354 : }
355 :
356 0 : HcclResult InsTempReduceMesh2D::SendFromInput(const u32 slice, const u32 axis, const ResLinks &tempLinks,
357 : std::vector<InsQuePtr> &axisTempInsQues)
358 : {
359 0 : HCCL_DEBUG("[InsTempReduceMesh2D] Send from input start");
360 :
361 0 : u64 sliceSize = sliceSize_[slice];
362 :
363 0 : RankId rmtRank = tempVTopo_.at(axis).at(axisRoot_[axis]);
364 0 : const LinkData &sendLink = tempLinks.at(rmtRank).at(0);
365 :
366 0 : DataSlice srcDataSlice(BufferType::INPUT, sliceInputBaseOffset_[slice], sliceSize);
367 0 : DataSlice dstDataSlice(BufferType::SCRATCH, sliceScratchBaseOffset_[slice] + axisRank_[axis] * sliceSize, sliceSize);
368 0 : SlicesList sendSlicesList({srcDataSlice}, {dstDataSlice});
369 0 : DataInfo sendInfo(sendLink, sendSlicesList);
370 :
371 0 : CHK_PRT_RET(Send(sendInfo, axisTempInsQues[0], 0, true, DmaMode::PUT),
372 : HCCL_ERROR("[InsTempReduceMesh2D] Send data failed"),
373 : HcclResult::HCCL_E_INTERNAL);
374 :
375 0 : return HcclResult::HCCL_SUCCESS;
376 0 : }
377 :
378 0 : HcclResult InsTempReduceMesh2D::SendFromScratch(const u32 slice, const u32 axis, const ResLinks &tempLinks,
379 : std::vector<InsQuePtr> &axisTempInsQues)
380 : {
381 0 : HCCL_DEBUG("[InsTempReduceMesh2D] Send from scratch start");
382 :
383 0 : u64 sliceSize = sliceSize_[slice];
384 0 : u64 sliceScratchBaseOffset = sliceScratchBaseOffset_[slice];
385 :
386 0 : RankId rmtRank = tempVTopo_.at(axis).at(axisRoot_[axis]);
387 0 : const LinkData &sendLink = tempLinks.at(rmtRank).at(0);
388 :
389 0 : DataSlice srcDataSlice(BufferType::SCRATCH, sliceScratchBaseOffset + axisRoot_[axis] * sliceSize, sliceSize);
390 0 : DataSlice dstDataSlice(BufferType::SCRATCH, sliceScratchBaseOffset + axisRank_[axis] * sliceSize, sliceSize);
391 0 : SlicesList sendSlicesList({srcDataSlice}, {dstDataSlice});
392 0 : DataInfo sendInfo(sendLink, sendSlicesList);
393 :
394 0 : CHK_PRT_RET(Send(sendInfo, axisTempInsQues[0], 0, true, DmaMode::PUT),
395 : HCCL_ERROR("[InsTempReduceMesh2D] Send data failed"),
396 : HcclResult::HCCL_E_INTERNAL);
397 :
398 0 : return HcclResult::HCCL_SUCCESS;
399 0 : }
400 :
401 0 : HcclResult InsTempReduceMesh2D::ReduceToScratch(const u32 slice, const u32 axis, std::vector<InsQuePtr> &axisTempInsQues)
402 : {
403 0 : HCCL_DEBUG("[InsTempReduceMesh2D] Reduce to scratch start");
404 :
405 0 : u64 sliceSize = sliceSize_[slice];
406 0 : u64 sliceScratchBaseOffset = sliceScratchBaseOffset_[slice];
407 :
408 : // 数据规约到scratch时,下一步会交换处理另一片数据,因此数据规约至axisRoot_[1-axis]的偏移位置,便于后续数据搬运
409 0 : DataSlice dstDataSlice(BufferType::SCRATCH, sliceScratchBaseOffset + axisRoot_[1-axis] * sliceSize, sliceSize);
410 :
411 : // 另一轴Root值大于等于当前轴的RankSize时,需要将数据规约至原本无数据的区域,需要先拷贝第一片数据
412 0 : bool needLocalCopy = axisRoot_[1-axis] >= axisRankSize_[axis];
413 0 : if (needLocalCopy) {
414 0 : DataSlice srcLocalSlice(BufferType::SCRATCH, sliceScratchBaseOffset, sliceSize);
415 0 : CHK_PRT_RET(LocalCopy(axisTempInsQues[0], srcLocalSlice, dstDataSlice),
416 : HCCL_ERROR("[InsTempReduceMesh2D] LocalCopy data failed"),
417 : HcclResult::HCCL_E_INTERNAL);
418 :
419 0 : for (u32 sliceId = 1; sliceId < axisRankSize_[axis]; ++sliceId) {
420 0 : DataSlice srcDataSlice(BufferType::SCRATCH, sliceScratchBaseOffset + sliceId * sliceSize, sliceSize);
421 0 : CHK_PRT_RET(LocalReduce(axisTempInsQues[0], srcDataSlice, dstDataSlice, dataType_, redOp_),
422 : HCCL_ERROR("[InsTempReduceMesh2D] Local reduce data failed"),
423 : HcclResult::HCCL_E_INTERNAL);
424 : }
425 :
426 0 : return HcclResult::HCCL_SUCCESS;
427 : }
428 :
429 : // 另一轴Root值小于当前轴RankSize时,按照数据片顺序逐个Reduce
430 0 : for (u32 sliceId = 0; sliceId < axisRankSize_[axis]; ++sliceId) {
431 0 : if (sliceId == axisRoot_[1-axis]) {
432 0 : continue;
433 : }
434 0 : DataSlice srcDataSlice(BufferType::SCRATCH, sliceScratchBaseOffset + sliceId * sliceSize, sliceSize);
435 0 : CHK_PRT_RET(LocalReduce(axisTempInsQues[0], srcDataSlice, dstDataSlice, dataType_, redOp_),
436 : HCCL_ERROR("[InsTempReduceMesh2D] Local reduce data failed"),
437 : HcclResult::HCCL_E_INTERNAL);
438 : }
439 :
440 0 : return HcclResult::HCCL_SUCCESS;
441 : }
442 :
443 0 : HcclResult InsTempReduceMesh2D::ReduceToOutput(const u32 slice, const u32 axis, std::vector<InsQuePtr> &axisTempInsQues)
444 : {
445 0 : HCCL_DEBUG("[InsTempReduceMesh2D] Reduce to output start");
446 :
447 0 : u64 sliceSize = sliceSize_[slice];
448 0 : u64 sliceScratchBaseOffset = sliceScratchBaseOffset_[slice];
449 :
450 0 : DataSlice dstDataSlice(BufferType::OUTPUT, sliceOutputBaseOffset_[slice], sliceSize);
451 :
452 0 : for (u32 sliceId = 0; sliceId < axisRankSize_[axis]; ++sliceId) {
453 0 : if (sliceId == axisRoot_[axis]) { // 跳过轴向root的数据片,这一片已经提前拷贝至Output
454 0 : continue;
455 : }
456 0 : DataSlice srcDataSlice(BufferType::SCRATCH, sliceScratchBaseOffset + sliceId * sliceSize, sliceSize);
457 0 : CHK_PRT_RET(LocalReduce(axisTempInsQues[0], srcDataSlice, dstDataSlice, dataType_, redOp_),
458 : HCCL_ERROR("[InsTempReduceMesh2D] Local reduce data failed"),
459 : HcclResult::HCCL_E_INTERNAL);
460 : }
461 :
462 0 : return HcclResult::HCCL_SUCCESS;
463 : }
464 :
465 : } // namespace Hccl
|