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

A temporary fix (#254)

Browse Source
This commit is contained in:
Awni Hannun 2023-12-21 17:59:15 -08:00 committed by GitHub
parent 1d053e0d1d
commit a002797d52
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 18 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
mlx/backend/metal/device.cpp
View File

@ -108,6 +108,16 @@ MTL::Library* load_library(
} }
} }
struct PoolHolder {
PoolHolder() {
p = NS::AutoreleasePool::alloc()->init();
}
~PoolHolder() {
p->release();
}
NS::AutoreleasePool* p;
};
} // namespace } // namespace
Device::Device() Device::Device()
@ -125,6 +135,12 @@ Device::~Device() {
for (auto& l : library_map_) { for (auto& l : library_map_) {
l.second->release(); l.second->release();
} }
for (auto& b : buffer_map_) {
b.second.second->release();
}
for (auto& e : encoder_map_) {
e.second->release();
}
device_->release(); device_->release();
pool_->release(); pool_->release();
} }
@ -282,9 +298,8 @@ Device& device(mlx::core::Device) {
} }
NS::AutoreleasePool*& thread_autorelease_pool() { NS::AutoreleasePool*& thread_autorelease_pool() {
static thread_local NS::AutoreleasePool* p = static thread_local PoolHolder pool{};
NS::AutoreleasePool::alloc()->init(); return pool.p;
return p;
} }
void new_stream(Stream stream) { void new_stream(Stream stream) {