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 "aicpu_reduce_scatter.h"
12 :
13 : namespace {
14 : template<typename T>
15 0 : inline T MathCeil(T num1, T num2) {
16 0 : if (num2 == 0) {
17 0 : return num1;
18 : }
19 0 : return (num1 + num2 - 1) / num2;
20 : }
21 :
22 : template<typename T>
23 0 : inline T AlignUp(T num1, T num2) {
24 0 : return MathCeil(num1, num2) * num2;
25 : }
26 : }
27 :
28 21 : HcclResult AicpuReduceScatter::RunAlgorithm(HcclReduceOp opType, void *sendBuffer, void *recvBuffer, u64 dataCount,
29 : HcclDataType dataType, u64 strideLen, AivAicpuOpParam * /* nextTask */)
30 : {
31 21 : CHK_PTR_NULL(ctx_);
32 : // dataCount 为tile的输入的数据量（scatter前）
33 21 : if (dataCount % rankNum_ != 0) { // 每个tile数据量必须能均分至每张卡
34 2 : HCCL_ERROR("Reduce scatter dataCount %lu max be multiple of rankNum_.", dataCount);
35 2 : return HCCL_E_NOT_SUPPORT;
36 : }
37 :
38 19 : switch (ctx_->commAlg) {
39 15 : case CommAlgType::COMM_ALG_FULL_MESH: {
40 15 : if (ctx_->determinism) {
41 6 : return RunDeterministicReduceScatterLocal(opType, sendBuffer, recvBuffer, dataCount, dataType, strideLen);
42 : }
43 9 : return RunReduceScatterWriteMode(opType, sendBuffer, recvBuffer, dataCount, dataType, strideLen);
44 : }
45 2 : case CommAlgType::COMM_ALG_DOUBLE_RING: {
46 2 : return RunDoubleRingReduceScatter(opType, reinterpret_cast<u64>(sendBuffer),
47 2 : reinterpret_cast<u64>(recvBuffer), dataCount, dataType);
48 : }
49 2 : case CommAlgType::COMM_ALG_SWITCH_WING: {
50 2 : return RunSwitchReduceScatter(opType, reinterpret_cast<u64>(sendBuffer), reinterpret_cast<u64>(recvBuffer),
51 2 : dataCount, dataType);
52 : }
53 0 : default: {
54 0 : HCCL_ERROR("CommAlg %d is not supported.", ctx_->commAlg);
55 0 : return HCCL_E_NOT_SUPPORT;
56 : }
57 : }
58 : }
59 :
60 6 : HcclResult AicpuReduceScatter::RunDeterministicReduceScatterLocal(HcclReduceOp opType, void *sendBuffer,
61 : void *recvBuffer, u64 dataCount, HcclDataType dataType, u64 strideLen)
62 : {
63 6 : u8 *curInputPtr = static_cast<u8 *>(sendBuffer);
64 6 : u8 *curOutputPtr = static_cast<u8 *>(recvBuffer);
65 6 : u64 windowSize = ctx_->windowSize;
66 6 : u64 countLeft = dataCount / rankNum_;
67 6 : u64 maxCountPerLoop = windowSize / unitSize_ / rankNum_;
68 6 : u64 inputOffset = 0;
69 6 : u64 outputOffset = 0;
70 6 : u32 loopIdx = 0;
71 6 : u64 displs[AC_MAX_RANK_NUM] = {0};
72 6 : u64 windowSlices = rankId_ * countLeft * unitSize_;
73 :
74 54 : for (u32 i = 0; i < ctx_->rankNum; i++) {
75 48 : displs[i] = strideLen * i * unitSize_;
76 : }
77 :
78 14 : while (countLeft > 0) {
79 8 : curInputPtr += inputOffset;
80 8 : curOutputPtr += outputOffset;
81 8 : u64 curCount = (countLeft > maxCountPerLoop) ? maxCountPerLoop : countLeft;
82 8 : u64 curSize = curCount * unitSize_; // 单位 byte
83 8 : windowSlices = rankId_ * curSize;
84 :
85 8 : HCCL_DEBUG("RunDeterministicReduceScatterLocal: countLeft = %llu, loop %u, curInputPtr[%p], curOutputPtr[%p],"
86 : "curCount[%llu], curSize[%llu], strideLen[%llu]", countLeft,
87 : loopIdx++, curInputPtr, curOutputPtr, curCount, curSize, strideLen);
88 : // 1. 片内数据拷贝：send->win
89 8 : TaskOrchestrator::SelfCpySnd2Win(curInputPtr, curSize, displs[rankId_], windowSlices, HCCL_REDUCE_RESERVED,
90 : dataType);
91 :
92 : // 2. 前同步
93 8 : TaskOrchestrator::DoPreSync();
94 :
95 : // 3. 跨片SDMA，send->对端win
96 8 : TaskOrchestrator::IpcCpySnd2Win(curInputPtr, curSize, displs, windowSlices,
97 : HCCL_REDUCE_RESERVED, dataType);
98 :
99 : // 4. 后同步
100 8 : TaskOrchestrator::DoPostSync();
101 :
102 : // 5. 折半计算
103 8 : TaskOrchestrator::SelfLocalReduce(curSize, opType, dataType);
104 :
105 : // 6. 片内数据拷贝 本端win->当前rcv buff
106 8 : TaskOrchestrator::SelfCpyWin2Rcv(curOutputPtr, curSize, 0, 0, HCCL_REDUCE_RESERVED,
107 : dataType);
108 :
109 8 : TaskOrchestrator::LaunchTasks();
110 :
111 8 : countLeft -= curCount;
112 8 : inputOffset = curSize;
113 8 : outputOffset = curSize;
114 : }
115 6 : return HCCL_SUCCESS;
116 : }
117 :
118 9 : HcclResult AicpuReduceScatter::RunReduceScatterWriteMode(HcclReduceOp opType, void *sendBuffer, void *recvBuffer,
119 : u64 dataCount, HcclDataType dataType, u64 strideLen)
120 : {
121 9 : dataCount /= rankNum_;
122 :
123 9 : u64 windowSize = ctx_->windowSize;
124 9 : u64 maxCountPerLoop = windowSize / unitSize_; // 中转内存单次最多能够接受的output count
125 :
126 9 : uint8_t *curInputPtr = static_cast<uint8_t *>(sendBuffer);
127 9 : uint8_t *curOutputPtr = static_cast<uint8_t *>(recvBuffer);
128 9 : u64 inputOffset = 0;
129 9 : u64 outputOffset = 0;
130 9 : u64 countLeft = dataCount;
131 :
132 9 : u64 displs[AC_MAX_RANK_NUM] = {0};
133 9 : u64 windowOffsets[AC_MAX_RANK_NUM] = {0};
134 81 : for (u32 i = 0; i < rankNum_; i++) {
135 72 : displs[i] = i * strideLen * unitSize_;
136 72 : windowOffsets[i] = (i * strideLen * unitSize_) % HCCL_COPY_ALIGN;
137 : }
138 :
139 9 : uint32_t loopIdx = 0;
140 17 : while (countLeft > 0) {
141 8 : curInputPtr += inputOffset;
142 8 : curOutputPtr += outputOffset;
143 8 : u64 curCount = (countLeft > maxCountPerLoop) ? maxCountPerLoop : countLeft;
144 8 : u64 curSize = curCount * unitSize_; // 单位 byte
145 :
146 8 : HCCL_DEBUG(
147 : "RunReduceScatterWriteMode: loop %u, curInputPtr[%p], curOutputPtr[%p], curCount[%llu], curSize[%llu], strideLen[%llu]",
148 : loopIdx++, curInputPtr, curOutputPtr, curCount, curSize, strideLen);
149 :
150 : // 1. 片内数据 snd->win
151 8 : CHK_RET(TaskOrchestrator::SelfCpySnd2Win(curInputPtr, curSize, displs[rankId_], windowOffsets[rankId_],
152 : HCCL_REDUCE_RESERVED, dataType)); // 1, 0
153 : // 2. 前同步
154 8 : CHK_RET(TaskOrchestrator::DoPreSync()); // 15 sqe, 35
155 :
156 : // 3. 跨片SDMA send->其他window
157 8 : CHK_RET(TaskOrchestrator::IpcCpySnd2Win(curInputPtr, curSize, displs, windowOffsets, opType,
158 : dataType)); // 0, 7
159 :
160 : // 4. 后同步
161 8 : CHK_RET(TaskOrchestrator::DoPostSync()); // 8, 21
162 :
163 : // 5. 片内数据 拷贝到recv
164 8 : CHK_RET(TaskOrchestrator::SelfCpyWin2Rcv(curOutputPtr, curSize, windowOffsets[rankId_], 0, HCCL_REDUCE_RESERVED,
165 : dataType)); // 1, 0
166 :
167 8 : CHK_RET(TaskOrchestrator::LaunchTasks()); // 25, 63
168 :
169 8 : countLeft -= curCount;
170 8 : inputOffset = curSize;
171 8 : outputOffset = curSize;
172 : }
173 :
174 9 : return HCCL_SUCCESS;
175 : }
176 :
177 0 : HcclResult AicpuReduceScatter::RunReduceScatterReadMode(HcclReduceOp opType, void *sendBuffer, void *recvBuffer,
178 : u64 dataCount, HcclDataType dataType)
179 : {
180 0 : uint32_t dataSize = dataCount * unitSize_; // 输入大小
181 0 : u64 scatterSize = dataSize / rankNum_; // 输出大小
182 :
183 0 : HCCL_DEBUG("RunReduceScatterReadMode: sendBuffer[%p], recvBuffer[%p], dataCount[%llu], dataSize[%llu]", sendBuffer,
184 : recvBuffer, dataCount, dataSize);
185 :
186 : // 1. 片内数据 snd->win 一次性全拷贝
187 0 : CHK_RET(TaskOrchestrator::SelfCpySnd2Win(sendBuffer, dataSize, 0, 0, HCCL_REDUCE_RESERVED, dataType));
188 : // 2. 前同步
189 0 : CHK_RET(TaskOrchestrator::DoPreSync()); // 15 sqe, 35
190 :
191 : // 3. 片内win->recv 实测放在和跨片sdma并行时性能更好
192 0 : CHK_RET(TaskOrchestrator::SelfCpyWin2Rcv(recvBuffer, scatterSize, rankId_ * scatterSize, 0, HCCL_REDUCE_RESERVED,
193 : dataType));
194 :
195 : // 3. 跨片SDMA 其他win->recv
196 0 : u64 winOffsets[AC_MAX_RANK_NUM] = {0};
197 0 : for (size_t i = 0; i < rankNum_; i++) {
198 0 : winOffsets[i] = rankId_ * scatterSize;
199 : }
200 0 : CHK_RET(TaskOrchestrator::IpcCpyWin2Rcv(recvBuffer, scatterSize, winOffsets, nullptr, opType, dataType));
201 :
202 : // 4. 后同步
203 0 : CHK_RET(TaskOrchestrator::DoPostSync()); // 8 sqe, 21
204 :
205 0 : CHK_RET(TaskOrchestrator::LaunchTasks()); // 25, 63
206 :
207 0 : return HCCL_SUCCESS;
208 : }
209 :
210 0 : std::vector<Slice> AicpuReduceScatter::PrepareMeshSlice(u64 dataSize, uint32_t rankNum_)
211 : {
212 0 : std::vector<Slice> meshSlices;
213 0 : uint32_t roundCnt = rankNum_ - 1;
214 :
215 0 : u64 sizePerRound = AlignUp<u64>((dataSize + roundCnt - 1) / roundCnt, HCCL_MIN_SLICE_ALIGN);
216 0 : int64_t rankResidueSize = static_cast<int64_t>(dataSize);
217 :
218 0 : while (rankResidueSize > 0) {
219 0 : Slice singleRoundSlice;
220 0 : singleRoundSlice.offset = dataSize - rankResidueSize;
221 0 : singleRoundSlice.size = std::min<u64>(sizePerRound, rankResidueSize);
222 0 : rankResidueSize -= singleRoundSlice.size;
223 0 : meshSlices.push_back(singleRoundSlice);
224 : }
225 :
226 0 : for (const auto &slice : meshSlices) {
227 0 : HCCL_DEBUG("Slice offset:%lu, size:%lu", slice.offset, slice.size);
228 : }
229 0 : return meshSlices;
230 0 : }
231 :
232 0 : HcclResult AicpuReduceScatter::RunDeterministicReduceScatter(HcclReduceOp opType, void *sendBuffer,
233 : void *recvBuffer, u64 dataCount, HcclDataType dataType, u64 strideCount)
234 : {
235 0 : HCCL_INFO("RunDeterministicReduceScatter strideCount:%u", strideCount);
236 0 : dataCount /= rankNum_;
237 :
238 0 : u64 maxCountPerLoop = ctx_->windowSize / (rankNum_ * unitSize_); // 中转内存单次最多能够接受的output count
239 0 : u64 curCount = 0;
240 0 : uint8_t *curInputPtr = static_cast<uint8_t *>(sendBuffer);
241 0 : uint8_t *curOutputPtr = static_cast<uint8_t *>(recvBuffer);
242 0 : for (u64 countLeft = dataCount, inputOffset = 0, outputOffset = 0; countLeft > 0; countLeft -= curCount) {
243 0 : curInputPtr += inputOffset;
244 0 : curOutputPtr += outputOffset;
245 0 : curCount = ((countLeft * unitSize_ * rankNum_) > ctx_->windowSize) ? maxCountPerLoop : countLeft;
246 0 : u64 curSize = curCount * unitSize_;
247 :
248 0 : CHK_RET(RunCurrentDeterministicReduceScatter(opType, curInputPtr, curOutputPtr, strideCount * unitSize_,
249 : curSize, dataType));
250 :
251 0 : inputOffset = curSize;
252 0 : outputOffset = curSize;
253 : }
254 :
255 0 : return HCCL_SUCCESS;
256 : }
257 :
258 0 : HcclResult AicpuReduceScatter::RunCurrentDeterministicReduceScatter(HcclReduceOp opType, uint8_t *curInputPtr,
259 : uint8_t *curOutputPtr, u64 strideSize, u64 curSize, HcclDataType dataType)
260 : {
261 0 : std::vector<Slice> slices = PrepareMeshSlice(curSize, rankNum_);
262 0 : std::vector<uint32_t> srcRankOrder;
263 0 : std::vector<uint32_t> dstRankOrder;
264 0 : for (uint32_t i = 1; i <= rankNum_ - 1; i++) {
265 0 : srcRankOrder.push_back((rankId_ + rankNum_ - i) % rankNum_);
266 0 : dstRankOrder.push_back((rankId_ + i) % rankNum_);
267 0 : HCCL_DEBUG("SrcRank:%u dstRank:%u", (rankId_ + rankNum_ - i) % rankNum_, (rankId_ + i) % rankNum_);
268 : }
269 : // 1. 片内send->recv
270 0 : AicpuDispatcher::CopyData(rankId_, static_cast<void *>(curInputPtr + rankId_ * strideSize), curOutputPtr, curSize,
271 : dataType, HCCL_REDUCE_RESERVED, rankId_);
272 0 : for (uint32_t round = 0; round < rankNum_ - 1; round++) {
273 : // 2. 片内send->win
274 0 : for (size_t i = 0; i < slices.size(); i++) {
275 0 : uint32_t idx = (round + i) % (rankNum_ - 1);
276 0 : u64 sendOff = dstRankOrder[idx] * strideSize + slices[i].offset;
277 0 : HCCL_DEBUG("Cpy send to win, dstRank:%u, srcOffset:%lu, dstOffset:%lu, size:%lu", dstRankOrder[idx],
278 : sendOff, slices[i].offset, slices[i].size);
279 0 : CHK_RET(TaskOrchestrator::SelfCpySnd2Win(curInputPtr, slices[i].size, sendOff, slices[i].offset,
280 : HCCL_REDUCE_RESERVED, dataType));
281 : }
282 : // 3. 前同步
283 0 : CHK_RET(TaskOrchestrator::DoPreSync());
284 : // 4. 跨片读 inline reduce
285 0 : for (size_t i = 0; i < slices.size(); i++) {
286 0 : uint32_t idx = (round + i) % (rankNum_ - 1);
287 0 : HCCL_DEBUG("Ipc read win to rcv, srcRank:%u, srcOffset:%lu, dstOffset:%lu, size:%lu", srcRankOrder[idx],
288 : slices[i].offset, slices[i].offset, slices[i].size);
289 0 : CHK_RET(TaskOrchestrator::IpcCpyWin2RcvP2P(curOutputPtr, srcRankOrder[idx], slices[i].size,
290 : slices[i].offset, slices[i].offset, opType, dataType));
291 : }
292 : // 5. 后同步
293 0 : CHK_RET(TaskOrchestrator::DoPostSync());
294 : }
295 0 : CHK_RET(TaskOrchestrator::LaunchTasks());
296 0 : return HCCL_SUCCESS;
297 0 : }
298 :
299 28 : HcclResult AicpuReduceScatter::GenRingTask(HcclReduceOp opType, u64 sndAddr, u64 rcvAddr, u64 curSize,
300 : u64 scatterSize, HcclDataType dataType, uint32_t streamId, bool isClockwise, uint32_t step) const
301 : {
302 28 : const u64 winIn = ctx_->rankInfo[rankId_].window + (isClockwise ? 0U : ctx_->windowSize / RING_NUM);
303 28 : const u64 winOut = ctx_->rankInfo[rankId_].windowOut + (isClockwise ? 0U : ctx_->windowSize / RING_NUM);
304 28 : const uint32_t preRankId = isClockwise ? (rankId_ + rankNum_ - 1U) % rankNum_ : (rankId_ + 1) % rankNum_;
305 28 : const uint32_t postRankId = isClockwise ? (rankId_ + 1) % rankNum_ : (rankId_ + rankNum_ - 1U) % rankNum_;
306 28 : const u64 preWinIn = ctx_->rankInfo[preRankId].window + (isClockwise ? 0U : ctx_->windowSize / RING_NUM);
307 28 : const u64 preWinOut = ctx_->rankInfo[preRankId].windowOut + (isClockwise ? 0U : ctx_->windowSize / RING_NUM);
308 28 : const bool evenStep = (step % 2U == 0U); // 2: 环上winIn winOut轮流接收, 奇数轮in->out, 偶数轮out->in
309 :
310 28 : HcclResult ret = HCCL_SUCCESS;
311 28 : if (step == 1U) { // 首轮send->winIn winOut
312 4 : ret = AicpuDispatcher::CopyData(streamId, sndAddr, winIn, curSize, dataType, HCCL_REDUCE_RESERVED, rankId_);
313 4 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u send to clock winIn failed", turn_, step), ret);
314 :
315 8 : ret = AicpuDispatcher::CopyData(streamId, sndAddr + scatterSize, winOut, curSize, dataType,
316 4 : HCCL_REDUCE_RESERVED, rankId_);
317 4 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u send to clock winOut failed", turn_, step), ret);
318 : } else {
319 24 : ret = AicpuDispatcher::CopyData(streamId, sndAddr, evenStep ? winIn : winOut, curSize, dataType,
320 24 : HCCL_REDUCE_RESERVED, rankId_);
321 24 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u send to clock winOut failed", turn_, step), ret);
322 : }
323 : // 片间同步 notify后卡 wait前卡
324 28 : ret = AicpuDispatcher::SignalRecord(streamId, postRankId, AicpuDispatcher::IPC, AicpuDispatcher::PRE_SYNC);
325 28 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u notify post rank failed", turn_, step), ret);
326 :
327 28 : ret = AicpuDispatcher::SignalWait(streamId, preRankId, AicpuDispatcher::IPC, AicpuDispatcher::PRE_SYNC);
328 28 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u wait pre rank failed", turn_, step), ret);
329 :
330 : // 片间memcpy 前卡window->window
331 28 : ret = AicpuDispatcher::CopyData(streamId, evenStep ? preWinOut : preWinIn, evenStep ? winIn : winOut, curSize,
332 : dataType, opType, preRankId);
333 28 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u cpy pre window failed", turn_, step), ret);
334 28 : if (step == rankNum_ - 1U) { // 最后一轮 win输出至recv
335 4 : ret = AicpuDispatcher::CopyData(streamId, evenStep ? winIn : winOut, rcvAddr, curSize, dataType,
336 4 : HCCL_REDUCE_RESERVED, rankId_);
337 4 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u window to recv failed", turn_, step), ret);
338 : }
339 28 : ret = AicpuDispatcher::SignalRecord(streamId, preRankId, AicpuDispatcher::IPC, AicpuDispatcher::POST_SYNC);
340 28 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u notify pre rank failed", turn_, step), ret);
341 : // 片间同步 notify前卡 wait后卡
342 28 : ret = AicpuDispatcher::SignalWait(streamId, postRankId, AicpuDispatcher::IPC, AicpuDispatcher::POST_SYNC);
343 28 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u wait post rank failed", turn_, step), ret);
344 28 : return HCCL_SUCCESS;
345 : }
346 :
347 2 : HcclResult AicpuReduceScatter::RunDoubleRingReduceScatter(HcclReduceOp opType, u64 sendBuffer, u64 recvBuffer,
348 : u64 dataCount, HcclDataType dataType) const
349 : {
350 2 : const u64 scatterSize = dataCount / rankNum_ * unitSize_;
351 2 : if (scatterSize > ctx_->windowSize / RING_NUM) {
352 0 : HCCL_INFO("Tile scatter size %lu max less than window size %lu/RING_NUM", scatterSize, ctx_->windowSize);
353 : }
354 : thread_local static u64 sndAddr[RING_NUM] = { 0UL };
355 : thread_local static u64 rcvAddr[RING_NUM] = { 0UL };
356 2 : if (turn_ == 0U) { // 首个tile时初始化snd rcv地址
357 1 : sndAddr[0] = sendBuffer;
358 1 : rcvAddr[0] = recvBuffer;
359 1 : rcvAddr[1] = recvBuffer + ctx_->totalCnt * unitSize_ / rankNum_;
360 : }
361 2 : sndAddr[1] = sndAddr[0] + 2U * scatterSize + scatterSize; // 2 首轮需要同时输出到winIn和winOut
362 :
363 2 : HCCL_INFO("DR reducescatter snd addr %p %p, rcv addr %p %p, size %lu", sndAddr[0], sndAddr[1], rcvAddr[0],
364 : rcvAddr[1], scatterSize);
365 2 : uint32_t mainStream = rankId_;
366 2 : uint32_t subStream = (rankId_ + 1U) % rankNum_;
367 2 : HcclResult ret = HCCL_SUCCESS;
368 :
369 2 : u64 maxCountPerLoop = ctx_->windowSize / RING_NUM / unitSize_; // 中转内存单次最多能够接受的output count
370 2 : u64 countLeft = dataCount / rankNum_;
371 2 : bool isWindowFirst = true;
372 2 : bool isWindowLast = false;
373 : // windowSize循环
374 2 : u32 loopIdx = 0;
375 4 : while (countLeft > 0) {
376 2 : u64 curCount = (countLeft > maxCountPerLoop) ? maxCountPerLoop : countLeft;
377 2 : u64 curSize = curCount * unitSize_; // 单位 byte
378 :
379 2 : HCCL_DEBUG("DR reducescatter: loop %u, snd addr[%p %p], rcv addr[%p %p], curCount[%llu], curSize[%llu]",
380 : loopIdx++, sndAddr[0], sndAddr[1], rcvAddr[0], rcvAddr[1], curCount, curSize);
381 :
382 2 : sndAddr[1] -= curSize; // 逆时针输入在执行前向上偏移
383 2 : rcvAddr[1] -= curSize; // 逆时针输出在执行前向上偏移
384 :
385 2 : isWindowLast = ((countLeft - curCount) == 0);
386 :
387 2 : u64 stepSndAddrClockwise = sndAddr[0];
388 2 : u64 stepSndAddrAnticlockwise = sndAddr[1];
389 16 : for (uint32_t step = 1U; step <= rankNum_ - 1U; step++) {
390 14 : if (isWindowFirst) {
391 : // ccore wait
392 14 : u64 waitAddr = ctx_->workSpaceAddr + ctx_->notifyOff + offsetof(AivAicpuOpParam, sendCnt);
393 14 : ret = AicpuDispatcher::AddCcoreWait(mainStream, waitAddr, turn_ * (rankNum_ - 1U) + step,
394 14 : (turn_ + 1U >= ctx_->totalTurnCnt) && (step == rankNum_ - 1U));
395 14 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u add ccore wait failed", turn_, step), ret);
396 : }
397 :
398 : // 主->从
399 14 : CHK_RET(
400 : AicpuDispatcher::SignalRecord(mainStream, subStream, AicpuDispatcher::NO_IPC, AicpuDispatcher::PRE_SYNC));
401 14 : CHK_RET(AicpuDispatcher::SignalWait(subStream, subStream, AicpuDispatcher::NO_IPC, AicpuDispatcher::PRE_SYNC));
402 :
403 : // 顺时针环
404 14 : CHK_RET(GenRingTask(opType, stepSndAddrClockwise, rcvAddr[0], curSize, scatterSize, dataType, mainStream, true, step));
405 14 : stepSndAddrClockwise += ((step == 1U) ? 2U * scatterSize * RING_NUM : scatterSize * RING_NUM); // 2 首轮需要同时输出到winIn和winOut
406 :
407 : // 逆时针环
408 14 : CHK_RET(GenRingTask(opType, stepSndAddrAnticlockwise, rcvAddr[1], curSize, scatterSize, dataType, subStream, false, step));
409 14 : stepSndAddrAnticlockwise += ((step == 1U) ? 2U * scatterSize + scatterSize : scatterSize * RING_NUM); // 2 首轮需要同时输出到winIn和winOut
410 :
411 : // 从->主
412 14 : CHK_RET(
413 : AicpuDispatcher::SignalRecord(subStream, subStream, AicpuDispatcher::NO_IPC, AicpuDispatcher::POST_SYNC));
414 14 : CHK_RET(
415 : AicpuDispatcher::SignalWait(mainStream, subStream, AicpuDispatcher::NO_IPC, AicpuDispatcher::POST_SYNC));
416 14 : if (isWindowLast) {
417 14 : ret = TaskOrchestrator::AddBarrier(mainStream, rankId_, rankNum_);
418 14 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u add barrier failed", turn_, step), ret);
419 : // ccore notify
420 14 : ret = AicpuDispatcher::AddCcoreNotify(mainStream, turn_ * (rankNum_ - 1U) + step);
421 14 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u add ccore notify failed", turn_, step), ret);
422 : }
423 : }
424 2 : isWindowFirst = false;
425 2 : countLeft -= curCount;
426 2 : sndAddr[0] += curSize; // 顺时针输入向下偏移
427 2 : rcvAddr[0] += curSize; // 顺时针输出向下偏移
428 : }
429 :
430 : // tile之间回退一个scatterSize再偏移
431 2 : sndAddr[0] = sndAddr[0] - scatterSize + rankNum_ * RING_NUM * scatterSize;
432 :
433 2 : ret = TaskOrchestrator::LaunchTasks();
434 2 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Launch tasks failed"), ret);
435 2 : return HCCL_SUCCESS;
436 : }
437 :
438 2 : HcclResult AicpuReduceScatter::RunSwitchReduceScatter(HcclReduceOp opType, u64 sendBuffer, u64 recvBuffer,
439 : u64 dataCount, HcclDataType dataType) const
440 : {
441 2 : const u64 scatterSize = dataCount / rankNum_ * unitSize_;
442 2 : if (scatterSize > ctx_->windowSize) {
443 0 : HCCL_INFO("Tile scatter size %lu max less than window size %lu", scatterSize, ctx_->windowSize);
444 : }
445 : thread_local static u64 sndAddr = 0UL;
446 : thread_local static u64 rcvAddr = 0UL;
447 2 : if (turn_ == 0U) { // 首个tile时初始化snd rcv地址
448 1 : sndAddr = sendBuffer;
449 1 : rcvAddr = recvBuffer;
450 : }
451 2 : uint32_t mainStream = rankId_;
452 2 : HCCL_INFO("SW reducescatter snd addr %p, rcv addr %p, size %lu", sndAddr, rcvAddr, scatterSize);
453 :
454 2 : HcclResult ret = HCCL_SUCCESS;
455 2 : u64 maxCountPerLoop = ctx_->windowSize / unitSize_; // 中转内存单次最多能够接受的output count
456 2 : u64 countLeft = dataCount / rankNum_;
457 2 : bool isWindowFirst = true;
458 2 : bool isWindowLast = false;
459 : // windowSize循环
460 4 : while (countLeft > 0) {
461 2 : u64 curCount = (countLeft > maxCountPerLoop) ? maxCountPerLoop : countLeft;
462 2 : u64 curSize = curCount * unitSize_; // 单位 byte
463 :
464 2 : isWindowLast = ((countLeft - curCount) == 0);
465 :
466 2 : u64 stepSndAddr = sndAddr;
467 18 : for (uint32_t step = 1U; step <= rankNum_; step++) {
468 16 : uint32_t preRankId = (rankId_ + rankNum_ - step) % rankNum_;
469 16 : uint32_t postRankId = (rankId_ + step) % rankNum_;
470 16 : if (isWindowFirst) {
471 : // ccore wait
472 16 : u64 waitAddr = ctx_->workSpaceAddr + ctx_->notifyOff + offsetof(AivAicpuOpParam, sendCnt);
473 16 : ret = AicpuDispatcher::AddCcoreWait(mainStream, waitAddr, turn_ * rankNum_ + step,
474 16 : (turn_ + 1u >= ctx_->totalTurnCnt) && (step == rankNum_));
475 16 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u add ccore wait failed", turn_, step), ret);
476 : }
477 :
478 16 : if (step != rankNum_) {
479 : // 片间同步 notify前卡 wait后卡
480 14 : ret = AicpuDispatcher::SignalRecord(mainStream, preRankId, AicpuDispatcher::IPC, AicpuDispatcher::PRE_SYNC);
481 14 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u notify pre rank failed", turn_, step), ret);
482 :
483 14 : ret = AicpuDispatcher::SignalWait(mainStream, postRankId, AicpuDispatcher::IPC, AicpuDispatcher::PRE_SYNC);
484 14 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u wait post rank failed", turn_, step), ret);
485 :
486 : // 片间memcpy snd->后序卡win 首轮采用覆盖
487 14 : ret = AicpuDispatcher::CopyData(mainStream, stepSndAddr, ctx_->rankInfo[postRankId].window, curSize,
488 : dataType, step == 1U ? HCCL_REDUCE_RESERVED : opType, postRankId);
489 14 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u cpy snd to post win failed", turn_, step),
490 : ret);
491 14 : stepSndAddr += scatterSize;
492 :
493 : // 片间同步 notify后卡 wait前卡
494 : ret =
495 14 : AicpuDispatcher::SignalRecord(mainStream, postRankId, AicpuDispatcher::IPC, AicpuDispatcher::POST_SYNC);
496 14 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u notify post rank failed", turn_, step), ret);
497 :
498 14 : ret = AicpuDispatcher::SignalWait(mainStream, preRankId, AicpuDispatcher::IPC, AicpuDispatcher::POST_SYNC);
499 14 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u wait pre rank failed", turn_, step), ret);
500 : } else { // 最后一轮仅作本卡内拷贝
501 : // snd->recv
502 4 : ret = AicpuDispatcher::CopyData(mainStream, stepSndAddr, rcvAddr, curSize, dataType,
503 2 : HCCL_REDUCE_RESERVED, rankId_);
504 2 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u cpy snd to rcv failed", turn_, step), ret);
505 2 : stepSndAddr += scatterSize;
506 :
507 : // win->recv
508 4 : ret = AicpuDispatcher::CopyData(mainStream, ctx_->rankInfo[rankId_].window, rcvAddr, curSize, dataType,
509 2 : opType, rankId_);
510 2 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u cpy win to rcv failed", turn_, step), ret);
511 : }
512 16 : if (isWindowLast) {
513 16 : ret = TaskOrchestrator::AddBarrier(mainStream, rankId_, rankNum_);
514 16 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u add barrier failed", turn_, step), ret);
515 :
516 : // ccore notify
517 16 : ret = AicpuDispatcher::AddCcoreNotify(mainStream, turn_ * rankNum_ + step);
518 16 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u add ccore notify failed", turn_, step), ret);
519 : }
520 : }
521 2 : isWindowFirst = false;
522 2 : countLeft -= curCount;
523 2 : sndAddr += curSize;
524 2 : rcvAddr += curSize;
525 : }
526 :
527 : // 回退一个scatterSize再偏移
528 2 : sndAddr = sndAddr - scatterSize + rankNum_ * scatterSize;
529 :
530 2 : ret = TaskOrchestrator::LaunchTasks();
531 2 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Launch tasks failed"), ret);
532 2 : return HCCL_SUCCESS;
533 : }
|
