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