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