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 ZERO_COPY_MEMORY_AGENT_H
12 : #define ZERO_COPY_MEMORY_AGENT_H
13 :
14 : #include <atomic>
15 : #include <thread>
16 : #include <unordered_map>
17 : #include <condition_variable>
18 : #include "topoinfo_struct.h"
19 : #include "hccl_socket_manager.h"
20 : #include "common.h"
21 : #include "coll_alg_param.h"
22 : #include "zero_copy_address_mgr.h"
23 :
24 : namespace hccl {
25 :
26 : enum class RequestType {
27 : SET_MEMORY_RANGE = 0,
28 : SET_MEMORY_RANGE_ACK,
29 : UNSET_MEMORY_RANGE,
30 : UNSET_MEMORY_RANGE_ACK,
31 : ACTIVATE_COMM_MEMORY,
32 : ACTIVATE_COMM_MEMORY_ACK,
33 : DEACTIVATE_COMM_MEMORY,
34 : DEACTIVATE_COMM_MEMORY_ACK,
35 : SET_REMOTE_BARE_TGID,
36 : SET_REMOTE_BARE_TGID_ACK,
37 : BARRIER_CLOSE,
38 : BARRIER_CLOSE_ACK,
39 : RESERVED
40 : };
41 :
42 : const std::map<RequestType, std::string> REQUEST_TYPE_STR {
43 : {RequestType::SET_MEMORY_RANGE, "SET_MEMORY_RANGE"},
44 : {RequestType::SET_MEMORY_RANGE_ACK, "SET_MEMORY_RANGE_ACK"},
45 : {RequestType::UNSET_MEMORY_RANGE, "UNSET_MEMORY_RANGE"},
46 : {RequestType::UNSET_MEMORY_RANGE_ACK, "UNSET_MEMORY_RANGE_ACK"},
47 : {RequestType::ACTIVATE_COMM_MEMORY, "ACTIVATE_COMM_MEMORY"},
48 : {RequestType::ACTIVATE_COMM_MEMORY_ACK, "ACTIVATE_COMM_MEMORY_ACK"},
49 : {RequestType::DEACTIVATE_COMM_MEMORY, "DEACTIVATE_COMM_MEMORY"},
50 : {RequestType::DEACTIVATE_COMM_MEMORY_ACK, "DEACTIVATE_COMM_MEMORY_ACK"},
51 : {RequestType::SET_REMOTE_BARE_TGID, "SET_REMOTE_BARE_TGID"},
52 : {RequestType::SET_REMOTE_BARE_TGID_ACK, "SET_REMOTE_BARE_TGID_ACK"},
53 : {RequestType::BARRIER_CLOSE, "BARRIER_CLOSE"},
54 : {RequestType::BARRIER_CLOSE_ACK, "BARRIER_CLOSE_ACK"},
55 : {RequestType::RESERVED, "RESERVED"}
56 : };
57 :
58 0 : inline const char *GetReadableRequestType(RequestType type) {
59 0 : auto it = REQUEST_TYPE_STR.find(type);
60 0 : return (it != REQUEST_TYPE_STR.end()) ? it->second.c_str() : "unknown type";
61 : }
62 :
63 0 : inline bool IsAckRequestType(RequestType type)
64 : {
65 0 : return (type == RequestType::SET_MEMORY_RANGE_ACK) || (type == RequestType::UNSET_MEMORY_RANGE_ACK) ||
66 0 : (type == RequestType::ACTIVATE_COMM_MEMORY_ACK) || (type == RequestType::DEACTIVATE_COMM_MEMORY_ACK) ||
67 0 : (type == RequestType::SET_REMOTE_BARE_TGID_ACK) || (type == RequestType::BARRIER_CLOSE_ACK);
68 : }
69 :
70 : constexpr u32 ZERO_COPY_MEMORY_AGENT_SEND_QUEUE_SIZE = 2;
71 : constexpr u32 ZERO_COPY_MEMORY_AGENT_RECV_QUEUE_SIZE = 2;
72 :
73 : struct ZeroCopyMemoryAgentSendMgr {
74 0 : ZeroCopyMemoryAgentSendMgr() {
75 0 : for (u32 i = 0; i < ZERO_COPY_MEMORY_AGENT_SEND_QUEUE_SIZE; i++) {
76 0 : reqDatas_[i] = nullptr;
77 0 : hasReq_[i] = false;
78 : }
79 0 : }
80 : ~ZeroCopyMemoryAgentSendMgr() = default;
81 :
82 0 : inline void AddRequest(bool isAck, const std::vector<u8>& req) {
83 0 : u32 index = isAck ? 0 : 1;
84 0 : reqDatas_[index] = &req;
85 0 : hasReq_[index] = true;
86 0 : }
87 :
88 : const std::vector<u8>* reqDatas_[ZERO_COPY_MEMORY_AGENT_SEND_QUEUE_SIZE]; // 发送队列，size为2， 0：ack， 1：request
89 : std::atomic<bool> hasReq_[ZERO_COPY_MEMORY_AGENT_SEND_QUEUE_SIZE]; // 发送队列对应位置是否有数据
90 : u32 reqDataSize_{0}; // 本次要发送数据的大小
91 : u32 currIndex_{0}; // 正在发送的req在waitingReq_的下标
92 : u64 sentSize_{0}; // 已发送的字节数， 用于断点续传
93 : u64 lastSendSize_{0}; // 前次异步发送的字节数
94 : void *lastSendHandle_{nullptr}; // 前次异步发送的句柄
95 : };
96 :
97 : struct ZeroCopyMemoryAgentRecvMgr {
98 : std::vector<std::vector<u8>> receivedData_; // 用于待处理的数据队列，size为2
99 : u32 recvIndex_{0}; // 下一个待接收req写入receivedData_的位置
100 : u32 praseIndex_{0}; // 解析线程当前解析的req在receivedData_的位置
101 : u64 receivedSize_{0}; // 已接收的字节数
102 : u64 lastRecvSize_{0}; // 前次异步接收的字节数
103 : void *lastRecvHandle_{nullptr}; // 前次异步接收的句柄
104 : };
105 :
106 : class ZeroCopyMemoryAgent {
107 : public:
108 : ZeroCopyMemoryAgent(const std::unique_ptr<HcclSocketManager> &socketManager, u32 devicePhyId,
109 : s32 deviceLogicId, const HcclIpAddress &localVnicIp, const std::vector<RankInfo> &rankInfoList, RankId userRank,
110 : bool useSuperPodMode, const std::string &identifier);
111 0 : virtual ~ZeroCopyMemoryAgent() = default;
112 :
113 : HcclResult Init();
114 : HcclResult DeInit();
115 :
116 : HcclResult SetMemoryRange(void *virPtr, size_t size, size_t alignment, uint64_t flags);
117 : HcclResult UnsetMemoryRange(void *virPtr);
118 :
119 : HcclResult ActivateCommMemory(void *virPtr, size_t size, size_t offset, void* memHandle, uint64_t flags);
120 : HcclResult DeactivateCommMemory(void *virPtr);
121 :
122 : HcclResult BarrierClose();
123 :
124 : static bool IsActivateCommMemoryAddr(void *virPtr, u64 length);
125 : static HcclResult GetRingBufferAddr(u64 &bufferPtr, u64 &headPtr, u64 &tailPtr);
126 : static bool IsAddressMgrInited();
127 :
128 : bool IsPaused() const;
129 : bool IsResumed() const;
130 :
131 : private:
132 : // member functions
133 : std::string GenerateSocketTag(u32 localRank, u32 remoteRank);
134 : HcclResult SendRequest(RequestType requestType, const std::vector<u8> &req, u32 remoteDevPhyId = INVALID_VALUE_RANKID);
135 : HcclResult SendRequestSync(RequestType requestType, const std::vector<u8>& req, u32 remoteDevPhyId);
136 :
137 : // main thread functions
138 : HcclResult SetRemoteTgid();
139 : HcclResult EstablishSockets();
140 : HcclResult InitInnerThread();
141 : HcclResult WaitForAllRemoteComplete(RequestType requestType);
142 :
143 : // sub thread functions
144 : void InnerThread();
145 : void RequestBatchSendAsync();
146 : void CheckBatchSendAsyncResult();
147 : void RequestBatchRecvAsync();
148 : void CheckBatchRecvAsyncResult();
149 : inline void RecvRequest(ZeroCopyMemoryAgentRecvMgr &recvMgr, u32 remoteDevicePhyId);
150 : void ParseReceivedRequests();
151 : void RequestBatchRecvSync();
152 :
153 : HcclResult ParseReceivedRequest(std::vector<u8>& receivedData, u32 remoteRank);
154 : HcclResult ParseSetMemoryRange(u8* &exchangeDataPtr, u32 &exchangeDataBlankSize);
155 : HcclResult ParseUnsetMemoryRange(u8* &exchangeDataPtr, u32 &exchangeDataBlankSize);
156 : HcclResult ParseBareTgid(u8* &exchangeDataPtr, u32 &exchangeDataBlankSize);
157 : HcclResult ParseBareTgidAck(u8* &exchangeDataPtr, u32 &exchangeDataBlankSize);
158 : HcclResult ParseActivateCommMemory(u8* &exchangeDataPtr, u32 &exchangeDataBlankSize);
159 : HcclResult ParseDeactivateCommMemory(u8* &exchangeDataPtr, u32 &exchangeDataBlankSize);
160 : HcclResult ParseSetMemoryRangeAck(u8* &exchangeDataPtr, u32 &exchangeDataBlankSize);
161 : HcclResult ParseBarrierClose(u8* &exchangeDataPtr, u32 &exchangeDataBlankSize);
162 : HcclResult ParseBarrierCloseAck(u8* &exchangeDataPtr, u32 &exchangeDataBlankSize);
163 : HcclResult SendAckAfterParse(RequestType requestType, RequestType ackType, u32 remoteDevicePhyId, void *extraData = nullptr, u64 extraDataLen = 0);
164 : HcclResult ParseRemoteAck(RequestType requestType, u32 remoteRank);
165 :
166 : // 生成维测日志
167 : std::string DumpFinishInfo(RequestType requestType);
168 :
169 : // 是否操作需要barrier
170 0 : bool NeedBarrier(RequestType type)
171 : {
172 0 : return type == RequestType::BARRIER_CLOSE || type == RequestType::BARRIER_CLOSE_ACK;
173 : }
174 :
175 : void CheckSnapshotStatus();
176 :
177 : // member variables
178 : bool initiated_;
179 : bool isSingleRank_{false};
180 : HcclNetDevCtx vnicPortCtx_{nullptr};
181 : const std::unique_ptr<HcclSocketManager> &socketManager_;
182 : u32 devicePhyId_;
183 : s32 deviceLogicId_;
184 : HcclIpAddress localVnicIp_;
185 : const std::vector<RankInfo> &rankInfoList_;
186 : RankId userRank_;
187 : u32 rankSize_;
188 : bool useSuperPodMode_;
189 : std::string identifier_{};
190 : std::vector<s32> remotePids_;
191 :
192 : std::unique_ptr<std::thread> innerThread_;
193 : std::mutex commRefCntLock_;
194 : std::unordered_map<u32, std::shared_ptr<HcclSocket> > mapDevPhyIdconnectedSockets_;
195 : std::unordered_set<u32> receivedBarrierCloseAck_;
196 : std::unordered_set<u32> receivedBarrierClose_{};
197 : std::unordered_map<u32, u32> mapDevPhyId2RankId_; // 维测信息使用
198 : std::vector<u8> exchangeDataForSend_;
199 : std::unordered_map<u32, std::vector<u8>> exchangeDataForAck_;
200 : std::atomic<bool> threadRun_{false};
201 : std::atomic<u32> reqMsgCounter_[static_cast<int>(RequestType::RESERVED)]{};
202 : std::mutex dfxMutex_;
203 : std::set<u32> reqMsgFinishedRanks_[static_cast<int>(RequestType::RESERVED)]{}; // 维测信息使用
204 : std::atomic<u32> reqMsgDeliverCnt_{};
205 : std::atomic<u32> reqMsgFinishCnt_{};
206 :
207 : std::condition_variable waitCompleteCv_;
208 : bool isSocketSupportAsync_{false};
209 : bool ioRecvWaiting_{false};
210 : bool hasReceivedRequest_{false};
211 : std::atomic<bool> hasSendRequest_{false};
212 : std::mutex sendMutex_;
213 : std::condition_variable sendCv_;
214 : std::unordered_map<u32, ZeroCopyMemoryAgentSendMgr> sendMgrs_;
215 : std::unordered_map<u32, ZeroCopyMemoryAgentRecvMgr> recvMgrs_;
216 :
217 : static std::unique_ptr<ZeroCopyAddressMgr> addressMgr_;
218 : bool isPaused_ { false }; // need to be paused when snapshot
219 : };
220 : }
221 :
222 : #endif // ZERO_COPY_MEMORY_AGENT_H
|
