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