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