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