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 "hccl_mem_v2.h"
11 : #include "log.h"
12 : #include "exchange_ub_buffer_dto.h"
13 : #include "local_ub_rma_buffer_manager.h"
14 : #include "remote_rma_buffer.h"
15 : #include "local_ub_rma_buffer.h"
16 :
17 : using namespace Hccl;
18 :
19 2 : HcclResult HcclMemRegV2(HcclNetDev netDev, const HcclMem *mem, HcclBuf *buf)
20 : {
21 2 : if (netDev == nullptr || mem == nullptr || buf == nullptr) {
22 0 : HCCL_ERROR("[%s] netDev[%p] or mem[%p] or buf[%p] is null", __func__, netDev, mem, buf);
23 0 : return HCCL_E_PTR;
24 : }
25 6 : HCCL_INFO("[%s] Begin, addr[%p], size[%llu], type[%d]", __func__, mem->addr, mem->size, mem->type);
26 : // 仅支持UB类型
27 2 : HcclNetDevice *hcclNetDevice = static_cast<HcclNetDevice *>(netDev);
28 2 : if (!hcclNetDevice->IsUB()) {
29 0 : HCCL_ERROR("[%s] only support UB", __func__);
30 0 : return HCCL_E_NOT_SUPPORT;
31 : }
32 :
33 : // 构造LocalUbRmaBuffer
34 2 : auto getBuffFunc = [&]() -> HcclResult {
35 : std::shared_ptr<Buffer> localBufferPtr
36 2 : = make_shared<Buffer>(reinterpret_cast<uintptr_t>(mem->addr), mem->size, mem->type);
37 : std::shared_ptr<LocalUbRmaBuffer> localUbRmaBuffer
38 2 : = make_shared<LocalUbRmaBuffer>(localBufferPtr, hcclNetDevice, false);
39 2 : LocalUbRmaBufferMgr *localRmaBufferMgr = LocalUbRmaBufferManager::GetInstance();
40 :
41 : // 注册到LocalUbRmaBuffer计数器
42 2 : BufferKey<uintptr_t, u64> tempKey(reinterpret_cast<uintptr_t>(mem->addr), mem->size);
43 2 : auto resultPair = localRmaBufferMgr->Add(tempKey, localUbRmaBuffer);
44 2 : if (resultPair.first == localRmaBufferMgr->End()) {
45 : // 若已注册内存有交叉，返回HCCL_E_INTERNAL
46 0 : HCCL_ERROR("[%s]The memory overlaps with the memory that has been registered.", __func__);
47 0 : return HCCL_E_INTERNAL;
48 : }
49 2 : buf->addr = mem->addr;
50 2 : buf->len = mem->size;
51 2 : buf->handle = resultPair.first->second.buffer.get();
52 2 : return HCCL_SUCCESS;
53 2 : };
54 2 : TRY_CATCH_RETURN(getBuffFunc());
55 :
56 6 : HCCL_INFO("[%s]End, addr[%p], size[%llu], handle[%p]", __func__, buf->addr, buf->len, buf->handle);
57 2 : return HCCL_SUCCESS;
58 : }
59 :
60 2 : HcclResult HcclMemDeregV2(const HcclBuf *buf)
61 : {
62 2 : if (buf == nullptr) {
63 0 : HCCL_ERROR("[%s]buf[%p] is null", __func__, buf);
64 0 : return HCCL_E_PTR;
65 : }
66 6 : HCCL_INFO("[%s] Begin, addr[%p], size[%llu], handle[%p]", __func__, buf->addr, buf->len, buf->handle);
67 : // 从LocalRamBuffer计数器删除HcclBuf
68 2 : LocalUbRmaBufferMgr *localRmaBufferMgr = LocalUbRmaBufferManager::GetInstance();
69 2 : BufferKey<uintptr_t, u64> tempKey(reinterpret_cast<uintptr_t>(buf->addr), buf->len);
70 : try {
71 2 : auto resultPair = localRmaBufferMgr->Del(tempKey);
72 : // 计数器大于1时，返回false，说明框架层有其它设备在使用这段内存，返回HCCL_E_AGAIN
73 2 : if (!resultPair) {
74 0 : HCCL_INFO("[HcclOneSidedService][DeregMem]Memory reference count is larger than 0"
75 : "(used by other RemoteRank), do not deregister memory.");
76 0 : return HCCL_E_AGAIN;
77 : }
78 2 : return HCCL_SUCCESS;
79 0 : } catch (const std::out_of_range &e) {
80 : // 若计数器内未找到buf，返回HCCL_E_NOT_FOUND
81 0 : HCCL_ERROR("[%s] %s", __func__, e.what());
82 0 : return HCCL_E_NOT_FOUND;
83 0 : }
84 : }
85 :
86 3 : HcclResult HcclMemExportV2(HcclBuf *buf, char **outDesc, uint64_t *outDescLen)
87 : {
88 3 : if (buf == nullptr || buf->handle == nullptr || outDesc == nullptr || outDescLen == nullptr) {
89 0 : HCCL_ERROR("[%s] buf[%p] or buf->handle or outDesc[%p] or outDescLen[%p] is null",
90 : __func__, buf, outDesc, outDescLen);
91 0 : return HCCL_E_PTR;
92 : }
93 9 : HCCL_INFO("[%s] Begin, addr[%p], size[%llu], handle[%p]", __func__, buf->addr, buf->len, buf->handle);
94 : // 获取序列化信息
95 3 : LocalUbRmaBuffer *localUbRmaBuffer = reinterpret_cast<LocalUbRmaBuffer *>(buf->handle);
96 3 : std::unique_ptr<Serializable> dto = localUbRmaBuffer->GetExchangeDto();
97 3 : BinaryStream localRdmaRmaBufferStream;
98 3 : dto->Serialize(localRdmaRmaBufferStream);
99 3 : std::vector<char> tempLocalMemDesc;
100 3 : localRdmaRmaBufferStream.Dump(tempLocalMemDesc);
101 9 : HCCL_DEBUG("[%s] dump data size [%zu]", __func__, tempLocalMemDesc.size());
102 : // 判断内存描述符是否正确导出
103 3 : if (tempLocalMemDesc.empty()) {
104 0 : HCCL_ERROR("[%s] tempLocalMemDesc export failed.", __func__);
105 0 : return HCCL_E_INTERNAL;
106 : }
107 :
108 : // 内存描述符拷贝
109 3 : *outDescLen = tempLocalMemDesc.size();
110 3 : if (memcpy_s(*outDesc, TRANSPORT_EMD_ESC_SIZE, tempLocalMemDesc.data(), tempLocalMemDesc.size()) != EOK) {
111 0 : HCCL_ERROR("[%s] tempLocalMemDesc copy error. aim size:[%zu]", __func__, tempLocalMemDesc.size());
112 0 : return HCCL_E_INTERNAL;
113 : }
114 :
115 9 : HCCL_INFO("[%s]End, outDescLen[%llu]", __func__, *outDescLen);
116 3 : return HCCL_SUCCESS;
117 3 : }
118 :
119 1 : HcclResult HcclMemImportV2(const char *description, uint64_t descLen, bool isRemote, HcclBuf *outBuf, HcclNetDev netDev)
120 : {
121 1 : if (description == nullptr || outBuf == nullptr || netDev == nullptr) {
122 0 : HCCL_ERROR("[%s] description[%p] or outBuf[%p] or netDev[%p] is null", __func__,
123 : description, outBuf, netDev);
124 0 : return HCCL_E_PTR;
125 : }
126 : (void)(isRemote);
127 3 : HCCL_INFO("[%s] Begin, descLen[%llu]", __func__, descLen);
128 : // 仅支持UB类型
129 1 : HcclNetDevice *hcclNetDevice = static_cast<HcclNetDevice *>(netDev);
130 1 : if (!hcclNetDevice->IsUB()) {
131 0 : HCCL_ERROR("[%s] only support UB", __func__);
132 0 : return HCCL_E_NOT_SUPPORT;
133 : }
134 :
135 : // 反序列化
136 1 : std::vector<char> tempDesc{};
137 1 : tempDesc.resize(TRANSPORT_EMD_ESC_SIZE);
138 1 : tempDesc.assign(description, description + descLen);
139 1 : ExchangeUbBufferDto dto;
140 1 : BinaryStream remoteRdmaRmaBufferStream(tempDesc);
141 1 : dto.Deserialize(remoteRdmaRmaBufferStream);
142 :
143 : // 构造RemoteUbRmaBuffer
144 1 : RemoteUbRmaBuffer *remoteUbRmaBuffer = new(std::nothrow) RemoteUbRmaBuffer(hcclNetDevice->GetRdmaHandle(), dto);
145 1 : if(remoteUbRmaBuffer == nullptr) {
146 0 : HCCL_ERROR("[%s] Failed to allocate RemoteUbRmaBuffer", __func__);
147 0 : return HCCL_E_PTR;
148 : }
149 :
150 : // 填充HcclBuf
151 1 : outBuf->addr = reinterpret_cast<void *>(remoteUbRmaBuffer->GetAddr());
152 1 : outBuf->len = remoteUbRmaBuffer->GetSize();
153 1 : outBuf->handle = static_cast<void *>(remoteUbRmaBuffer);
154 3 : HCCL_INFO("[%s]End, addr[%p], size[%llu], handle[%p]", __func__, outBuf->addr, outBuf->len, outBuf->handle);
155 1 : return HCCL_SUCCESS;
156 1 : }
157 :
158 1 : HcclResult HcclMemCloseV2(HcclBuf *buf)
159 : {
160 1 : if (buf == nullptr || buf->handle == nullptr) {
161 0 : HCCL_ERROR("[%s] buf[%p] or buf->handle is null", __func__, buf);
162 0 : return HCCL_E_PTR;
163 : }
164 3 : HCCL_INFO("[%s] Begin, addr[%p], size[%llu], handle[%p]", __func__, buf->addr, buf->len, buf->handle);
165 : // 仅支持UB类型
166 1 : RemoteRmaBuffer *remoteRmaBuffer = static_cast<RemoteRmaBuffer *>(buf->handle);
167 1 : if (remoteRmaBuffer->GetRmaType() != RmaType::UB) {
168 0 : HCCL_ERROR("[%s] only support UB", __func__);
169 0 : return HCCL_E_NOT_SUPPORT;
170 : }
171 :
172 : // 删除RemoteUbRmaBuffer
173 3 : HCCL_INFO("[HcclMemCloseV2][Ub] CloseMem");
174 1 : RemoteUbRmaBuffer *remoteUbRmaBuffer = static_cast<RemoteUbRmaBuffer *>(remoteRmaBuffer);
175 1 : delete remoteUbRmaBuffer;
176 1 : return HCCL_SUCCESS;
177 : }
|
