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 : #include "endpoint_pair.h"
11 : #include "log.h"
12 : #include "roce_mem.h"
13 : #include "exchange_rdma_buffer_dto.h"
14 : #include "local_rdma_rma_buffer.h"
15 : #include "hccl_one_sided_data.h"
16 :
17 : namespace hcomm {
18 :
19 16 : RoceRegedMemMgr::RoceRegedMemMgr()
20 : {
21 16 : localRdmaRmaBufferMgr_ = std::make_unique<LocalRdmaRmaBufferMgr>();
22 16 : }
23 :
24 20 : HcclResult RoceRegedMemMgr::RegisterMemory(HcommMem mem, const char *memTag, void **memHandle)
25 : {
26 20 : HCCL_INFO("[%s] Begin", __FUNCTION__);
27 20 : CHK_PTR_NULL(localRdmaRmaBufferMgr_);
28 20 : std::lock_guard<std::mutex> lock(memMtx_);
29 40 : return RegisterMemoryImpl(mem, memTag, memHandle,
30 20 : localRdmaRmaBufferMgr_, allRegisteredBuffers_, &handlesRecords_, "RoceRegedMemMgr",
31 0 : [&](auto& bufPtr, auto& parent) {
32 9 : return std::make_shared<Hccl::LocalRdmaRmaBuffer>(bufPtr, rdmaHandle_,
33 18 : parent->GetLkey(), parent->GetRkey(), parent->GetMrHandle());
34 : },
35 20 : [&](auto& bufPtr) {
36 10 : return std::make_shared<Hccl::LocalRdmaRmaBuffer>(bufPtr, rdmaHandle_);
37 20 : });
38 20 : }
39 :
40 13 : HcclResult RoceRegedMemMgr::UnregisterMemory(void* memHandle)
41 : {
42 13 : HCCL_INFO("[%s] Begin", __FUNCTION__);
43 13 : CHK_PTR_NULL(localRdmaRmaBufferMgr_);
44 13 : std::lock_guard<std::mutex> lock(memMtx_);
45 13 : return UnregisterMemoryImpl(memHandle, localRdmaRmaBufferMgr_, allRegisteredBuffers_, &handlesRecords_,
46 12 : [](auto* b) { return b->GetLkey(); },
47 19 : [](auto a, auto b) { return a == b; });
48 13 : }
49 :
50 0 : HcclResult RoceRegedMemMgr::GetMemDesc(const EndpointDesc endpointDesc, Hccl::LocalRdmaRmaBuffer *localRdmaRmaBuffer)
51 : {
52 0 : auto dto = localRdmaRmaBuffer->GetExchangeDto();
53 0 : Hccl::BinaryStream localRdmaRmaBufferStream;
54 0 : dto->Serialize(localRdmaRmaBufferStream);
55 0 : std::vector<char> tempLocalMemDesc;
56 0 : localRdmaRmaBufferStream.Dump(tempLocalMemDesc);
57 0 : HCCL_DEBUG("[RoceRegedMemMgr][GetMemDesc] [%s] dump data size [%u]", __func__, tempLocalMemDesc.size());
58 : // 判断内存描述符是否正确导出
59 0 : if (tempLocalMemDesc.empty()) {
60 0 : HCCL_ERROR("[RoceRegedMemMgr][GetMemDesc] [%s] tempLocalMemDesc export failed.", __func__);
61 0 : return HCCL_E_INTERNAL;
62 : }
63 :
64 0 : std::vector<char> tempLocalEndpointDesc;
65 0 : tempLocalEndpointDesc.resize(sizeof(EndpointDesc));
66 0 : if(memcpy_s(tempLocalEndpointDesc.data(), sizeof(EndpointDesc), &endpointDesc, sizeof(EndpointDesc)) != EOK) {
67 0 : HCCL_ERROR("[RoceRegedMemMgr][GetMemDesc] [%s] endpointDesc memcpy_s failed.", __func__);
68 0 : return HCCL_E_INTERNAL;
69 : }
70 :
71 0 : tempLocalMemDesc.insert(tempLocalMemDesc.end(),
72 : tempLocalEndpointDesc.begin(),
73 : tempLocalEndpointDesc.end());
74 :
75 : // 内存描述符拷贝
76 0 : localRdmaRmaBuffer->Desc = std::move(tempLocalMemDesc);
77 0 : return HCCL_SUCCESS;
78 0 : }
79 :
80 1 : HcclResult RoceRegedMemMgr::MemoryExport(const EndpointDesc endpointDesc, void *memHandle, void **memDesc, uint32_t *memDescLen)
81 : {
82 1 : HCCL_INFO("[%s] Begin", __FUNCTION__);
83 1 : CHK_PTR_NULL(memHandle);
84 1 : CHK_PTR_NULL(memDesc);
85 1 : CHK_PTR_NULL(memDescLen);
86 1 : std::lock_guard<std::mutex> lock(memMtx_);
87 :
88 1 : Hccl::LocalRdmaRmaBuffer *localRdmaRmaBuffer = nullptr;
89 1 : CHK_RET(ValidateMemExportHandle(memHandle, allRegisteredBuffers_, localRdmaRmaBuffer));
90 :
91 : // 获取序列化信息
92 0 : CHK_RET(GetMemDesc(endpointDesc, localRdmaRmaBuffer));
93 :
94 0 : *memDescLen = static_cast<uint32_t>(localRdmaRmaBuffer->Desc.size());
95 0 : *memDesc = static_cast<void *>(localRdmaRmaBuffer->Desc.data());
96 0 : return HCCL_SUCCESS;
97 1 : }
98 :
99 2 : HcclResult RoceRegedMemMgr::GetParamsFromMemDesc(const void *memDesc, uint32_t descLen,
100 : EndpointDesc &endpointDesc, Hccl::ExchangeRdmaBufferDto &dto)
101 : {
102 2 : const char *description = static_cast<const char *>(memDesc);
103 :
104 2 : if (descLen < sizeof(EndpointDesc)) {
105 1 : HCCL_ERROR("[RoceRegedMemMgr][GetParamsFromMemDesc] [%s] descLen[%u] is too small. aim size:[%llu]", __func__, descLen, sizeof(EndpointDesc));
106 1 : return HCCL_E_INTERNAL;
107 : }
108 : // 从memDesc末尾提取EndpointDesc
109 1 : if (memcpy_s(&endpointDesc, sizeof(EndpointDesc), description + descLen - sizeof(EndpointDesc), sizeof(EndpointDesc)) != EOK) {
110 0 : HCCL_ERROR("[RoceRegedMemMgr][GetParamsFromMemDesc] [%s] endpointDesc copy error. aim size:[%llu]", __func__, sizeof(EndpointDesc));
111 0 : return HCCL_E_INTERNAL;
112 : }
113 :
114 : // 反序列化
115 1 : std::vector<char> tempDesc{};
116 1 : tempDesc.resize(TRANSPORT_EMD_ESC_SIZE);
117 1 : tempDesc.assign(description, description + descLen - sizeof(EndpointDesc));
118 1 : Hccl::BinaryStream remoteRdmaRmaBufferStream(tempDesc);
119 1 : dto.Deserialize(remoteRdmaRmaBufferStream);
120 1 : return HCCL_SUCCESS;
121 1 : }
122 :
123 0 : HcclResult RoceRegedMemMgr::MemoryImport(const void *memDesc, uint32_t descLen, HcommMem *outMem)
124 : {
125 0 : HCCL_INFO("[%s] Begin", __FUNCTION__);
126 0 : std::lock_guard<std::mutex> lock(memMtx_);
127 :
128 : EndpointDesc endpointDesc;
129 0 : Hccl::ExchangeRdmaBufferDto dto;
130 0 : CHK_RET(GetParamsFromMemDesc(memDesc, descLen, endpointDesc, dto));
131 :
132 : // 构造RemoteRdmaRmaBuffer
133 0 : std::shared_ptr<Hccl::RemoteRdmaRmaBuffer> remoteRdmaRmaBuffer;
134 0 : EXCEPTION_CATCH(
135 : remoteRdmaRmaBuffer = std::make_shared<Hccl::RemoteRdmaRmaBuffer>(rdmaHandle_, dto),
136 : return HCCL_E_PTR;
137 : );
138 :
139 : // 放到RemoteRdmaRmaBufferMgr_
140 0 : hccl::BufferKey<uintptr_t, u64> tempKey(static_cast<uintptr_t>(dto.addr), dto.size);
141 0 : if(remoteRdmaRmaBufferMgrs_.find(endpointDesc) == remoteRdmaRmaBufferMgrs_.end()) {
142 0 : std::unique_ptr<RemoteRdmaRmaBufferMgr> remoteRdmaRmaBufferMgr;
143 0 : EXCEPTION_CATCH((remoteRdmaRmaBufferMgr = std::make_unique<RemoteRdmaRmaBufferMgr>()),
144 : return HCCL_E_PTR);
145 0 : CHK_SMART_PTR_NULL(remoteRdmaRmaBufferMgr);
146 0 : remoteRdmaRmaBufferMgrs_[endpointDesc] = std::move(remoteRdmaRmaBufferMgr);
147 0 : HCCL_INFO("remoteRdmaRmaBufferMgrs_ add remoteRdmaRmaBufferMgr successfully!");
148 0 : }
149 :
150 0 : auto resultPair = remoteRdmaRmaBufferMgrs_[endpointDesc]->Add(tempKey, remoteRdmaRmaBuffer);
151 0 : if(!resultPair.second) {
152 0 : HCCL_ERROR("[RoceRegedMemMgr][MemoryImport] This memDesc has already been imported!");
153 0 : return HCCL_E_AGAIN;
154 : }
155 :
156 0 : outMem->addr = reinterpret_cast<void *>(remoteRdmaRmaBuffer->GetAddr());
157 0 : outMem->size = remoteRdmaRmaBuffer->GetSize();
158 :
159 0 : return HCCL_SUCCESS;
160 0 : }
161 :
162 0 : HcclResult RoceRegedMemMgr::MemoryUnimport(const void *memDesc, uint32_t descLen)
163 : {
164 0 : HCCL_INFO("[%s] Begin", __FUNCTION__);
165 0 : std::lock_guard<std::mutex> lock(memMtx_);
166 :
167 : EndpointDesc endpointDesc;
168 0 : Hccl::ExchangeRdmaBufferDto dto;
169 0 : CHK_RET(GetParamsFromMemDesc(memDesc, descLen, endpointDesc, dto));
170 :
171 0 : if(remoteRdmaRmaBufferMgrs_.find(endpointDesc) == remoteRdmaRmaBufferMgrs_.end()) {
172 0 : HCCL_ERROR("[RoceRegedMemMgr][MemoryUnimport] Remote buffer manager Not Found.");
173 0 : return HCCL_E_NOT_FOUND;
174 : }
175 :
176 : // 删除RemoteRdmaRmaBuffer
177 0 : HCCL_INFO("[MemoryUnimport][Rdma] MemoryUnimport");
178 0 : hccl::BufferKey<uintptr_t, u64> tempKey(static_cast<uintptr_t>(dto.addr), dto.size);
179 :
180 0 : bool resultPair = false;
181 0 : EXCEPTION_CATCH(resultPair = remoteRdmaRmaBufferMgrs_[endpointDesc]->Del(tempKey), return HCCL_E_NOT_FOUND);
182 : // 计数器大于1时，返回false，说明框架层有其它设备在使用这段内存，返回HCCL_E_AGAIN
183 0 : if (!resultPair) {
184 0 : HCCL_INFO("[RoceRegedMemMgr][[MemoryUnimport] Memory reference count is larger than 0"
185 : "(used by other RemoteRank).");
186 0 : return HCCL_E_AGAIN;
187 : }
188 0 : if (!remoteRdmaRmaBufferMgrs_[endpointDesc]->size()) {
189 0 : remoteRdmaRmaBufferMgrs_.erase(endpointDesc);
190 : }
191 0 : return HCCL_SUCCESS;
192 0 : }
193 :
194 5 : HcclResult RoceRegedMemMgr::GetAllMemHandles(void **memHandles, uint32_t *memHandleNum)
195 : {
196 5 : HCCL_INFO("[%s] Begin", __FUNCTION__);
197 5 : std::lock_guard<std::mutex> lock(memMtx_);
198 5 : CHK_PTR_NULL(memHandles);
199 5 : CHK_PTR_NULL(memHandleNum);
200 5 : *memHandleNum = static_cast<uint32_t>(handlesRecords_.size());
201 5 : *memHandles = handlesRecords_.empty() ? nullptr : static_cast<void *>(handlesRecords_.data());
202 5 : HCCL_INFO("[RoceRegedMemMgr][GetAllMemHandles] memHandleNum[%u]", *memHandleNum);
203 5 : return HCCL_SUCCESS;
204 5 : }
205 :
206 : }
|
