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