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