zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-22 02:58:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix multistream GPU deadlock (#1969)

Browse Source 
* fix multistream GPU deadlock

* comments
This commit is contained in:
Awni Hannun
2025-03-20 07:19:47 -07:00
committed by GitHub
parent 95e335db7b
commit 3c164fca8c
6 changed files with 31 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -19,9 +19,6 @@ namespace mlx::core::metal {
namespace {
// TODO nicer way to set this or possibly expose as an environment variable
constexpr int MAX_BUFFERS_PER_QUEUE = 12;
constexpr const char* default_mtllib_path = METAL_PATH;
auto get_metal_version() {
@@ -256,7 +253,7 @@ Device::~Device() {
void Device::new_queue(int index) {
auto thread_pool = metal::new_scoped_memory_pool();
auto q = device_->newCommandQueue(MAX_BUFFERS_PER_QUEUE);
auto q = device_->newCommandQueue();
debug_set_stream_queue_label(q, index);
if (!q) {
throw std::runtime_error(

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

@@ -75,11 +75,7 @@ void finalize(Stream s) {
auto& d = metal::device(s.device);
auto cb = d.get_command_buffer(s.index);
d.end_encoding(s.index);
scheduler::notify_new_task(s);
cb->addCompletedHandler([s](MTL::CommandBuffer* cbuf) {
scheduler::notify_task_completion(s);
check_error(cbuf);
});
cb->addCompletedHandler([s](MTL::CommandBuffer* cbuf) { check_error(cbuf); });
d.commit_command_buffer(s.index);
d.get_command_buffer(s.index);
}