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 : #include "endpoint_mgr.h"
12 : #include <algorithm>
13 : #include "hcomm_c_adpt.h"
14 :
15 : namespace hcomm {
16 :
17 32 : TaggedMemMap::~TaggedMemMap()
18 : {
19 16 : if (handle_ == nullptr) {
20 0 : return;
21 : }
22 43 : for (const auto& kv : tagToHandle_) {
23 27 : HcommResult ret = HcommMemUnreg(handle_, kv.second);
24 27 : if (ret != HCCL_SUCCESS) {
25 12 : HCCL_ERROR(
26 : "[TaggedMemMap::~TaggedMemMap] HcommMemUnreg failed, handle[%p] tag[%s] ret[%d]", handle_,
27 : kv.first.c_str(), ret);
28 : }
29 : }
30 16 : }
31 :
32 24 : MemHandle TaggedMemMap::FindHandle(const std::string& tag) const
33 : {
34 24 : auto it = tagToHandle_.find(tag);
35 24 : return it != tagToHandle_.end() ? it->second : nullptr;
36 : }
37 :
38 33 : bool TaggedMemMap::HasTag(const std::string& tag) const { return tagToHandle_.find(tag) != tagToHandle_.end(); }
39 :
40 31 : void TaggedMemMap::EmplaceHandle(const std::string& tag, MemHandle handle) { tagToHandle_.emplace(tag, handle); }
41 :
42 4 : MemHandle TaggedMemMap::RemoveTag(const std::string& tag)
43 : {
44 4 : auto it = tagToHandle_.find(tag);
45 4 : if (it == tagToHandle_.end()) {
46 0 : return nullptr;
47 : }
48 4 : MemHandle handle = it->second;
49 4 : tagToHandle_.erase(it);
50 4 : return handle;
51 : }
52 :
53 218 : EndpointMgr::~EndpointMgr()
54 : {
55 218 : endpointTagMemMap_.clear();
56 230 : for (const auto& kv : endpointMap_) {
57 12 : const EndpointHandle& endpointHandle = kv.second;
58 12 : (void)HcommEndpointDestroy(endpointHandle);
59 : }
60 : // 销毁共享 jetty 场景按 tag 创建的独立 Endpoint
61 221 : for (const auto& kv : taggedEndpointMap_) {
62 3 : (void)HcommEndpointDestroy(kv.second);
63 : }
64 218 : taggedEndpointMap_.clear();
65 218 : }
66 :
67 37 : HcclResult EndpointMgr::Get(EndpointDesc epDesc, EndpointHandle& handle)
68 : {
69 37 : std::lock_guard<std::mutex> lock(mutex_);
70 37 : auto iterPtr = endpointMap_.find(epDesc);
71 37 : if (iterPtr != endpointMap_.end()) {
72 23 : handle = iterPtr->second;
73 23 : return HCCL_SUCCESS;
74 : }
75 14 : HCCL_INFO("[EndpointMgr::Get] create Endpoint");
76 14 : CHK_RET(static_cast<HcclResult>(HcommEndpointCreate(&epDesc, &handle)));
77 :
78 12 : endpointMap_.emplace(epDesc, handle);
79 12 : return HCCL_SUCCESS;
80 37 : }
81 :
82 5 : HcclResult EndpointMgr::GetWithTag(EndpointDesc epDesc, const std::string& sharedQueueTag, EndpointHandle& handle)
83 : {
84 : // tag 为空：退化为默认 Get，兼容非共享路径或无 tag 场景
85 5 : if (sharedQueueTag.empty()) {
86 1 : return Get(epDesc, handle);
87 : }
88 :
89 4 : EndpointDescTagKey key{epDesc, sharedQueueTag};
90 :
91 : // 快路径：持锁查缓存，命中直接返回
92 : {
93 4 : std::lock_guard<std::mutex> lock(mutex_);
94 4 : auto iter = taggedEndpointMap_.find(key);
95 4 : if (iter != taggedEndpointMap_.end()) {
96 1 : handle = iter->second;
97 1 : return HCCL_SUCCESS;
98 : }
99 4 : }
100 :
101 : // 慢路径：持锁创建 + 二次检查。
102 : // 不采用"无锁创建+失败销毁"乐观模式：HcommEndpointCreate 涉及 device context 分配等重操作，
103 : // 高并发同 key 多线程重复创建+销毁的代价高于锁内串行等待；且 create/destroy 非严格幂等时可能残留状态。
104 3 : std::lock_guard<std::mutex> lock(mutex_);
105 : // 二次检查：另一线程可能已在快路径后、本线程拿锁前完成创建
106 3 : auto iter = taggedEndpointMap_.find(key);
107 3 : if (iter != taggedEndpointMap_.end()) {
108 0 : handle = iter->second;
109 0 : return HCCL_SUCCESS;
110 : }
111 : // 锁内创建：同一 key 不会有并发的重复创建
112 3 : CHK_RET(static_cast<HcclResult>(HcommEndpointCreate(&epDesc, &handle)));
113 3 : taggedEndpointMap_.emplace(std::move(key), handle);
114 3 : HCCL_INFO("[EndpointMgr::GetWithTag] create tagged Endpoint, tag[%s], handle[%p].", sharedQueueTag.c_str(), handle);
115 3 : return HCCL_SUCCESS;
116 4 : }
117 :
118 21 : HcclResult EndpointMgr::RegisterMemory(
119 : EndpointHandle epHandle, const std::vector<std::string>& memTag, const std::vector<HcclMem>& memVec,
120 : uint64_t commMemsVersion)
121 : {
122 21 : std::lock_guard<std::mutex> lock(mutex_);
123 21 : auto& taggedMap = endpointTagMemMap_.try_emplace(epHandle, epHandle).first->second;
124 :
125 : // 版本一致，CommMems 无变更，跳过注册
126 21 : if (taggedMap.GetVersion() == commMemsVersion) {
127 6 : HCCL_INFO(
128 : "[%s]commMemsVersion[%llu] unchanged, skip registration, epHandle[%p]", __FUNCTION__, commMemsVersion,
129 : epHandle);
130 6 : return HCCL_SUCCESS;
131 : }
132 15 : CHK_PRT_RET(
133 : memTag.size() < memVec.size(),
134 : HCCL_ERROR("[%s] memTag.size()[%zu] < memVec.size()[%zu]", __FUNCTION__, memTag.size(), memVec.size()),
135 : HCCL_E_PARA);
136 :
137 15 : size_t index = 0;
138 48 : for (const auto& mem : memVec) {
139 33 : const std::string& tag = memTag[index];
140 33 : index++;
141 : // 检查tag是否已注册，避免重复注册
142 33 : if (taggedMap.HasTag(tag)) {
143 2 : HCCL_INFO("[%s]tag already registered, reuse existing handle, tag=%s", __FUNCTION__, tag.c_str());
144 2 : continue;
145 : }
146 31 : MemHandle memHandle = nullptr;
147 31 : CommMem commMem{static_cast<CommMemType>(mem.type), mem.addr, mem.size};
148 31 : HcclResult ret = static_cast<HcclResult>(HcommMemReg(epHandle, tag.c_str(), &commMem, &memHandle));
149 31 : if (ret != HCCL_SUCCESS && ret != HCCL_E_AGAIN) {
150 0 : HCCL_ERROR("[%s]call trace: hcclRet -> %d", __FUNCTION__, ret);
151 0 : return ret;
152 : }
153 31 : CHK_PTR_NULL(memHandle);
154 31 : taggedMap.EmplaceHandle(tag, memHandle); // 记录到tag映射，后续相同tag直接命中
155 31 : if (ret == HCCL_E_AGAIN) {
156 0 : HCCL_WARNING("This mem has already been registered, addr=%p, size=%llu", mem.addr, mem.size);
157 : }
158 : }
159 :
160 15 : taggedMap.SetVersion(commMemsVersion);
161 15 : return HCCL_SUCCESS;
162 21 : }
163 :
164 14 : HcclResult EndpointMgr::GetMemHandlesByTags(
165 : EndpointHandle epHandle, const std::vector<std::string>& memTags, std::vector<MemHandle>& memHandleVec)
166 : {
167 14 : std::lock_guard<std::mutex> lock(mutex_);
168 14 : memHandleVec.clear();
169 14 : auto it = endpointTagMemMap_.find(epHandle);
170 14 : if (it == endpointTagMemMap_.end()) {
171 0 : HCCL_ERROR("[%s] epHandle[%p] not found in endpointTagMemMap_", __FUNCTION__, epHandle);
172 0 : return HCCL_E_MEMORY;
173 : }
174 14 : const auto& taggedMap = it->second;
175 32 : for (const auto& tag : memTags) {
176 19 : MemHandle handle = taggedMap.FindHandle(tag);
177 19 : if (handle == nullptr) {
178 1 : HCCL_ERROR(
179 : "[%s] tag[%s] not found in endpoint[%p], registration may have been skipped", __FUNCTION__, tag.c_str(),
180 : epHandle);
181 1 : return HCCL_E_NOT_FOUND;
182 : }
183 18 : memHandleVec.push_back(handle);
184 : }
185 13 : return HCCL_SUCCESS;
186 14 : }
187 :
188 3 : HcclResult EndpointMgr::UnregMemByTag(const std::string& tag)
189 : {
190 3 : std::lock_guard<std::mutex> lock(mutex_);
191 3 : HcclResult lastErr = HCCL_SUCCESS;
192 8 : for (auto& kv : endpointTagMemMap_) {
193 5 : MemHandle handle = kv.second.FindHandle(tag);
194 5 : if (handle == nullptr) {
195 1 : continue;
196 : }
197 4 : HcommResult ret = HcommMemUnreg(kv.first, handle);
198 4 : if (ret != HCCL_SUCCESS) {
199 0 : HCCL_ERROR(
200 : "[%s] HcommMemUnreg failed, epHandle[%p] tag[%s] ret[%d]", __FUNCTION__, kv.first, tag.c_str(), ret);
201 0 : lastErr = static_cast<HcclResult>(ret);
202 0 : continue;
203 : }
204 4 : kv.second.RemoveTag(tag);
205 : }
206 3 : return lastErr;
207 3 : }
208 :
209 1 : bool EndpointMgr::IsDescExist(EndpointDesc epDesc) { return endpointMap_.find(epDesc) != endpointMap_.end(); }
210 :
211 : } // namespace hcomm
|
