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              : #ifndef __AICPU_ALLREDUCE_H__
      12              : #define __AICPU_ALLREDUCE_H__
      13              : 
      14              : #include "aicpu_algorithm.h"
      15              : 
      16              : class AicpuAllreduce : public AicpuAlgorithm {
      17              : public:
      18           40 :     explicit AicpuAllreduce(AicpuComContext* ctx) : AicpuAlgorithm(ctx) {}
      19           39 :     ~AicpuAllreduce() override = default;
      20              : 
      21              :     HcclResult RunAlgorithm(
      22              :         HcclReduceOp opType, void* sendBuffer, void* recvBuffer, u64 dataCount, HcclDataType dataType,
      23              :         u64 strideLen = 0, AivAicpuOpParam* nextTask = nullptr) override;
      24              : 
      25              : private:
      26              :     int64_t RoundUpWithDivisor(u64 value, u64 divisor) const;
      27              :     HcclResult PrepareSlice(u64 dataCount, HcclDataType dataType, u32 sliceNum, std::vector<Slice>& dataSlice) const;
      28              :     void GetDataSizes16K(std::vector<u64>& dataSizes, u64 allDataSize) const;
      29              :     void RunAllReduceSlice(
      30              :         u8* curOutputPtr, u8* curInputPtr, u64* sliceSize, u64* dataSlice, HcclReduceOp opType,
      31              :         HcclDataType dataType) const;
      32              :     void RunAllReduceSliceWin2Win(
      33              :         u8* curOutputPtr, u64* sliceSize, u64* dataSlice, HcclReduceOp opType, HcclDataType dataType) const;
      34              :     HcclResult RunAllReduceAlign(
      35              :         HcclReduceOp opType, void* sendBuffer, void* recvBuffer, u64 dataCount, HcclDataType dataType) const;
      36              :     HcclResult
      37              :     RunAllReduceAlignWin2Win(HcclReduceOp opType, void* recvBuffer, u64 dataCount, HcclDataType dataType) const;
      38              :     HcclResult
      39              :     RunAllReduce(HcclReduceOp opType, void* sendBuffer, void* recvBuffer, u64 dataCount, HcclDataType dataType) const;
      40              :     HcclResult RunAllReduceOneShot4Stream(
      41              :         HcclReduceOp opType, void* sendBuffer, void* recvBuffer, u64 dataSize, HcclDataType dataType) const;
      42              :     HcclResult RunAllReduceReduceBcast(
      43              :         HcclReduceOp opType, void* sendBuffer, void* recvBuffer, u64 dataSize, HcclDataType dataType) const;
      44              :     HcclResult RunReduceBcastOnMainSq(
      45              :         u32 mainRankId, u32 maxStreamNum, u32 /* startRank */, u32 /* end_rank */, void* sendBuffer, void* recvBuffer,
      46              :         u64 dataSize, HcclDataType dataType) const;
      47              :     HcclResult RunReduceBcastOnOtherSq(
      48              :         HcclReduceOp opType, u32 mainRankId, u32 maxStreamNum, void* sendBuffer, void* recvBuffer, u64 dataSize,
      49              :         HcclDataType dataType) const;
      50              :     HcclResult RunAllReduceOneShot1Stream(
      51              :         HcclReduceOp opType, void* sendBuffer, void* recvBuffer, u64 dataSize, HcclDataType dataType) const;
      52              :     HcclResult RunAllReduceTwoShot1Stream(
      53              :         HcclReduceOp opType, void* sendBuffer, void* recvBuffer, u64 dataCount, HcclDataType dataType) const;
      54              :     u32 GetHdPeer(const u32 hdRound, const u32 curRank) const;
      55              :     HcclResult RunAllReduceOneshotHD(
      56              :         HcclReduceOp opType, void* sendBuffer, void* recvBuffer, u64 dataSize, HcclDataType dataType) const;
      57              :     HcclResult RunAllReduceRing(
      58              :         HcclReduceOp opType, void* sendBuffer, void* recvBuffer, u64 dataCount, HcclDataType dataType) const;
      59              :     u64 AlignWith(u64 oriValue, u64 alignValue) const;
      60              :     HcclResult GetBurstDataCounts(u64 windowSize, u64 dataCount, std::vector<u64>& burstDataCounts) const;
      61              :     std::vector<std::vector<u32>> GetRingOrders() const;
      62              :     HcclResult reorderRingSlice(
      63              :         const std::vector<u32>& ringOrder, const std::vector<Slice>& ringSlices,
      64              :         std::vector<Slice>& orderedRingSlices) const;
      65              :     HcclResult PrepareRingSlice(
      66              :         const std::vector<std::vector<u32>>& ringOrders, u64 dataCount, HcclDataType dataType,
      67              :         std::vector<std::vector<Slice>>& orderedAllRingSlice) const;
      68              :     HcclResult RingIPCPreSync(const u32 stream, const u32 prevRank, const u32 nextRank) const;
      69              :     HcclResult RingIPCPostSync(const u32 stream, const u32 prevRank, const u32 nextRank) const;
      70              :     HcclResult GetPrevRankList(const std::vector<u32>& ringOrder, std::vector<u32>& previousRankList) const;
      71              :     size_t FindNextRank(const std::vector<u32>& previousRankList, const u32 localRank) const;
      72              :     Slice* GetNextRingSlice(
      73              :         const std::vector<u32>& previousRankList, std::vector<Slice>& orderedRingSlices, u32& curSliceIdx) const;
      74              :     HcclResult RunAllReduceRingSingleBurst(
      75              :         HcclReduceOp opType, void* sendBuffer, void* recvBuffer, u64 dataCount, HcclDataType dataType,
      76              :         std::vector<std::vector<u32>>& ringOrders) const;
      77              :     HcclResult RunAllReduceRingAlg(
      78              :         HcclReduceOp opType, void* sendBuffer, void* recvBuffer, std::vector<Slice>& orderedRingSlices,
      79              :         std::vector<u32>& ringOrder, HcclDataType dataType) const;
      80              : };
      81              : 
      82              : #endif