diff --git a/mlx/array.cpp b/mlx/array.cpp index 07d3acd616..e6dbbe7851 100644 --- a/mlx/array.cpp +++ b/mlx/array.cpp @@ -208,34 +208,33 @@ void array::copy_shared_buffer(const array& other) { copy_shared_buffer(other, other.strides(), other.flags(), other.data_size()); } -array::~array() { - if (array_desc_ == nullptr) { +void array::release(std::shared_ptr desc) { + if (desc == nullptr) { return; } // Detached/detaching - if (array_desc_->primitive == nullptr) { + if (desc->primitive == nullptr) { return; } // Break circular reference for non-detached arrays with siblings - if (auto n = siblings().size(); n > 0) { + if (auto n = desc->siblings.size(); n > 0) { bool do_detach = true; // If all siblings have siblings.size() references except - // the one we are currently destroying (which has siblings.size() + 1) + // the one we are currently releasing (which has siblings.size() + 1) // then there are no more external references - do_detach &= (array_desc_.use_count() == (n + 1)); - for (auto& s : siblings()) { + do_detach &= (desc.use_count() == (n + 1)); + for (auto& s : desc->siblings) { do_detach &= (s.array_desc_.use_count() == n); if (!do_detach) { break; } } if (do_detach) { - for (auto& s : siblings()) { + for (auto& s : desc->siblings) { for (auto& ss : s.siblings()) { - // Set to null here to avoid descending into array destructor - // for siblings + // Set to null here to avoid descending into reset() for siblings ss.array_desc_ = nullptr; } s.array_desc_->siblings.clear(); diff --git a/mlx/array.h b/mlx/array.h index 3f45e9cb9d..cc29290d2c 100644 --- a/mlx/array.h +++ b/mlx/array.h @@ -5,6 +5,7 @@ #include #include #include +#include #include #include "mlx/allocator.h" @@ -80,13 +81,22 @@ class MLX_API array { array& operator=(array&& other) && = delete; /** Default copy and move constructors otherwise. */ - array& operator=(array&& other) & = default; array(const array& other) = default; array(array&& other) = default; + /** Assignment goes through release() to break sibling reference cycle. */ array& operator=(const array& other) & { - if (this->id() != other.id()) { - this->array_desc_ = other.array_desc_; + if (array_desc_ != other.array_desc_) { + // Hold the old descriptor while assigning, and then release it. Otherwise + // when |this| and |other| are siblings they would be both destroyed. + release(std::exchange(array_desc_, other.array_desc_)); + } + return *this; + } + + array& operator=(array&& other) & noexcept { + if (array_desc_ != other.array_desc_) { + release(std::exchange(array_desc_, std::move(other.array_desc_))); } return *this; } @@ -466,10 +476,12 @@ class MLX_API array { void copy_shared_buffer(const array& other); void overwrite_descriptor(const array& other) { - array_desc_ = other.array_desc_; + release(std::exchange(array_desc_, other.array_desc_)); } - ~array(); + ~array() { + release(std::move(array_desc_)); + } private: // Initialize the arrays data @@ -532,6 +544,12 @@ class MLX_API array { // the primitive which knows how to compute the array's data from its inputs // and the list of array's inputs for the primitive. std::shared_ptr array_desc_; + + /** + * Lets go of an ArrayDesc, breaking the sibling reference cycle of a + * multi-output primitive. + */ + static void release(std::shared_ptr desc); }; template diff --git a/tests/array_tests.cpp b/tests/array_tests.cpp index 68a4bed3a8..cc0cf52395 100644 --- a/tests/array_tests.cpp +++ b/tests/array_tests.cpp @@ -690,3 +690,60 @@ TEST_CASE("test siblings circular references without eval") { fun(); CHECK(tracker.expired()); } + +// https://github.com/ml-explore/mlx/pull/4453 +TEST_CASE("test siblings circular references released by assignment") { + std::weak_ptr tracker; + auto fun = [&]() { + array sink({3, 4}); + array key({1, 2}); + auto splits = split(key, 2); + { + // Set fake data as a tracker for ArrayDesc's lifetime. + splits[0].set_data(allocator::malloc(0)); + tracker = splits[0].data_shared_ptr(); + } + // Drop both outputs by assigning over them rather than destroying + // them. The siblings hold each other, so unless assignment runs the + // same check the destructor does, nothing ever releases them. + splits[0] = sink; // copy assignment + splits[1] = array({5, 6}); // move assignment + }; + fun(); + CHECK(tracker.expired()); +} + +TEST_CASE("test siblings circular references released by overwrite") { + std::weak_ptr tracker; + auto fun = [&]() { + array key({1, 2}); + auto splits = split(key, 2); + { + // Set fake data as a tracker for ArrayDesc's lifetime. + splits[0].set_data(allocator::malloc(0)); + tracker = splits[0].data_shared_ptr(); + } + // Leave one output as the last reference, then replace it in place. + // This is the path __setitem__ takes, with no compile involved. + array last = splits[0]; + splits.clear(); + last.overwrite_descriptor(array({7, 8})); + }; + fun(); + CHECK(tracker.expired()); +} + +TEST_CASE("test assigning a sibling over the last reference keeps the cycle") { + array key({1, 2}); + auto splits = split(key, 2); + // Leave one output as the last outside reference to the pair, then assign + // its sibling over it. The sibling is reachable only through the array + // being replaced, so it has to be held before the old descriptor is + // checked: checked first, the pair looks unreferenced, its cycle is broken + // and the sibling list the new value lives in is freed before it is read. + array last = splits[0]; + splits.clear(); + last = last.siblings()[0]; + REQUIRE_EQ(last.siblings().size(), 1); + CHECK(array_equal(last, array({2})).item()); +} diff --git a/tests/autograd_tests.cpp b/tests/autograd_tests.cpp index 9fad4376fc..1ff203e1e9 100644 --- a/tests/autograd_tests.cpp +++ b/tests/autograd_tests.cpp @@ -749,7 +749,7 @@ TEST_CASE("test op vjps") { TEST_CASE("test gather and take grads") { // Check linear takes auto linear_f = [](array indices) { - auto fun_linear = [&indices](array input) { return take(input, indices); }; + auto fun_linear = [indices](array input) { return take(input, indices); }; return fun_linear; };