zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 17:31:16 +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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 6 deletions
Show Stats 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)};
} }