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

Fix leak with multi-output primitives (#1274)

Browse Source 
* fix leak with multi-output primitives

* hopefully an actual fix
This commit is contained in:
Awni Hannun
2024-07-23 06:34:18 -07:00
committed by GitHub
parent df124e018a
commit 1fba87b0df
4 changed files with 29 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
mlx/array.cpp
View File

@@ -175,10 +175,11 @@ array::~array() {
return;
}
// Ignore arrays that will be detached
if (status() != array::Status::unscheduled) {
// Ignore arrays that might be detached during eval
if (status() == array::Status::scheduled) {
return;
}
// Break circular reference for non-detached arrays with siblings
if (auto n = siblings().size(); n > 0) {
bool do_detach = true;

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

@@ -75,6 +75,9 @@ std::function<void()> make_task(array arr, bool signal) {
if (!arr.is_tracer()) {
arr.detach();
}
for (auto& out : outputs) {
out.set_status(array::Status::available);
}
if (signal || d.get_command_buffer_ops(s.index) >= MAX_OPS_PER_BUFFER) {
d.end_encoding(s.index);

9
mlx/transforms.cpp
View File

@@ -141,10 +141,9 @@ array eval_impl(std::vector<array> outputs, bool async) {
tape.pop();
// Set the status of the array and siblings.
auto status = async ? array::Status::scheduled : array::Status::available;
arr.set_status(status);
arr.set_status(array::Status::scheduled);
for (auto& s : arr.siblings()) {
s.set_status(status);
s.set_status(array::Status::scheduled);
}
auto stream = arr.primitive().stream();
@@ -170,6 +169,10 @@ array eval_impl(std::vector<array> outputs, bool async) {
if (!arr.is_tracer()) {
arr.detach();
}
for (auto& out : outputs) {
out.set_status(array::Status::available);
}
if (signal) {
arr.event().signal();
}