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