zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-04 06:44:40 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Make LRUCache more like a normal container

Browse Source
This commit is contained in:
Cheng
2025-07-20 16:20:53 -07:00
parent 3d16cb5071
commit 4c0dc7745f
1 changed files with 12 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
mlx/backend/cuda/lru_cache.h
View File

@@ -16,8 +16,9 @@ class LRUCache {
public: public:
using value_type = std::pair<K, V>; using value_type = std::pair<K, V>;
using list_type = std::list<value_type>; using list_type = std::list<value_type>;
using list_iter = typename list_type::iterator; using iterator = typename list_type::iterator;
using map_type = M<K, list_iter>; using const_iterator = typename list_type::const_iterator;
using map_type = M<K, iterator>;
explicit LRUCache(size_t capacity) : capacity_(capacity) {} explicit LRUCache(size_t capacity) : capacity_(capacity) {}
@@ -36,16 +37,16 @@ class LRUCache {
trim(); trim();
} }
auto begin() { iterator begin() {
return vlist_.begin(); return vlist_.begin();
} }
auto begin() const { const_iterator begin() const {
return vlist_.begin(); return vlist_.begin();
} }
auto end() { iterator end() {
return vlist_.end(); return vlist_.end();
} }
auto end() const { const_iterator end() const {
return vlist_.end(); return vlist_.end();
} }
@@ -54,7 +55,7 @@ class LRUCache {
vlist_.clear(); vlist_.clear();
} }
list_iter find(const K& key) { iterator find(const K& key) {
auto it = map_.find(key); auto it = map_.find(key);
if (it == map_.end()) if (it == map_.end())
return end(); return end();
@@ -62,14 +63,15 @@ class LRUCache {
return it->second; return it->second;
} }
std::pair<list_iter, bool> emplace(const K& key, V value) { template <typename U>
std::pair<iterator, bool> emplace(const K& key, U&& value) {
auto it = map_.find(key); auto it = map_.find(key);
if (it != map_.end()) { if (it != map_.end()) {
vlist_.splice(vlist_.begin(), vlist_, it->second); vlist_.splice(vlist_.begin(), vlist_, it->second);
return {it->second, false}; return {it->second, false};
} }
vlist_.emplace_front(key, std::move(value)); vlist_.emplace_front(key, std::forward<U>(value));
map_[key] = vlist_.begin(); map_[key] = vlist_.begin();
trim(); trim();
@@ -77,7 +79,7 @@ class LRUCache {
return {vlist_.begin(), true}; return {vlist_.begin(), true};
} }
list_iter erase(list_iter pos) { iterator erase(iterator pos) {
map_.erase(pos->first); map_.erase(pos->first);
return vlist_.erase(pos); return vlist_.erase(pos);
} }