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

change fit crtieria

Browse Source
This commit is contained in:
Awni Hannun 2023-12-25 15:35:35 -08:00
parent 679d69220e
commit f3c7a7c1e5
1 changed files with 8 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -52,8 +52,8 @@ MTL::Buffer* BufferCache::reuse_from_cache(size_t size) {
MTL::Buffer* pbuf = nullptr; MTL::Buffer* pbuf = nullptr;
// Make sure we use > 50% of the available memory // Make sure we use > 50% of the available memory
if (auto it = buffer_pool_.lower_bound(size); if (auto it = buffer_pool_.lower_bound(size); it != buffer_pool_.end() &&
it != buffer_pool_.end() && it->first < 2 * size) { it->first < std::min(2 * size, size + vm_page_size)) {
// Collect from the cache // Collect from the cache
pbuf = it->second->buf; pbuf = it->second->buf;
// Remove from cache // Remove from cache
@ -150,17 +150,18 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
size = vm_page_size * ((size + vm_page_size - 1) / vm_page_size); size = vm_page_size * ((size + vm_page_size - 1) / vm_page_size);
} }
// Try the cache
MTL::Buffer* buf = buffer_cache_.reuse_from_cache(size); MTL::Buffer* buf = buffer_cache_.reuse_from_cache(size);
// Prepare to allocate new memory as needed
if (!buf) { if (!buf) {
// If there is still too much memory pressure, fail (likely causes a wait). // If we have memory pressure, first check if we can reclaim some memory
if (auto new_size = device_->currentAllocatedSize() + size; // from the cache
new_size >= block_limit_) { if (device_->currentAllocatedSize() + size >= block_limit_) {
buffer_cache_.clear(); buffer_cache_.clear();
} }
if (device_->currentAllocatedSize() >= block_limit_) { // If there is still too much memory pressure, fail (likely causes a wait).
if (!allow_swap && device_->currentAllocatedSize() + size >= block_limit_) {
return Buffer{nullptr}; return Buffer{nullptr};
} }