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