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