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 HCCL_ONE_SIDED_CONN_H
      12              : #define HCCL_ONE_SIDED_CONN_H
      13              : 
      14              : #include <hccl/hccl_types.h>
      15              : #include <hccl/base.h>
      16              : #include "hccl_socket_manager.h"
      17              : #include "hccl_network_pub.h"
      18              : #include <hccl/hccl_one_sided_services.h>
      19              : #include "notify_pool.h"
      20              : #include "transport_mem.h"
      21              : #include "exception_handler.h"
      22              : #include "rma_buffer_mgr.h"
      23              : #include "hccl_mem.h"
      24              : #include "aicpu_operator_pub.h"
      25              : 
      26              : namespace hccl {
      27              : constexpr u32 MAX_REMOTE_MEM_NUM = 256;
      28              : 
      29              : using RemoteRmaBufferMgr = RmaBufferMgr<BufferKey<uintptr_t, u64>, void*>; // (addr, size) handle 
      30              : class HcclOneSidedConn {
      31              : public:
      32              :     struct ProcessInfo {
      33              :         s32 pid;
      34              :         u32 sdid;
      35              :         u32 serverId;
      36              :     };
      37              : 
      38              :     struct RmaMemDesc {
      39              :         u32 localRankId;
      40              :         u32 remoteRankId;
      41              :         char memDesc[TRANSPORT_EMD_ESC_SIZE];
      42              :     };
      43              : 
      44              :     // 参数超过5个，最终交付前完成优化
      45              :     HcclOneSidedConn(const HcclNetDevCtx &netDevCtx, const HcclRankLinkInfo &localRankInfo,
      46              :         const HcclRankLinkInfo &remoteRankInfo, std::unique_ptr<HcclSocketManager> &socketManager,
      47              :         std::unique_ptr<NotifyPool> &notifyPool, const HcclDispatcher &dispatcher, const bool &useRdma, u32 sdid,
      48              :         u32 serverId, u32 trafficClass = HCCL_COMM_TRAFFIC_CLASS_CONFIG_NOT_SET,
      49              :         u32 serviceLevel = HCCL_COMM_SERVICE_LEVEL_CONFIG_NOT_SET, bool aicpuUnfoldMode = false,
      50              :         bool isStandardCard = false, bool isNeedEnableP2P = false);
      51              : 
      52              :     ~HcclOneSidedConn();
      53              : 
      54              :     HcclResult Connect(const std::string &commIdentifier, s32 timeoutSec);
      55              :     HcclResult ExchangeIpcProcessInfo(const ProcessInfo &localProcess, ProcessInfo &remoteProcess);
      56              :     HcclResult ExchangeMemDesc(const HcclMemDescs &localMemDescs, HcclMemDescs &remoteMemDescs, u32 &actualNumOfRemote);
      57              : 
      58              :     void EnableMemAccess(const HcclMemDesc &remoteMemDesc, HcclMem &remoteMem);
      59              :     void DisableMemAccess(const HcclMemDesc &remoteMemDesc);
      60              : 
      61              :     void BatchWrite(const HcclOneSideOpDesc* oneSideDescs, u32 descNum, const rtStream_t& stream);
      62              :     void BatchRead(const HcclOneSideOpDesc* oneSideDescs, u32 descNum, const rtStream_t& stream);
      63              : 
      64              :     HcclResult GetTransInfo(HcclOneSideOpDescParam* descParam, const HcclOneSideOpDesc* desc, u32 descNum,
      65              :         u64 &transportDataAddr, u64 &transportDataSize);
      66              :     HcclResult WaitOpFence(const rtStream_t &stream);
      67              : 
      68              :     HcclResult ConnectWithRemote(const std::string &commIdentifier, ProcessInfo localProcess, s32 timeoutSec);
      69              :     HcclResult GetRemoteProcessInfo(ProcessInfo& remoteProcess);
      70              : 
      71              :     HcclResult ExchangeMemDesc(const HcclMemDescs &localMemDescs);
      72              :     HcclResult EnableMemAccess();
      73              :     HcclResult DisableMemAccess();
      74              :     void CleanSocketResource(const std::string &commIdentifier);
      75              : 
      76              : private:
      77            0 :     std::string RmaMemDescCopyToStr(const RmaMemDesc &rmaMemDesc) const
      78              :     {
      79            0 :         return std::string(rmaMemDesc.memDesc, TRANSPORT_EMD_ESC_SIZE);
      80              :     }
      81              :     HcclResult GetMemType(const char *description, RmaMemType &memType);
      82              :     HcclNetDevCtx netDevCtx_{};
      83              : 
      84              :     const HcclRankLinkInfo &localRankInfo_;
      85              :     HcclRankLinkInfo remoteRankInfo_{};
      86              :     std::unique_ptr<HcclSocketManager> &socketManager_;
      87              : 
      88              :     std::shared_ptr<HcclSocket> socket_{};
      89              :     std::shared_ptr<HcclSocket> rdmaSocket_{};
      90              : 
      91              :     std::unique_ptr<NotifyPool> &notifyPool_;
      92              : 
      93              :     std::shared_ptr<TransportMem> transportMemPtr_{};
      94              : 
      95              :     RemoteRmaBufferMgr remoteRmaBufferMgr_{};
      96              :     std::unordered_map <std::string, HcclBuf> memDescMap_;
      97              :     bool useRdma_{true};
      98              : 
      99              :     ProcessInfo remoteProcess_{};
     100              :     std::vector<TransportMem::RmaMemDesc> remoteMemDescsVec_{};
     101              :     u32 actualNumOfRemote_;
     102              : 
     103              :     bool aicpuUnfoldMode_{false};
     104              :     TransportDeviceNormalData transportData_;
     105              :     DeviceMem transportDataDevice_;
     106              :     bool isStandardCard_{false};
     107              :     bool isNeedEnableP2P_{false};
     108              :     std::vector<u32> enableP2PDevices_;
     109              : };
     110              : }
     111              : #endif