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 "local_ub_rma_buffer.h"
12 :
13 : #include "null_ptr_exception.h"
14 : #include "invalid_params_exception.h"
15 : #include "exchange_ub_buffer_dto.h"
16 : #include "rdma_handle_manager.h"
17 :
18 : namespace Hccl {
19 :
20 : constexpr u32 TEN_MILLISECOND_OF_USLEEP = 10000;
21 :
22 525 : LocalUbRmaBuffer::LocalUbRmaBuffer(std::shared_ptr<Buffer> buf, RdmaHandle rdmaHandle)
23 : : LocalRmaBuffer(buf, RmaType::UB),
24 525 : rdmaHandle(rdmaHandle)
25 : {
26 525 : if (rdmaHandle == nullptr) {
27 1 : THROW<NullPtrException>("LocalUbRmaBuffer's rdmaHandle is nullptr");
28 : }
29 524 : std::pair<u64, u64> alignBuf = BufAlign(buf->GetAddr(), buf->GetSize());
30 :
31 524 : bufKey_ = BufferKey<uintptr_t, u64>{alignBuf.first, alignBuf.second};
32 524 : const auto& tokenIdInfoPair = RdmaHandleManager::GetInstance().GetTokenIdInfo(rdmaHandle, bufKey_);
33 524 : tokenIdHandle = tokenIdInfoPair.first;
34 524 : tokenId = tokenIdInfoPair.second;
35 524 : tokenValue = GetUbToken();
36 524 : HrtRaUbLocMemRegParam lmemReg{alignBuf.first, alignBuf.second, tokenValue, tokenIdHandle, 1};
37 524 : reqReg = HrtRaUbLocalMemReg(rdmaHandle, lmemReg);
38 524 : auto ret = memcpy_s(key, HRT_UB_MEM_KEY_MAX_LEN, reqReg.key, HRT_UB_MEM_KEY_MAX_LEN);
39 524 : if (ret != EOK) {
40 0 : HCCL_ERROR("[LocalUbRmaBuffer::%s] copy key failed, ret[%d].", __func__, ret);
41 0 : THROW<InvalidParamsException>("LocalUbRmaBuffer copy key failed, ret[%d]", ret);
42 : }
43 :
44 1572 : HCCL_INFO(
45 : "[LocalUbRmaBuffer::%s] end, rdmaHandle[%p], lmemHandle[0x%llx], reqReg.keySize[%u]", __func__, rdmaHandle,
46 : reqReg.handle, reqReg.keySize);
47 526 : }
48 :
49 0 : LocalUbRmaBuffer::LocalUbRmaBuffer(std::shared_ptr<Buffer> buf, RdmaHandle rdmaHandle, const LocalUbRmaBuffer& parent)
50 : : LocalRmaBuffer(buf, RmaType::UB, true),
51 0 : rdmaHandle(rdmaHandle),
52 0 : tokenValue(parent.tokenValue),
53 0 : tokenId(parent.tokenId),
54 0 : tokenIdHandle(parent.tokenIdHandle),
55 0 : reqReg(parent.reqReg)
56 : {
57 0 : if (rdmaHandle == nullptr) {
58 0 : THROW<NullPtrException>("LocalUbRmaBuffer alias rdmaHandle is nullptr");
59 : }
60 :
61 0 : auto ret = memcpy_s(key, HRT_UB_MEM_KEY_MAX_LEN, parent.key, HRT_UB_MEM_KEY_MAX_LEN);
62 0 : if (ret != EOK) {
63 0 : HCCL_ERROR("[LocalUbRmaBuffer::%s] alias copy key failed, ret[%d].", __func__, ret);
64 0 : THROW<InvalidParamsException>("LocalUbRmaBuffer alias copy key failed, ret[%d]", ret);
65 : }
66 :
67 0 : HCCL_INFO(
68 : "[LocalUbRmaBuffer::%s] alias, rdmaHandle[%p], lmemHandle[0x%llx], keySize[%u]", __func__, rdmaHandle,
69 : reqReg.handle, reqReg.keySize);
70 0 : }
71 :
72 2 : LocalUbRmaBuffer::LocalUbRmaBuffer(std::shared_ptr<Buffer> buf, void* netDevice, bool flag)
73 2 : : LocalRmaBuffer(buf, RmaType::UB)
74 : {
75 : (void)flag;
76 2 : if (netDevice == nullptr) {
77 0 : THROW<NullPtrException>("LocalUbRmaBuffer's netDevice is nullptr");
78 : }
79 2 : tokenValue = GetUbToken();
80 2 : netDev = reinterpret_cast<HcclNetDevice*>(netDevice);
81 2 : rdmaHandle = netDev->GetRdmaHandle();
82 :
83 2 : std::pair<u64, u64> alignBuf = BufAlign(buf->GetAddr(), buf->GetSize());
84 :
85 2 : bufKey_ = BufferKey<uintptr_t, u64>{alignBuf.first, alignBuf.second};
86 2 : const auto& tokenIdInfoPair = netDev->GetTokenIdInfo(bufKey_);
87 2 : tokenIdHandle = tokenIdInfoPair.first;
88 2 : tokenId = tokenIdInfoPair.second;
89 2 : tokenValue = GetUbToken();
90 2 : HrtRaUbLocMemRegParam lmemReg{alignBuf.first, alignBuf.second, tokenValue, tokenIdHandle, 1};
91 2 : reqReg = HrtRaUbLocalMemReg(rdmaHandle, lmemReg);
92 2 : auto ret = memcpy_s(key, HRT_UB_MEM_KEY_MAX_LEN, reqReg.key, HRT_UB_MEM_KEY_MAX_LEN);
93 2 : if (ret != EOK) {
94 0 : HCCL_ERROR("[LocalUbRmaBuffer::%s] netDevice copy key failed, ret[%d].", __func__, ret);
95 0 : THROW<InvalidParamsException>("LocalUbRmaBuffer netDevice copy key failed, ret[%d]", ret);
96 : }
97 6 : HCCL_INFO(
98 : "[LocalUbRmaBuffer::%s] end, rdmaHandle[%p], lmemHandle[0x%llx], reqReg.keySize[%u]", __func__, rdmaHandle,
99 : reqReg.handle, reqReg.keySize);
100 2 : }
101 :
102 29 : LocalUbRmaBuffer::LocalUbRmaBuffer(std::shared_ptr<Buffer> buf) : LocalRmaBuffer(buf, RmaType::UB), rdmaHandle(nullptr)
103 : {
104 : rtMemUbTokenInfo info;
105 29 : info.va = buf->GetAddr();
106 29 : info.size = buf->GetSize();
107 29 : HrtUbDevQueryInfo(QUERY_PROCESS_TOKEN, &info);
108 29 : tokenId = info.tokenId;
109 29 : tokenValue = info.tokenValue; // 未处理tokenIdHandle
110 87 : HCCL_INFO("LocalUbRmaBuffer Construct: buf=[%s]", buf->Describe().c_str());
111 29 : }
112 :
113 72 : string LocalUbRmaBuffer::Describe() const
114 : {
115 : return StringFormat(
116 144 : "LocalUbRmaBuffer[rdmaHandle=%p, buf=%s, reqReg.handle=0x%llx]", rdmaHandle, buf->Describe().c_str(),
117 216 : static_cast<unsigned long long>(reqReg.handle));
118 : }
119 :
120 12 : std::unique_ptr<Serializable> LocalUbRmaBuffer::GetExchangeDto()
121 : {
122 : std::unique_ptr<ExchangeUbBufferDto> dto = make_unique<ExchangeUbBufferDto>(
123 24 : buf->GetAddr(), buf->GetSize(), buf->GetMemType(), buf->GetMemInfo().c_str(), tokenValue, tokenId,
124 24 : reqReg.keySize);
125 12 : (void)memcpy_s(dto->key, HRT_UB_MEM_KEY_MAX_LEN, key, HRT_UB_MEM_KEY_MAX_LEN);
126 12 : dto->segVa = reqReg.targetSegVa;
127 24 : return std::unique_ptr<Serializable>(dto.release());
128 12 : }
129 :
130 1038 : LocalUbRmaBuffer::~LocalUbRmaBuffer()
131 : {
132 555 : if (isAlias_) {
133 0 : return;
134 : }
135 555 : if (netDev != nullptr && reqReg.handle != 0) {
136 0 : RdmaHandle h = netDev->GetRdmaHandle();
137 0 : const bool ctxValid = h != nullptr && RdmaHandleManager::GetInstance().IsHandleValid(h);
138 0 : if (!ctxValid) {
139 0 : HCCL_WARNING(
140 : "[LocalUbRmaBuffer][%s] skip HrtRaUbLocalMemUnreg (netDev), "
141 : "rdmaHandle=%p invalid, lmemHandle=0x%llx",
142 : __func__, h, static_cast<unsigned long long>(reqReg.handle));
143 : } else {
144 0 : DECTOR_TRY_CATCH("LocalUbRmaBuffer", HrtRaUbLocalMemUnreg(h, reqReg.handle));
145 0 : netDev->PutTokenIdInfo(bufKey_, tokenIdHandle);
146 : }
147 0 : reqReg.handle = 0;
148 555 : } else if (rdmaHandle != nullptr && reqReg.handle != 0) {
149 10 : const bool ctxValid = RdmaHandleManager::GetInstance().IsHandleValid(rdmaHandle);
150 10 : if (!ctxValid) {
151 30 : HCCL_WARNING(
152 : "[LocalUbRmaBuffer][%s] skip HrtRaUbLocalMemUnreg, "
153 : "rdmaHandle=%p invalid, lmemHandle=0x%llx",
154 : __func__, rdmaHandle, static_cast<unsigned long long>(reqReg.handle));
155 : } else {
156 0 : HCCL_INFO(
157 : "[LocalUbRmaBuffer::%s] rdmaHandle[%p], lmemHandle[0x%llx]", __func__, rdmaHandle,
158 : static_cast<unsigned long long>(reqReg.handle));
159 0 : DECTOR_TRY_CATCH("LocalUbRmaBuffer", HrtRaUbLocalMemUnreg(rdmaHandle, reqReg.handle));
160 0 : RdmaHandleManager::GetInstance().PutTokenIdInfo(rdmaHandle, bufKey_, tokenIdHandle);
161 : }
162 10 : reqReg.handle = 0;
163 555 : } else if (reqReg.handle != 0) {
164 0 : HCCL_WARNING(
165 : "[LocalUbRmaBuffer::%s] reqReg.handle[0x%llx] is non-zero but no valid cleanup path "
166 : "(netDev[%p], rdmaHandle[%p])",
167 : __func__, reqReg.handle, netDev, rdmaHandle);
168 : }
169 1038 : }
170 :
171 62 : u32 LocalUbRmaBuffer::GetTokenId() const { return tokenId; }
172 :
173 88 : u32 LocalUbRmaBuffer::GetTokenValue() const { return tokenValue; }
174 :
175 0 : TokenIdHandle LocalUbRmaBuffer::GetTokenIdHandle() const { return tokenIdHandle; }
176 :
177 : static bool isInitialized = false; // 标记是否已经初始化
178 : static u32 token = 0; // 存储生成的随机数
179 : static std::mutex ubTokenMutex;
180 1670 : u32 GetUbToken()
181 : {
182 1670 : std::lock_guard<std::mutex> lock(ubTokenMutex);
183 1670 : if (!isInitialized) {
184 1 : s32 devLogicId = HrtGetDevice();
185 1 : u32 devPhyId = HrtGetDevicePhyIdByIndex(devLogicId);
186 1 : HrtRaGetSecRandom(&token, devPhyId);
187 1 : isInitialized = true;
188 : }
189 1670 : return token;
190 1670 : }
191 :
192 : } // namespace Hccl
|
