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