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

try alternative gc

Browse Source
This commit is contained in:
Awni Hannun 2023-12-24 11:09:51 -08:00
parent 295ce9db09
commit b900e60972
2 changed files with 17 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
mlx/backend/metal/allocator.cpp
View File

@ -26,11 +26,7 @@ namespace metal {
namespace { namespace {
BufferCache::BufferCache(MTL::Device* device) BufferCache::BufferCache(MTL::Device* device)
: device_(device), : device_(device), head_(nullptr), tail_(nullptr), pool_size_(0) {}
head_(nullptr),
tail_(nullptr),
pool_size_(0),
gc_limit_(0.95 * device_->recommendedMaxWorkingSetSize()) {}
BufferCache::~BufferCache() { BufferCache::~BufferCache() {
clear(); clear();
@ -54,10 +50,10 @@ MTL::Buffer* BufferCache::reuse_from_cache(size_t size) {
// Find the closest buffer in pool // Find the closest buffer in pool
MTL::Buffer* pbuf = nullptr; MTL::Buffer* pbuf = nullptr;
auto it = buffer_pool_.lower_bound(size);
// Make sure we use > 50% of the available memory // Make sure we use > 50% of the available memory
while (!pbuf && it != buffer_pool_.end() && it->first < 2 * size) { if (auto it = buffer_pool_.lower_bound(size);
it != buffer_pool_.end() && it->first < 2 * size) {
// Collect from the cache // Collect from the cache
pbuf = it->second->buf; pbuf = it->second->buf;
// Remove from cache // Remove from cache
@ -85,17 +81,12 @@ void BufferCache::recycle_to_cache(MTL::Buffer* buf) {
} }
} }
size_t BufferCache::release_cached_buffers(size_t min_bytes_to_free) { void BufferCache::release_cached_buffers(size_t min_bytes_to_free) {
min_bytes_to_free += device_->currentAllocatedSize() - gc_limit_;
if (min_bytes_to_free >= 0.9 * pool_size_) { if (min_bytes_to_free >= 0.9 * pool_size_) {
size_t old_pool_size = pool_size_;
clear(); clear();
return old_pool_size;
} else { } else {
std::lock_guard<std::mutex> lk(cache_mutex_); std::lock_guard<std::mutex> lk(cache_mutex_);
size_t total_bytes_freed = 0; size_t total_bytes_freed = 0;
while (tail_ && (total_bytes_freed < min_bytes_to_free)) { while (tail_ && (total_bytes_freed < min_bytes_to_free)) {
if (tail_->buf) { if (tail_->buf) {
total_bytes_freed += tail_->buf->length(); total_bytes_freed += tail_->buf->length();
@ -104,9 +95,7 @@ size_t BufferCache::release_cached_buffers(size_t min_bytes_to_free) {
} }
remove_from_list(tail_); remove_from_list(tail_);
} }
pool_size_ -= total_bytes_freed; pool_size_ -= total_bytes_freed;
return total_bytes_freed;
} }
} }
@ -165,14 +154,17 @@ Buffer MetalAllocator::malloc(size_t size) {
// Prepare to allocate new memory as needed // Prepare to allocate new memory as needed
if (!buf) { if (!buf) {
// If we are under very high memory pressure, we don't allocate further // First check if the cache is big but nothing fits, garbage collect
if (device_->currentAllocatedSize() >= block_limit_) { // if so
return Buffer{nullptr}; // TODO maybe block limit and gc limit should be different
if (buffer_cache_.size() >= block_limit_) {
buffer_cache_.release_cached_buffers(
std::max(buffer_cache_.size() - block_limit_, size));
} }
// If we are still under memory pressure, try cleaning cache // If there is still too much memory pressure, fail (likely causes a wait).
if (buffer_cache_.can_garbage_collect()) { if (device_->currentAllocatedSize() >= block_limit_) {
buffer_cache_.release_cached_buffers(size); return Buffer{nullptr};
} }
// Allocate new buffer if needed // Allocate new buffer if needed

8
mlx/backend/metal/allocator.h
View File

@ -23,10 +23,11 @@ class BufferCache {
MTL::Buffer* reuse_from_cache(size_t size); MTL::Buffer* reuse_from_cache(size_t size);
void recycle_to_cache(MTL::Buffer* buf); void recycle_to_cache(MTL::Buffer* buf);
size_t release_cached_buffers(size_t min_bytes_to_free); void release_cached_buffers(size_t min_bytes_to_free);
bool can_garbage_collect() { // Returnt he size in bytes of cached memory
return pool_size_ > 0 && device_->currentAllocatedSize() > gc_limit_; size_t size() {
return pool_size_;
} }
private: private:
@ -49,7 +50,6 @@ class BufferCache {
BufferHolder* head_; BufferHolder* head_;
BufferHolder* tail_; BufferHolder* tail_;
size_t pool_size_; size_t pool_size_;
size_t gc_limit_;
}; };
} // namespace } // namespace