Line data Source code
1 : /**
2 : * Copyright (c) 2026 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 : #include "aicpu_ts_roce_endpoint.h"
11 : #include "log.h"
12 : #include "hccl_net_dev.h"
13 : #include "aicpu_ts_roce_mem.h"
14 : #include "adapter_rts_common.h"
15 : #include "hccl_network.h"
16 : #include "network_manager_pub.h"
17 : #include "hccl_socket.h"
18 : #include <exception>
19 :
20 : namespace hcomm {
21 : namespace {
22 : constexpr uint32_t kDefaultRocePort = 16666;
23 : }
24 :
25 21 : AicpuTsRoceEndpoint::AicpuTsRoceEndpoint(const EndpointDesc &endpointDesc)
26 21 : : Endpoint(endpointDesc)
27 : {
28 21 : }
29 :
30 22 : AicpuTsRoceEndpoint::~AicpuTsRoceEndpoint()
31 : {
32 21 : regedMemMgr_.reset();
33 21 : ctxHandle_ = nullptr;
34 21 : ReleaseListenSocketRefs();
35 21 : ReleaseSharedNetDev();
36 22 : }
37 :
38 21 : void AicpuTsRoceEndpoint::ReleaseListenSocketRefs()
39 : {
40 21 : std::lock_guard<std::mutex> lk(ListenSocketMapMutex());
41 21 : HCCL_INFO("[ReleaseListenSocketRefs] netDevRefPhyId_[%u], listenRefKeys_.size[%zu]",
42 : netDevRefPhyId_, listenRefKeys_.size());
43 :
44 21 : std::vector<SocketMapKey> keys = std::move(listenRefKeys_);
45 21 : auto &sockMap = GetServerSocketMap();
46 23 : for (const auto &key : keys) {
47 2 : auto it = sockMap.find(key);
48 2 : if (it == sockMap.end()) {
49 0 : HCCL_INFO("[ReleaseListenSocketRefs] key[dev=%u,port=%u] not found in sockMap",
50 : key.devicePhyId, key.port);
51 0 : continue;
52 : }
53 2 : HCCL_INFO("[ReleaseListenSocketRefs] key[dev=%u,port=%u] refCount[%u] before decrement",
54 : key.devicePhyId, key.port, it->second.refCount);
55 2 : if (it->second.refCount > 0U) {
56 2 : it->second.refCount--;
57 : }
58 2 : HCCL_INFO("[ReleaseListenSocketRefs] key[dev=%u,port=%u] refCount[%u] after decrement, socket shared_ptr use_count[%ld]",
59 : key.devicePhyId, key.port, it->second.refCount, it->second.socket.use_count());
60 2 : if (it->second.refCount == 0U) {
61 1 : HCCL_INFO("[ReleaseListenSocketRefs] erasing key[dev=%u,port=%u] from sockMap",
62 : key.devicePhyId, key.port);
63 1 : (void)sockMap.erase(it);
64 : }
65 : }
66 21 : }
67 :
68 13 : std::mutex &AicpuTsRoceEndpoint::NetDevMapMutex()
69 : {
70 : static std::mutex mutex;
71 13 : return mutex;
72 : }
73 :
74 35 : std::unordered_map<uint32_t, AicpuTsNetDevSlot> &AicpuTsRoceEndpoint::GetNetDevMap()
75 : {
76 35 : static std::unordered_map<uint32_t, AicpuTsNetDevSlot> netDevMap;
77 35 : return netDevMap;
78 : }
79 :
80 4 : HcclResult AicpuTsRoceEndpoint::AcquireSharedNetDev(uint32_t devicePhyId, const HcclNetDevInfos &info)
81 : {
82 4 : std::lock_guard<std::mutex> lk(NetDevMapMutex());
83 4 : auto &netDevMap = GetNetDevMap();
84 4 : const auto it = netDevMap.find(devicePhyId);
85 4 : if (it != netDevMap.end()) {
86 1 : it->second.refCount++;
87 1 : netDev_ = it->second.netDev;
88 1 : netDevRefPhyId_ = devicePhyId;
89 1 : HCCL_INFO("[AicpuTsRoceEndpoint][%s] reuse HcclNetDev for devicePhyId[%u], ref[%u]", __func__, devicePhyId,
90 : it->second.refCount);
91 1 : return HCCL_SUCCESS;
92 : }
93 :
94 3 : HcclNetDev netDev = nullptr;
95 3 : const HcclResult ret = HcclNetDevOpen(&info, &netDev);
96 3 : if (ret != HCCL_SUCCESS) {
97 0 : HCCL_ERROR("[AicpuTsRoceEndpoint][%s] HcclNetDevOpen failed, ret[%d]", __func__, ret);
98 0 : return ret;
99 : }
100 3 : netDevMap[devicePhyId] = AicpuTsNetDevSlot{ netDev, 1U };
101 3 : netDev_ = netDev;
102 3 : netDevRefPhyId_ = devicePhyId;
103 3 : return HCCL_SUCCESS;
104 4 : }
105 :
106 21 : void AicpuTsRoceEndpoint::ReleaseSharedNetDev()
107 : {
108 21 : if (netDevRefPhyId_ == UINT32_MAX) {
109 12 : return;
110 : }
111 9 : const uint32_t key = netDevRefPhyId_;
112 9 : netDevRefPhyId_ = UINT32_MAX;
113 9 : HcclNetDev toClose = nullptr;
114 : {
115 9 : std::lock_guard<std::mutex> lk(NetDevMapMutex());
116 9 : auto &netDevMap = GetNetDevMap();
117 9 : const auto it = netDevMap.find(key);
118 9 : if (it == netDevMap.end()) {
119 5 : HCCL_ERROR("[AicpuTsRoceEndpoint][ReleaseSharedNetDev] missing slot for devicePhyId[%u]", key);
120 : } else {
121 4 : if (it->second.refCount > 0U) {
122 4 : it->second.refCount--;
123 : }
124 4 : if (it->second.refCount == 0U) {
125 3 : toClose = it->second.netDev;
126 3 : (void)netDevMap.erase(it);
127 : }
128 : }
129 9 : }
130 9 : netDev_ = nullptr;
131 9 : if (toClose != nullptr) {
132 3 : if (!hasListenSocketRef_) {
133 2 : ReleaseNicSocketHandle(toClose);
134 : }
135 3 : HCCL_INFO("[AicpuTsRoceEndpoint][ReleaseSharedNetDev] closing HcclNetDev for devicePhyId[%u]", key);
136 3 : const HcclResult ret = HcclNetDevClose(toClose);
137 3 : if (ret != HCCL_SUCCESS) {
138 0 : HCCL_ERROR("[AicpuTsRoceEndpoint][ReleaseSharedNetDev] HcclNetDevClose failed, ret[%d]", ret);
139 : }
140 : }
141 : }
142 :
143 2 : void AicpuTsRoceEndpoint::ReleaseNicSocketHandle(HcclNetDev netDev)
144 : {
145 2 : auto *netDevCtx = static_cast<hccl::NetDevContext *>(netDev);
146 2 : if (netDevCtx == nullptr) {
147 0 : return;
148 : }
149 2 : const hccl::HcclIpAddress localIp = netDevCtx->GetLocalIp();
150 2 : const HcclResult ret = hccl::NetworkManager::GetInstance(netDevCtx->GetLogicId()).StopNicSocketHandle(localIp);
151 2 : if (ret != HCCL_SUCCESS) {
152 2 : HCCL_WARNING("[AicpuTsRoceEndpoint][%s] StopNicSocketHandle failed, ip[%s], ret[%d]",
153 : __func__, localIp.GetReadableAddress(), ret);
154 : }
155 2 : }
156 :
157 4 : HcclResult AicpuTsRoceEndpoint::AcquireRdmaContext(uint32_t devPhyId, const EndpointDesc &endpointDesc)
158 : {
159 4 : HcclNetDevInfos info;
160 4 : info.addr.protoType = HCCL_PROTO_TYPE_ROCE;
161 4 : CHK_RET(CommAddrTypeToHcclAddressType(endpointDesc.commAddr.type, info.addr.type));
162 4 : if (endpointDesc.commAddr.type == COMM_ADDR_TYPE_IP_V4) {
163 4 : info.addr.addr = endpointDesc.commAddr.addr;
164 : } else {
165 0 : info.addr.addr6 = endpointDesc.commAddr.addr6;
166 : }
167 4 : info.netdevDeployment = HCCL_NETDEV_DEPLOYMENT_DEVICE;
168 4 : info.devicePhyId = static_cast<int32_t>(devPhyId);
169 4 : HcclResult ret = AcquireSharedNetDev(devPhyId, info);
170 4 : if (ret != HCCL_SUCCESS) {
171 0 : return ret;
172 : }
173 :
174 4 : auto *netDevCtx = static_cast<hccl::NetDevContext *>(netDev_);
175 4 : if (netDevCtx == nullptr) {
176 0 : ReleaseSharedNetDev();
177 0 : return HCCL_E_PTR;
178 : }
179 4 : const hccl::HcclIpAddress ipAddr = netDevCtx->GetLocalIp();
180 4 : RdmaHandle rdmaHandle = nullptr;
181 4 : ret = hccl::NetworkManager::GetInstance(netDevCtx->GetLogicId()).GetRdmaHandleByIpAddr(ipAddr, rdmaHandle);
182 4 : if (ret != HCCL_SUCCESS) {
183 0 : HCCL_ERROR("[%s]call trace: hcclRet -> %d", __func__, ret);
184 0 : ReleaseSharedNetDev();
185 0 : return ret;
186 : }
187 4 : ctxHandle_ = rdmaHandle;
188 4 : if (ctxHandle_ == nullptr) {
189 0 : HCCL_ERROR("[%s]errNo[0x%016llx]ptr [ctxHandle_] is nullptr, return HCCL_E_PTR",
190 : __func__, HCCL_ERROR_CODE(HCCL_E_PTR));
191 0 : ReleaseSharedNetDev();
192 0 : return HCCL_E_PTR;
193 : }
194 4 : HCCL_INFO("AicpuTsRoceEndpoint::%s success, devId[%u], ipAddr[%s], ctxHandle[%p]",
195 : __func__, devPhyId, ipAddr.GetReadableAddress(), ctxHandle_);
196 4 : return HCCL_SUCCESS;
197 4 : }
198 :
199 5 : HcclResult AicpuTsRoceEndpoint::Init()
200 : {
201 5 : HCCL_INFO("[%s] localEndpoint protocol[%d]", __func__, endpointDesc_.protocol);
202 :
203 5 : if (endpointDesc_.loc.locType != ENDPOINT_LOC_TYPE_DEVICE) {
204 1 : HCCL_INFO("[AicpuTsRoceEndpoint][%s] AicpuTsRoceEndpoint not support host", __func__);
205 1 : return HCCL_E_NOT_SUPPORT;
206 : }
207 :
208 4 : s32 devId = 0;
209 4 : CHK_RET(hrtGetDevice(&devId));
210 4 : u32 devPhyId = 0;
211 4 : CHK_RET(hrtGetDevicePhyIdByIndex(devId, devPhyId));
212 :
213 4 : HcclResult ret = AcquireRdmaContext(devPhyId, endpointDesc_);
214 4 : if (ret != HCCL_SUCCESS) {
215 0 : return ret;
216 : }
217 :
218 : try {
219 4 : regedMemMgr_ = std::make_shared<AicpuTsRoceRegedMemMgr>(netDev_, ctxHandle_);
220 0 : } catch (std::exception &e) {
221 0 : HCCL_ERROR("[%s]Failed, exception caught:%s", __func__, e.what());
222 0 : ctxHandle_ = nullptr;
223 0 : ReleaseSharedNetDev();
224 0 : return HCCL_E_PTR;
225 0 : }
226 4 : this->regedMemMgr_->rdmaHandle_ = this->ctxHandle_;
227 :
228 4 : return HCCL_SUCCESS;
229 : }
230 :
231 2 : HcclResult AicpuTsRoceEndpoint::ServerSocketListen(const uint32_t port)
232 : {
233 2 : const uint32_t listenPort = (port != 0U) ? port : kDefaultRocePort;
234 2 : const SocketMapKey key{netDevRefPhyId_, listenPort};
235 2 : std::lock_guard<std::mutex> lk(ListenSocketMapMutex());
236 2 : if (ReuseListenSocketIfExist(key, "reuse serverSocket")) {
237 1 : return HCCL_SUCCESS;
238 : }
239 :
240 1 : std::shared_ptr<hccl::HcclSocket> newServerSocket = nullptr;
241 1 : EXCEPTION_CATCH(newServerSocket = std::make_shared<hccl::HcclSocket>(static_cast<HcclNetDevCtx>(netDev_),
242 : listenPort),
243 : return HCCL_E_PTR);
244 1 : CHK_SMART_PTR_NULL(newServerSocket);
245 :
246 1 : HcclResult ret = newServerSocket->Init();
247 1 : if (ret != HCCL_SUCCESS) {
248 0 : HCCL_ERROR("[AicpuTsRoceEndpoint][%s] HcclSocket Init failed, ret[%d]", __func__, ret);
249 0 : return ret;
250 : }
251 :
252 1 : ret = newServerSocket->Listen();
253 1 : if (ret != HCCL_SUCCESS) {
254 0 : HCCL_ERROR("[AicpuTsRoceEndpoint][%s] HcclSocket Listen failed, ret[%d]", __func__, ret);
255 0 : return ret;
256 : }
257 :
258 1 : auto &serverSocketMap = GetServerSocketMap();
259 1 : serverSocketMap[key] = AicpuTsListenSocketSlot{newServerSocket, 1U};
260 1 : listenRefKeys_.push_back(key);
261 1 : hasListenSocketRef_ = true;
262 1 : HCCL_INFO("[AicpuTsRoceEndpoint][%s] listen on key[dev=%u,port=%u] success",
263 : __func__, key.devicePhyId, key.port);
264 1 : return HCCL_SUCCESS;
265 2 : }
266 :
267 2 : bool AicpuTsRoceEndpoint::ReuseListenSocketIfExist(const SocketMapKey &key, const char *logPrefix)
268 : {
269 2 : auto &serverSocketMap = GetServerSocketMap();
270 2 : auto it = serverSocketMap.find(key);
271 2 : if (it == serverSocketMap.end() || it->second.socket == nullptr) {
272 1 : return false;
273 : }
274 1 : it->second.refCount++;
275 1 : listenRefKeys_.push_back(key);
276 1 : hasListenSocketRef_ = true;
277 1 : HCCL_INFO("[AicpuTsRoceEndpoint::%s] %s key[dev=%u,port=%u], ref[%u]",
278 : __func__, logPrefix, key.devicePhyId, key.port, it->second.refCount);
279 1 : return true;
280 : }
281 :
282 27 : std::mutex &AicpuTsRoceEndpoint::ListenSocketMapMutex()
283 : {
284 : static std::mutex mutex;
285 27 : return mutex;
286 : }
287 :
288 54 : std::unordered_map<SocketMapKey, AicpuTsListenSocketSlot, SocketMapKeyHash> &AicpuTsRoceEndpoint::GetServerSocketMap()
289 : {
290 54 : static std::unordered_map<SocketMapKey, AicpuTsListenSocketSlot, SocketMapKeyHash> serverSocketMap;
291 54 : return serverSocketMap;
292 : }
293 :
294 3 : HcclResult AicpuTsRoceEndpoint::AddListenSocketWhiteList(uint32_t port, const std::vector<SocketWlistInfo> &wlistInfos)
295 : {
296 3 : if (wlistInfos.empty()) {
297 1 : HCCL_ERROR("[AicpuTsRoceEndpoint][%s] empty whitelist", __func__);
298 1 : return HCCL_E_PARA;
299 : }
300 2 : std::lock_guard<std::mutex> lk(ListenSocketMapMutex());
301 2 : auto &sockMap = GetServerSocketMap();
302 2 : const uint32_t listenPort = (port != 0U) ? port : kDefaultRocePort;
303 2 : const SocketMapKey key{netDevRefPhyId_, listenPort};
304 2 : auto it = sockMap.find(key);
305 2 : if (it == sockMap.end() || it->second.socket == nullptr) {
306 1 : HCCL_ERROR("[AicpuTsRoceEndpoint][%s] no listen socket for key[dev=%u,port=%u]",
307 : __func__, key.devicePhyId, key.port);
308 1 : return HCCL_E_NOT_FOUND;
309 : }
310 1 : std::vector<SocketWlistInfo> mutableCopy = wlistInfos;
311 1 : return it->second.socket->AddWhiteList(mutableCopy);
312 2 : }
313 :
314 0 : HcclResult AicpuTsRoceEndpoint::GetSocket(uint32_t port, const std::string &tag,
315 : std::shared_ptr<hccl::HcclSocket> &outConnected)
316 : {
317 0 : EXCEPTION_CATCH((outConnected = std::make_shared<hccl::HcclSocket>(tag,
318 : static_cast<HcclNetDevCtx>(netDev_), hccl::HcclIpAddress(), 0, hccl::HcclSocketRole::SOCKET_ROLE_SERVER)), return HCCL_E_PTR);
319 0 : CHK_SMART_PTR_NULL(outConnected);
320 0 : CHK_RET(outConnected->Init());
321 :
322 0 : return HCCL_SUCCESS;
323 : }
324 :
325 2 : HcclResult AicpuTsRoceEndpoint::AcceptDataSocket(uint32_t port, const std::string &tag,
326 : std::shared_ptr<hccl::HcclSocket> &outConnected, uint32_t acceptTimeoutMs)
327 : {
328 2 : std::lock_guard<std::mutex> lk(ListenSocketMapMutex());
329 2 : auto &map = GetServerSocketMap();
330 2 : const uint32_t listenPort = (port != 0U) ? port : kDefaultRocePort;
331 2 : const SocketMapKey key{netDevRefPhyId_, listenPort};
332 2 : auto it = map.find(key);
333 2 : if (it == map.end() || it->second.socket == nullptr) {
334 1 : HCCL_ERROR("[AicpuTsRoceEndpoint][%s] no listen socket for key[dev=%u,port=%u]",
335 : __func__, key.devicePhyId, key.port);
336 1 : return HCCL_E_NOT_FOUND;
337 : }
338 1 : return it->second.socket->Accept(tag, outConnected, acceptTimeoutMs);
339 2 : }
340 :
341 2 : HcclResult AicpuTsRoceEndpoint::RegisterMemory(HcommMem mem, const char *memTag, void **memHandle)
342 : {
343 2 : CHK_RET(this->regedMemMgr_->RegisterMemory(mem, memTag, memHandle));
344 1 : return HCCL_SUCCESS;
345 : }
346 :
347 2 : HcclResult AicpuTsRoceEndpoint::UnregisterMemory(void* memHandle)
348 : {
349 2 : CHK_RET(this->regedMemMgr_->UnregisterMemory(memHandle));
350 1 : return HCCL_SUCCESS;
351 : }
352 :
353 1 : HcclResult AicpuTsRoceEndpoint::MemoryExport(void *memHandle, void **memDesc, uint32_t *memDescLen)
354 : {
355 1 : CHK_RET(this->regedMemMgr_->MemoryExport(this->endpointDesc_, memHandle, memDesc, memDescLen));
356 1 : return HCCL_SUCCESS;
357 : }
358 :
359 1 : HcclResult AicpuTsRoceEndpoint::MemoryImport(const void *memDesc, uint32_t descLen, HcommMem *outMem)
360 : {
361 1 : CHK_RET(this->regedMemMgr_->MemoryImport(memDesc, descLen, outMem));
362 1 : return HCCL_SUCCESS;
363 : }
364 :
365 1 : HcclResult AicpuTsRoceEndpoint::MemoryUnimport(const void *memDesc, uint32_t descLen)
366 : {
367 1 : CHK_RET(this->regedMemMgr_->MemoryUnimport(memDesc, descLen));
368 1 : return HCCL_SUCCESS;
369 : }
370 :
371 1 : HcclResult AicpuTsRoceEndpoint::GetAllMemHandles(void **memHandles, uint32_t *memHandleNum)
372 : {
373 1 : CHK_RET(this->regedMemMgr_->GetAllMemHandles(memHandles, memHandleNum));
374 1 : return HCCL_SUCCESS;
375 : }
376 : }
|
