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

Revert "More buffer donation in some cases (#1858)" (#1863)

Browse Source 
This reverts commit d274ae77f2.
This commit is contained in:
Awni Hannun
2025-02-13 14:21:44 -08:00
committed by GitHub
parent 5cd97f7ffe
commit 428f589364
2 changed files with 4 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,5 +1,6 @@
// Copyright © 2023-2024 Apple Inc.
#include <memory>
#include "mlx/backend/metal/device.h"
#include "mlx/backend/metal/event.h"
#include "mlx/backend/metal/utils.h"
@@ -53,34 +54,16 @@ std::function<void()> make_task(array arr, bool signal) {
abort_with_exception(error);
}
}
// Use a set to de-dup input buffers
std::unordered_set<std::shared_ptr<array::Data>> buffers;
std::vector<std::shared_ptr<array::Data>> buffers;
for (auto& in : arr.inputs()) {
if (in.data_shared_ptr() != nullptr) {
buffers.insert(in.data_shared_ptr());
}
buffers.push_back(in.data_shared_ptr());
}
for (auto& s : arr.siblings()) {
buffers.insert(s.data_shared_ptr());
buffers.push_back(s.data_shared_ptr());
}
if (!arr.is_tracer()) {
arr.detach();
}
// Erase any input buffers that have a use count > 1
// to keep them elligible for donation
for (auto it = buffers.begin(); it != buffers.end();) {
// Always hold buffers which could belong to outputs
// TODO: this shouldn't be necessary, but metal validation
// complains if buffers get released before the command buffer is
// finished, even if the ready signal has fired
if (!signal && it->use_count() > 1) {
it = buffers.erase(it);
} else {
++it;
}
}
for (auto& out : outputs) {
out.set_status(array::Status::evaluated);
}