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