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