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