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 : #ifndef UB_MEM_TRANSPORT_H
11 : #define UB_MEM_TRANSPORT_H
12 : #include "base_mem_transport.h"
13 : #include "local_cnt_notify.h"
14 : #include "local_ub_rma_buffer.h"
15 : #include "task_param.h"
16 : #include "virtual_topo.h"
17 : #include "ub_local_notify.h"
18 :
19 : namespace Hccl {
20 :
21 : class UbMemTransport : public BaseMemTransport {
22 : public:
23 : UbMemTransport(
24 : CommonLocRes& commonLocRes, Attribution& attr, const LinkData& linkData, const Socket& socket,
25 : RdmaHandle rdmaHandle1, LocCntNotifyRes& locCntNotifyRes1, bool isRecvFirst);
26 :
27 : UbMemTransport(
28 : CommonLocRes& commonLocRes, Attribution& attr, const LinkData& linkData, const Socket& socket,
29 : RdmaHandle rdmaHandle1, LocCntNotifyRes& locCntNotifyRes1,
30 : std::function<void(u32 streamId, u32 taskId, const TaskParam& taskParam)> callback);
31 :
32 : std::string Describe() const override;
33 : HcclResult Describe(std::string& dfxMsg);
34 :
35 : HcclResult StatusMachine();
36 : TransportStatus GetStatus() override;
37 :
38 : std::vector<char> GetUniqueId() override;
39 :
40 : std::vector<char> GetUniqueIdV2();
41 : std::vector<char> PackConnData();
42 :
43 4 : vector<char>& GetRmtCntNotifyDesc() override // 仅UB 支持
44 : {
45 4 : return rmtCntNotifyDesc;
46 : }
47 :
48 2 : void SetConnVec(std::vector<RmaConnection*>& connectVec) override { commonLocRes.connVec = connectVec; }
49 :
50 : void Post(u32 index, const Stream& stream) override;
51 :
52 : void Wait(u32 index, const Stream& stream, u32 timeout) override;
53 :
54 : void Read(const RmaBufferSlice& locSlice, const RmtRmaBufferSlice& rmtSlice, const Stream& stream) override;
55 :
56 : void ReadReduce(
57 : const RmaBufferSlice& locSlice, const RmtRmaBufferSlice& rmtSlice, const ReduceIn& reduceIn,
58 : const Stream& stream) override;
59 :
60 : void Write(const RmaBufferSlice& locSlice, const RmtRmaBufferSlice& rmtSlice, const Stream& stream) override;
61 :
62 : void WriteReduce(
63 : const RmaBufferSlice& locSlice, const RmtRmaBufferSlice& rmtSlice, const ReduceIn& reduceIn,
64 : const Stream& stream) override;
65 :
66 : void WriteWithNotify(
67 : const RmaBufferSlice& locSlice, const RmtRmaBufferSlice& rmtSlice, const WithNotifyIn& withNotify,
68 : const Stream& stream) override;
69 :
70 : void WriteReduceWithNotify(
71 : const RmaBufferSlice& locSlice, const RmtRmaBufferSlice& rmtSlice, const ReduceIn& reduceIn,
72 : const WithNotifyIn& withNotify, const Stream& stream) override;
73 :
74 : u32 GetCurrentStatus() { return static_cast<u32>(baseStatus); }
75 :
76 : HcclResult GetRemoteMems(uint32_t* memNum, CommMem** remoteMem, char*** memInfos);
77 : HcclResult CheckSocketStatus(std::string socketOpreator);
78 : HcclResult UpdateMemInfo(std::vector<LocalRmaBuffer*>& bufferVecTemp);
79 :
80 : // hostUb使用
81 : HcclResult GetRemoteSeg(const void* addr, u64 len, u64* seg);
82 :
83 : HcclResult Init();
84 : HcclResult DeInit() const;
85 :
86 : private:
87 : RdmaHandle rdmaHandle;
88 :
89 : LocCntNotifyRes locCntNotifyRes;
90 :
91 : MemoryBuffer GetLocMemBuffer(const RmaBufferSlice& locSlice) const;
92 : MemoryBuffer GetRmtMemBuffer(const RmtRmaBufferSlice& rmtSlice) const;
93 : MemoryBuffer GetRmtNotifyMemBuffer(u32 index);
94 : MemoryBuffer GetRmtCntNotifyMemBuffer(const WithNotifyIn& withNotify);
95 :
96 : static constexpr u64 NORMAL_NOTIFY_VAL = 1;
97 :
98 325 : MAKE_ENUM(UbStatus, INIT, RECV_SIZE, SEND_DATA, RECV_DATA, SEND_FIN, RECV_FIN, PROCESS_DATA, SET_READY, READY)
99 : UbStatus ubStatus{UbStatus::INIT};
100 : bool isRecvFirst_{false};
101 :
102 : u32 cntNotifyNum{0};
103 : u32 cntNotifyDescSize{0};
104 : vector<char> rmtCntNotifyDesc;
105 :
106 : using RemoteBufferVec = std::vector<std::unique_ptr<RemoteUbRmaBuffer>>;
107 : using LocalBufferVec = std::vector<LocalUbRmaBuffer*>;
108 :
109 243 : MAKE_ENUM(UbRmtBufType, NOTIFY, BUFFER, CNT_NOTIFY)
110 :
111 : std::mutex remoteMemsMutex_; // 远端内存列表互斥锁
112 : RemoteBufferVec rmtNotifyVec; // 远端普通 notify
113 : RemoteBufferVec rmtBufferVec; // 远端 buffer
114 : RemoteBufferVec rmtCntNotifyVec; // 远端 cnt Notify
115 : LocalBufferVec locBufferVec; // 本端 buffer
116 : bool cacheValid_ = false; // 当前缓存是否有效
117 : std::vector<CommMem> remoteUserMems_; // 内存基本信息缓存
118 : std::vector<std::string> memInfoCopies_; // 储存 Tag 字符串副本
119 : std::vector<char*> memInfoPointers_; // Tag 缓存
120 : std::unique_ptr<Hccl::BaseLocalNotify> drainNotify_; // 本端drain notify
121 : std::unique_ptr<Hccl::LocalUbRmaBuffer> drainBuffer_; // 本端常量1 buffer
122 : std::unique_ptr<Hccl::RemoteUbRmaBuffer> rmtDrainBuffer_; // 对端常量1 buffer
123 :
124 : HcclResult BuildDrainResource();
125 :
126 : HcclResult SendAll();
127 : HcclResult RecvDataSize();
128 : HcclResult SendExchangeData();
129 : HcclResult RecvExchangeData();
130 :
131 : HcclResult SendFinish();
132 : HcclResult RecvFinish();
133 :
134 : void BufferVecPack(BinaryStream& binaryStream, std::vector<LocalRmaBuffer*>& bufferVec);
135 : void CntNotifyVecPack(BinaryStream& binaryStream);
136 :
137 : HcclResult DrainBufPack(BinaryStream& binaryStream);
138 : HcclResult DrainBufUnpack(BinaryStream& binaryStream);
139 :
140 : void CntNotifyDescPack(BinaryStream& binaryStream);
141 : HcclResult CntNotifyDescUnpack(BinaryStream& binaryStream);
142 :
143 : HcclResult
144 : RmtBufferVecUnpackProc(u32 locNum, BinaryStream& binaryStream, RemoteBufferVec& bufferVec, UbRmtBufType type);
145 : HcclResult ConnVecUnpackProc(BinaryStream& binaryStream, bool& needSendFinish);
146 :
147 : void FillRmtRmaBufferVec(RemoteRmaBuffer* rmaBuffer, UbRmtBufType type);
148 :
149 : void SubmitNotify(const MemoryBuffer& rmtNotify, u64 data, const Stream& stream);
150 :
151 : void SubmitWriteEmptyWithNotify(const WithNotifyIn& withNotify, const Stream& stream);
152 :
153 : void SubmitWriteWithNotify(
154 : const MemoryBuffer& rmt, const MemoryBuffer& loc, u64 data, const MemoryBuffer& rmtNotify,
155 : const Stream& stream);
156 :
157 : void SubmitWriteReduceWithNotify(
158 : const MemoryBuffer& rmt, const MemoryBuffer& loc, const ReduceIn& reduceIn, u64 data,
159 : const MemoryBuffer& rmtNotify, const Stream& stream);
160 :
161 : std::vector<char> GetNotifyUniqueIds();
162 : std::vector<char> GetRmtBufferUniqueIds(RemoteBufferVec& bufferVec, UbRmtBufType type) const;
163 : std::vector<char> GetLocBufferUniqueIds(LocalBufferVec& bufferVec, UbRmtBufType type) const;
164 : std::vector<char> GetSingleRmtBufferUniqueId(u64 addr, u64 size, u32 tokenId, u32 tokenValue, u32 notifyId) const;
165 : std::vector<char> GetDrainUniqueIds() const;
166 : std::vector<char> GetConnUniqueIds();
167 :
168 : bool IsResReady();
169 : bool IsConnsReady();
170 : HcclResult RecvDataProcess(bool& needSendFinish);
171 :
172 : HcclResult HandleInitStatus();
173 : HcclResult HandleSendAllStatus();
174 : HcclResult HandleRecvSizeStatus();
175 : HcclResult HandleRecvDataStatus();
176 : HcclResult HandleProcessDataStatus();
177 : HcclResult HandleSendFinStatus();
178 : HcclResult HandleRecvFinStatus();
179 : HcclResult HandleSetReadyStatus();
180 : vector<char> recvData{};
181 : vector<char> recvFinishMsg{};
182 : vector<char> sendData{};
183 : vector<char> sendDataPack_{};
184 : vector<char> sendFinishMsg{};
185 :
186 : void SaveDfxTaskInfo(const TaskParam& taskParam);
187 : };
188 : } // namespace Hccl
189 : #endif
|
