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 : #ifndef RMA_BUFFER_MGR_H
12 : #define RMA_BUFFER_MGR_H
13 :
14 : #include <map>
15 : #include <mutex>
16 : #include <shared_mutex>
17 : #include <utility>
18 : #include "buffer_key.h"
19 : #include "log.h"
20 :
21 : namespace hcomm {
22 : template <typename KeyType, typename BufferType, template <typename...> class M = std::map, typename... MapArgs>
23 : class RmaBufferMgr {
24 : public:
25 : struct BufferWithRef {
26 : // BufferType可以是指针的类型
27 : BufferType buffer{};
28 : uint64_t ref{}; // 引用计数
29 :
30 151 : BufferWithRef(BufferType buf, uint64_t r) : buffer(buf), ref(r) {}
31 : };
32 :
33 : using AddrType = typename KeyType::AddrType;
34 : using SizeType = typename KeyType::SizeType;
35 :
36 : using MapType = M<KeyType, BufferWithRef, MapArgs...>;
37 : using Iterator = typename MapType::iterator;
38 : using ConstIterator = typename MapType::const_iterator;
39 :
40 : template <typename... BufferArgs>
41 : // 禁止直接调用,请使用加锁版本的 AddWithoutCheck 替代
42 151 : std::pair<Iterator, bool> AddToTree(const KeyType& key, BufferArgs&&... bufferArgs)
43 : {
44 151 : auto result = intervalTree_.emplace(
45 151 : std::piecewise_construct, std::forward_as_tuple(key),
46 302 : std::forward_as_tuple(BufferWithRef{BufferType{std::forward<BufferArgs>(bufferArgs)...}, 1}));
47 151 : if (!result.second) {
48 31 : result.first->second.ref++;
49 : // 翻转
50 31 : if (result.first->second.ref == 0) {
51 0 : HCCL_ERROR("Error: ref = 0, ref++ flipped");
52 0 : throw std::logic_error("ref++ = 0, ref++ flipped");
53 31 : } else if (result.first->second.ref > 1) {
54 31 : HCCL_INFO(
55 : "Memory is already registered, just increase the reference count, "
56 : "current memory reference count[%llu], %s.",
57 : result.first->second.ref, key.ToString().c_str());
58 : }
59 : }
60 :
61 151 : return result;
62 : }
63 :
64 : template <typename... BufferArgs>
65 127 : std::pair<Iterator, bool> AddWithoutCheck(const KeyType& key, BufferArgs&&... bufferArgs)
66 : {
67 127 : std::unique_lock<std::shared_mutex> lock(mtx_);
68 254 : return AddToTree(key, std::forward<BufferArgs>(bufferArgs)...);
69 127 : }
70 :
71 : // 1.添加成功：输入key是表中某一最相近key的空集。 计数+1，返回添加成功的迭代器，及true
72 : // 2.添加已存在：输入key是表中某一最相近key的全集。 计数+1，返回添加该key的迭代器，及false
73 : // 3.添加失败：输入key是表中某一个最相近key的交集、子集、超集。返回空迭代器，及false
74 : template <typename... BufferArgs>
75 25 : std::pair<Iterator, bool> Add(const KeyType& key, BufferArgs&&... bufferArgs)
76 : {
77 25 : std::unique_lock<std::shared_mutex> lock(mtx_);
78 25 : auto overlapResult = CheckOverlap(key);
79 25 : if (overlapResult.second) {
80 3 : HCCL_ERROR("Error: Buffer key overlaps with existing buffer key.");
81 1 : return std::make_pair(intervalTree_.end(), false);
82 : }
83 24 : return AddToTree(key, std::forward<BufferArgs>(bufferArgs)...);
84 25 : }
85 :
86 : // 1.查询成功：输入key是表中某一最相近key的子集、全集。 返回true，最相近key的bufferType
87 : // 2.查询失败：输入key是表中某一个最相近key的空集、交集。返回false，空bufferType
88 150 : std::pair<bool, BufferType> Find(const KeyType& key) const
89 : {
90 150 : std::shared_lock<std::shared_mutex> lock(mtx_);
91 150 : auto it = intervalTree_.lower_bound(key);
92 150 : if (it != intervalTree_.end() && (it->first == key || it->first.IsSuperset(key))) {
93 80 : return std::make_pair(true, it->second.buffer);
94 : }
95 :
96 70 : if (it != intervalTree_.begin()) {
97 7 : auto prevIt = std::prev(it);
98 7 : if (prevIt->first.IsSuperset(key)) {
99 1 : return std::make_pair(true, prevIt->second.buffer);
100 : }
101 6 : if (it != intervalTree_.end()) {
102 2 : HCCL_WARNING(
103 : "Key[%s] not found. The near key is [%s] or [%s].", key.ToString().c_str(),
104 : it->first.ToString().c_str(), prevIt->first.ToString().c_str());
105 : } else {
106 4 : HCCL_WARNING(
107 : "Key[%s] not found. The near key is [%s]", key.ToString().c_str(),
108 : prevIt->first.ToString().c_str());
109 : }
110 : } else {
111 63 : if (it != intervalTree_.end()) {
112 0 : HCCL_WARNING(
113 : "Key[%s] not found. The near key is [%s]", key.ToString().c_str(), it->first.ToString().c_str());
114 : } else {
115 63 : HCCL_WARNING("Key[%s] not found. There is no key in table.", key.ToString().c_str());
116 : }
117 : }
118 :
119 69 : return std::make_pair(false, BufferType{}); // 未找到
120 150 : }
121 :
122 : // 1.删除成功：输入key是表中某一最相近key的全集。 计数-1且之后为0。 返回true
123 : // 2.删除引用数-1但未删除：输入key是表中某一最相近key的全集。 计数-1且之后大于0。 返回false
124 : // 3.删除失败：输入key是表中某一个最相近key的交集、子集、超集、空集。——抛出NOT_FOUND异常
125 93 : bool Del(const KeyType& key)
126 : {
127 93 : std::unique_lock<std::shared_mutex> lock(mtx_);
128 93 : auto it = intervalTree_.find(key);
129 93 : if (it == intervalTree_.end()) {
130 9 : HCCL_ERROR("Error: Buffer key not found.");
131 9 : throw std::out_of_range("Del NOT_FOUND");
132 : }
133 :
134 84 : if (--(it->second.ref) == 0) {
135 56 : intervalTree_.erase(it);
136 56 : return true;
137 : }
138 : // 引用计数大于0，不删除
139 28 : HCCL_RUN_INFO(
140 : "Memory reference count is larger than 0, (used by other RemoteRank), do not deregister memory."
141 : "current memory reference count[%llu], %s.",
142 : it->second.ref, key.ToString().c_str());
143 28 : return false;
144 93 : }
145 :
146 77 : bool IsInTree(const KeyType& key)
147 : {
148 77 : std::shared_lock<std::shared_mutex> lock(mtx_);
149 77 : auto it = intervalTree_.find(key);
150 77 : if (it == intervalTree_.end()) {
151 52 : return false;
152 : }
153 25 : return true;
154 77 : }
155 :
156 69 : ConstIterator Begin() { return intervalTree_.begin(); }
157 :
158 6 : ConstIterator Next(ConstIterator it) { return std::next(it); }
159 :
160 85 : ConstIterator End() { return intervalTree_.end(); }
161 13 : size_t size() const { return intervalTree_.size(); }
162 : void PrintContents() const
163 : {
164 : std::shared_lock<std::shared_mutex> lock(mtx_);
165 : for (const auto& pair : intervalTree_) {
166 : HCCL_INFO("Key: %s, Value: %p", pair.first.ToString().c_str(), pair.second.buffer.get());
167 : }
168 : }
169 :
170 : template <typename Fn>
171 1 : void ForEach(Fn&& fn) const
172 : {
173 1 : std::shared_lock<std::shared_mutex> lock(mtx_);
174 2 : for (const auto& pair : intervalTree_) {
175 1 : std::forward<Fn>(fn)(pair.first, pair.second.buffer);
176 : }
177 1 : }
178 :
179 : private:
180 : MapType intervalTree_;
181 : mutable std::shared_mutex mtx_{};
182 :
183 25 : std::pair<Iterator, bool> CheckOverlap(const KeyType& key)
184 : {
185 25 : auto it = intervalTree_.lower_bound(key);
186 25 : if (it != intervalTree_.end()) {
187 : // 情况1：addr_ == it->first.addr_ && size_ == it->first.size_
188 0 : if (it->first == key) {
189 0 : return std::make_pair(it, false);
190 : }
191 :
192 : // 情况2：addr_ == it->first.addr_ && size_ < it->first.size_。it->first.IsSubset(key)非必须
193 : // 情况3：addr_ < it->first.addr_
194 0 : if (it->first.IsSuperset(key) || it->first.IsIntersect(key)) {
195 0 : return std::make_pair(it, true);
196 : }
197 : }
198 :
199 : // 剩下的是空集
200 25 : if (it != intervalTree_.begin()) {
201 2 : auto prevIt = std::prev(it);
202 : // 情况4：addr_ > prevIt->first.addr_
203 : // 情况5： 1) addr_ > prevIt->first.addr_的子集情况；
204 : // 2) addr_ == prevIt->first.addr_，size_ > prevIt->first.size
205 2 : if (prevIt->first.IsIntersect(key) || prevIt->first.IsSubset(key) || prevIt->first.IsSuperset(key)) {
206 1 : return std::make_pair(prevIt, true);
207 : }
208 :
209 : // 6. 剩下的是空集
210 1 : return std::make_pair(prevIt, false);
211 : }
212 :
213 : // 剩下的是空集
214 23 : return std::make_pair(it, false);
215 : }
216 : };
217 : } // namespace hcomm
218 :
219 : #endif
|
