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