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