@@ -0,0 +1,151 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef AICPU_TS_UBOE_CHANNEL_H +#define AICPU_TS_UBOE_CHANNEL_H + +#include "channel.h" +#include "socket_mgr.h" + +// Orion +#include "../../../../../../legacy/unified_platform/resource/socket/socket.h" +#include "../../../../../../legacy/framework/resource_manager/socket/socket_manager.h" +#include "../../../../../../legacy/unified_platform/pub_inc/buffer_key.h" +#include "rma_connection.h" +#include "ub_mem_transport.h" +#include "dev_ub_connection.h" +#include "ub_local_notify.h" + +namespace hcomm { + +constexpr u32 FINISH_MSG_SIZE = 128; +constexpr char_t FINISH_MSG[FINISH_MSG_SIZE] = "Uboe Comm Pipe ready!"; + +class AicpuTsUboeChannel : public Channel { +public: + AicpuTsUboeChannel(EndpointHandle endpointHandle, const HcommChannelDesc &channelDesc); + + HcclResult Init() override; + HcclResult GetNotifyNum(uint32_t *notifyNum) const override; + HcclResult GetRemoteMem(HcclMem **remoteMem, uint32_t *memNum, char **memTags) override; + ChannelStatus GetStatus() override; + HcclResult GetUserRemoteMem(CommMem **remoteMem, char ***memTag, uint32_t *memNum) override; + + HcclResult H2DResPack(std::vector& buffer); + + virtual HcclResult Clean() override; + virtual HcclResult Resume() override; + +private: + MAKE_ENUM(UboeRmtBufType, NOTIFY, BUFFER) + using RemoteBufferVec = std::vector>; + using LocalBufferVec = std::vector; + + HcclResult Makebufs(HcommMemHandle *memHandles, uint32_t memHandleNum, + std::vector> &bufs); + HcclResult ParseInputParam(); + HcclResult BuildConnection(); + HcclResult BuildNotify(); + HcclResult BuildBuffer(std::vector> &bufs); + HcclResult BuildUbMemTransport(); + HcclResult BuildSocket(); + + std::vector GetUniqueIdV2(); + HcclResult PackOpData(std::vector &data); + + HcclResult FillTagVec(std::vector &bufferVec, + std::vector> &localUserMemTag); + bool IsSocketReady(); + bool IsResReady(); + bool IsConnsReady(); + bool RecvDataProcess(); + void RecvEidDataProcess(); + void SendDataSize(); + void RecvDataSize(); + void SendEidData(); + void SendExchangeData(); + void RecvExchangeData(); + void RecvEidData(); + void SendFinish(); + void RecvFinish(); + + void HandleProcessData(); + void ProcessUboeState(); + + void EidPack(); + void NotifyVecPack(Hccl::BinaryStream &binaryStream); + void BufferVecPack(Hccl::BinaryStream &binaryStream, std::vector &bufferVec, + std::vector> &tagVec); + void ConnVecPack(Hccl::BinaryStream &binaryStream); + void RmtEidUnpackProc(Hccl::IpAddress& rmtAddr); + void RmtBufferVecUnpackProc(u32 locNum, Hccl::BinaryStream &binaryStream, + RemoteBufferVec &bufferVec, UboeRmtBufType type); + bool ConnVecUnpackProc(Hccl::BinaryStream &binaryStream); + void BuildConn(); + + std::vector GetNotifyUniqueIds(); + std::vector GetRmtBufferUniqueIds(RemoteBufferVec &bufferVec, UboeRmtBufType type) const; + std::vector GetLocBufferUniqueIds(LocalBufferVec &bufferVec, UboeRmtBufType type) const; + std::vector GetSingleRmtBufferUniqueId(u64 addr, u64 size, u32 tokenId, u32 tokenValue) const; + std::vector GetConnUniqueIds(); +private: + // --------------------- 入参 --------------------- + EndpointHandle endpointHandle_; + HcommChannelDesc channelDesc_; + + // TODO: 成员变量全部初始化 + // --------------------- 转换参数 --------------------- + EndpointDesc localEp_{}; + EndpointDesc remoteEp_{}; + uint32_t notifyNum_{0}; + std::vector> bufs_{}; + std::vector> bufsTemp{}; // channel 复用时暂存新增 buffer + + // --------------------- 具体成员 --------------------- + Hccl::Socket* socket_{nullptr}; + RdmaHandle rdmaHandle_{nullptr}; + std::unique_ptr memTransport_{nullptr}; + Hccl::BaseMemTransport::CommonLocRes commonRes_{}; + std::vector bufferVecTemp_; // channel 复用时暂存新增 rmaBuffer + std::vector> connections_{}; + std::vector> localRmaBuffers_{}; + std::vector> localNotifies_{}; + std::unique_ptr serverSocket_; + std::unique_ptr socketMgr_{nullptr}; + + ChannelStatus channelStatus{ChannelStatus::INIT}; + MAKE_ENUM(UboeStatus, INIT, SEND_EID, RECV_EID, PROCESS_EID_DATA, BUILD_CONN, SEND_SIZE, RECV_SIZE, SEND_DATA, + RECV_DATA, SEND_FIN, RECV_FIN, PROCESS_DATA, SET_READY, READY) + UboeStatus uboeStatus{UboeStatus::INIT}; + + u32 bufferNum_{0}; + u32 connNum_{0}; + u32 recvDataSize_{0}; + std::vector> localUserMemTag_{}; + std::vector> remoteUserMemTag_{}; + std::vector> rmtMemTagTemp_{}; + + RemoteBufferVec rmtNotifyVec_; // 远端 notify + RemoteBufferVec rmtBufferVec_; // 远端 buffer + LocalBufferVec locBufferVec_; // 本端 buffer + std::vector recvData_{}; + std::vector recvFinishMsg_{}; + std::vector recvEidData_{}; + std::vector sendData_{}; + std::vector sendFinishMsg_{}; + std::vector sendEidData_{}; + bool isRecvFirst_{false}; + + Hccl::IpAddress locAddr_; + Hccl::IpAddress rmtAddr_; +}; + +} // namespace hcomm + +#endif // AICPU_TS_UBOE_CHANNEL_H \ No newline at end of file