Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2026 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_TS_ROCE_CHANNEL_H
      12              : #define AICPU_TS_ROCE_CHANNEL_H
      13              : 
      14              : #include <cstddef>
      15              : #include <cstdint>
      16              : #include <memory>
      17              : #include <string>
      18              : #include <vector>
      19              : #include "../channel.h"
      20              : #include "aicpu_ts_channel_helper.h"
      21              : #include "hccl_dispatcher_ctx.h"
      22              : #include "hccl_socket.h"
      23              : #include "transport_pub.h"
      24              : 
      25              : struct HcommRoceChannelRes;
      26              : 
      27              : namespace hcomm {
      28              : class AicpuTsRoceEndpoint;
      29              : 
      30              : class AicpuTsRoceChannel : public Channel {
      31              : public:
      32              :     explicit AicpuTsRoceChannel(EndpointHandle endpointHandle, const HcommChannelDesc& channelDesc);
      33              :     ~AicpuTsRoceChannel() override;
      34              : 
      35              :     HcclResult Init() override;
      36              :     HcclResult GetNotifyNum(uint32_t* notifyNum) const override;
      37              :     HcclResult GetRemoteMems(uint32_t* memNum, CommMem** remoteMem, char*** memInfos) override;
      38              :     ChannelStatus GetStatus() override;
      39              :     HcclResult Serialize(std::shared_ptr<hccl::DeviceMem>& out) override;
      40              :     HcommChannelKind GetChannelKind() const override;
      41            0 :     const HcommChannelDesc& GetChannelDesc() const override { return channelDesc_; }
      42              :     HcclResult Clean() override;
      43              :     HcclResult Resume() override;
      44              : 
      45              :     // 数据面接口
      46              :     HcclResult NotifyRecord(const uint32_t remoteNotifyIdx) override;
      47              :     HcclResult NotifyWait(const uint32_t localNotifyIdx, const uint32_t timeout) override;
      48              :     HcclResult WriteWithNotify(void* dst, const void* src, const uint64_t len, uint32_t remoteNotifyIdx) override;
      49              :     HcclResult Write(void* dst, const void* src, uint64_t len) override;
      50              :     HcclResult Read(void* dst, const void* src, uint64_t len) override;
      51              :     HcclResult ChannelFence() override;
      52              : 
      53            0 :     AicpuTsChannelHelper* GetAicpuTsHelper() override { return &aicpuTsHelper_; }
      54              : 
      55              : private:
      56              :     AicpuTsChannelHelper aicpuTsHelper_;
      57              :     /** Owns res / local+remote RoceMemDetails arrays as separate device allocations for AICPU kernel blob. */
      58              :     struct AicpuTsRoceChannelMem {
      59              :         hccl::DeviceMem resAlloc{};
      60              :         hccl::DeviceMem localAlloc{};
      61              :         hccl::DeviceMem remoteAlloc{};
      62              :     };
      63              : 
      64              :     enum class RoceStatus {
      65              :         INIT,              // 初始状态
      66              :         SOCKET_CONNECTING, // Socket正在连接
      67              :         SOCKET_OK,         // Socket建立连接
      68              :         READY,             // Dispatcher + Transport初始化完成
      69              :         FAILED             // 建链失败
      70              :     };
      71              : 
      72              :     HcclResult ParseInputParam();
      73              :     HcclResult BuildDataSocket();
      74              :     HcclResult BuildClientDataSocket(
      75              :         HcclNetDevCtx netDevCtx, const hccl::HcclIpAddress& remoteIp, uint32_t port, const std::string& socketTag);
      76              :     HcclResult BuildServerDataSocket(
      77              :         AicpuTsRoceEndpoint* roceEp, const hccl::HcclIpAddress& remoteIp, uint32_t port, const std::string& socketTag);
      78              :     HcclResult BuildDispatcherAndTransport();
      79              :     HcclResult AssignDispatcherCommId();
      80              :     HcclResult EnsureDispatcherCtx(u32 devPhyId);
      81              :     HcclResult ConfigureMachineParaForTransport();
      82              :     void ConfigureTransportParaForRoce();
      83              :     HcclResult CreateAndInitTransport(HcclDispatcher dispatcher);
      84              :     HcclResult BuildSocketTagName(std::string& outTag) const;
      85              :     HcclResult ValidateSerializeParams(u32 qpNum, size_t localMemCount, size_t remoteMemCount) const;
      86              :     HcclResult InitSerializeRoceChannelRes(
      87              :         HcommRoceChannelRes& res, size_t localMemCount, size_t remoteMemCount, void* localMem, void* remoteMem,
      88              :         const std::vector<HcclQpInfoV2>& aiQpInfos, u32 qpNum) const;
      89              :     HcclResult BuildSerializeChannelMem(
      90              :         AicpuTsRoceChannelMem& bundle, const std::vector<RoceMemDetails>& localMd,
      91              :         const std::vector<RoceMemDetails>& remoteMd, const std::vector<HcclQpInfoV2>& aiQpInfos, u32 qpNum);
      92              : 
      93           75 :     const char* SocketRoleTag() const noexcept { return isLocalIpClient_ ? "client" : "server"; }
      94              :     HcclResult SerializeDrainNotifyInfo(HcommRoceChannelRes& res) const;
      95              : 
      96              :     EndpointHandle endpointHandle_{};
      97              :     HcommChannelDesc channelDesc_{};
      98              : 
      99              :     EndpointDesc localEp_{};
     100              :     EndpointDesc remoteEp_{};
     101              :     bool isLocalIpClient_{false};
     102              :     uint32_t notifyNum_{0};
     103              :     RdmaHandle rdmaHandle_{nullptr};
     104              : 
     105              :     std::shared_ptr<hccl::HcclSocket> dataSocket_{};
     106              :     std::string dispatcherCommId_{};
     107              :     DispatcherCtxPtr dispatcherCtx_{nullptr};
     108              :     bool ownsDispatcherCtx_{false};
     109              : 
     110              :     std::unique_ptr<hccl::NotifyPool> notifyPool_{};
     111              :     hccl::MachinePara machinePara_{};
     112              :     hccl::TransportPara transportPara_{};
     113              :     std::unique_ptr<hccl::Transport> transport_{};
     114              : 
     115              :     bool inited_{false};
     116              :     RoceStatus roceStatus_{RoceStatus::INIT};
     117              : };
     118              : 
     119              : } // namespace hcomm
     120              : 
     121              : #endif // AICPU_TS_ROCE_CHANNEL_H