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