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