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(CommonLocRes &commonLocRes, Attribution &attr, const LinkData &linkData, const Socket &socket,
24 : RdmaHandle rdmaHandle1, LocCntNotifyRes &locCntNotifyRes1, bool isRecvFirst);
25 :
26 : UbMemTransport(CommonLocRes &commonLocRes, Attribution &attr, const LinkData &linkData, const Socket &socket,
27 : RdmaHandle rdmaHandle1, LocCntNotifyRes &locCntNotifyRes1,
28 : std::function<void(u32 streamId, u32 taskId, const TaskParam &taskParam)> callback);
29 :
30 : std::string Describe() const override;
31 : HcclResult Describe(std::string &dfxMsg);
32 :
33 : HcclResult StatusMachine();
34 : TransportStatus GetStatus() override;
35 :
36 : std::vector<char> GetUniqueId() override;
37 :
38 : std::vector<char> GetUniqueIdV2();
39 : std::vector<char> PackConnData();
40 :
41 4 : vector<char> &GetRmtCntNotifyDesc() override // 仅UB 支持
42 : {
43 4 : return rmtCntNotifyDesc;
44 : }
45 :
46 2 : void SetConnVec(std::vector<RmaConnection *> &connectVec) override
47 : {
48 2 : commonLocRes.connVec = connectVec;
49 2 : }
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(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(const RmaBufferSlice &locSlice, const RmtRmaBufferSlice &rmtSlice, const ReduceIn &reduceIn,
63 : const Stream &stream) override;
64 :
65 : void WriteWithNotify(const RmaBufferSlice &locSlice, const RmtRmaBufferSlice &rmtSlice,
66 : const WithNotifyIn &withNotify, const Stream &stream) override;
67 :
68 : void WriteReduceWithNotify(const RmaBufferSlice &locSlice, const RmtRmaBufferSlice &rmtSlice,
69 : const ReduceIn &reduceIn, const WithNotifyIn &withNotify, const Stream &stream) override;
70 :
71 : u32 GetCurrentStatus()
72 : {
73 : return static_cast<u32>(baseStatus);
74 : }
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 321 : 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 227 : 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 RmtBufferVecUnpackProc(u32 locNum, BinaryStream &binaryStream, RemoteBufferVec &bufferVec, UbRmtBufType type);
144 : HcclResult ConnVecUnpackProc(BinaryStream &binaryStream, bool &needSendFinish);
145 :
146 : void FillRmtRmaBufferVec(RemoteRmaBuffer *rmaBuffer, UbRmtBufType type);
147 :
148 : void SubmitNotify(const MemoryBuffer &rmtNotify, u64 data, const Stream &stream);
149 :
150 : void SubmitWriteEmptyWithNotify(const WithNotifyIn &withNotify, const Stream &stream);
151 :
152 : void SubmitWriteWithNotify(const MemoryBuffer &rmt, const MemoryBuffer &loc, u64 data,
153 : const MemoryBuffer &rmtNotify, const Stream &stream);
154 :
155 : void SubmitWriteReduceWithNotify(const MemoryBuffer &rmt, const MemoryBuffer &loc, const ReduceIn &reduceIn,
156 : u64 data, const MemoryBuffer &rmtNotify, const Stream &stream);
157 :
158 : std::vector<char> GetNotifyUniqueIds();
159 : std::vector<char> GetRmtBufferUniqueIds(RemoteBufferVec &bufferVec, UbRmtBufType type) const;
160 : std::vector<char> GetLocBufferUniqueIds(LocalBufferVec &bufferVec, UbRmtBufType type) const;
161 : std::vector<char> GetSingleRmtBufferUniqueId(u64 addr, u64 size, u32 tokenId, u32 tokenValue, u32 notifyId) const;
162 : std::vector<char> GetDrainUniqueIds() const;
163 : std::vector<char> GetConnUniqueIds();
164 :
165 : bool IsResReady();
166 : bool IsConnsReady();
167 : HcclResult RecvDataProcess(bool &needSendFinish);
168 :
169 : HcclResult HandleInitStatus();
170 : HcclResult HandleSendAllStatus();
171 : HcclResult HandleRecvSizeStatus();
172 : HcclResult HandleRecvDataStatus();
173 : HcclResult HandleProcessDataStatus();
174 : HcclResult HandleSendFinStatus();
175 : HcclResult HandleRecvFinStatus();
176 : HcclResult HandleSetReadyStatus();
177 : vector<char> recvData{};
178 : vector<char> recvFinishMsg{};
179 : vector<char> sendData{};
180 : vector<char> sendDataPack_{};
181 : vector<char> sendFinishMsg{};
182 :
183 : void SaveDfxTaskInfo(const TaskParam &taskParam);
184 : };
185 : } // namespace Hccl
186 : #endif
|
