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

new_memory_pool -> new_scoped_memory_pool

Browse Source
This commit is contained in:
Ronan Collobert 2023-12-22 10:51:49 -08:00
parent 5ccc1fb314
commit d9478d0eb0
4 changed files with 7 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -20,7 +20,7 @@ namespace mlx::core::metal {
namespace {
// Catch things related to the main-thread static variables
static std::shared_ptr<void> global_memory_pool = new_memory_pool();
static std::shared_ptr<void> global_memory_pool = new_scoped_memory_pool();
// TODO nicer way to set this or possibly expose as an environment variable
static constexpr int MAX_BUFFERS_PER_QUEUE = 12;
@ -114,7 +114,7 @@ MTL::Library* load_library(
} // namespace
Device::Device() {
auto pool = new_memory_pool();
auto pool = new_scoped_memory_pool();
device_ = load_device();
library_map_ = {{"mlx", load_library(device_)}};
}
@ -244,7 +244,7 @@ void Device::register_library(
MTL::ComputePipelineState* Device::get_kernel(
const std::string& name,
const std::string& lib_name /* = "mlx" */) {
auto pool = new_memory_pool();
auto pool = new_scoped_memory_pool();
// Look for cached kernel
if (auto it = kernel_map_.find(name); it != kernel_map_.end()) {
return it->second;
@ -291,7 +291,7 @@ Device& device(mlx::core::Device) {
return metal_device;
}
std::shared_ptr<void> new_memory_pool() {
std::shared_ptr<void> new_scoped_memory_pool() {
auto dtor = [](void* ptr) {
static_cast<NS::AutoreleasePool*>(ptr)->release();
};

2
mlx/backend/metal/metal.cpp
View File

@ -50,7 +50,7 @@ std::function<void()> make_task(
bool retain_graph) {
auto task =
[retain_graph, arr, deps = std::move(deps), p = std::move(p)]() mutable {
auto pool = new_memory_pool();
auto pool = new_scoped_memory_pool();
for (auto& d : deps) {
d.wait();
}

2
mlx/backend/metal/metal.h
View File

@ -20,7 +20,7 @@ constexpr bool is_available() {
}
void new_stream(Stream stream);
std::shared_ptr<void> new_memory_pool();
std::shared_ptr<void> new_scoped_memory_pool();
std::function<void()> make_task(
array& arr,

2
mlx/scheduler.h
View File

@ -35,7 +35,7 @@ struct StreamThread {
}
void thread_fn() {
auto thread_pool = metal::new_memory_pool();
auto thread_pool = metal::new_scoped_memory_pool();
metal::new_stream(stream);
while (true) {
std::function<void()> task;