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