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_allgather.h"
12 : #include "common/aicpu_hccl_common.h"
13 :
14 : bool AicpuAllgather::isMCFirstCall = true;
15 :
16 32 : HcclResult AicpuAllgather::RunAlgorithm(HcclReduceOp opType, void *sendBuffer, void *recvBuffer, u64 dataCount,
17 : HcclDataType dataType, u64 strideLen, AivAicpuOpParam *nextTask)
18 : {
19 32 : CHK_PTR_NULL(ctx_);
20 32 : switch (ctx_->commAlg) {
21 24 : case CommAlgType::COMM_ALG_FULL_MESH: {
22 24 : if (ctx_->commLen < ctx_->windowSize / AC_DEFAULT_WINDOW_DIM) {
23 22 : if (nextTask == nullptr ||
24 6 : DataUnitSize(nextTask->hcclDataType) * nextTask->count > ctx_->windowSize / AC_DEFAULT_WINDOW_DIM) {
25 10 : return RunAllGathervMC(opType, sendBuffer, recvBuffer, dataCount, dataType, strideLen, nullptr);
26 : } else {
27 6 : return RunAllGathervMC(opType, sendBuffer, recvBuffer, dataCount, dataType, strideLen, nextTask);
28 : }
29 : }
30 8 : return RunAllGatherv(opType, sendBuffer, recvBuffer, dataCount, dataType, strideLen);
31 : }
32 8 : case CommAlgType::COMM_ALG_DOUBLE_RING: {
33 8 : return RunDoubleRingAllGather(opType, reinterpret_cast<u64>(sendBuffer), reinterpret_cast<u64>(recvBuffer),
34 8 : dataCount, dataType);
35 : }
36 0 : default: {
37 0 : HCCL_ERROR("CommAlg %d is not supported.", ctx_->commAlg);
38 0 : return HCCL_E_NOT_SUPPORT;
39 : }
40 : }
41 : }
42 :
43 8 : HcclResult AicpuAllgather::RunAllGatherv(HcclReduceOp opType, void *sendBuffer, void *recvBuffer, u64 dataCount,
44 : HcclDataType dataType, u64 strideCnt)
45 : {
46 8 : CHK_PTR_NULL(ctx_);
47 :
48 : // dataCount 为 gather前的数据量
49 8 : u64 windowSize = ctx_->windowSize;
50 8 : u32 unitSize = ctx_->unitSize;
51 8 : u64 maxCountPerLoop = windowSize / unitSize; // 中转内存单次最多能够接受的output count
52 :
53 8 : u8 *curInputPtr = static_cast<u8 *>(sendBuffer);
54 8 : u8 *curOutputPtr = static_cast<u8 *>(recvBuffer);
55 8 : u64 inputOffset = 0;
56 8 : u64 outputOffset = 0;
57 8 : u64 countLeft = dataCount;
58 :
59 8 : u64 displs[AC_MAX_RANK_NUM] = {0};
60 72 : for (u32 i = 0; i < ctx_->rankNum; i++) {
61 64 : displs[i] = i * strideCnt * ctx_->unitSize;
62 : }
63 :
64 8 : u32 loopIdx = 0;
65 16 : while (countLeft > 0) {
66 8 : curInputPtr += inputOffset;
67 8 : curOutputPtr += outputOffset;
68 8 : u64 curCount = (countLeft > maxCountPerLoop) ? maxCountPerLoop : countLeft;
69 8 : u64 curSize = curCount * unitSize; // 单位 byte
70 :
71 8 : HCCL_DEBUG("RunAllGatherv: loop %u, curInputPtr[%p], curOutputPtr[%p], curCount[%llu], curSize[%llu]",
72 : loopIdx++, curInputPtr, curOutputPtr, curCount, curSize);
73 :
74 : // 1. 片内数据拷贝 snd->win
75 8 : CHK_RET(TaskOrchestrator::SelfCpySnd2Win(curInputPtr, curSize, 0, 0, HCCL_REDUCE_RESERVED, dataType));
76 : // 2. 前同步
77 8 : CHK_RET(TaskOrchestrator::DoPreSync());
78 :
79 : // 缺省情况allgather时op_type没用，融合时aic设置op_type为1时表示不需要本卡数据。
80 8 : if (opType != HCCL_REDUCE_PROD) {
81 0 : CHK_RET(TaskOrchestrator::SelfCpyWin2Rcv(curOutputPtr, curSize, 0, displs[ctx_->rankId],
82 : HCCL_REDUCE_RESERVED, dataType));
83 : }
84 :
85 : // 3. 跨片SDMA，其他win->当前rcv buff
86 8 : CHK_RET(
87 : TaskOrchestrator::IpcCpyWin2Rcv(curOutputPtr, curSize, nullptr, displs, HCCL_REDUCE_RESERVED, dataType));
88 :
89 : // 4. 后同步
90 8 : CHK_RET(TaskOrchestrator::DoPostSync());
91 :
92 8 : CHK_RET(TaskOrchestrator::LaunchTasks());
93 :
94 8 : countLeft -= curCount;
95 8 : inputOffset = curSize;
96 8 : outputOffset = curSize;
97 : }
98 8 : return HCCL_SUCCESS;
99 : }
100 :
101 22 : u64 AicpuAllgather::GetWindowOffset(u32 curTurnCnt, u64 curSize, u64 strideCnt, u64 recvBuffer)
102 : {
103 : (void) recvBuffer;
104 22 : u64 windowOffset = 0;
105 22 : u64 bufferFlag = static_cast<uint64_t>(curTurnCnt % AC_DEFAULT_WINDOW_DIM);
106 22 : u64 recvOffset = (static_cast<uint64_t>(ctx_->rankId) * strideCnt * static_cast<uint64_t>(ctx_->unitSize)) % HCCL_COPY_ALIGN;
107 22 : u64 windowOffsetTmp = bufferFlag * (ctx_->windowSize / AC_DEFAULT_WINDOW_DIM / HCCL_COPY_ALIGN + 1) *
108 : HCCL_COPY_ALIGN + recvOffset;
109 22 : if ((bufferFlag == 0 && (recvOffset + curSize) < ctx_->windowSize / AC_DEFAULT_WINDOW_DIM) ||
110 13 : (bufferFlag == 1 && (windowOffsetTmp + curSize) < ctx_->windowSize)) {
111 22 : windowOffset = windowOffsetTmp;
112 : } else {
113 0 : windowOffset = bufferFlag * (ctx_->windowSize / AC_DEFAULT_WINDOW_DIM);
114 : }
115 22 : return windowOffset;
116 : }
117 :
118 16 : HcclResult AicpuAllgather::RunAllGathervMC(HcclReduceOp opType, void *sendBuffer, void *recvBuffer, u64 dataCount,
119 : HcclDataType dataType, u64 strideCnt, AivAicpuOpParam *nextTask)
120 : {
121 16 : u64 displs[AC_MAX_RANK_NUM] = {0};
122 144 : for (u32 i = 0; i < ctx_->rankNum; i++) {
123 128 : displs[i] = i * strideCnt * ctx_->unitSize;
124 : }
125 :
126 16 : u64 curSize = ctx_->unitSize * dataCount;
127 16 : u64 windowOffset = 0;
128 16 : if (curSize < ctx_->windowSize / AC_DEFAULT_WINDOW_DIM) {
129 16 : windowOffset = GetWindowOffset(ctx_->curTurnCnt, curSize, strideCnt, reinterpret_cast<uint64_t>(recvBuffer));
130 : }
131 16 : HCCL_INFO("current task: windowSize[%llu], windowOffset[%llu], dataCount[%llu], "
132 : "curSize[%llu], strideCnt[%llu], gatherOut [%#llx]",
133 : ctx_->windowSize, windowOffset, dataCount, curSize, strideCnt, ctx_->gatherOut);
134 :
135 : // 1. 片内数据拷贝 snd->win
136 16 : if (isMCFirstCall) {
137 10 : CHK_RET(TaskOrchestrator::SelfCpySnd2Win(sendBuffer, curSize, 0, windowOffset, HCCL_REDUCE_RESERVED, dataType));
138 10 : isMCFirstCall = false;
139 : }
140 :
141 : // 2. 前同步
142 16 : CHK_RET(TaskOrchestrator::DoPreSync());
143 :
144 : // 缺省情况allgather时op_type没用，融合时aic设置op_type为1时表示不需要本卡数据。
145 16 : if (opType != HCCL_REDUCE_PROD) {
146 0 : CHK_RET(TaskOrchestrator::SelfCpyWin2Rcv(recvBuffer, curSize, windowOffset, displs[ctx_->rankId],
147 : HCCL_REDUCE_RESERVED, dataType));
148 : }
149 :
150 : // 3. 跨片SDMA，其他win->当前rcv buff
151 16 : u64 winOffsets[AC_MAX_RANK_NUM] = {0};
152 144 : for (u32 i = 0; i < ctx_->rankNum; i++) {
153 128 : winOffsets[i] = (ctx_->curTurnCnt % AC_DEFAULT_WINDOW_DIM) *
154 128 : (ctx_->windowSize / AC_DEFAULT_WINDOW_DIM / HCCL_COPY_ALIGN + 1) * HCCL_COPY_ALIGN +
155 128 : (i * strideCnt * ctx_->unitSize) % HCCL_COPY_ALIGN;
156 : }
157 16 : CHK_RET(TaskOrchestrator::IpcCpyWin2Rcv(recvBuffer, curSize, winOffsets, displs, HCCL_REDUCE_RESERVED, dataType));
158 :
159 16 : if (nextTask != nullptr) {
160 : // 1. 片内数据拷贝 snd->win
161 6 : u64 nextSize = DataUnitSize(nextTask->hcclDataType) * nextTask->count;
162 6 : u64 nextWindowOffset = 0;
163 6 : if (curSize < ctx_->windowSize / AC_DEFAULT_WINDOW_DIM) {
164 6 : nextWindowOffset = GetWindowOffset(ctx_->curTurnCnt + 1, nextSize, strideCnt, nextTask->recvBuffer);
165 : }
166 :
167 6 : HCCL_INFO("curTurnCnt[%u], windowSize[%llu], windowOffset[%llu], dataCount[%llu], nextSize[%llu]",
168 : ctx_->curTurnCnt, ctx_->windowSize, nextWindowOffset, nextTask->count, nextSize);
169 :
170 6 : CHK_RET(TaskOrchestrator::SelfCpySnd2Win(reinterpret_cast<void *>(nextTask->sendBuffer), nextSize, 0,
171 : nextWindowOffset, HCCL_REDUCE_RESERVED, nextTask->hcclDataType));
172 : } else {
173 10 : isMCFirstCall = true;
174 : }
175 :
176 : // 4. 后同步
177 16 : CHK_RET(TaskOrchestrator::DoPostSync());
178 :
179 16 : CHK_RET(TaskOrchestrator::LaunchTasks());
180 :
181 16 : return HCCL_SUCCESS;
182 : }
183 :
184 112 : HcclResult AicpuAllgather::GenRingTask(HcclReduceOp opType, u64 sndAddr, u64 rcvAddr, u64 gatherSize,
185 : HcclDataType dataType, uint32_t streamId, bool isClockwise, uint32_t step, bool isWindowLast) const
186 : {
187 112 : const u64 winIn = ctx_->rankInfo[rankId_].window + (isClockwise ? 0U : ctx_->windowSize / RING_NUM);
188 112 : const u64 winOut = ctx_->rankInfo[rankId_].windowOut + (isClockwise ? 0U : ctx_->windowSize / RING_NUM);
189 112 : const uint32_t preRankId = isClockwise ? (rankId_ + rankNum_ - 1U) % rankNum_ : (rankId_ + 1U) % rankNum_;
190 112 : const uint32_t postRankId = isClockwise ? (rankId_ + 1U) % rankNum_ : (rankId_ + rankNum_ - 1U) % rankNum_;
191 112 : const u64 preWinIn = ctx_->rankInfo[preRankId].window + (isClockwise ? 0U : ctx_->windowSize / RING_NUM);
192 112 : const u64 preWinOut = ctx_->rankInfo[preRankId].windowOut + (isClockwise ? 0U : ctx_->windowSize / RING_NUM);
193 112 : const bool evenStep = (step % 2U == 0U); // 2: 环上winIn winOut轮流接收, 奇数轮in->out, 偶数轮out->in
194 :
195 112 : HcclResult ret = HCCL_SUCCESS;
196 112 : if (step == 1U) { // 首轮send->winIn
197 16 : ret = AicpuDispatcher::CopyData(streamId, sndAddr, winIn, gatherSize, dataType, HCCL_REDUCE_RESERVED, rankId_);
198 16 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u send to clock winIn failed", turn_, step), ret);
199 : }
200 : // 片间同步 notify后卡 wait前卡
201 112 : ret = AicpuDispatcher::SignalRecord(streamId, postRankId, AicpuDispatcher::IPC, AicpuDispatcher::PRE_SYNC);
202 112 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u notify post rank failed", turn_, step), ret);
203 112 : ret = AicpuDispatcher::SignalWait(streamId, preRankId, AicpuDispatcher::IPC, AicpuDispatcher::PRE_SYNC);
204 112 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u wait pre rank failed", turn_, step), ret);
205 :
206 : // 片间memcpy 前卡window->recv
207 112 : auto recvBufAddr = rcvAddr + (isClockwise ? (rankId_ + rankNum_ - step) % rankNum_ : (rankId_ + step) % rankNum_) *
208 112 : ctx_->totalCnt * unitSize_;
209 112 : ret = AicpuDispatcher::CopyData(streamId, evenStep ? preWinOut : preWinIn, recvBufAddr, gatherSize,
210 : dataType, HCCL_REDUCE_RESERVED, preRankId);
211 112 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u cpy pre window failed", turn_, step), ret);
212 :
213 112 : if (isClockwise) {
214 56 : if (!((step == rankNum_ - 1U) && isWindowLast)) {
215 48 : CHK_RET(AicpuDispatcher::SignalWait(streamId, (rankId_ + 1U) % rankNum_, AicpuDispatcher::NO_IPC,
216 : AicpuDispatcher::POST_SYNC));
217 : }
218 : // ccore notify
219 56 : if ((step < rankNum_ - 1U) && isWindowLast) {
220 48 : ret = AicpuDispatcher::AddCcoreNotify(streamId, turn_ * (rankNum_ - 1U) + step);
221 48 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u add ccore notify failed", turn_, step), ret);
222 : }
223 : } else {
224 56 : if (!((step == rankNum_ - 1U) && isWindowLast)) {
225 48 : CHK_RET(AicpuDispatcher::SignalRecord(streamId, (rankId_ + 1U) % rankNum_, AicpuDispatcher::NO_IPC,
226 : AicpuDispatcher::POST_SYNC));
227 : }
228 : }
229 : // recv->windowout
230 112 : if (step < rankNum_ - 1U){
231 96 : ret = AicpuDispatcher::CopyData(streamId, recvBufAddr, evenStep ? winIn : winOut, gatherSize, dataType,
232 96 : HCCL_REDUCE_RESERVED, rankId_);
233 96 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u cpy pre window failed", turn_, step), ret);
234 : }
235 :
236 : // 片间同步 notify前卡 wait后卡
237 112 : ret = AicpuDispatcher::SignalRecord(streamId, preRankId, AicpuDispatcher::IPC, AicpuDispatcher::POST_SYNC);
238 112 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u notify pre rank failed", turn_, step), ret);
239 112 : ret = AicpuDispatcher::SignalWait(streamId, postRankId, AicpuDispatcher::IPC, AicpuDispatcher::POST_SYNC);
240 112 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u wait post rank failed", turn_, step), ret);
241 112 : return HCCL_SUCCESS;
242 : }
243 :
244 8 : HcclResult AicpuAllgather::RunDoubleRingAllGather(HcclReduceOp opType, u64 sendBuffer, u64 recvBuffer,
245 : u64 dataCount, HcclDataType dataType) const
246 : {
247 8 : const u64 gatherSize = dataCount * unitSize_;
248 8 : if (gatherSize > ctx_->windowSize / RING_NUM) {
249 0 : HCCL_INFO("Tile gather size %lu max less than window size %lu/RING_NUM", gatherSize, ctx_->windowSize);
250 : }
251 8 : const u64 rankDataSize = ctx_->totalCnt * unitSize_;
252 : thread_local static u64 sndAddr[RING_NUM] = { 0UL };
253 : thread_local static u64 rcvAddr[RING_NUM] = { 0UL };
254 8 : if (turn_ == 0U) { // 首个tile时初始化snd rcv地址
255 1 : sndAddr[0] = sendBuffer;
256 1 : sndAddr[1] = sendBuffer + rankDataSize;
257 1 : rcvAddr[0] = recvBuffer;
258 1 : rcvAddr[1] = recvBuffer + rankDataSize;
259 : }
260 :
261 8 : HCCL_INFO("DR AllGather snd addr %p %p, rcv addr %p %p, size %lu", sndAddr[0], sndAddr[1], rcvAddr[0], rcvAddr[1],
262 : gatherSize);
263 8 : uint32_t mainStream = rankId_;
264 8 : uint32_t subStream = (rankId_ + 1U) % rankNum_;
265 8 : HcclResult ret = HCCL_SUCCESS;
266 :
267 8 : u64 maxCountPerLoop = ctx_->windowSize / RING_NUM / unitSize_; // 中转内存单次最多能够接受的output count
268 8 : u64 countLeft = dataCount;
269 8 : bool isWindowFirst = true;
270 8 : bool isWindowLast = false;
271 : // windowSize循环
272 8 : u32 loopIdx = 0;
273 16 : while (countLeft > 0) {
274 8 : u64 curCount = (countLeft > maxCountPerLoop) ? maxCountPerLoop : countLeft;
275 8 : u64 curSize = curCount * unitSize_; // 单位 byte
276 :
277 8 : HCCL_DEBUG("DR AllGather: loop %u, snd addr[%u %u], rcv addr[%u %u], curCount[%llu], curSize[%llu]",
278 : loopIdx++, sndAddr[0], sndAddr[1], rcvAddr[0], rcvAddr[1], curCount, curSize);
279 :
280 8 : sndAddr[1] -= curSize;
281 8 : rcvAddr[1] -= curSize; // 逆时针输出在执行前向上偏移
282 :
283 8 : isWindowLast = ((countLeft - curCount) == 0);
284 :
285 : // 顺、逆时针各平移ranknum-1次
286 64 : for (uint32_t step = 1U; step <= rankNum_ - 1U; step++) {
287 56 : if (isWindowFirst) {
288 : // ccore wait
289 56 : u64 waitAddr = ctx_->workSpaceAddr + ctx_->notifyOff + offsetof(AivAicpuOpParam, sendCnt);
290 56 : ret = AicpuDispatcher::AddCcoreWait(mainStream, waitAddr, turn_ * (rankNum_ - 1U) + step,
291 56 : (turn_ + 1U >= ctx_->totalTurnCnt) && (step == rankNum_ - 1U));
292 56 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u add ccore wait failed", turn_, step), ret);
293 : }
294 :
295 : // 主->从
296 56 : CHK_RET(
297 : AicpuDispatcher::SignalRecord(mainStream, subStream, AicpuDispatcher::NO_IPC, AicpuDispatcher::PRE_SYNC));
298 56 : CHK_RET(AicpuDispatcher::SignalWait(subStream, subStream, AicpuDispatcher::NO_IPC, AicpuDispatcher::PRE_SYNC));
299 :
300 : // 顺时针环
301 56 : CHK_RET(GenRingTask(opType, sndAddr[0], rcvAddr[0], curSize, dataType, mainStream, true, step, isWindowLast));
302 :
303 : // 逆时针环
304 56 : CHK_RET(GenRingTask(opType, sndAddr[1], rcvAddr[1], curSize, dataType, subStream, false, step, isWindowLast));
305 :
306 : // 缺省情况allgather时op_type没用，融合时aic设置op_type为1时表示不需要本卡数据。
307 56 : if ((step == rankNum_ - 1U) && isWindowLast) { // 拷贝全量本卡直接从send->recv
308 8 : if (opType != HCCL_REDUCE_PROD) {
309 0 : ret = AicpuDispatcher::CopyData(mainStream, sendBuffer, recvBuffer + rankId_ * rankDataSize, rankDataSize,
310 0 : dataType, HCCL_REDUCE_RESERVED, rankId_);
311 0 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u send to recv winIn failed", turn_, step), ret);
312 : }
313 :
314 : // 从->主
315 8 : CHK_RET(
316 : AicpuDispatcher::SignalRecord(subStream, subStream, AicpuDispatcher::NO_IPC, AicpuDispatcher::POST_SYNC));
317 8 : CHK_RET(
318 : AicpuDispatcher::SignalWait(mainStream, subStream, AicpuDispatcher::NO_IPC, AicpuDispatcher::POST_SYNC));
319 :
320 8 : ret = TaskOrchestrator::AddBarrier(mainStream, rankId_, rankNum_);
321 8 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u add barrier failed.", turn_, step), ret);
322 :
323 : // ccore notify 非step7由GenRingTask处理AddCcoreNotify
324 8 : ret = AicpuDispatcher::AddCcoreNotify(mainStream, turn_ * (rankNum_ - 1U) + step);
325 8 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Turn %u step %u add ccore notify failed", turn_, step), ret);
326 : }
327 : }
328 8 : isWindowFirst = false;
329 8 : countLeft -= curCount;
330 8 : sndAddr[0] += curSize;
331 8 : rcvAddr[0] += curSize; // 顺时针输出在执行后向下偏移
332 : }
333 :
334 8 : ret = TaskOrchestrator::LaunchTasks();
335 8 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("Launch tasks failed"), ret);
336 8 : return HCCL_SUCCESS;
337 : }
|
