Line data Source code
1 : /**
2 : * Copyright (c) 2026 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 "aicpu_ts_roce_mem.h"
11 : #include <algorithm>
12 : #include <mutex>
13 : #include <unordered_map>
14 : #include "securec.h"
15 : #include "adapter_hccp.h"
16 : #include "hccl_network.h"
17 : #include "log.h"
18 : #include "rma_buffer.h"
19 :
20 : namespace {
21 : using LocalRdmaRmaBufferMgr = hccl::NetDevContext::LocalRdmaRmaBufferMgr;
22 : struct LocalBufferMgrCtx {
23 : std::shared_ptr<std::mutex> mu;
24 : std::shared_ptr<LocalRdmaRmaBufferMgr> mgr;
25 : };
26 :
27 : std::mutex g_phyLocalRdmaBundleMapMu;
28 : std::unordered_map<s32, std::shared_ptr<LocalBufferMgrCtx>> g_phyIdToLocalBufferMgrCtx;
29 :
30 28 : std::shared_ptr<LocalBufferMgrCtx> GetOrCreateLocalBufferMgr(s32 devicePhyId)
31 : {
32 28 : std::lock_guard<std::mutex> mapLock(g_phyLocalRdmaBundleMapMu);
33 28 : std::shared_ptr<LocalBufferMgrCtx> &slot = g_phyIdToLocalBufferMgrCtx[devicePhyId];
34 28 : if (slot == nullptr) {
35 1 : std::shared_ptr<LocalRdmaRmaBufferMgr> mgr;
36 1 : EXCEPTION_CATCH((mgr = std::make_shared<LocalRdmaRmaBufferMgr>()), return nullptr);
37 1 : slot = std::make_shared<LocalBufferMgrCtx>();
38 1 : slot->mu = std::make_shared<std::mutex>();
39 1 : slot->mgr = std::move(mgr);
40 1 : }
41 28 : return slot;
42 28 : }
43 : } // namespace
44 :
45 : namespace hcomm {
46 25 : AicpuTsRoceRegedMemMgr::AicpuTsRoceRegedMemMgr(HcclNetDev netDev, RdmaHandle rdmaHandle)
47 25 : : netDev_(netDev)
48 : {
49 25 : rdmaHandle_ = rdmaHandle;
50 25 : if (netDev_ != nullptr) {
51 10 : auto *netDevCtx = static_cast<hccl::NetDevContext *>(netDev_);
52 10 : const s32 phyId = netDevCtx->GetPhyId();
53 10 : std::shared_ptr<LocalBufferMgrCtx> ctx = GetOrCreateLocalBufferMgr(phyId);
54 10 : if (ctx != nullptr) {
55 10 : localRdmaRmaBufferMgr_ = ctx->mgr;
56 : }
57 10 : HCCL_INFO(
58 : "[AicpuTsRoceRegedMemMgr] ctor netDev[%p] phyId[%d] process-local localRdmaRmaBufferMgr[%p] (not NetDev embed)",
59 : static_cast<void *>(netDev_), static_cast<int>(phyId), static_cast<void *>(localRdmaRmaBufferMgr_.get()));
60 10 : } else {
61 15 : HCCL_INFO("[AicpuTsRoceRegedMemMgr] ctor netDev is null, local mgr unset");
62 : }
63 25 : }
64 :
65 8 : void AicpuTsRoceRegedMemMgr::TrackRegisteredBuffer(const std::shared_ptr<hccl::LocalRdmaRmaBuffer> &localBuffer)
66 : {
67 8 : void *const handlePtr = static_cast<void *>(localBuffer.get());
68 8 : const bool alreadyListed = std::any_of(allRegisteredBuffers_.begin(), allRegisteredBuffers_.end(),
69 4 : [handlePtr](const auto &entry) { return static_cast<void *>(entry.first.get()) == handlePtr; });
70 8 : if (alreadyListed) {
71 0 : return;
72 : }
73 8 : allRegisteredBuffers_.emplace_back(localBuffer, false);
74 8 : HcclBuf rec{};
75 8 : rec.addr = localBuffer->GetAddr();
76 8 : rec.len = localBuffer->GetSize();
77 8 : auto *rma = dynamic_cast<hccl::RmaBuffer *>(localBuffer.get());
78 8 : if (rma == nullptr) {
79 0 : HCCL_ERROR("[AicpuTsRoceRegedMemMgr][TrackRegisteredBuffer] rma is nullptr");
80 0 : return;
81 : }
82 8 : rec.handle = static_cast<void *>(rma);
83 8 : hcclBufRecords_.push_back(rec);
84 : }
85 :
86 11 : HcclResult AicpuTsRoceRegedMemMgr::RegisterMemory(HcommMem mem, const char *memTag, void **memHandle)
87 : {
88 : (void)memTag;
89 11 : HCCL_INFO("[%s] Begin", __FUNCTION__);
90 11 : CHK_PTR_NULL(netDev_);
91 9 : CHK_PTR_NULL(memHandle);
92 9 : CHK_PTR_NULL(localRdmaRmaBufferMgr_);
93 :
94 9 : auto *netDevCtx = static_cast<hccl::NetDevContext *>(netDev_);
95 :
96 9 : std::shared_ptr<LocalBufferMgrCtx> ctx = GetOrCreateLocalBufferMgr(netDevCtx->GetPhyId());
97 9 : CHK_PTR_NULL(ctx);
98 9 : std::lock_guard<std::mutex> phyLocalLock(*ctx->mu);
99 :
100 9 : hccl::RmaMemType memType = static_cast<hccl::RmaMemType>(mem.type);
101 9 : hccl::BufferKey<uintptr_t, u64> tempKey(reinterpret_cast<uintptr_t>(mem.addr), static_cast<u64>(mem.size));
102 9 : auto findPair = localRdmaRmaBufferMgr_->Find(tempKey);
103 :
104 9 : std::shared_ptr<hccl::LocalRdmaRmaBuffer> localRdmaRmaBuffer;
105 9 : if (findPair.first) {
106 4 : auto parentBuffer = findPair.second;
107 4 : EXCEPTION_CATCH((localRdmaRmaBuffer = std::make_shared<hccl::LocalRdmaRmaBuffer>(
108 : netDevCtx, mem.addr, static_cast<u64>(mem.size), memType, *parentBuffer)),
109 : return HCCL_E_PTR);
110 4 : CHK_RET(AddBuffer(localRdmaRmaBufferMgr_, parentBuffer));
111 4 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][RegisterMemory] alias created, key {%p, %llu}", mem.addr, mem.size);
112 4 : } else {
113 5 : EXCEPTION_CATCH((localRdmaRmaBuffer = std::make_shared<hccl::LocalRdmaRmaBuffer>(
114 : netDevCtx, mem.addr, static_cast<u64>(mem.size), memType)),
115 : return HCCL_E_PTR);
116 :
117 5 : HcclResult ret = localRdmaRmaBuffer->Init();
118 5 : if (ret != HCCL_SUCCESS) {
119 1 : HCCL_ERROR("[AicpuTsRoceRegedMemMgr][RegisterMemory] Init failed, ret[%d]", ret);
120 1 : return ret;
121 : }
122 :
123 4 : CHK_RET(AddBuffer(localRdmaRmaBufferMgr_, localRdmaRmaBuffer));
124 4 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][RegisterMemory] success, key {%p, %llu}", mem.addr, mem.size);
125 : }
126 :
127 8 : *memHandle = static_cast<void *>(localRdmaRmaBuffer.get());
128 8 : TrackRegisteredBuffer(localRdmaRmaBuffer);
129 8 : return HCCL_SUCCESS;
130 9 : }
131 :
132 10 : HcclResult AicpuTsRoceRegedMemMgr::UnregisterMemory(void *memHandle)
133 : {
134 10 : HCCL_INFO("[%s] Begin", __FUNCTION__);
135 10 : CHK_PTR_NULL(netDev_);
136 8 : CHK_PTR_NULL(memHandle);
137 8 : CHK_PTR_NULL(localRdmaRmaBufferMgr_);
138 :
139 8 : auto *netDevCtx = static_cast<hccl::NetDevContext *>(netDev_);
140 8 : std::shared_ptr<LocalBufferMgrCtx> ctx = GetOrCreateLocalBufferMgr(netDevCtx->GetPhyId());
141 8 : CHK_PTR_NULL(ctx);
142 8 : std::lock_guard<std::mutex> phyLocalLock(*ctx->mu);
143 :
144 8 : auto *buffer = static_cast<hccl::LocalRdmaRmaBuffer *>(memHandle);
145 :
146 : // IsAlias() 直接区分父子buffer：
147 : // - 父buffer (IsAlias()=false): 自己的key在tree中 → Del(ownKey)
148 : // - 子buffer (IsAlias()=true): 自己的key不在tree中 → 通过GetKey找父key做Del
149 8 : hccl::BufferKey<uintptr_t, u64> ownKey(reinterpret_cast<uintptr_t>(buffer->GetAddr()), buffer->GetSize());
150 8 : hccl::LocalRdmaRmaBuffer* refBuffer = buffer;
151 8 : if (buffer->IsAlias()) {
152 4 : refBuffer = ResolveAliasParent(localRdmaRmaBufferMgr_, ownKey, buffer, allRegisteredBuffers_,
153 8 : [](auto* b) { return b->GetKey(); },
154 4 : [](auto a, auto b) { return a == b; });
155 4 : if (refBuffer == nullptr) {
156 0 : HCCL_ERROR("[AicpuTsRoceRegedMemMgr][UnregisterMemory] alias parent not found");
157 0 : return HCCL_E_NOT_FOUND;
158 : }
159 : }
160 :
161 : hccl::BufferKey<uintptr_t, u64> tempKey(
162 8 : reinterpret_cast<uintptr_t>(refBuffer->GetAddr()), refBuffer->GetSize());
163 :
164 8 : bool delOk = false;
165 8 : EXCEPTION_CATCH(delOk = localRdmaRmaBufferMgr_->Del(tempKey), return HCCL_E_NOT_FOUND);
166 8 : if (!delOk) {
167 4 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][UnregisterMemory] ref count > 0");
168 : }
169 :
170 8 : exportDescByBuffer_.erase(buffer);
171 :
172 8 : auto it = std::find_if(allRegisteredBuffers_.begin(), allRegisteredBuffers_.end(),
173 12 : [memHandle](const auto &entry) { return entry.first.get() == memHandle; });
174 8 : if (it != allRegisteredBuffers_.end()) {
175 : // IsInTree判断tree中是否还有该key的引用
176 8 : if (!localRdmaRmaBufferMgr_->IsInTree(ownKey)) {
177 4 : allRegisteredBuffers_.erase(it);
178 : } else {
179 4 : it->second = true;
180 : }
181 : }
182 :
183 8 : hcclBufRecords_.erase(std::remove_if(hcclBufRecords_.begin(), hcclBufRecords_.end(),
184 12 : [memHandle](const HcclBuf &b) { return b.handle == memHandle; }),
185 8 : hcclBufRecords_.end());
186 8 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][UnregisterMemory] success, memHandle[%p] key {%p, %llu}", memHandle,
187 : buffer->GetAddr(), static_cast<unsigned long long>(buffer->GetSize()));
188 8 : return HCCL_SUCCESS;
189 8 : }
190 :
191 3 : HcclResult AicpuTsRoceRegedMemMgr::MemoryExport(const EndpointDesc endpointDesc, void *memHandle, void **memDesc,
192 : uint32_t *memDescLen)
193 : {
194 3 : HCCL_INFO("[%s] Begin", __FUNCTION__);
195 3 : CHK_PTR_NULL(memHandle);
196 2 : CHK_PTR_NULL(memDesc);
197 1 : CHK_PTR_NULL(memDescLen);
198 1 : CHK_PTR_NULL(netDev_);
199 1 : CHK_PTR_NULL(localRdmaRmaBufferMgr_);
200 :
201 1 : auto *netDevCtx = static_cast<hccl::NetDevContext *>(netDev_);
202 1 : std::shared_ptr<LocalBufferMgrCtx> ctx = GetOrCreateLocalBufferMgr(netDevCtx->GetPhyId());
203 1 : CHK_PTR_NULL(ctx);
204 1 : std::lock_guard<std::mutex> phyLocalLock(*ctx->mu);
205 :
206 1 : hccl::LocalRdmaRmaBuffer *buf = nullptr;
207 1 : CHK_RET(ValidateMemExportHandle(memHandle, allRegisteredBuffers_, buf));
208 0 : std::string &ser = buf->Serialize();
209 0 : if (ser.empty()) {
210 0 : HCCL_ERROR("[AicpuTsRoceRegedMemMgr][MemoryExport] Serialize empty");
211 0 : return HCCL_E_INTERNAL;
212 : }
213 :
214 0 : std::vector<char> &blob = exportDescByBuffer_[buf];
215 0 : blob.clear();
216 0 : blob.reserve(ser.size() + sizeof(EndpointDesc));
217 0 : blob.insert(blob.end(), ser.begin(), ser.end());
218 :
219 0 : std::vector<char> ep(sizeof(EndpointDesc));
220 0 : if (memcpy_s(ep.data(), sizeof(EndpointDesc), &endpointDesc, sizeof(EndpointDesc)) != EOK) {
221 0 : HCCL_ERROR("[AicpuTsRoceRegedMemMgr][MemoryExport] endpointDesc memcpy_s failed");
222 0 : return HCCL_E_INTERNAL;
223 : }
224 0 : blob.insert(blob.end(), ep.begin(), ep.end());
225 :
226 0 : *memDesc = static_cast<void *>(blob.data());
227 0 : *memDescLen = static_cast<uint32_t>(blob.size());
228 0 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][MemoryExport] success memHandle[%p] memDescLen[%u] rdmaSerLen[%zu]", memHandle,
229 : *memDescLen, ser.size());
230 0 : return HCCL_SUCCESS;
231 1 : }
232 :
233 4 : HcclResult AicpuTsRoceRegedMemMgr::GetParamsFromMemDesc(const void *memDesc, uint32_t descLen,
234 : EndpointDesc &endpointDesc, std::string &rdmaBlob)
235 : {
236 4 : CHK_PTR_NULL(memDesc);
237 3 : if (descLen < sizeof(EndpointDesc)) {
238 1 : HCCL_ERROR("[AicpuTsRoceRegedMemMgr][GetParamsFromMemDesc] descLen[%u] too small", descLen);
239 1 : return HCCL_E_PARA;
240 : }
241 2 : const auto *base = static_cast<const char *>(memDesc);
242 2 : if (memcpy_s(&endpointDesc, sizeof(EndpointDesc), base + descLen - sizeof(EndpointDesc), sizeof(EndpointDesc)) != EOK) {
243 0 : HCCL_ERROR("[AicpuTsRoceRegedMemMgr][GetParamsFromMemDesc] endpointDesc copy failed");
244 0 : return HCCL_E_INTERNAL;
245 : }
246 2 : rdmaBlob.assign(base, base + descLen - sizeof(EndpointDesc));
247 2 : return HCCL_SUCCESS;
248 : }
249 :
250 2 : HcclResult AicpuTsRoceRegedMemMgr::MemoryImport(const void *memDesc, uint32_t descLen, HcommMem *outMem)
251 : {
252 2 : HCCL_INFO("[%s] Begin", __FUNCTION__);
253 2 : CHK_PTR_NULL(outMem);
254 :
255 1 : EndpointDesc endpointDesc{};
256 1 : std::string rdmaBlob;
257 1 : CHK_RET(GetParamsFromMemDesc(memDesc, descLen, endpointDesc, rdmaBlob));
258 :
259 0 : std::shared_ptr<hccl::RemoteRdmaRmaBuffer> remoteBuf;
260 0 : EXCEPTION_CATCH(remoteBuf = std::make_shared<hccl::RemoteRdmaRmaBuffer>(), return HCCL_E_PTR);
261 0 : CHK_RET(remoteBuf->Deserialize(rdmaBlob));
262 :
263 0 : if (remoteRdmaRmaBufferMgrs_.find(endpointDesc) == remoteRdmaRmaBufferMgrs_.end()) {
264 0 : std::unique_ptr<RemoteRdmaRmaBufferMgr> mgr;
265 0 : EXCEPTION_CATCH(mgr = std::make_unique<RemoteRdmaRmaBufferMgr>(), return HCCL_E_PTR);
266 0 : CHK_SMART_PTR_NULL(mgr);
267 0 : remoteRdmaRmaBufferMgrs_[endpointDesc] = std::move(mgr);
268 0 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][MemoryImport] created RemoteRdmaRmaBufferMgr for new endpoint, mgrCnt[%zu]",
269 : remoteRdmaRmaBufferMgrs_.size());
270 0 : }
271 :
272 0 : hccl::BufferKey<uintptr_t, u64> tempKey(reinterpret_cast<uintptr_t>(remoteBuf->GetAddr()), remoteBuf->GetSize());
273 0 : auto resultPair = remoteRdmaRmaBufferMgrs_[endpointDesc]->Add(tempKey, remoteBuf);
274 0 : if (!resultPair.second) {
275 0 : HCCL_ERROR("[AicpuTsRoceRegedMemMgr][MemoryImport] memDesc already imported");
276 0 : return HCCL_E_AGAIN;
277 : }
278 :
279 0 : outMem->addr = remoteBuf->GetAddr();
280 0 : outMem->size = remoteBuf->GetSize();
281 0 : outMem->type = COMM_MEM_TYPE_DEVICE;
282 0 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][MemoryImport] success descLen[%u] outMem addr[%p] size[%llu], mgrCnt[%zu]",
283 : descLen, outMem->addr, static_cast<unsigned long long>(outMem->size), remoteRdmaRmaBufferMgrs_.size());
284 0 : return HCCL_SUCCESS;
285 1 : }
286 :
287 2 : HcclResult AicpuTsRoceRegedMemMgr::MemoryUnimport(const void *memDesc, uint32_t descLen)
288 : {
289 2 : HCCL_INFO("[%s] Begin", __FUNCTION__);
290 :
291 2 : EndpointDesc endpointDesc{};
292 2 : std::string rdmaBlob;
293 2 : CHK_RET(GetParamsFromMemDesc(memDesc, descLen, endpointDesc, rdmaBlob));
294 :
295 1 : auto mgrIt = remoteRdmaRmaBufferMgrs_.find(endpointDesc);
296 1 : if (mgrIt == remoteRdmaRmaBufferMgrs_.end()) {
297 1 : HCCL_ERROR("[AicpuTsRoceRegedMemMgr][MemoryUnimport] remote mgr not found");
298 1 : return HCCL_E_NOT_FOUND;
299 : }
300 :
301 0 : std::shared_ptr<hccl::RemoteRdmaRmaBuffer> probe;
302 0 : EXCEPTION_CATCH(probe = std::make_shared<hccl::RemoteRdmaRmaBuffer>(), return HCCL_E_PTR);
303 0 : CHK_RET(probe->Deserialize(rdmaBlob));
304 :
305 0 : hccl::BufferKey<uintptr_t, u64> tempKey(reinterpret_cast<uintptr_t>(probe->GetAddr()), probe->GetSize());
306 0 : bool delOk = false;
307 0 : EXCEPTION_CATCH(delOk = mgrIt->second->Del(tempKey), return HCCL_E_NOT_FOUND);
308 0 : if (!delOk) {
309 0 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][MemoryUnimport] ref count > 0");
310 0 : return HCCL_E_AGAIN;
311 : }
312 0 : if (mgrIt->second->size() == 0) {
313 0 : remoteRdmaRmaBufferMgrs_.erase(mgrIt);
314 0 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][MemoryUnimport] erased empty remote mgr, descLen[%u] remainingMgrCnt[%zu]",
315 : descLen, remoteRdmaRmaBufferMgrs_.size());
316 : }
317 0 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][MemoryUnimport] success descLen[%u] key {%p, %llu}", descLen, probe->GetAddr(),
318 : static_cast<unsigned long long>(probe->GetSize()));
319 0 : return HCCL_SUCCESS;
320 2 : }
321 :
322 7 : HcclResult AicpuTsRoceRegedMemMgr::GetAllMemHandles(void **memHandles, uint32_t *memHandleNum)
323 : {
324 7 : HCCL_INFO("[%s] Begin", __FUNCTION__);
325 7 : CHK_PTR_NULL(memHandleNum);
326 :
327 6 : *memHandleNum = static_cast<uint32_t>(hcclBufRecords_.size());
328 6 : if (*memHandleNum == 0U) {
329 3 : *memHandles = nullptr;
330 3 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][GetAllMemHandles] no records, memHandleNum[0]");
331 3 : return HCCL_SUCCESS;
332 : }
333 3 : *memHandles = static_cast<void *>(hcclBufRecords_.data());
334 3 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][GetAllMemHandles] memHandleNum[%u] hcclBufRecords[%p]", *memHandleNum,
335 : static_cast<void *>(hcclBufRecords_.data()));
336 3 : return HCCL_SUCCESS;
337 : }
338 :
339 1 : HcclResult AicpuTsRoceRegedMemMgr::GatherLocalMemDetails(std::vector<RoceMemDetails> &localOut) const
340 : {
341 1 : localOut.reserve(allRegisteredBuffers_.size());
342 1 : for (const auto &entry : allRegisteredBuffers_) {
343 0 : if (entry.second) {
344 0 : continue;
345 : }
346 0 : const auto &buf = entry.first;
347 0 : if (buf == nullptr) {
348 0 : continue;
349 : }
350 0 : auto *rma = static_cast<hccl::RmaBuffer *>(buf.get());
351 0 : CHK_PTR_NULL(rma->GetDevAddr());
352 0 : RoceMemDetails r{};
353 0 : r.addr = static_cast<u64>(reinterpret_cast<uintptr_t>(rma->GetAddr()));
354 0 : r.devAddr = static_cast<u64>(reinterpret_cast<uintptr_t>(rma->GetDevAddr()));
355 0 : r.size = buf->GetSize();
356 0 : r.key = buf->GetKey();
357 0 : localOut.push_back(r);
358 0 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][GetAllMemDetails][local][%zu] addr[0x%llx] devAddr[0x%llx] size[%llu] key[%u]",
359 : localOut.size() - 1U, static_cast<unsigned long long>(r.addr), static_cast<unsigned long long>(r.devAddr),
360 : static_cast<unsigned long long>(r.size), r.key);
361 : }
362 1 : return HCCL_SUCCESS;
363 : }
364 :
365 1 : HcclResult AicpuTsRoceRegedMemMgr::AppendLocalNotifyMemDetails(std::vector<RoceMemDetails> &localOut) const
366 : {
367 1 : CHK_PTR_NULL(rdmaHandle_);
368 0 : CHK_PTR_NULL(netDev_);
369 0 : auto *netCtx = static_cast<hccl::NetDevContext *>(netDev_);
370 0 : struct MrInfoT mrInfo{};
371 0 : CHK_RET(HrtRaGetNotifyMrInfo(static_cast<u32>(netCtx->GetPhyId()), rdmaHandle_, &mrInfo));
372 0 : CHK_PTR_NULL(mrInfo.addr);
373 0 : RoceMemDetails notifyMd{};
374 0 : notifyMd.addr = static_cast<u64>(reinterpret_cast<uintptr_t>(mrInfo.addr));
375 0 : notifyMd.devAddr = notifyMd.addr;
376 0 : notifyMd.size = static_cast<u64>(mrInfo.size);
377 0 : notifyMd.key = mrInfo.lkey;
378 0 : localOut.push_back(notifyMd);
379 0 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][GetAllMemDetails][local][notify] addr[0x%llx] devAddr[0x%llx] size[%llu] key[%u]",
380 : static_cast<unsigned long long>(notifyMd.addr), static_cast<unsigned long long>(notifyMd.devAddr),
381 : static_cast<unsigned long long>(notifyMd.size), notifyMd.key);
382 0 : return HCCL_SUCCESS;
383 : }
384 :
385 0 : void AicpuTsRoceRegedMemMgr::GatherRemoteMemDetails(std::vector<RoceMemDetails> &remoteOut) const
386 : {
387 0 : for (const auto &epMgr : remoteRdmaRmaBufferMgrs_) {
388 0 : const auto &mgr = epMgr.second;
389 0 : if (mgr == nullptr) {
390 0 : continue;
391 : }
392 0 : mgr->ForEach([&remoteOut](const hccl::BufferKey<uintptr_t, u64> &, const std::shared_ptr<hccl::RemoteRdmaRmaBuffer> &rb) {
393 0 : if (rb == nullptr) {
394 0 : return;
395 : }
396 0 : auto *rma = static_cast<hccl::RmaBuffer *>(rb.get());
397 0 : if (rma->GetDevAddr() == nullptr) {
398 0 : return;
399 : }
400 0 : RoceMemDetails r{};
401 0 : r.addr = static_cast<u64>(reinterpret_cast<uintptr_t>(rma->GetAddr()));
402 0 : r.devAddr = static_cast<u64>(reinterpret_cast<uintptr_t>(rma->GetDevAddr()));
403 0 : r.size = rb->GetSize();
404 0 : r.key = rb->GetKey();
405 0 : remoteOut.push_back(r);
406 0 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][GetAllMemDetails][remote][%zu] addr[0x%llx] devAddr[0x%llx] size[%llu] key[%u]",
407 : remoteOut.size() - 1U, static_cast<unsigned long long>(r.addr), static_cast<unsigned long long>(r.devAddr),
408 : static_cast<unsigned long long>(r.size), r.key);
409 : });
410 : }
411 0 : }
412 :
413 1 : HcclResult AicpuTsRoceRegedMemMgr::GetAllMemDetails(std::vector<RoceMemDetails> &localOut,
414 : std::vector<RoceMemDetails> &remoteOut) const
415 : {
416 1 : localOut.clear();
417 1 : remoteOut.clear();
418 1 : CHK_RET(GatherLocalMemDetails(localOut));
419 1 : CHK_RET(AppendLocalNotifyMemDetails(localOut));
420 0 : GatherRemoteMemDetails(remoteOut);
421 0 : HCCL_INFO("[AicpuTsRoceRegedMemMgr][GetAllMemDetails] summary localCnt[%zu] remoteCnt[%zu] remoteMgrCnt[%zu]",
422 : localOut.size(), remoteOut.size(), remoteRdmaRmaBufferMgrs_.size());
423 0 : return HCCL_SUCCESS;
424 : }
425 :
426 : } // namespace hcomm
|
