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 : #ifndef REGED_MEM_MGR_H
11 : #define REGED_MEM_MGR_H
12 :
13 : #include <algorithm>
14 : #include <cstdint>
15 : #include <memory>
16 : #include <utility>
17 : #include <vector>
18 : #include "hcomm_c_adpt.h"
19 : #include "log.h"
20 : #include "buffer_key.h"
21 : #include "buffer.h"
22 :
23 : using RdmaHandle = void *;
24 :
25 : namespace hcomm {
26 : /**
27 : * @note 职责：用于通信设备EndPoint的注册内存信息管理，支持基于RmaBufferMgr类的重叠内存的检测报错等。
28 : */
29 : class RegedMemMgr {
30 : public:
31 237 : RegedMemMgr() = default;
32 208 : virtual ~RegedMemMgr() = default;
33 :
34 : // 注册内存
35 : virtual HcclResult RegisterMemory(HcommMem mem, const char *memTag, void **memHandle) = 0;
36 :
37 : // 注销内存
38 : virtual HcclResult UnregisterMemory(void* memHandle) = 0;
39 :
40 : // 导出指定内存描述，用于交换
41 : virtual HcclResult MemoryExport(const EndpointDesc endpointDesc, void *memHandle, void **memDesc, uint32_t *memDescLen) = 0;
42 :
43 : // 基于内存描述，导入获得内存
44 : virtual HcclResult MemoryImport(const void *memDesc, uint32_t descLen, HcommMem *outMem) = 0;
45 :
46 : // 关闭内存
47 : virtual HcclResult MemoryUnimport(const void *memDesc, uint32_t descLen) = 0;
48 :
49 : virtual HcclResult GetAllMemHandles(void **memHandles, uint32_t *memHandleNum) = 0;
50 :
51 : // 授权
52 0 : virtual HcclResult MemoryGrant(const HcommMemGrantInfo *remoteGrantInfo)
53 : {
54 0 : return HCCL_SUCCESS;
55 : }
56 :
57 : RdmaHandle rdmaHandle_{nullptr};
58 :
59 : protected:
60 : template <typename RmaBuffer>
61 : using RegedBufferEntry = std::pair<std::shared_ptr<RmaBuffer>, bool>;
62 :
63 58 : static HcclResult ValidateMemParams(HcommMem mem, void **memHandle)
64 : {
65 58 : CHK_PTR_NULL(memHandle);
66 57 : CHK_PTR_NULL(mem.addr);
67 54 : CHK_PRT_RET(mem.size == 0, HCCL_ERROR("[%s] mem size is zero", __func__), HCCL_E_PARA);
68 51 : CHK_PRT_RET(mem.type == COMM_MEM_TYPE_INVALID,
69 : HCCL_ERROR("[%s] invalid mem type [%d]", __func__, mem.type), HCCL_E_PARA);
70 48 : return HCCL_SUCCESS;
71 : }
72 :
73 : // MemoryExport: 从allBuffers中校验memHandle并获取buffer指针
74 : template <typename RmaBuffer>
75 6 : static HcclResult ValidateMemExportHandle(void *memHandle,
76 : const std::vector<RegedBufferEntry<RmaBuffer>>& allBuffers, RmaBuffer*& outBuffer)
77 : {
78 6 : auto it = std::find_if(allBuffers.begin(), allBuffers.end(),
79 2 : [memHandle](const auto &entry) {
80 2 : return entry.first != nullptr && entry.first.get() == memHandle;
81 : });
82 6 : if (it == allBuffers.end()) {
83 4 : HCCL_ERROR("[RegedMemMgr][MemoryExport] memHandle[%p] is not registered.", memHandle);
84 4 : return HCCL_E_NOT_FOUND;
85 : }
86 2 : outBuffer = it->first.get();
87 2 : return HCCL_SUCCESS;
88 : }
89 :
90 : // GetAllMemHandles: 收集所有未标记删除的buffer指针
91 : template <typename RmaBuffer>
92 0 : static HcclResult GetAllMemHandlesImpl(std::vector<RegedBufferEntry<RmaBuffer>>& allBuffers,
93 : std::vector<std::shared_ptr<RmaBuffer>>& activeHandles,
94 : void **memHandles, uint32_t *memHandleNum, const char* logTag)
95 : {
96 0 : CHK_PTR_NULL(memHandles);
97 0 : CHK_PTR_NULL(memHandleNum);
98 :
99 0 : activeHandles.clear();
100 0 : for (auto& entry : allBuffers) {
101 0 : if (!entry.second) {
102 0 : activeHandles.push_back(entry.first);
103 : }
104 : }
105 :
106 0 : *memHandleNum = static_cast<uint32_t>(activeHandles.size());
107 0 : *memHandles = activeHandles.empty() ? nullptr : static_cast<void *>(activeHandles.data());
108 0 : HCCL_INFO("[%s][GetAllMemHandles] memHandleNum[%u]", logTag, *memHandleNum);
109 0 : return HCCL_SUCCESS;
110 : }
111 :
112 : // ---- 底层：tree 操作 ----
113 :
114 : // 用实际注册后的addr/size做key，对buffer增加引用计数
115 : template <typename Mgr, typename BufferPtr>
116 51 : static HcclResult AddBuffer(Mgr& mgr, const BufferPtr& buffer)
117 : {
118 51 : hccl::BufferKey<uintptr_t, u64> actualRegKey(
119 51 : reinterpret_cast<uintptr_t>(buffer->GetAddr()), static_cast<uint64_t>(buffer->GetSize()));
120 51 : EXCEPTION_CATCH((void)mgr->AddWithoutCheck(actualRegKey, buffer), return HCCL_E_INTERNAL);
121 51 : return HCCL_SUCCESS;
122 : }
123 :
124 : // UnregisterMemory: IsAlias为true时，通过硬件句柄定位父buffer
125 : template <typename Mgr, typename BufferPtr, typename BufferVec, typename HwHandleFn, typename EqualFn>
126 15 : static BufferPtr ResolveAliasParent(Mgr& mgr, const hccl::BufferKey<uintptr_t, u64>& ownKey,
127 : BufferPtr buffer, BufferVec& allBuffers, HwHandleFn&& hwHandleGetter, EqualFn&& tokenEqual)
128 : {
129 15 : auto token = hwHandleGetter(buffer);
130 15 : auto findResult = mgr->Find(ownKey);
131 15 : if (findResult.first && tokenEqual(hwHandleGetter(findResult.second.get()), token)) {
132 15 : return findResult.second.get();
133 : }
134 0 : for (auto& entry : allBuffers) {
135 0 : const auto& ptr = entry.first;
136 0 : if (ptr.get() == buffer) {
137 0 : continue;
138 : }
139 0 : if (tokenEqual(hwHandleGetter(ptr.get()), token) && !ptr->IsAlias()) {
140 0 : return ptr.get();
141 : }
142 : }
143 0 : return nullptr;
144 15 : }
145 :
146 : // ---- 中层：注册/注销 核心逻辑 ----
147 :
148 : // Find命中则基于父buffer构造别名，未命中则构造新buffer并注册
149 : template <typename Mgr, typename FindResult, typename BufferPtr, typename MakeAlias, typename MakeNew>
150 30 : static HcclResult RegisterOrAlias(Mgr& mgr, const FindResult& findPair, BufferPtr& buffer,
151 : MakeAlias&& makeAlias, MakeNew&& makeNew)
152 : {
153 30 : if (findPair.first) {
154 11 : auto parentBuffer = findPair.second;
155 11 : EXCEPTION_CATCH((buffer = makeAlias(parentBuffer)), return HCCL_E_PTR);
156 11 : CHK_RET(AddBuffer(mgr, parentBuffer));
157 11 : } else {
158 19 : EXCEPTION_CATCH((buffer = makeNew()), return HCCL_E_PTR);
159 19 : CHK_RET(AddBuffer(mgr, buffer));
160 : }
161 30 : return HCCL_SUCCESS;
162 : }
163 :
164 : // ---- 上层：RegisterMemory / UnregisterMemory 通用实现 ----
165 :
166 : // 构造基类Buffer → RegisterOrAlias → 写回memHandle
167 : // makeAlias(bufPtr, parent) — 基于父buffer构造别名
168 : // makeNew(bufPtr) — 构造新buffer
169 : template <typename Mgr, typename RmaBuffer, typename MakeAlias, typename MakeNew>
170 37 : static HcclResult RegisterMemoryImpl(HcommMem mem, const char *memTag, void **memHandle,
171 : Mgr& mgr, std::vector<RegedBufferEntry<RmaBuffer>>& allBuffers,
172 : const char* logTag, MakeAlias&& makeAlias, MakeNew&& makeNew)
173 : {
174 37 : CHK_RET(ValidateMemParams(mem, memHandle));
175 :
176 30 : hccl::BufferKey<uintptr_t, u64> tempKey(reinterpret_cast<uintptr_t>(mem.addr), mem.size);
177 30 : auto findPair = mgr->Find(tempKey);
178 :
179 30 : std::shared_ptr<Hccl::Buffer> localBufferPtr = nullptr;
180 30 : EXCEPTION_CATCH((localBufferPtr = std::make_shared<Hccl::Buffer>(reinterpret_cast<uintptr_t>(mem.addr),
181 : mem.size, static_cast<HcclMemType>(mem.type), memTag)),
182 : return HCCL_E_PTR);
183 :
184 30 : std::shared_ptr<RmaBuffer> rmaBuffer;
185 60 : CHK_RET(RegisterOrAlias(mgr, findPair, rmaBuffer,
186 : [&](auto& parent) { return makeAlias(localBufferPtr, parent); },
187 : [&]() { return makeNew(localBufferPtr); }));
188 :
189 30 : HCCL_INFO("[%s][RegisterMemory] success, key {%p, %llu}", logTag, mem.addr, mem.size);
190 30 : *memHandle = static_cast<void *>(rmaBuffer.get());
191 30 : allBuffers.emplace_back(rmaBuffer, false);
192 30 : return HCCL_SUCCESS;
193 30 : }
194 :
195 : // IsAlias → ResolveAliasParent → Del → erase
196 : // hwHandleGetter(buffer) — 获取硬件句柄用于ResolveAliasParent
197 : // tokenEqual(lhs, rhs) — 比较两个句柄是否相等
198 : template <typename Mgr, typename RmaBuffer, typename HwHandleFn, typename EqualFn>
199 38 : static HcclResult UnregisterMemoryImpl(void* memHandle, Mgr& mgr,
200 : std::vector<RegedBufferEntry<RmaBuffer>>& allBuffers, HwHandleFn&& hwHandleGetter, EqualFn&& tokenEqual)
201 : {
202 38 : CHK_PTR_NULL(memHandle);
203 37 : RmaBuffer* buffer = static_cast<RmaBuffer*>(memHandle);
204 37 : CHK_PTR_NULL(buffer);
205 37 : auto bufferInfo = buffer->GetBufferInfo();
206 :
207 37 : hccl::BufferKey<uintptr_t, u64> ownKey(bufferInfo.first, bufferInfo.second);
208 37 : RmaBuffer* refBuffer = buffer;
209 37 : if (buffer->IsAlias()) {
210 13 : refBuffer = ResolveAliasParent(mgr, ownKey, buffer, allBuffers,
211 : std::forward<HwHandleFn>(hwHandleGetter), std::forward<EqualFn>(tokenEqual));
212 13 : if (refBuffer == nullptr) {
213 0 : HCCL_ERROR("[UnregisterMemory] alias parent not found");
214 0 : return HCCL_E_NOT_FOUND;
215 : }
216 : }
217 :
218 37 : auto refBufferInfo = refBuffer->GetBufferInfo();
219 37 : hccl::BufferKey<uintptr_t, u64> tempKey(refBufferInfo.first, refBufferInfo.second);
220 : // Del returns false when ref remains nonzero; local unregister still succeeds after erasing this handle.
221 37 : EXCEPTION_CATCH((void)mgr->Del(tempKey), return HCCL_E_NOT_FOUND);
222 : // IsInTree判断tree中是否还有该key的引用
223 30 : auto it = std::find_if(allBuffers.begin(), allBuffers.end(),
224 46 : [buffer](const RegedBufferEntry<RmaBuffer>& entry) { return entry.first.get() == buffer; });
225 30 : if (it != allBuffers.end()) {
226 30 : if(!mgr->IsInTree(ownKey)) {
227 22 : allBuffers.erase(it);
228 : } else {
229 8 : it->second = true;
230 : }
231 : }
232 30 : return HCCL_SUCCESS;
233 : }
234 :
235 : };
236 : }
237 : #endif // REGED_MEM_MGR_H
|
