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 TRANSPORT_MEM_H
      12              : #define TRANSPORT_MEM_H
      13              : 
      14              : #include <hccl/hccl_types.h>
      15              : #include <atomic>
      16              : #include "dispatcher.h"
      17              : #include "notify_pool.h"
      18              : #include "hccl_socket.h"
      19              : #include "hccl_network_pub.h"
      20              : #include "hccl_common.h"
      21              : #include "hccl_mem.h"
      22              : #include "transport_pub.h"
      23              : 
      24              : namespace hccl {
      25              : 
      26              : enum class RmaMemType : int {
      27              :     DEVICE = 0,  // device侧内存
      28              :     HOST = 1,    // host侧内存
      29              :     TYPE_NUM
      30              : };
      31              : 
      32              : constexpr size_t TRANSPORT_EMD_ESC_SIZE = 512U - (sizeof(u32) * 2);
      33              : 
      34              : class TransportMem {
      35              : public:
      36              :     enum class TpType : int {
      37              :         IPC = 0,
      38              :         ROCE = 1,
      39              :         ROCE_DEVICE,
      40              :         TYPE_NUM
      41              :     };
      42              : 
      43              :     struct AttrInfo {
      44              :         u32 localRankId{INVALID_VALUE_RANKID};
      45              :         u32 remoteRankId{INVALID_VALUE_RANKID};
      46              :         u32 sdid{INVALID_UINT};      // 本端所属超节点
      47              :         u32 serverId{INVALID_UINT};  // 本端所属server
      48              :         u32 trafficClass{HCCL_COMM_TRAFFIC_CLASS_CONFIG_NOT_SET};
      49              :         u32 serviceLevel{HCCL_COMM_SERVICE_LEVEL_CONFIG_NOT_SET};
      50              :         u32 timeout{INVALID_UINT};  // 传输超时时间
      51              :     };
      52              : 
      53              :     struct RmaMemDesc {
      54              :         u32 localRankId;
      55              :         u32 remoteRankId;
      56              :         char memDesc[TRANSPORT_EMD_ESC_SIZE];
      57              :     };
      58              : 
      59              :     struct RmaMemDescs {
      60              :         RmaMemDesc *array;
      61              :         u32 arrayLength;
      62              :     };
      63              : 
      64              :     struct RmaOpMem {
      65              :         void *addr;
      66              :         u64 size;
      67              :     };
      68              : 
      69              :     struct RmaMem {
      70              :         RmaMemType type;  // segment的内存类型
      71              :         void *addr;       // segment的虚拟地址
      72              :         u64 size;         // segment的size
      73              :     };
      74              : 
      75              :     static std::shared_ptr<TransportMem> Create(TpType tpType,
      76              :         const std::unique_ptr<NotifyPool> &notifyPool, const HcclNetDevCtx &netDevCtx, const HcclDispatcher &dispatcher,
      77              :         AttrInfo &attrInfo);
      78              :     static std::shared_ptr<TransportMem> Create(TpType tpType, const std::unique_ptr<NotifyPool> &notifyPool,
      79              :         const HcclNetDevCtx &netDevCtx, const HcclDispatcher &dispatcher, AttrInfo &attrInfo,
      80              :         bool aicpuUnfoldMode);
      81              :     // AICPU侧创建Transport
      82              :     static std::shared_ptr<TransportMem> Create(TpType tpType, const HcclQpInfoV2 &qpInfo,
      83              :         const HcclDispatcher &dispatcher, AttrInfo &attrInfo);
      84              : 
      85              :     explicit TransportMem(const std::unique_ptr<NotifyPool> &notifyPool, const HcclNetDevCtx &netDevCtx,
      86              :         const HcclDispatcher &dispatcher, AttrInfo &attrInfo);
      87              :     TransportMem(const std::unique_ptr<NotifyPool> &notifyPool, const HcclNetDevCtx &netDevCtx,
      88              :         const HcclDispatcher &dispatcher, AttrInfo &attrInfo, bool aicpuUnfoldMode);
      89              :     virtual ~TransportMem();
      90              :     virtual HcclResult ExchangeMemDesc(
      91              :         const RmaMemDescs &localMemDescs, RmaMemDescs &remoteMemDescs, u32 &actualNumOfRemote) = 0;
      92              :     virtual HcclResult EnableMemAccess(const RmaMemDesc &remoteMemDesc, RmaMem &remoteMem) = 0;
      93              :     virtual HcclResult DisableMemAccess(const RmaMemDesc &remoteMemDesc) = 0;
      94              :     virtual HcclResult SetDataSocket(const std::shared_ptr<HcclSocket> &socket);
      95              : 
      96              :     virtual HcclResult SetSocket(const std::shared_ptr<HcclSocket> &socket) = 0;
      97              :     virtual HcclResult Connect(s32 timeoutSec) = 0;
      98              :     virtual HcclResult Write(const HcclBuf &remoteMem, const HcclBuf &localMem, const rtStream_t &stream) = 0;
      99              :     virtual HcclResult Read(const HcclBuf &localMem, const HcclBuf &remoteMem, const rtStream_t &stream) = 0;
     100              :     /**
     101              :     * @brief 旧版Write
     102              :     * @deprecated 参数优化，改用 `Write(const HcclBuf &remoteMem, const HcclBuf &localMem, const rtStream_t &stream)`。
     103              :     */
     104              :     virtual HcclResult Write(const RmaOpMem &remoteMem, const RmaOpMem &localMem, const rtStream_t &stream) = 0;
     105              :     /**
     106              :     * @brief 旧版Read
     107              :     * @deprecated 参数优化，改用 `Read(const HcclBuf &localMem, const HcclBuf &remoteMem, const rtStream_t &stream)`。
     108              :     */
     109              :     virtual HcclResult Read(const RmaOpMem &localMem, const RmaOpMem &remoteMem, const rtStream_t &stream) = 0;
     110              :     virtual HcclResult AddOpFence(const rtStream_t &stream) = 0;
     111              : 
     112              :     virtual HcclResult GetTransInfo(HcclQpInfoV2 &qpInfo, u32 *lkey, u32 *rkey, HcclBuf *localMem, HcclBuf *remoteMem,
     113              :         u32 num) = 0;
     114              :     virtual HcclResult WaitOpFence(const rtStream_t &stream) = 0;
     115              : 
     116              :     // AICPU侧批量下发读、写操作，下发wr后敲Doorbell
     117              :     virtual HcclResult BatchWrite(const std::vector<MemDetails> &remoteMems, const std::vector<MemDetails> &localMems,
     118              :         Stream &stream) = 0;
     119              :     virtual HcclResult BatchRead(const std::vector<MemDetails> &localMems, const std::vector<MemDetails> &remoteMems,
     120              :         Stream &stream) = 0;
     121              :     virtual HcclResult AddOpFence(const MemDetails &localFenceMem, const MemDetails &remoteFenceMem,
     122              :         Stream &stream) = 0;
     123              : 
     124              : protected:
     125              :     // 从 string 拷贝到 memDesc
     126              :     HcclResult RmaMemDescCopyFromStr(RmaMemDesc &rmaMemDesc, const std::string &memDescStr) const
     127              :     {
     128              :         if (memcpy_s(rmaMemDesc.memDesc, TRANSPORT_EMD_ESC_SIZE, memDescStr.c_str(), memDescStr.size() + 1) != EOK) {
     129              :             return HCCL_E_INTERNAL;
     130              :         }
     131              :         return HCCL_SUCCESS;
     132              :     }
     133              : 
     134              :     // 从 memDesc 转换为 string
     135            0 :     std::string RmaMemDescCopyToStr(const RmaMemDesc &rmaMemDesc) const
     136              :     {
     137            0 :         return std::string(rmaMemDesc.memDesc, TRANSPORT_EMD_ESC_SIZE);
     138              :     }
     139              : 
     140              :     HcclResult DoExchangeMemDesc(const RmaMemDescs &localMemDescs, RmaMemDescs &remoteMemDescs, u32 &actualNumOfRemote);
     141              :     HcclResult SendLocalMemDesc(const RmaMemDescs &localMemDescs);
     142              :     HcclResult ReceiveRemoteMemDesc(RmaMemDescs &remoteMemDescs, u32 &actualNumOfRemote);
     143              : 
     144              :     const std::unique_ptr<NotifyPool> &notifyPool_;
     145              :     HcclNetDevCtx netDevCtx_{nullptr};
     146              :     HcclDispatcher dispatcher_{nullptr};
     147              : 
     148              :     u32 localRankId_{0};
     149              :     u32 remoteRankId_{0};
     150              :     std::shared_ptr<HcclSocket> socket_{nullptr};
     151              : 
     152              :     std::shared_ptr<HcclSocket> dataSocket_{nullptr};
     153              : 
     154              :     bool aicpuUnfoldMode_{false};
     155              : };
     156              : }  // namespace hccl
     157              : #endif