zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-25 01:41:17 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix (#255)

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

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

@ -17,6 +17,8 @@ namespace fs = std::filesystem;
namespace mlx::core::metal { namespace mlx::core::metal {
static Device metal_device_;
namespace { namespace {
// TODO nicer way to set this or possibly expose as an environment variable // TODO nicer way to set this or possibly expose as an environment variable
@ -108,16 +110,6 @@ MTL::Library* load_library(
} }
} }
struct PoolHolder {
PoolHolder() {
p = NS::AutoreleasePool::alloc()->init();
}
~PoolHolder() {
p->release();
}
NS::AutoreleasePool* p;
};
} // namespace } // namespace
Device::Device() Device::Device()
@ -293,13 +285,13 @@ MTL::ComputePipelineState* Device::get_kernel(
} }
Device& device(mlx::core::Device) { Device& device(mlx::core::Device) {
static Device metal_device_;
return metal_device_; return metal_device_;
} }
NS::AutoreleasePool*& thread_autorelease_pool() { NS::AutoreleasePool*& thread_autorelease_pool() {
static thread_local PoolHolder pool{}; static thread_local NS::AutoreleasePool* p =
return pool.p; NS::AutoreleasePool::alloc()->init();
return p;
} }
void new_stream(Stream stream) { void new_stream(Stream stream) {