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 1038 : const auto &tokenIdInfoPair = RdmaHandleManager::GetInstance().GetTokenIdInfo(rdmaHandle,
31 519 : BufferKey<uintptr_t, u64>{alignBuf.first, alignBuf.second});
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 : const auto &tokenIdInfoPair = netDev->GetTokenIdInfo(BufferKey<uintptr_t, u64>{alignBuf.first, alignBuf.second});
79 2 : tokenIdHandle = tokenIdInfoPair.first;
80 2 : tokenId = tokenIdInfoPair.second;
81 2 : tokenValue = GetUbToken();
82 2 : HrtRaUbLocMemRegParam lmemReg{alignBuf.first, alignBuf.second, tokenValue, tokenIdHandle, 1};
83 2 : reqReg = HrtRaUbLocalMemReg(rdmaHandle, lmemReg);
84 2 : memcpy_s(key, HRT_UB_MEM_KEY_MAX_LEN, reqReg.key, HRT_UB_MEM_KEY_MAX_LEN);
85 6 : HCCL_INFO("[LocalUbRmaBuffer::%s] end, rdmaHandle[%p], lmemHandle[0x%llx], reqReg.keySize[%u]", __func__, rdmaHandle,
86 : reqReg.handle, reqReg.keySize);
87 2 : }
88 :
89 28 : LocalUbRmaBuffer::LocalUbRmaBuffer(std::shared_ptr<Buffer> buf) : LocalRmaBuffer(buf, RmaType::UB), rdmaHandle(nullptr)
90 : {
91 : rtMemUbTokenInfo info;
92 28 : info.va = buf->GetAddr();
93 28 : info.size = buf->GetSize();
94 28 : HrtUbDevQueryInfo(QUERY_PROCESS_TOKEN, &info);
95 28 : tokenId = info.tokenId;
96 28 : tokenValue = info.tokenValue; // 未处理tokenIdHandle
97 84 : HCCL_INFO("LocalUbRmaBuffer Construct: buf=[%s]", buf->Describe().c_str());
98 28 : }
99 :
100 67 : string LocalUbRmaBuffer::Describe() const
101 : {
102 : return StringFormat("LocalUbRmaBuffer[rdmaHandle=%p, buf=%s, reqReg.handle=0x%llx]",
103 134 : rdmaHandle, buf->Describe().c_str(),
104 201 : static_cast<unsigned long long>(reqReg.handle));
105 : }
106 :
107 12 : std::unique_ptr<Serializable> LocalUbRmaBuffer::GetExchangeDto()
108 : {
109 12 : std::unique_ptr<ExchangeUbBufferDto> dto = make_unique<ExchangeUbBufferDto>(buf->GetAddr(),
110 12 : buf->GetSize(),
111 12 : buf->GetMemType(),
112 24 : buf->GetMemInfo().c_str(),
113 12 : tokenValue,
114 12 : 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 1030 : LocalUbRmaBuffer::~LocalUbRmaBuffer()
122 : {
123 549 : if (rdmaHandle != nullptr && reqReg.handle != 0 && !isAlias_) {
124 30 : HCCL_INFO("[LocalUbRmaBuffer::%s] rdmaHandle[%p], lmemHandle[0x%llx]", __func__, rdmaHandle, reqReg.handle);
125 10 : DECTOR_TRY_CATCH("LocalUbRmaBuffer", HrtRaUbLocalMemUnreg(rdmaHandle, reqReg.handle));
126 : }
127 1030 : }
128 :
129 62 : u32 LocalUbRmaBuffer::GetTokenId() const
130 : {
131 62 : return tokenId;
132 : }
133 :
134 88 : u32 LocalUbRmaBuffer::GetTokenValue() const
135 : {
136 88 : return tokenValue;
137 : }
138 :
139 26 : TokenIdHandle LocalUbRmaBuffer::GetTokenIdHandle() const
140 : {
141 26 : return tokenIdHandle;
142 : }
143 :
144 : static bool isInitialized = false; // 标记是否已经初始化
145 : static u32 token = 0; // 存储生成的随机数
146 : static std::mutex ubTokenMutex;
147 1639 : u32 GetUbToken()
148 : {
149 1639 : std::lock_guard<std::mutex> lock(ubTokenMutex);
150 1639 : if (!isInitialized) {
151 1 : s32 devLogicId = HrtGetDevice();
152 1 : u32 devPhyId = HrtGetDevicePhyIdByIndex(devLogicId);
153 1 : HrtRaGetSecRandom(&token, devPhyId);
154 1 : isInitialized = true;
155 : }
156 1639 : return token;
157 1639 : }
158 :
159 :
160 : } // namespace Hccl
|
