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