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 : #include "endpoint_mgr.h"
11 : #include <algorithm>
12 : #include "hcomm_c_adpt.h"
13 :
14 : namespace hcomm {
15 :
16 32 : TaggedMemMap::~TaggedMemMap()
17 : {
18 16 : if (handle_ == nullptr) {
19 0 : return;
20 : }
21 43 : for (const auto& kv : tagToHandle_) {
22 27 : HcommResult ret = HcommMemUnreg(handle_, kv.second);
23 27 : if (ret != HCCL_SUCCESS) {
24 12 : HCCL_ERROR(
25 : "[TaggedMemMap::~TaggedMemMap] HcommMemUnreg failed, handle[%p] tag[%s] ret[%d]", handle_,
26 : kv.first.c_str(), ret);
27 : }
28 : }
29 16 : }
30 :
31 24 : MemHandle TaggedMemMap::FindHandle(const std::string& tag) const
32 : {
33 24 : auto it = tagToHandle_.find(tag);
34 24 : return it != tagToHandle_.end() ? it->second : nullptr;
35 : }
36 :
37 33 : bool TaggedMemMap::HasTag(const std::string& tag) const { return tagToHandle_.find(tag) != tagToHandle_.end(); }
38 :
39 31 : void TaggedMemMap::EmplaceHandle(const std::string& tag, MemHandle handle) { tagToHandle_.emplace(tag, handle); }
40 :
41 4 : MemHandle TaggedMemMap::RemoveTag(const std::string& tag)
42 : {
43 4 : auto it = tagToHandle_.find(tag);
44 4 : if (it == tagToHandle_.end()) {
45 0 : return nullptr;
46 : }
47 4 : MemHandle handle = it->second;
48 4 : tagToHandle_.erase(it);
49 4 : return handle;
50 : }
51 :
52 173 : EndpointMgr::~EndpointMgr()
53 : {
54 173 : endpointTagMemMap_.clear();
55 183 : for (const auto& kv : endpointMap_) {
56 10 : const EndpointHandle& endpointHandle = kv.second;
57 10 : (void)HcommEndpointDestroy(endpointHandle);
58 : }
59 173 : }
60 :
61 34 : HcclResult EndpointMgr::Get(EndpointDesc epDesc, EndpointHandle& handle)
62 : {
63 34 : auto iterPtr = endpointMap_.find(epDesc);
64 34 : if (iterPtr != endpointMap_.end()) {
65 23 : handle = iterPtr->second;
66 23 : return HCCL_SUCCESS;
67 : }
68 11 : HCCL_INFO("[EndpointMgr::Get] create Endpoint");
69 11 : CHK_RET(static_cast<HcclResult>(HcommEndpointCreate(&epDesc, &handle)));
70 :
71 10 : endpointMap_.emplace(epDesc, handle);
72 10 : return HCCL_SUCCESS;
73 : }
74 :
75 21 : HcclResult EndpointMgr::RegisterMemory(
76 : EndpointHandle epHandle, const std::vector<std::string>& memTag, const std::vector<HcclMem>& memVec,
77 : uint64_t commMemsVersion)
78 : {
79 21 : std::lock_guard<std::mutex> lock(mutex_);
80 21 : auto& taggedMap = endpointTagMemMap_.try_emplace(epHandle, epHandle).first->second;
81 :
82 : // 版本一致,CommMems 无变更,跳过注册
83 21 : if (taggedMap.GetVersion() == commMemsVersion) {
84 6 : HCCL_INFO(
85 : "[%s]commMemsVersion[%llu] unchanged, skip registration, epHandle[%p]", __FUNCTION__, commMemsVersion,
86 : epHandle);
87 6 : return HCCL_SUCCESS;
88 : }
89 15 : CHK_PRT_RET(
90 : memTag.size() < memVec.size(),
91 : HCCL_ERROR("[%s] memTag.size()[%zu] < memVec.size()[%zu]", __FUNCTION__, memTag.size(), memVec.size()),
92 : HCCL_E_PARA);
93 :
94 15 : size_t index = 0;
95 48 : for (const auto& mem : memVec) {
96 33 : const std::string& tag = memTag[index];
97 33 : index++;
98 : // 检查tag是否已注册,避免重复注册
99 33 : if (taggedMap.HasTag(tag)) {
100 2 : HCCL_INFO("[%s]tag already registered, reuse existing handle, tag=%s", __FUNCTION__, tag.c_str());
101 2 : continue;
102 : }
103 31 : MemHandle memHandle = nullptr;
104 31 : CommMem commMem{static_cast<CommMemType>(mem.type), mem.addr, mem.size};
105 31 : HcclResult ret = static_cast<HcclResult>(HcommMemReg(epHandle, tag.c_str(), &commMem, &memHandle));
106 31 : if (ret != HCCL_SUCCESS && ret != HCCL_E_AGAIN) {
107 0 : HCCL_ERROR("[%s]call trace: hcclRet -> %d", __FUNCTION__, ret);
108 0 : return ret;
109 : }
110 31 : CHK_PTR_NULL(memHandle);
111 31 : taggedMap.EmplaceHandle(tag, memHandle); // 记录到tag映射,后续相同tag直接命中
112 31 : if (ret == HCCL_E_AGAIN) {
113 0 : HCCL_WARNING("This mem has already been registered, addr=%p, size=%llu", mem.addr, mem.size);
114 : }
115 : }
116 :
117 15 : taggedMap.SetVersion(commMemsVersion);
118 15 : return HCCL_SUCCESS;
119 21 : }
120 :
121 14 : HcclResult EndpointMgr::GetMemHandlesByTags(
122 : EndpointHandle epHandle, const std::vector<std::string>& memTags, std::vector<MemHandle>& memHandleVec)
123 : {
124 14 : std::lock_guard<std::mutex> lock(mutex_);
125 14 : memHandleVec.clear();
126 14 : auto it = endpointTagMemMap_.find(epHandle);
127 14 : if (it == endpointTagMemMap_.end()) {
128 0 : HCCL_ERROR("[%s] epHandle[%p] not found in endpointTagMemMap_", __FUNCTION__, epHandle);
129 0 : return HCCL_E_MEMORY;
130 : }
131 14 : const auto& taggedMap = it->second;
132 32 : for (const auto& tag : memTags) {
133 19 : MemHandle handle = taggedMap.FindHandle(tag);
134 19 : if (handle == nullptr) {
135 1 : HCCL_ERROR(
136 : "[%s] tag[%s] not found in endpoint[%p], registration may have been skipped", __FUNCTION__, tag.c_str(),
137 : epHandle);
138 1 : return HCCL_E_NOT_FOUND;
139 : }
140 18 : memHandleVec.push_back(handle);
141 : }
142 13 : return HCCL_SUCCESS;
143 14 : }
144 :
145 3 : HcclResult EndpointMgr::UnregMemByTag(const std::string& tag)
146 : {
147 3 : std::lock_guard<std::mutex> lock(mutex_);
148 3 : HcclResult lastErr = HCCL_SUCCESS;
149 8 : for (auto& kv : endpointTagMemMap_) {
150 5 : MemHandle handle = kv.second.FindHandle(tag);
151 5 : if (handle == nullptr) {
152 1 : continue;
153 : }
154 4 : HcommResult ret = HcommMemUnreg(kv.first, handle);
155 4 : if (ret != HCCL_SUCCESS) {
156 0 : HCCL_ERROR(
157 : "[%s] HcommMemUnreg failed, epHandle[%p] tag[%s] ret[%d]", __FUNCTION__, kv.first, tag.c_str(), ret);
158 0 : lastErr = static_cast<HcclResult>(ret);
159 0 : continue;
160 : }
161 4 : kv.second.RemoveTag(tag);
162 : }
163 3 : return lastErr;
164 3 : }
165 :
166 1 : bool EndpointMgr::IsDescExist(EndpointDesc epDesc) { return endpointMap_.find(epDesc) != endpointMap_.end(); }
167 :
168 : } // namespace hcomm
|