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 "comm_mem_manager.h"
12 : namespace hccl {
13 :
14 235 : void CommMemMgr::CommSetHcclBufferManager(CCLBufferManager& bufferManager) { bufferManager_ = &bufferManager; }
15 :
16 7 : HcclResult CommMemMgr::GetHcclBuffer(CommBuffer* buffer)
17 : {
18 7 : CHK_PTR_NULL(buffer);
19 7 : CHK_PTR_NULL(bufferManager_);
20 7 : std::lock_guard<std::mutex> lock(bufferMutex_);
21 7 : void* temp = nullptr;
22 7 : uint64_t tempSize = 0;
23 7 : HcclResult ret = bufferManager_->GetIndependentOpCCLbuffer(temp, tempSize);
24 7 : CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[GetHcclBuffer] GetHcclBuffer failed"), ret);
25 7 : buffer->addr = temp;
26 7 : buffer->size = tempSize;
27 7 : return HCCL_SUCCESS;
28 7 : }
29 :
30 : // 绑定：opTag -> 句柄（幂等）
31 9 : HcclResult CommMemMgr::CommRegMem(const std::string& memTag, const HcclMem& mem, HcclRegMemAttr attr, void** memHandle)
32 : {
33 9 : CHK_PRT_RET(
34 : memHandle == nullptr, HCCL_ERROR("[CommRegMem] memHandle is null. tag[%s]", memTag.c_str()), HCCL_E_PARA);
35 9 : CHK_PRT_RET(
36 : mem.addr == nullptr || mem.size == 0,
37 : HCCL_ERROR("[CommRegMem] invalid mem. addr[%p] size[%llu]", mem.addr, mem.size), HCCL_E_PARA);
38 :
39 : // 组装句柄（仅域内管理，无进程级注册）
40 9 : Handle h;
41 9 : EXCEPTION_CATCH(h = std::make_shared<HcclMemoryHandle>(), return HCCL_E_PTR);
42 9 : h->addr = mem.addr;
43 9 : h->size = static_cast<uint64_t>(mem.size);
44 9 : h->memType = static_cast<HcclMemType>(mem.type);
45 9 : h->attr = attr;
46 :
47 9 : const auto key = MakeKey(mem.addr, static_cast<size_t>(mem.size));
48 :
49 9 : std::lock_guard<std::mutex> addLock(memMutex_);
50 9 : auto& reg = tagRegs_[memTag];
51 :
52 : // 同tag内做区间冲突/幂等复用
53 9 : auto res = reg.table.Add(key, h);
54 9 : if (!res.second) {
55 : // 只能用 Find 的返回值来判定：
56 : // - 等于(全集命中)：Find(key).first == true（允许，Add 内已 ref）
57 : // - 子集/超集/交集：Find(key).first 可能为 true(子) 或 false(交/超/空)，但都属于冲突！
58 0 : auto f = reg.table.Find(key);
59 0 : if (!f.first || !(f.second && f.second->addr == mem.addr && f.second->size == mem.size)) {
60 0 : HCCL_ERROR("[CommRegMem] overlap in tag[%s], key=%s", memTag.c_str(), key.ToString().c_str());
61 0 : return HCCL_E_PARA;
62 : }
63 : // HcclRegMemAttr不同时更新
64 0 : if (f.second->attr.value != attr.value) {
65 0 : HCCL_WARNING("[CommRegMem] inconsistent attr for same mem. tag[%s]", memTag.c_str());
66 0 : f.second->attr.value = attr.value;
67 : }
68 : // 复用已有句柄：直接用 Find 返回的 buffer，避免解引用 res.first（可能是 end()）
69 0 : h = f.second;
70 0 : }
71 :
72 : // 幂等加入绑定列表（同memHandle不重复）
73 9 : auto& vec = opBindings_[memTag];
74 9 : bool exists = std::any_of(vec.begin(), vec.end(), [&h](const Handle& x) {
75 1 : return x && (x.get() == h.get());
76 : });
77 9 : if (!exists)
78 9 : vec.emplace_back(h);
79 :
80 9 : *memHandle = h.get();
81 9 : HCCL_INFO("[CommRegMem] ok. tag[%s] memHandle[%p] size[%llu]", memTag.c_str(), *memHandle, h->size);
82 9 : return HCCL_SUCCESS;
83 9 : }
84 :
85 : // 解绑：在该通信域实例内，移除“指定算子(memTag)”下的该句柄
86 0 : HcclResult CommMemMgr::CommUnregMem(const std::string& memTag, const void* memHandle)
87 : {
88 0 : CHK_PRT_RET(memHandle == nullptr, HCCL_ERROR("[CommUnregMem] memHandle is null"), HCCL_E_PARA);
89 0 : CHK_PRT_RET(memTag.empty(), HCCL_ERROR("[CommUnregMem] memTag is null or empty"), HCCL_E_PARA);
90 :
91 0 : std::lock_guard<std::mutex> addLock(memMutex_);
92 :
93 0 : auto itTag = opBindings_.find(memTag);
94 0 : CHK_PRT_RET(
95 : itTag == opBindings_.end(), HCCL_WARNING("[CommUnregMem] tag[%s] not found in bindings", memTag.c_str()),
96 : HCCL_E_NOT_FOUND);
97 :
98 0 : auto& vec = itTag->second; // vector<Handle> under this tag
99 0 : auto& reg = tagRegs_[itTag->first]; // TagRegistry for this tag
100 0 : size_t unboundCount = 0; // 本次解绑命中的句柄个数（即便 Del 未真正擦除也计数）
101 0 : size_t erasedCount = 0; // RmaBufferMgr::Del 返回 true 的次数（ref 归零而“擦除”）
102 :
103 0 : vec.erase(
104 0 : std::remove_if(
105 : vec.begin(), vec.end(),
106 0 : [&](const Handle& h) {
107 0 : if (!h || h.get() != memHandle)
108 0 : return false;
109 0 : const auto key = MakeKey(h->addr, static_cast<size_t>(h->size));
110 : try {
111 0 : if (reg.table.Del(key)) {
112 0 : ++erasedCount; // 该 key 的引用归零并从表中移除
113 : }
114 0 : } catch (const std::out_of_range&) {
115 0 : HCCL_ERROR(
116 : "[CommUnregMem] tag[%s] key not found on Del (maybe already removed)", itTag->first.c_str());
117 0 : }
118 :
119 0 : ++unboundCount; // 从绑定列表移除，无论 Del 是否真正擦除
120 0 : return true; // erase-remove：删除该 handle
121 : }),
122 0 : vec.end());
123 :
124 : // 若该 tag 已无绑定，可按需清理映射条目（以及空表）
125 0 : if (vec.empty()) {
126 0 : opBindings_.erase(itTag);
127 0 : if (reg.table.size() == 0) {
128 0 : tagRegs_.erase(std::string(memTag));
129 : }
130 : }
131 :
132 0 : CHK_PRT_RET(
133 : unboundCount == 0, HCCL_WARNING("[CommUnregMem] tag[%s] memHandle[%p] not found", memTag.c_str(), memHandle),
134 : HCCL_E_NOT_FOUND);
135 :
136 0 : HCCL_INFO(
137 : "[CommUnregMem] tag[%s] memHandle[%p] unbound=%zu, erased=%zu", memTag.c_str(), memHandle, unboundCount,
138 : erasedCount);
139 0 : return HCCL_SUCCESS;
140 0 : }
141 :
142 0 : HcclResult CommMemMgr::CommGetLocalRegMemByTag(const std::string& tag, std::vector<HcclMem>& memVec)
143 : {
144 0 : std::lock_guard<std::mutex> lock(memMutex_);
145 0 : auto it = opBindings_.find(tag);
146 0 : if (it == opBindings_.end()) {
147 0 : HCCL_INFO("[CommMemMgr] tag[%s] key not found", tag.c_str());
148 0 : return HCCL_SUCCESS;
149 : }
150 :
151 0 : const auto& vec = it->second;
152 0 : memVec.reserve(vec.size());
153 0 : for (const auto& handle : vec) {
154 : HcclMem mem;
155 0 : mem.addr = handle->addr;
156 0 : mem.size = handle->size;
157 0 : mem.type = handle->memType;
158 0 : memVec.push_back(mem);
159 : }
160 0 : return HCCL_SUCCESS;
161 0 : }
162 :
163 8 : HcclResult CommMemMgr::CommGetLocalRegMemByHandles(
164 : const HcclMemHandle* memHandles, uint32_t memHandleNum, std::vector<HcclMem>& memVec)
165 : {
166 8 : if (memHandleNum == 0) {
167 1 : return HCCL_SUCCESS;
168 : }
169 7 : CHK_PTR_NULL(memHandles);
170 :
171 6 : std::lock_guard<std::mutex> lock(memMutex_);
172 6 : memVec.clear();
173 6 : memVec.reserve(memHandleNum);
174 :
175 14 : for (uint32_t i = 0; i < memHandleNum; ++i) {
176 10 : bool found = false;
177 13 : for (const auto& tagBinding : opBindings_) {
178 15 : for (const auto& handle : tagBinding.second) {
179 12 : if (handle != nullptr && handle.get() == memHandles[i]) {
180 : HcclMem mem;
181 8 : mem.addr = handle->addr;
182 8 : mem.size = handle->size;
183 8 : mem.type = handle->memType;
184 8 : memVec.push_back(mem);
185 8 : found = true;
186 8 : break;
187 : }
188 : }
189 11 : if (found) {
190 8 : break;
191 : }
192 : }
193 10 : if (!found) {
194 2 : HCCL_ERROR("[CommGetLocalRegMemByHandles] memHandle[%p] not found in any tag", memHandles[i]);
195 2 : return HCCL_E_NOT_FOUND;
196 : }
197 : }
198 4 : return HCCL_SUCCESS;
199 6 : }
200 : } // namespace hccl
|
