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_alltoall.h"
      12              : 
      13            8 : HcclResult AicpuAllToAll::RunAlgorithm(HcclReduceOp opType, void *sendBuffer, void *recvBuffer, u64 dataCount,
      14              :                                        HcclDataType dataType, u64 strideLen, AivAicpuOpParam *nextTask)
      15              : {
      16            8 :     CHK_PTR_NULL(ctx_);
      17            8 :     if (ctx_->useBufferType == MC2_BUFFER_TYPE_WINDOW_IN) {
      18            3 :         return RunAllToAllWinIn(recvBuffer, dataCount, dataType, strideLen);
      19              :     } else {
      20            5 :         return RunAllToAll(sendBuffer, recvBuffer, dataCount, dataType, strideLen);
      21              :     }
      22              : }
      23              : 
      24            3 : HcclResult AicpuAllToAll::RunAllToAllWinIn(void *recvBuffer, u64 dataCount, HcclDataType dataType, u64 strideLen)
      25              : {
      26            3 :     u8 *curOutputPtr = static_cast<u8 *>(recvBuffer);
      27            3 :     u64 averageSize = dataCount * unitSize_; // 单位 byte
      28            3 :     u64 displs[AC_MAX_RANK_NUM] = {0};
      29              : 
      30           27 :     for (u32 i = 0; i < rankNum_; i++) {
      31           24 :         displs[i] = i * strideLen * unitSize_;
      32              :     }
      33            3 :     u64 winOffset = ctx_->winOffset + displs[rankId_];
      34            3 :     HCCL_INFO("dataCount %lu, strideLen %lu, winOffset %lu, unitSize_ %u", dataCount, strideLen, winOffset, unitSize_);
      35              : 
      36              :     // 1. 前同步
      37            3 :     CHK_RET(TaskOrchestrator::DoPreSync());
      38              : 
      39              :     // 2. 片内win->rcv, skipLocalDataCopy为true表示不需要本卡数据。
      40            3 :     if (!ctx_->skipLocalDataCopy) {
      41            1 :         CHK_RET(TaskOrchestrator::SelfCpyWin2Rcv(curOutputPtr, averageSize, winOffset, displs[rankId_],
      42              :                                                 HCCL_REDUCE_RESERVED, dataType));
      43              :     }
      44              : 
      45              :     // 3. 跨片SDMA，其他win->当前rcv buff
      46            3 :     CHK_RET(
      47              :         TaskOrchestrator::IpcCpyWin2Rcv(curOutputPtr, averageSize, winOffset, displs, HCCL_REDUCE_RESERVED, dataType));
      48              : 
      49              :     // 4. 后同步
      50            3 :     CHK_RET(TaskOrchestrator::DoPostSync());
      51              : 
      52            3 :     CHK_RET(TaskOrchestrator::LaunchTasks());
      53              : 
      54            3 :     return HCCL_SUCCESS;
      55              : }
      56              : 
      57            5 : HcclResult AicpuAllToAll::RunAllToAll(void *sendBuffer, void *recvBuffer, u64 dataCount, HcclDataType dataType,
      58              :                                       u64 strideLen)
      59              : {
      60            5 :     u64 windowSize = ctx_->windowSize;
      61            5 :     u64 maxCountPerLoop = windowSize / unitSize_; // 中转内存单次最多能够接受的output count
      62              : 
      63            5 :     u8 *curInputPtr = static_cast<u8 *>(sendBuffer);
      64            5 :     u8 *curOutputPtr = static_cast<u8 *>(recvBuffer);
      65            5 :     u64 inputOffset = 0;
      66            5 :     u64 outputOffset = 0;
      67            5 :     u64 countLeft = dataCount * rankNum_;
      68              : 
      69            5 :     u64 displs[AC_MAX_RANK_NUM] = {0};
      70           45 :     for (u32 i = 0; i < rankNum_; i++) {
      71           40 :         displs[i] = i * strideLen * unitSize_;
      72              :     }
      73            5 :     HCCL_INFO("windowSize %lu, maxCountPerLoop %lu, strideLen %lu, unitSize_ %u", windowSize, maxCountPerLoop,
      74              :               strideLen, unitSize_);
      75              : 
      76           13 :     while (countLeft > 0) {
      77            8 :         curInputPtr += inputOffset;
      78            8 :         curOutputPtr += outputOffset;
      79            8 :         u64 curCount = (countLeft > maxCountPerLoop) ? maxCountPerLoop : countLeft;
      80            8 :         u64 curSize = curCount * unitSize_; // 单位 byte
      81            8 :         u64 averageSize = curSize / rankNum_;
      82            8 :         u64 winOffsets[AC_MAX_RANK_NUM] = {0};
      83           72 :         for (u32 i = 0; i < rankNum_; i++) {
      84           64 :             winOffsets[i] = i * averageSize;
      85              :         }
      86            8 :         HCCL_INFO("inputOffset %lu, outputOffset %lu, countLeft %lu, averageSize %lu", inputOffset,
      87              :                   outputOffset, countLeft, averageSize);
      88              : 
      89              :         // 1. 片内数据拷贝 snd->win
      90           72 :         for (u32 i = 0; i < rankNum_; i++) {
      91           64 :             if ((ctx_->skipLocalDataCopy) && (i == rankId_)) {
      92            2 :                 continue;
      93              :             }
      94              : 
      95           62 :             CHK_RET(TaskOrchestrator::SelfCpySnd2Win(curInputPtr, averageSize, displs[i], winOffsets[i],
      96              :                                                     HCCL_REDUCE_RESERVED, dataType));
      97              :         }
      98              :              
      99              :         // 2. 前同步
     100            8 :         CHK_RET(TaskOrchestrator::DoPreSync());
     101              : 
     102              :         // 3. 片内win->rcv, skipLocalDataCopy为true表示不需要本卡数据。
     103            8 :         if (!ctx_->skipLocalDataCopy) {
     104            6 :             CHK_RET(TaskOrchestrator::SelfCpyWin2Rcv(curOutputPtr, averageSize, winOffsets[rankId_], displs[rankId_],
     105              :                                                     HCCL_REDUCE_RESERVED, dataType));
     106              :         }
     107              : 
     108              :         // 4. 跨片SDMA，其他win->当前rcv buff
     109            8 :         CHK_RET(TaskOrchestrator::IpcCpyWin2Rcv(curOutputPtr, averageSize, winOffsets[rankId_], displs,
     110              :                                                HCCL_REDUCE_RESERVED, dataType));
     111              : 
     112              :         // 5. 后同步
     113            8 :         CHK_RET(TaskOrchestrator::DoPostSync());
     114              : 
     115            8 :         CHK_RET(TaskOrchestrator::LaunchTasks());
     116              : 
     117            8 :         countLeft -= curCount;
     118            8 :         inputOffset = averageSize;
     119            8 :         outputOffset = averageSize;
     120              :     }
     121              : 
     122            5 :     return HCCL_SUCCESS;
     123              : }