From 713f2b4a8f464023a6c82ebf8125d0c0cfd1d697 Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Sun, 21 Jun 2026 15:02:08 +0400 Subject: [PATCH 1/2] Add load tests --- python/tests/test_load.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/python/tests/test_load.py b/python/tests/test_load.py index 10fb63ea63..5ab7d6b68b 100644 --- a/python/tests/test_load.py +++ b/python/tests/test_load.py @@ -88,6 +88,34 @@ def test_load_npy_dtype(self): with self.assertRaises(Exception): out = mx.load(save_file, stream=mx.cpu) + def test_load_npy_read_error(self): + save_file = os.path.join(self.test_dir, "truncated.npy") + expected = np.arange(16, dtype=np.float32) + np.save(save_file, expected) + with open(save_file, "r+b") as f: + f.truncate(os.path.getsize(save_file) - expected.nbytes) + + out = mx.load(save_file, stream=mx.cpu) + with self.assertRaises(RuntimeError): + mx.eval(out) + + def test_async_load_npy_read_error_across_streams(self): + save_file = os.path.join(self.test_dir, "truncated_async.npy") + expected = np.arange(16, dtype=np.float32) + np.save(save_file, expected) + with open(save_file, "r+b") as f: + f.truncate(os.path.getsize(save_file) - expected.nbytes) + + producer_stream = mx.new_stream(mx.cpu) + consumer_stream = mx.new_stream(mx.cpu) + out = mx.add( + mx.load(save_file, stream=producer_stream), + 1.0, + stream=consumer_stream, + ) + with self.assertRaises(RuntimeError): + mx.eval(out) + def test_save_and_load_safetensors(self): test_file = os.path.join(self.test_dir, "test.safetensors") with self.assertRaises(Exception): From b71f0ec7e2cd0b7f3f94de4c9196729a0c9d5e7c Mon Sep 17 00:00:00 2001 From: Cheng Date: Mon, 22 Jun 2026 11:04:47 +0900 Subject: [PATCH 2/2] Propagate CPU errors to events --- mlx/array.h | 1 + mlx/backend/common/load.cpp | 11 +- mlx/backend/cuda/event.cu | 190 +++++++++++++++++------------------ mlx/backend/cuda/event.h | 20 +++- mlx/backend/cuda/fence.cpp | 20 ++-- mlx/backend/metal/device.cpp | 23 ++--- mlx/backend/metal/device.h | 11 +- mlx/backend/metal/event.cpp | 37 +++---- mlx/backend/metal/event.h | 9 +- mlx/backend/metal/fence.cpp | 3 +- mlx/backend/no_gpu/event.cpp | 48 +++++---- mlx/backend/no_gpu/fence.cpp | 50 +++------ mlx/event.h | 31 ++++++ mlx/fence.h | 7 +- mlx/scheduler.cpp | 52 ++++++++++ mlx/scheduler.h | 33 +++++- mlx/transforms.cpp | 3 + 17 files changed, 331 insertions(+), 218 deletions(-) diff --git a/mlx/array.h b/mlx/array.h index 8e14ca4726..3f45e9cb9d 100644 --- a/mlx/array.h +++ b/mlx/array.h @@ -426,6 +426,7 @@ class MLX_API array { } void detach_event() const { + array_desc_->event.check_error(); array_desc_->event = Event{}; } diff --git a/mlx/backend/common/load.cpp b/mlx/backend/common/load.cpp index ce41963de7..d664c7411c 100644 --- a/mlx/backend/common/load.cpp +++ b/mlx/backend/common/load.cpp @@ -3,6 +3,8 @@ #include #include +#include + #include "mlx/primitives.h" #include "mlx/scheduler.h" @@ -51,7 +53,14 @@ void Load::eval_cpu(const std::vector& inputs, array& out) { } }; auto fut = io::thread_pool().enqueue(std::move(read_task)).share(); - scheduler::enqueue(stream(), [fut = std::move(fut)]() { fut.wait(); }); + auto s = stream(); + scheduler::enqueue(s, [s, fut = std::move(fut)]() { + try { + fut.get(); + } catch (const std::exception& error) { + scheduler::set_error(s, fmt::format("[Load::eval_cpu] {}", error.what())); + } + }); } } // namespace mlx::core diff --git a/mlx/backend/cuda/event.cu b/mlx/backend/cuda/event.cu index b73937ec38..9278697a58 100644 --- a/mlx/backend/cuda/event.cu +++ b/mlx/backend/cuda/event.cu @@ -113,10 +113,7 @@ void CudaEvent::init_pool() { cuda_event_pool(); } -// Wraps CudaEvent with a few features: -// 1. The class can be copied. -// 2. Make wait/record work with CPU streams. -// 3. Add checks for waiting on un-recorded event. +// Wraps CudaEvent so it can be copied. class CopyableCudaEvent { public: explicit CopyableCudaEvent(Device& d) @@ -126,32 +123,24 @@ class CopyableCudaEvent { cudaEventDisableTiming | cudaEventBlockingSync)) {} void wait() { + check_recorded(); event_->wait(); } void wait(Stream s) { - if (s.device == mlx::core::Device::cpu) { - scheduler::enqueue(s, [*this]() mutable { - check_recorded(); - event_->wait(); - }); - } else { - check_recorded(); - auto& encoder = cu::get_command_encoder(s); - encoder.commit(); - event_->wait(encoder.stream()); - } + assert(s.device == mlx::core::Device::gpu); + check_recorded(); + auto& encoder = cu::get_command_encoder(s); + encoder.commit(); + event_->wait(encoder.stream()); } void record(Stream s) { - if (s.device == mlx::core::Device::cpu) { - throw std::runtime_error("CudaEvent can not wait on CPU stream."); - } else { - auto& encoder = cu::get_command_encoder(s); - encoder.commit(); - event_->record(encoder.stream()); - recorded_ = true; - } + assert(s.device == mlx::core::Device::gpu); + auto& encoder = cu::get_command_encoder(s); + encoder.commit(); + event_->record(encoder.stream()); + recorded_ = true; } bool is_signaled() const { @@ -213,6 +202,11 @@ auto check_gpu_coherency() { return coherency; } +const CudaStream& signal_stream() { + static CudaStream stream(device(0)); + return stream; +} + AtomicEvent::AtomicEvent(Device& d) { void* buf; cudaError_t (*cuda_free)(void*); @@ -264,14 +258,11 @@ void AtomicEvent::wait(cudaStream_t stream, uint32_t value) { void AtomicEvent::wait(Stream s, uint32_t value) { nvtx3::scoped_range r("cu::AtomicEvent::wait(s)"); - if (s.device == mlx::core::Device::cpu) { - scheduler::enqueue(s, [*this, value]() mutable { wait(value); }); - } else { - auto& encoder = get_command_encoder(s); - encoder.commit(); - wait(encoder.stream(), value); - encoder.add_completed_handler([buf = buf_]() {}); - } + assert(s.device == mlx::core::Device::gpu); + auto& encoder = get_command_encoder(s); + encoder.commit(); + wait(encoder.stream(), value); + encoder.add_completed_handler([buf = buf_]() {}); } void AtomicEvent::signal(uint32_t value) { @@ -289,17 +280,11 @@ void AtomicEvent::signal(cudaStream_t stream, uint32_t value) { void AtomicEvent::signal(Stream s, uint32_t value) { nvtx3::scoped_range r("cu::AtomicEvent::signal(s)"); - if (s.device == mlx::core::Device::cpu) { - // Signal through a GPU stream so the atomic is updated in GPU - updating - // the atomic in CPU sometimes does not get GPU notified. - scheduler::enqueue( - s, [*this, value]() mutable { signal(signal_stream(), value); }); - } else { - auto& encoder = get_command_encoder(s); - encoder.commit(); - signal(encoder.stream(), value); - encoder.add_completed_handler([buf = buf_]() {}); - } + assert(s.device == mlx::core::Device::gpu); + auto& encoder = get_command_encoder(s); + encoder.commit(); + signal(encoder.stream(), value); + encoder.add_completed_handler([buf = buf_]() {}); } bool AtomicEvent::is_signaled(uint32_t val) const { @@ -319,9 +304,21 @@ uint32_t AtomicEvent::value() const { } } -const CudaStream& AtomicEvent::signal_stream() { - static CudaStream stream(device(0)); - return stream; +/////////////////////////////////////////////////////////////////////////////// +// EventImpl implementations +/////////////////////////////////////////////////////////////////////////////// + +void EventImpl::ensure_created(Stream s, uint64_t signal_value) { + if (is_created()) { + return; + } + auto& d = cu::device(s.device); + if (s.device == mlx::core::Device::cpu || signal_value > 1) { + nvtx3::mark("Using slow AtomicEvent"); + atomic = std::make_unique(d); + } else { + cuda = std::make_unique(d); + } } } // namespace cu @@ -330,86 +327,85 @@ const CudaStream& AtomicEvent::signal_stream() { // Event implementations /////////////////////////////////////////////////////////////////////////////// -namespace { - -struct EventImpl { - // CudaEvent is preferred when possible because it is fast, however we have - // to fallback to AtomicEvent in following cases: - // 1. the event is used to wait/signal a cpu stream; - // 2. signal value other than 1 has been specified. - std::unique_ptr cuda; - std::unique_ptr atomic; - - bool is_created() const { - return cuda || atomic; - } - - void ensure_created(Stream s, uint64_t signal_value) { - if (is_created()) { - return; - } - auto& d = cu::device(s.device); - if (s.device == mlx::core::Device::cpu || signal_value > 1) { - nvtx3::mark("Using slow AtomicEvent"); - atomic = std::make_unique(d); - } else { - cuda = std::make_unique(d); - } - } -}; - -} // namespace - Event::Event(Stream s) : stream_(s) { - event_ = std::shared_ptr( - new EventImpl(), [](void* ptr) { delete static_cast(ptr); }); + event_ = std::make_shared(); } void Event::wait() { - auto* event = static_cast(event_.get()); - assert(event->is_created()); - if (event->cuda) { + check_error(); + auto& event = cast(); + assert(event.is_created()); + if (event.cuda) { assert(value() == 1); - event->cuda->wait(); + event.cuda->wait(); } else { - event->atomic->wait(value()); + event.atomic->wait(value()); } CHECK_CUDA_ERROR(cudaPeekAtLastError()); + check_error(); } void Event::wait(Stream s) { - auto* event = static_cast(event_.get()); - assert(event->is_created()); - if (event->cuda) { + auto& event = cast(); + assert(event.is_created()); + if (event.cuda) { assert(value() == 1); - event->cuda->wait(s); + if (s.device == mlx::core::Device::cpu) { + scheduler::wait_event(s, *this, [value = value()](Event& self) { + self.cast().cuda->wait(); + }); + } else { + event.cuda->wait(s); + } } else { - event->atomic->wait(s, value()); + if (s.device == mlx::core::Device::cpu) { + scheduler::wait_event(s, *this, [value = value()](Event& self) { + self.cast().atomic->wait(value); + }); + } else { + event.atomic->wait(s, value()); + } } } void Event::signal(Stream s) { - auto* event = static_cast(event_.get()); - event->ensure_created(s, value()); - if (event->cuda) { + auto& event = cast(); + event.ensure_created(s, value()); + if (event.cuda) { assert(value() == 1); - event->cuda->record(s); + if (s.device == mlx::core::Device::cpu) { + throw std::runtime_error("CudaEvent can not wait on CPU stream."); + } else { + event.cuda->record(s); + } } else { - event->atomic->signal(s, value()); + if (s.device == mlx::core::Device::cpu) { + // Signal through a GPU stream so the atomic is updated in GPU - updating + // the atomic in CPU sometimes does not get GPU notified. + scheduler::signal_event(s, *this, [value = value()](Event& self) { + self.cast().atomic->signal(cu::signal_stream(), value); + }); + } else { + event.atomic->signal(s, value()); + } } } bool Event::is_signaled() const { - auto* event = static_cast(event_.get()); - if (!event->is_created()) { + auto& event = cast(); + if (!event.is_created()) { return false; } - if (event->cuda) { + if (event.cuda) { assert(value() == 1); - return event->cuda->is_signaled(); + return event.cuda->is_signaled(); } else { - return event->atomic->is_signaled(value()); + return event.atomic->is_signaled(value()); } } +Event::Error& Event::error() { + return cast().error; +} + } // namespace mlx::core diff --git a/mlx/backend/cuda/event.h b/mlx/backend/cuda/event.h index 53afeb0117..39d710fa7f 100644 --- a/mlx/backend/cuda/event.h +++ b/mlx/backend/cuda/event.h @@ -13,6 +13,7 @@ namespace mlx::core::cu { +class CopyableCudaEvent; class Device; // RAII-managed move-only wrapper of cudaEvent_t. @@ -66,8 +67,6 @@ class AtomicEvent { uint32_t value() const; private: - const CudaStream& signal_stream(); - uint32_t* ptr() const { return static_cast(buf_.get()); } @@ -76,4 +75,21 @@ class AtomicEvent { std::shared_ptr buf_; }; +struct EventImpl { + Event::Error error; + + // CudaEvent is preferred when possible because it is fast, however we have + // to fallback to AtomicEvent in following cases: + // 1. the event is used to wait/signal a cpu stream; + // 2. signal value other than 1 has been specified. + std::unique_ptr cuda; + std::unique_ptr atomic; + + bool is_created() const { + return cuda || atomic; + } + + void ensure_created(Stream s, uint64_t signal_value); +}; + } // namespace mlx::core::cu diff --git a/mlx/backend/cuda/fence.cpp b/mlx/backend/cuda/fence.cpp index c6a41f0e60..3a3acdba09 100644 --- a/mlx/backend/cuda/fence.cpp +++ b/mlx/backend/cuda/fence.cpp @@ -9,22 +9,23 @@ namespace mlx::core { struct FenceImpl { uint32_t count; - cu::AtomicEvent event; + Event event; + + FenceImpl(uint32_t count, Stream s) : count(count), event(s) {} }; Fence::Fence(Stream s) { - fence_ = std::shared_ptr( - new FenceImpl{0, cu::device(s.device)}, - [](void* ptr) { delete static_cast(ptr); }); + fence_ = std::make_shared(0, s); + // Ensure that we use AtomicEvent. + cast().event.cast().ensure_created(s, 2); } void Fence::wait(Stream s, const array&) { - auto* fence = static_cast(fence_.get()); - fence->event.wait(fence->count); + cast().event.wait(); } void Fence::update(Stream s, const array& a, bool cross_device) { - auto* fence = static_cast(fence_.get()); + auto& f = cast(); if (cross_device) { // Move to managed memory if there is a device switch auto& cbuf = @@ -35,8 +36,9 @@ void Fence::update(Stream s, const array& a, bool cross_device) { cu::allocator().move_to_unified_memory(cbuf, encoder.stream()); } } - fence->count++; - fence->event.signal(s, fence->count); + f.count++; + f.event.set_value(f.count); + f.event.signal(s); } } // namespace mlx::core diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 2a8e15afd7..46102ef260 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -465,19 +465,15 @@ void CommandEncoder::end_encoding() { all_inputs_.clear(); } -void CommandEncoder::signal_event( - std::shared_ptr event, - uint64_t value) { +void CommandEncoder::signal_event(Event event, uint64_t value) { end_encoding(); - buffer_->encodeSignalEvent(event->mtl_event(), value); + buffer_->encodeSignalEvent(event.cast().mtl_event(), value); signal_events_.push_back({std::move(event), value}); } -void CommandEncoder::wait_event( - std::shared_ptr event, - uint64_t value) { +void CommandEncoder::wait_event(Event event, uint64_t value) { end_encoding(); - buffer_->encodeWait(event->mtl_event(), value); + buffer_->encodeWait(event.cast().mtl_event(), value); wait_events_.push_back(std::move(event)); } @@ -491,14 +487,15 @@ void CommandEncoder::commit(std::function completion) { [&error_ = error_, wait_events = std::move(wait_events_), signal_events = std::move(signal_events_), - completion = std::move(completion)](MTL::CommandBuffer* cbuf) { + completion = std::move(completion)](MTL::CommandBuffer* cbuf) mutable { if (completion) { completion(); } // If any of the waited event has error in it, poison the encoder. for (auto& event : wait_events) { - if (event->error()) { - error_ = event->error(); + auto err = event.load_error(); + if (err) { + error_ = std::move(err); break; } } @@ -512,14 +509,14 @@ void CommandEncoder::commit(std::function completion) { // Poison all the signaled events when error happened. if (error_) { for (auto& [event, value] : signal_events) { - event->set_error(error_); + event.set_error(error_); } } // Metal won't signal the events for us on error, manually signal them // to avoid infinite waiting. if (cbuf->status() == MTL::CommandBufferStatusError) { for (auto& [event, value] : signal_events) { - event->signal(value); + event.cast().signal(value); } } }); diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h index bed0cd636e..63c5c83777 100644 --- a/mlx/backend/metal/device.h +++ b/mlx/backend/metal/device.h @@ -20,7 +20,6 @@ using MTLFCList = std::vector>; class Device; -class EventImpl; class MLX_API CommandEncoder { public: @@ -91,8 +90,8 @@ class MLX_API CommandEncoder { void barrier(); void end_encoding(); - void wait_event(std::shared_ptr event, uint64_t value); - void signal_event(std::shared_ptr event, uint64_t value); + void wait_event(Event event, uint64_t value); + void signal_event(Event event, uint64_t value); bool needs_commit() const; void commit(std::function completion = nullptr); void synchronize(); @@ -114,11 +113,11 @@ class MLX_API CommandEncoder { size_t buffer_sizes_{0}; // The events hooked to current command buffer. - std::vector> wait_events_; - std::vector, uint64_t>> signal_events_; + std::vector wait_events_; + std::vector> signal_events_; // Error from previous commited command buffer. - std::shared_ptr error_; + Event::Error error_; // Encoder for issuing GPU commands. // The members are used within a single ComputeCommandEncoder and will be diff --git a/mlx/backend/metal/event.cpp b/mlx/backend/metal/event.cpp index 77f48f0838..73113cf19e 100644 --- a/mlx/backend/metal/event.cpp +++ b/mlx/backend/metal/event.cpp @@ -26,26 +26,13 @@ EventImpl::~EventImpl() { } void EventImpl::wait(uint64_t value) { - check_error(); mtl_event_->waitUntilSignaledValue(value, -1); // never times out - check_error(); } void EventImpl::signal(uint64_t value) { mtl_event_->setSignaledValue(value); } -void EventImpl::set_error(std::shared_ptr error) { - std::atomic_store(&error_, std::move(error)); -} - -void EventImpl::check_error() { - auto error = std::atomic_exchange(&error_, {}); - if (error) { - throw std::runtime_error(*error); - } -} - } // namespace metal /////////////////////////////////////////////////////////////////////////////// @@ -57,36 +44,40 @@ Event::Event(Stream stream) : stream_(stream) { } void Event::wait() { - static_cast(event_.get())->wait(value()); + check_error(); + cast().wait(value()); + check_error(); } void Event::wait(Stream stream) { - auto impl = std::static_pointer_cast(event_); if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [impl = std::move(impl), value = value()]() { - impl->wait(value); + scheduler::wait_event(stream, *this, [value = value()](Event& self) { + self.cast().wait(value); }); } else { auto& encoder = metal::get_command_encoder(stream); - encoder.wait_event(std::move(impl), value()); + encoder.wait_event(*this, value()); } } void Event::signal(Stream stream) { - auto impl = std::static_pointer_cast(event_); if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [impl = std::move(impl), value = value()]() { - impl->signal(value); + scheduler::signal_event(stream, *this, [value = value()](Event& self) { + self.cast().signal(value); }); } else { auto& encoder = metal::get_command_encoder(stream); - encoder.signal_event(std::move(impl), value()); + encoder.signal_event(*this, value()); } } bool Event::is_signaled() const { - auto* mtl_event = static_cast(event_.get())->mtl_event(); + auto* mtl_event = cast().mtl_event(); return mtl_event->signaledValue() >= value(); } +Event::Error& Event::error() { + return cast().error(); +} + } // namespace mlx::core diff --git a/mlx/backend/metal/event.h b/mlx/backend/metal/event.h index c5c82a7cd3..fec91530c3 100644 --- a/mlx/backend/metal/event.h +++ b/mlx/backend/metal/event.h @@ -12,20 +12,17 @@ class EventImpl { void wait(uint64_t value); void signal(uint64_t value); - void set_error(std::shared_ptr error); - void check_error(); - const auto& error() const { + Event::Error& error() { return error_; } - auto* mtl_event() { + auto* mtl_event() const { return mtl_event_.get(); } private: - // TODO: Use std::atomic when it gets supported in Xcode. - std::shared_ptr error_; + Event::Error error_; NS::SharedPtr mtl_event_; }; diff --git a/mlx/backend/metal/fence.cpp b/mlx/backend/metal/fence.cpp index 6fdd57a5f6..70dd0e33bd 100644 --- a/mlx/backend/metal/fence.cpp +++ b/mlx/backend/metal/fence.cpp @@ -41,8 +41,7 @@ struct FenceImpl { }; Fence::Fence(Stream stream) { - auto dtor = [](void* ptr) { delete static_cast(ptr); }; - fence_ = std::shared_ptr(new FenceImpl(stream), dtor); + fence_ = std::make_shared(stream); } void Fence::wait(Stream stream, const array& x) { diff --git a/mlx/backend/no_gpu/event.cpp b/mlx/backend/no_gpu/event.cpp index 6dde047ab4..52cd390a02 100644 --- a/mlx/backend/no_gpu/event.cpp +++ b/mlx/backend/no_gpu/event.cpp @@ -12,42 +12,52 @@ struct EventCounter { uint64_t value{0}; std::mutex mtx; std::condition_variable cv; + Event::Error error; + + void wait(uint64_t val) { + std::unique_lock lk(mtx); + if (value >= val) { + return; + } + cv.wait(lk, [this, val] { return value >= val; }); + } }; Event::Event(Stream stream) : stream_(stream) { - auto dtor = [](void* ptr) { delete static_cast(ptr); }; - event_ = std::shared_ptr(new EventCounter{}, dtor); + event_ = std::make_shared(); } void Event::wait() { - auto ec = static_cast(event_.get()); - std::unique_lock lk(ec->mtx); - if (ec->value >= value()) { - return; - } - ec->cv.wait(lk, [value = value(), ec] { return ec->value >= value; }); + check_error(); + cast().wait(value()); + check_error(); } void Event::wait(Stream stream) { - scheduler::enqueue(stream, [*this]() mutable { wait(); }); + scheduler::wait_event(stream, *this, [value = value()](Event& self) { + self.cast().wait(value); + }); } void Event::signal(Stream stream) { - scheduler::enqueue(stream, [*this]() mutable { - auto ec = static_cast(event_.get()); + scheduler::signal_event(stream, *this, [value = value()](Event& self) { + auto& ec = self.cast(); { - std::lock_guard lk(ec->mtx); - ec->value = value(); + std::lock_guard lk(ec.mtx); + ec.value = value; } - ec->cv.notify_all(); + ec.cv.notify_all(); }); } bool Event::is_signaled() const { - auto ec = static_cast(event_.get()); - { - std::lock_guard lk(ec->mtx); - return (ec->value >= value()); - } + auto& ec = cast(); + std::lock_guard lk(ec.mtx); + return ec.value >= value(); +} + +Event::Error& Event::error() { + return cast().error; } + } // namespace mlx::core diff --git a/mlx/backend/no_gpu/fence.cpp b/mlx/backend/no_gpu/fence.cpp index cd66d23cfe..05852c860b 100644 --- a/mlx/backend/no_gpu/fence.cpp +++ b/mlx/backend/no_gpu/fence.cpp @@ -1,54 +1,30 @@ // Copyright © 2024 Apple Inc. -#include -#include - #include "mlx/fence.h" -#include "mlx/scheduler.h" +#include "mlx/event.h" namespace mlx::core { struct FenceImpl { - uint32_t count{0}; - uint32_t value{0}; - std::mutex mtx; - std::condition_variable cv; + uint32_t count; + Event event; + + FenceImpl(uint32_t count, Stream s) : count(count), event(s) {} }; -Fence::Fence(Stream) { - auto dtor = [](void* ptr) { delete static_cast(ptr); }; - fence_ = std::shared_ptr(new FenceImpl{}, dtor); +Fence::Fence(Stream s) { + fence_ = std::make_shared(0, s); } -void Fence::wait(Stream stream, const array&) { - auto& f = *static_cast(fence_.get()); - if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [count = f.count, fence_ = fence_]() mutable { - auto& f = *static_cast(fence_.get()); - std::unique_lock lk(f.mtx); - if (f.value >= count) { - return; - } - f.cv.wait(lk, [&f, count] { return f.value >= count; }); - }); - } else { - throw std::runtime_error("[Fence::wait] Invalid stream."); - } +void Fence::wait(Stream s, const array&) { + cast().event.wait(s); } -void Fence::update(Stream stream, const array&, bool) { - auto& f = *static_cast(fence_.get()); +void Fence::update(Stream s, const array&, bool) { + auto& f = cast(); f.count++; - if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [count = f.count, fence_ = fence_]() mutable { - auto& f = *static_cast(fence_.get()); - std::unique_lock lk(f.mtx); - f.value = count; - f.cv.notify_all(); - }); - } else { - throw std::runtime_error("[Fence::update] Invalid stream."); - } + f.event.set_value(f.count); + f.event.signal(s); } } // namespace mlx::core diff --git a/mlx/event.h b/mlx/event.h index 66a6a75df5..543785c5a2 100644 --- a/mlx/event.h +++ b/mlx/event.h @@ -11,6 +11,9 @@ namespace mlx::core { class Event { public: + // TODO: Use std::atomic when it gets supported in Xcode. + using Error = std::shared_ptr; + Event() {}; explicit Event(Stream stream); @@ -26,6 +29,27 @@ class Event { // Check if the event has been signaled at its current value bool is_signaled() const; + // Associate an error to the event + void set_error(Error err) { + std::atomic_store(&error(), std::move(err)); + } + + // Do an atomic load of the associated error. + Error load_error() { + return std::atomic_load(&error()); + } + + // Throw and clear the associated error + void check_error() { + if (!event_) { + return; + } + auto err = std::atomic_exchange(&error(), {}); + if (err) { + throw std::runtime_error(*err); + } + } + // Check if the event is valid bool valid() const { return event_ != nullptr; @@ -47,7 +71,14 @@ class Event { return stream_; } + template + auto& cast() const { + return *static_cast(event_.get()); + } + private: + Error& error(); + // Default constructed stream should never be used // since the event is not yet valid Stream stream_{0, Device::cpu}; diff --git a/mlx/fence.h b/mlx/fence.h index 0ececdb6d7..3fd5da333b 100644 --- a/mlx/fence.h +++ b/mlx/fence.h @@ -32,8 +32,13 @@ class Fence { void update(Stream stream, const array& x, bool cross_device); void wait(Stream stream, const array& x); + template + auto& cast() const { + return *static_cast(fence_.get()); + } + private: - std::shared_ptr fence_{nullptr}; + std::shared_ptr fence_; }; } // namespace mlx::core diff --git a/mlx/scheduler.cpp b/mlx/scheduler.cpp index 7507917f5b..60991229d2 100644 --- a/mlx/scheduler.cpp +++ b/mlx/scheduler.cpp @@ -60,6 +60,58 @@ void Scheduler::enqueue(Stream s, std::function task) { st->enqueue(std::move(task)); } +void Scheduler::wait_event( + Stream s, + Event event, + std::function task) { + assert(s.device == Device::cpu); + enqueue( + s, [this, s, event = std::move(event), task = std::move(task)]() mutable { + task(event); + // Poison current stream if the waited event has error. + auto err = event.load_error(); + if (err) { + set_error(s, std::move(err)); + } + }); +} + +void Scheduler::signal_event( + Stream s, + Event event, + std::function task) { + assert(s.device == Device::cpu); + enqueue( + s, [this, s, event = std::move(event), task = std::move(task)]() mutable { + { + // Poison the signal event if current stream has error. + std::unique_lock lock(error_mtx_); + auto it = errors_.find(s.index); + if (it != errors_.end()) { + event.set_error(it->second); + } + } + task(event); + }); +} + +void Scheduler::set_error(Stream s, Event::Error error) { + assert(s.device == Device::cpu); + // Set error only when no error happended before, to preserve the + // earliest error. + std::unique_lock lock(error_mtx_); + errors_.try_emplace(s.index, std::move(error)); +} + +void Scheduler::finalize(Stream s) { + assert(s.device == Device::cpu); + // Clear error in the end of graph. + enqueue(s, [this, s]() { + std::unique_lock lock(error_mtx_); + errors_.erase(s.index); + }); +} + // Leak the scheduler singleton on all platforms. During static destruction, // worker threads may still be executing JIT-compiled code that has been // unmapped, causing SIGSEGV (macOS/Linux) or join() deadlocks (Windows/MSVC diff --git a/mlx/scheduler.h b/mlx/scheduler.h index c84ab62855..68d3334a1e 100644 --- a/mlx/scheduler.h +++ b/mlx/scheduler.h @@ -4,6 +4,7 @@ #include #include +#include #include #include #include @@ -76,6 +77,10 @@ class MLX_API Scheduler { Scheduler& operator=(Scheduler&&) = delete; void enqueue(Stream s, std::function task); + void wait_event(Stream s, Event event, std::function task); + void signal_event(Stream s, Event event, std::function task); + void set_error(Stream s, Event::Error error); + void finalize(Stream s); void notify_new_task(const Stream& stream) { { @@ -115,13 +120,37 @@ class MLX_API Scheduler { std::shared_mutex threads_mtx_; std::condition_variable completion_cv; std::mutex mtx; + + std::unordered_map errors_; + std::mutex error_mtx_; }; MLX_API Scheduler& scheduler(); template -void enqueue(const Stream& stream, F&& f) { - scheduler().enqueue(stream, std::forward(f)); +inline void enqueue(Stream s, F&& f) { + scheduler().enqueue(s, std::forward(f)); +} + +// Like enqueue but the task is used for processing the passed event. +template +inline void wait_event(Stream s, Event event, F&& f) { + scheduler().wait_event(s, std::move(event), std::forward(f)); +} + +template +inline void signal_event(Stream s, Event event, F&& f) { + scheduler().signal_event(s, std::move(event), std::forward(f)); +} + +// Set error in all the pending events associated with the stream. +inline void set_error(Stream s, const std::string& error) { + scheduler().set_error(s, std::make_shared(error)); +} + +// Mark the end of eval. +inline void finalize(Stream s) { + scheduler().finalize(s); } inline int n_active_tasks() { diff --git a/mlx/transforms.cpp b/mlx/transforms.cpp index 9a8207339e..90533323d6 100644 --- a/mlx/transforms.cpp +++ b/mlx/transforms.cpp @@ -311,6 +311,9 @@ array eval_impl(std::vector outputs, bool async) { if (auto e = events.find(s.index); e != events.end()) { e->second.signal(s); } + if (s.device == Device::cpu) { + scheduler::finalize(s); + } if (s.device == Device::gpu) { gpu::finalize(s); }