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

only add to residency set once (#2049)

Browse Source
This commit is contained in:
Awni Hannun
2025-04-06 17:38:25 -07:00
committed by GitHub
parent ba09f01ce8
commit 1a28b69ee2
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -272,9 +272,13 @@ Buffer MetalAllocator::malloc(size_t size) {
if (!buf) { if (!buf) {
buf = device_->newBuffer(size, resource_options); buf = device_->newBuffer(size, resource_options);
} }
if (!buf) {
return Buffer{nullptr};
}
lk.lock(); lk.lock();
if (buf) {
num_resources_++; num_resources_++;
if (!buf->heap()) {
residency_set_.insert(buf);
} }
} }
@@ -288,10 +292,6 @@ Buffer MetalAllocator::malloc(size_t size) {
get_cache_memory() - max_pool_size_); get_cache_memory() - max_pool_size_);
} }
if (!buf->heap()) {
residency_set_.insert(buf);
}
return Buffer{static_cast<void*>(buf)}; return Buffer{static_cast<void*>(buf)};
} }