From d4b222b6d399eaf3002495b661900e4ebe49e16a Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Wed, 27 Nov 2024 20:01:20 -0800 Subject: [PATCH] Fix some leaks and races (#1629) * fix leak and fix potential race * more leak fixes * fix one more --- mlx/backend/metal/allocator.cpp | 12 ++++++++---- mlx/backend/metal/device.cpp | 30 ++++++++++++++++++------------ mlx/backend/metal/metal.cpp | 1 + mlx/backend/metal/resident.cpp | 3 +++ 4 files changed, 30 insertions(+), 16 deletions(-) diff --git a/mlx/backend/metal/allocator.cpp b/mlx/backend/metal/allocator.cpp index cfbc82943..d36aec82f 100644 --- a/mlx/backend/metal/allocator.cpp +++ b/mlx/backend/metal/allocator.cpp @@ -30,7 +30,7 @@ BufferCache::BufferCache(MTL::Device* device) : device_(device), head_(nullptr), tail_(nullptr), pool_size_(0) {} BufferCache::~BufferCache() { - auto thread_pool = metal::new_scoped_memory_pool(); + auto pool = metal::new_scoped_memory_pool(); clear(); } @@ -155,11 +155,13 @@ MetalAllocator::MetalAllocator() } size_t MetalAllocator::set_cache_limit(size_t limit) { + std::unique_lock lk(mutex_); std::swap(limit, max_pool_size_); return limit; }; size_t MetalAllocator::set_memory_limit(size_t limit, bool relaxed) { + std::unique_lock lk(mutex_); std::swap(limit, block_limit_); relaxed_ = relaxed; gc_limit_ = std::min( @@ -169,6 +171,7 @@ size_t MetalAllocator::set_memory_limit(size_t limit, bool relaxed) { }; size_t MetalAllocator::set_wired_limit(size_t limit) { + std::unique_lock lk(mutex_); std::swap(limit, wired_limit_); residency_set_.resize(wired_limit_); return limit; @@ -205,7 +208,7 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) { return Buffer{nullptr}; } - auto thread_pool = metal::new_scoped_memory_pool(); + auto pool = metal::new_scoped_memory_pool(); // If we have a lot of memory pressure or are over the maximum cache size, // try to reclaim memory from the cache @@ -226,7 +229,7 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) { // Maintain the cache below the requested limit if (get_cache_memory() >= max_pool_size_) { - auto thread_pool = metal::new_scoped_memory_pool(); + auto pool = metal::new_scoped_memory_pool(); buffer_cache_.release_cached_buffers(get_cache_memory() - max_pool_size_); } @@ -237,6 +240,7 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) { void MetalAllocator::clear_cache() { std::unique_lock lk(mutex_); + auto pool = metal::new_scoped_memory_pool(); buffer_cache_.clear(); } @@ -252,7 +256,7 @@ void MetalAllocator::free(Buffer buffer) { buffer_cache_.recycle_to_cache(buf); } else { lk.unlock(); - auto thread_pool = metal::new_scoped_memory_pool(); + auto pool = metal::new_scoped_memory_pool(); buf->release(); } } diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 8f8a4468d..44a1ae9eb 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -645,21 +645,27 @@ void new_stream(Stream stream) { std::unordered_map> device_info() { - auto raw_device = device(default_device()).mtl_device(); - auto arch = std::string(raw_device->architecture()->name()->utf8String()); + auto init_device_info = []() + -> std::unordered_map> { + auto pool = new_scoped_memory_pool(); + auto raw_device = device(default_device()).mtl_device(); + auto arch = std::string(raw_device->architecture()->name()->utf8String()); - int mib[] = {CTL_HW, HW_MEMSIZE}; - size_t memsize = 0; - size_t length = sizeof(memsize); + int mib[] = {CTL_HW, HW_MEMSIZE}; + size_t memsize = 0; + size_t length = sizeof(memsize); - sysctl(mib, 2, &memsize, &length, NULL, 0); + sysctl(mib, 2, &memsize, &length, NULL, 0); - return { - {"architecture", arch}, - {"max_buffer_length", raw_device->maxBufferLength()}, - {"max_recommended_working_set_size", - raw_device->recommendedMaxWorkingSetSize()}, - {"memory_size", memsize}}; + return { + {"architecture", arch}, + {"max_buffer_length", raw_device->maxBufferLength()}, + {"max_recommended_working_set_size", + raw_device->recommendedMaxWorkingSetSize()}, + {"memory_size", memsize}}; + }; + static auto device_info_ = init_device_info(); + return device_info_; } } // namespace mlx::core::metal diff --git a/mlx/backend/metal/metal.cpp b/mlx/backend/metal/metal.cpp index 4b662bb36..661985c3f 100644 --- a/mlx/backend/metal/metal.cpp +++ b/mlx/backend/metal/metal.cpp @@ -94,6 +94,7 @@ std::function make_synchronize_task( Stream s, std::shared_ptr> p) { return [s, p = std::move(p)]() { + auto pool = new_scoped_memory_pool(); auto& d = metal::device(s.device); auto cb = d.get_command_buffer(s.index); cb->retain(); diff --git a/mlx/backend/metal/resident.cpp b/mlx/backend/metal/resident.cpp index ec2560aa4..545f67e49 100644 --- a/mlx/backend/metal/resident.cpp +++ b/mlx/backend/metal/resident.cpp @@ -63,6 +63,7 @@ void ResidencySet::resize(size_t size) { size_t current_size = wired_set_->allocatedSize(); if (current_size < size) { + auto pool = new_scoped_memory_pool(); // Add unwired allocations to the set for (auto it = unwired_set_.begin(); it != unwired_set_.end();) { auto buf_size = (*it)->allocatedSize(); @@ -77,6 +78,7 @@ void ResidencySet::resize(size_t size) { wired_set_->commit(); wired_set_->requestResidency(); } else if (current_size > size) { + auto pool = new_scoped_memory_pool(); // Remove wired allocations until under capacity auto allocations = wired_set_->allAllocations(); auto num_allocations = wired_set_->allocationCount(); @@ -92,6 +94,7 @@ void ResidencySet::resize(size_t size) { ResidencySet::~ResidencySet() { if (wired_set_) { + auto pool = new_scoped_memory_pool(); wired_set_->release(); } }