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 UNIVERSAL_CONCURRENT_MAP_H
12 : #define UNIVERSAL_CONCURRENT_MAP_H
13 :
14 : #include <mutex>
15 : #include <unordered_map>
16 : #include <map>
17 : #include <shared_mutex>
18 :
19 : namespace hccl {
20 :
21 : template <typename K, typename V, template <typename...> class M = std::unordered_map, typename... MapArgs>
22 : class UniversalConcurrentMap {
23 : public:
24 178 : UniversalConcurrentMap() = default;
25 175 : ~UniversalConcurrentMap() = default;
26 :
27 : using MapType = M<K, V, MapArgs...>;
28 : using Iterator = typename MapType::iterator;
29 : using ConstIterator = typename MapType::const_iterator;
30 : using SizeType = typename MapType::size_type;
31 :
32 : // true -> valid
33 1060 : inline std::pair<Iterator, bool> Find(const K& k)
34 : {
35 1060 : std::shared_lock<std::shared_timed_mutex> lock(mapMtx_);
36 1060 : Iterator it = map_.find(k);
37 1060 : if (it != map_.end()) {
38 984 : return { it, true };
39 : }
40 :
41 76 : return { map_.end(), false };
42 1060 : }
43 :
44 : // true -> valid
45 : inline std::pair<ConstIterator, bool> Find(const K& k) const
46 : {
47 : std::shared_lock<std::shared_timed_mutex> lock(mapMtx_);
48 : ConstIterator it = map_.find(k);
49 : if (it != map_.end()) {
50 : return { it, true };
51 : }
52 :
53 : return { map_.end(), false };
54 : }
55 :
56 : // true -> 新插入
57 : template <class... Args>
58 57 : inline std::pair<Iterator, bool> Emplace(Args&&... args)
59 : {
60 57 : std::lock_guard<std::shared_timed_mutex> lock(mapMtx_);
61 :
62 114 : return map_.emplace(std::forward<Args>(args)...);
63 57 : }
64 :
65 : // true -> 新插入，可能抛异常
66 : template<typename Func, typename... Args>
67 : inline std::pair<Iterator, bool> EmplaceIfNotExist(const K& k, Func func, Args&&... args)
68 : {
69 : std::lock_guard<std::shared_timed_mutex> lock(mapMtx_);
70 : Iterator it = map_.find(k);
71 : if (it == map_.end()) {
72 : return map_.emplace(k, func(std::forward<Args>(args)...));
73 : }
74 :
75 : return { it, false };
76 : }
77 :
78 : // 可能抛异常
79 : template<typename Func, typename... Args>
80 43 : inline std::pair<Iterator, bool> EmplaceAndUpdate(const K& k, Func func, Args&&... args)
81 : {
82 43 : std::lock_guard<std::shared_timed_mutex> lock(mapMtx_);
83 :
84 43 : std::pair<Iterator, bool> it = map_.emplace(k, V());
85 43 : func(it.first->second, std::forward<Args>(args)...);
86 :
87 43 : return it;
88 43 : }
89 :
90 25 : inline V& operator[] (K&& k)
91 : {
92 25 : std::lock_guard<std::shared_timed_mutex> lock(mapMtx_);
93 50 : return map_[std::forward<K>(k)];
94 25 : }
95 :
96 12 : inline V& operator[] (const K& k)
97 : {
98 12 : std::lock_guard<std::shared_timed_mutex> lock(mapMtx_);
99 24 : return map_[k];
100 12 : }
101 :
102 : V& At(const K& k)
103 : {
104 : std::lock_guard<std::shared_timed_mutex> lock(mapMtx_);
105 : return map_.at(k);
106 : }
107 :
108 : const V& At(const K& k) const
109 : {
110 : std::shared_lock<std::shared_timed_mutex> lock(mapMtx_);
111 : return map_.at(k);
112 : }
113 :
114 : // 可能抛异常
115 : template<typename Func, typename... Args>
116 : inline void EraseAll(Func func, Args&&... args)
117 : {
118 : std::lock_guard<std::shared_timed_mutex> lock(mapMtx_);
119 : for (auto it = map_.begin(); it != map_.end();) {
120 : func(it->second, std::forward<Args>(args)...);
121 : it = map_.erase(it);
122 : }
123 : }
124 :
125 18 : inline SizeType Size() const
126 : {
127 18 : std::shared_lock<std::shared_timed_mutex> lock(mapMtx_);
128 36 : return map_.size();
129 18 : }
130 :
131 30 : inline void Clear()
132 : {
133 30 : std::lock_guard<std::shared_timed_mutex> lock(mapMtx_);
134 30 : map_.clear();
135 30 : }
136 :
137 6 : inline SizeType Erase(const K& k)
138 : {
139 6 : std::lock_guard<std::shared_timed_mutex> lock(mapMtx_);
140 12 : return map_.erase(k);
141 6 : }
142 :
143 : // 尽量少使用LockFree结尾的函数
144 : inline SizeType EraseLockFree(const K& k)
145 : {
146 : return map_.erase(k);
147 : }
148 :
149 5 : inline std::shared_timed_mutex &GetMtx()
150 : {
151 5 : return mapMtx_;
152 : }
153 :
154 6 : inline Iterator FindLockFree(const K& k)
155 : {
156 6 : return map_.find(k);
157 : }
158 :
159 6 : inline Iterator EndLockFree()
160 : {
161 6 : return map_.end();
162 : }
163 :
164 : template <class... Args>
165 4 : inline std::pair<Iterator, bool> EmplaceLockFree(Args&&... args)
166 : {
167 4 : return map_.emplace(std::forward<Args>(args)...);
168 : }
169 :
170 : private:
171 : mutable std::shared_timed_mutex mapMtx_{};
172 : MapType map_{};
173 : };
174 : }
175 :
176 : // 兼容旧版本，后续可以直接使用hccl::UniversalConcurrentMap替换原来的ConcurrentMap
177 : namespace Hccl {
178 : template <typename K, typename V, template <typename...> class M = std::unordered_map, typename... MapArgs>
179 : using UniversalConcurrentMap = hccl::UniversalConcurrentMap<K, V, M, MapArgs...>;
180 : }
181 :
182 : #endif
|
