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_KFC_RPC_SERVER_H
      12              : #define AICPU_KFC_RPC_SERVER_H
      13              : 
      14              : #include "hccl_msg.h"
      15              : #include "common/aicpu_hccl_def.h"
      16              : #include "common/aicpu_kfc_def.h"
      17              : 
      18              : constexpr uint32_t AC_MAX_AIV = 64U; // 最多有64个AIV
      19              : 
      20              : enum RANK_MSG_TYPE { RANK_ADDR = 1, RANK_WORK = 2, RANK_ADD_AND_WORK = 3, RANK_TAIL_TIME = 6, RANK_MSG_END };
      21              : 
      22              : class AicpuKfcRpcServer {
      23              : public:
      24           10 :     AicpuKfcRpcServer() = default;
      25              :     ~AicpuKfcRpcServer() = default;
      26              : 
      27              :     void Init(u64 workSpaceAddr, uint32_t notifyOff, uint16_t notifyBeginCnt, KFCTask* taskParam);
      28              :     void Init(u64 workSpaceAddr);
      29              :     bool RcvMsg(AivAicpuOpParam* rMsg, uint32_t aivID, uint8_t msgType);
      30              :     bool CheckRcvAddrMsg(AivAicpuOpParam* rMsg, uint32_t aivID);
      31              :     bool CheckRcvAddrMsg(HcclApi::HcclMsg* hcclMsg, uint32_t msgPos);
      32              :     bool CheckRcvWorkMsg(AivAicpuOpParam* rMsg, uint32_t aivID, uint32_t curTurnCnt);
      33              :     bool PostMsg(AivAicpuOpParam* rMsg, uint32_t aivID);
      34              :     bool ReadWorkMsg(AivAicpuOpParam* rMsg, uint32_t aivID, uint32_t curTurnCnt);
      35              :     bool ReadAddrMsg(AivAicpuOpParam* rMsg, uint32_t aivID);
      36              :     bool ReadAddrMsg(HcclApi::HcclMsg* hcclMsg, uint32_t msgPos);
      37              :     bool ReadApiValidMsg(HcclApi::HcclMsg* rMsg, HcclApi::HcclMsg* msg, bool reset);
      38              :     bool PostMsg(uint32_t curTurnCnt) const;
      39              :     bool CheckAivIsEnd(uint32_t aivId);
      40              :     bool NeedAutoGenMsg();
      41              :     uint8_t GetWaitPolicy();
      42              :     uint8_t GetTaskType() const;
      43              :     uint8_t GetRspPolicy();
      44              :     uint8_t GetGenTaskNum();
      45              :     TASK_PREPARE_POSITION GetPreparePosition() const;
      46              :     void ClearWorkMsg() const;
      47              :     void WriteFinishWhenAllFinalize(uint32_t msgPos);
      48              :     void HcclMsg2AicAicpuOpParam(CommonHcclMsg* hcclMsg, AivAicpuOpParam* opMsg);
      49              :     void WriteTurnCnt(uint32_t msgPos);
      50              :     void PrintAllHcclMsgArea();
      51              :     void PrintAllHcclMsgAreaData();
      52              :     void PrintMsg(HcclApi::HcclMsg* hcclMsg, uint32_t msgPos);
      53              : 
      54              : private:
      55              :     template <typename T>
      56              :     bool ReadValidMsg(T* rMsg, T* msg, uint8_t msgType, bool reset);
      57              :     bool GenMsgIsLastMsg();
      58              :     std::string GetMsgTypeString(uint8_t msgType);
      59              :     void GenMsgByTaskParam(AivAicpuOpParam* outMsg);
      60              :     u64 GetSendOff() const;
      61              :     u64 GetRecvOff() const;
      62              :     void CalcAllgatherBuffer(AivAicpuOpParam* outMsg) const;
      63              :     void CalcAllreduceBuffer(AivAicpuOpParam* outMsg) const;
      64              :     void CalcReduceScatterBuffer(AivAicpuOpParam* outMsg) const;
      65              :     bool CheckDebugMode(HcclApi::HcclMsg* rMsg);
      66              : 
      67              : private:
      68              :     struct RpcMsgBody {
      69              :         // Rank* aiv * MsgSize * sizeof(消息)
      70              :         AivAicpuOpParam msgRcvArea[AC_MAX_AIV][HcclApi::HCCL_MSG_CNT];
      71              :         AivAicpuOpParam msgSndArea[AC_MAX_AIV][HcclApi::HCCL_MSG_CNT];
      72              :     };
      73              :     RpcMsgBody* msgBody_ = nullptr;
      74              : 
      75              :     HcclApi::HcclMsgArea* hcclMsgArea_ = nullptr;
      76              : 
      77              :     AivAicpuOpParam* msgSndWorkArea_ = nullptr;
      78              :     AivAicpuOpParam* msgRcvRspArea_ = nullptr;
      79              :     // 记录每个消息队列的位置
      80              :     uint32_t rcvMsgPos_[AC_MAX_AIV] = {0}; // 接收队列的初始化位置
      81              :     uint32_t sndMsgPos_[AC_MAX_AIV] = {0};
      82              :     uint32_t aivState_[AC_MAX_AIV] = {0};
      83              : 
      84              :     uint64_t genTaskNum_;
      85              :     KFCTask* genTaskParam_;
      86              :     HcclKFCTilingData* tilingData_;
      87              : };
      88              : 
      89              : #endif // __AICPU_RPC_SERVER_HPP__