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