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(
90 : "[%s] buf[%p] or buf->handle or outDesc[%p] or outDescLen[%p] is null", __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__, description, outBuf, netDev);
123 0 : return HCCL_E_PTR;
124 : }
125 : (void)(isRemote);
126 3 : HCCL_INFO("[%s] Begin, descLen[%llu]", __func__, descLen);
127 : // 仅支持UB类型
128 1 : HcclNetDevice* hcclNetDevice = static_cast<HcclNetDevice*>(netDev);
129 1 : if (!hcclNetDevice->IsUB()) {
130 0 : HCCL_ERROR("[%s] only support UB", __func__);
131 0 : return HCCL_E_NOT_SUPPORT;
132 : }
133 :
134 : // 反序列化
135 1 : std::vector<char> tempDesc{};
136 1 : tempDesc.resize(TRANSPORT_EMD_ESC_SIZE);
137 1 : tempDesc.assign(description, description + descLen);
138 1 : ExchangeUbBufferDto dto;
139 1 : BinaryStream remoteRdmaRmaBufferStream(tempDesc);
140 1 : dto.Deserialize(remoteRdmaRmaBufferStream);
141 :
142 : // 构造RemoteUbRmaBuffer
143 1 : RemoteUbRmaBuffer* remoteUbRmaBuffer = new (std::nothrow) RemoteUbRmaBuffer(hcclNetDevice->GetRdmaHandle(), dto);
144 1 : if (remoteUbRmaBuffer == nullptr) {
145 0 : HCCL_ERROR("[%s] Failed to allocate RemoteUbRmaBuffer", __func__);
146 0 : return HCCL_E_PTR;
147 : }
148 :
149 : // 填充HcclBuf
150 1 : outBuf->addr = reinterpret_cast<void*>(remoteUbRmaBuffer->GetAddr());
151 1 : outBuf->len = remoteUbRmaBuffer->GetSize();
152 1 : outBuf->handle = static_cast<void*>(remoteUbRmaBuffer);
153 3 : HCCL_INFO("[%s]End, addr[%p], size[%llu], handle[%p]", __func__, outBuf->addr, outBuf->len, outBuf->handle);
154 1 : return HCCL_SUCCESS;
155 1 : }
156 :
157 1 : HcclResult HcclMemCloseV2(HcclBuf* buf)
158 : {
159 1 : if (buf == nullptr || buf->handle == nullptr) {
160 0 : HCCL_ERROR("[%s] buf[%p] or buf->handle is null", __func__, buf);
161 0 : return HCCL_E_PTR;
162 : }
163 3 : HCCL_INFO("[%s] Begin, addr[%p], size[%llu], handle[%p]", __func__, buf->addr, buf->len, buf->handle);
164 : // 仅支持UB类型
165 1 : RemoteRmaBuffer* remoteRmaBuffer = static_cast<RemoteRmaBuffer*>(buf->handle);
166 1 : if (remoteRmaBuffer->GetRmaType() != RmaType::UB) {
167 0 : HCCL_ERROR("[%s] only support UB", __func__);
168 0 : return HCCL_E_NOT_SUPPORT;
169 : }
170 :
171 : // 删除RemoteUbRmaBuffer
172 3 : HCCL_INFO("[HcclMemCloseV2][Ub] CloseMem");
173 1 : RemoteUbRmaBuffer* remoteUbRmaBuffer = static_cast<RemoteUbRmaBuffer*>(remoteRmaBuffer);
174 1 : delete remoteUbRmaBuffer;
175 1 : return HCCL_SUCCESS;
176 : }
|
