diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp index 1cab6b4081..707624ea5f 100644 --- a/mlx/backend/cuda/allocator.cpp +++ b/mlx/backend/cuda/allocator.cpp @@ -67,11 +67,12 @@ inline void* unified_malloc(size_t size) { return data; } +// Errors are ignored since memory is freed in destructors, which can not throw. inline void unified_free(void* data) { if (supports_managed_memory()) { - CHECK_CUDA_ERROR(cudaFree(data)); + cudaFree(data); } else { - CHECK_CUDA_ERROR(cudaFreeHost(data)); + cudaFreeHost(data); } } @@ -319,14 +320,15 @@ void CudaAllocator::free_async(CudaBuffer& buf, cudaStream_t stream) { if (buf.device == -1) { unified_free(buf.data); } else { - // Free asynchronously when memory pools is supported. + // Free asynchronously when memory pools is supported. Errors are ignored + // since memory is freed in destructors, which can not throw. if (mem_pools_[buf.device]) { if (!stream) { stream = free_streams_[buf.device]; } - CHECK_CUDA_ERROR(cudaFreeAsync(buf.data, stream)); + cudaFreeAsync(buf.data, stream); } else { - CHECK_CUDA_ERROR(cudaFree(buf.data)); + cudaFree(buf.data); } } } diff --git a/mlx/backend/cuda/cuda_utils.h b/mlx/backend/cuda/cuda_utils.h index f8a234ee65..95fb8f2dd2 100644 --- a/mlx/backend/cuda/cuda_utils.h +++ b/mlx/backend/cuda/cuda_utils.h @@ -26,11 +26,11 @@ class CudaHandle { } ~CudaHandle() { - // Skip if there was an error to avoid throwing in the destructors - if (cudaPeekAtLastError() != cudaSuccess) { - return; + // Errors are ignored since destructors can not throw, and destroying can + // fail when the CUDA runtime is shutting down. + if (handle_ != nullptr) { + Destroy(handle_); } - reset(); } CudaHandle(const CudaHandle&) = delete; diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp index e7d8f2620d..7662fb2583 100644 --- a/mlx/backend/cuda/device.cpp +++ b/mlx/backend/cuda/device.cpp @@ -206,15 +206,17 @@ CommandEncoder::CommandEncoder(Device& d) : device_(d), stream_(d), graph_(d), - worker_(std::make_shared(d)), + worker_(std::make_unique(d)), graph_cache_("MLX_CUDA_GRAPH_CACHE_SIZE", /* default_capacity */ 400) { std::tie(max_ops_per_graph_, max_mb_per_graph_) = get_graph_limits(d); - worker_->start(); } CommandEncoder::~CommandEncoder() { - synchronize(); - worker_->stop(); + try { + synchronize(); + } catch (...) { + // Synchronizing can fail when the CUDA runtime is shutting down. + } } void CommandEncoder::add_completed_handler(std::function task) { @@ -490,6 +492,8 @@ void CommandEncoder::commit() { void CommandEncoder::synchronize() { CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_)); + // Commit first so the handler below runs after all pending tasks. + commit(); auto p = std::make_shared>(); std::future f = p->get_future(); add_completed_handler([p = std::move(p)]() { p->set_value(); }); diff --git a/mlx/backend/cuda/device.h b/mlx/backend/cuda/device.h index 198f0b5ad8..cf525e97ed 100644 --- a/mlx/backend/cuda/device.h +++ b/mlx/backend/cuda/device.h @@ -146,7 +146,7 @@ class CommandEncoder { Device& device_; CudaStream stream_; CudaGraph graph_; - std::shared_ptr worker_; + std::unique_ptr worker_; int node_count_{0}; bool in_concurrent_{false}; std::vector from_nodes_; diff --git a/mlx/backend/cuda/event.cu b/mlx/backend/cuda/event.cu index d3b6f97f5d..d1f27e1230 100644 --- a/mlx/backend/cuda/event.cu +++ b/mlx/backend/cuda/event.cu @@ -232,8 +232,8 @@ AtomicEvent::AtomicEvent(Device& d) { cuda_free = cudaFree; coherent_ = false; } - buf_ = std::shared_ptr( - buf, [cuda_free](void* buf) { CHECK_CUDA_ERROR(cuda_free(buf)); }); + // Errors are ignored since the deleter may run after CUDA is shut down. + buf_ = std::shared_ptr(buf, [cuda_free](void* buf) { cuda_free(buf); }); if (coherent_) { *ptr() = 0; } else { diff --git a/mlx/backend/cuda/worker.cpp b/mlx/backend/cuda/worker.cpp index 6266d6099b..86e8a4cc55 100644 --- a/mlx/backend/cuda/worker.cpp +++ b/mlx/backend/cuda/worker.cpp @@ -3,74 +3,73 @@ #include "mlx/backend/cuda/worker.h" #include "mlx/backend/cuda/device.h" +#include + namespace mlx::core::cu { Worker::Worker(Device& d) : signal_stream_(d), - signal_event_(d, cudaEventDisableTiming | cudaEventBlockingSync) {} - -Worker::~Worker() = default; - -void Worker::start() { - // Note that |shared_from_this| can not be called in constructor. - worker_ = std::thread(&Worker::thread_fn, shared_from_this()); - // Detach the thread and let it free itself after finishing tasks. + signal_event_(d, cudaEventDisableTiming | cudaEventBlockingSync), + state_(std::make_shared()) { + // Detach the thread and let it free the state after finishing tasks. // This is to avoid deadlock when joining threads on exit on Windows: // https://developercommunity.visualstudio.com/t/1654756 - worker_.detach(); + std::thread(&State::thread_fn, state_).detach(); } -void Worker::stop() { +Worker::~Worker() { { - std::lock_guard lock(mtx_); - stop_ = true; + std::lock_guard lock(state_->mtx); + state_->stop = true; } - cond_.notify_one(); + state_->cond.notify_one(); } void Worker::add_task(std::function task) { pending_tasks_.push_back(std::move(task)); } -void Worker::signal(void* data) { - auto w = static_cast(data); - { - std::lock_guard lock(w->mtx_); - w->signaled_batch_++; - } - w->cond_.notify_one(); -} - void Worker::commit(cudaStream_t stream) { // Move pending tasks into tasks if (pending_tasks_.empty()) { return; } { - std::lock_guard lock(mtx_); + std::lock_guard lock(state_->mtx); // Move pending tasks into ready tasks - worker_tasks_[++committed_batch_] = std::move(pending_tasks_); + state_->worker_tasks[++committed_batch_] = std::move(pending_tasks_); } signal_event_.record(stream); signal_event_.wait(signal_stream_); - CHECK_CUDA_ERROR(cudaLaunchHostFunc(signal_stream_, signal, this)); + CHECK_CUDA_ERROR( + cudaLaunchHostFunc(signal_stream_, State::signal, state_.get())); +} + +// static +void Worker::State::signal(void* data) { + auto state = static_cast(data); + { + std::lock_guard lock(state->mtx); + state->signaled_batch++; + } + state->cond.notify_one(); } -void Worker::thread_fn() { +void Worker::State::thread_fn() { uint64_t current_batch = 0; while (true) { Tasks tasks; { - std::unique_lock lk(mtx_); - cond_.wait(lk, [this, current_batch] { - return this->signaled_batch_ > current_batch || this->stop_; + std::unique_lock lk(mtx); + cond.wait(lk, [this, current_batch] { + return this->signaled_batch > current_batch || this->stop; }); - if (stop_) { + if (stop) { return; } - current_batch = signaled_batch_; - auto end = worker_tasks_.upper_bound(current_batch); - for (auto it = worker_tasks_.begin(); it != end; ++it) { + current_batch = signaled_batch; + auto end = worker_tasks.upper_bound(current_batch); + for (auto it = worker_tasks.begin(); it != end; ++it) { if (tasks.empty()) { tasks = std::move(it->second); } else { @@ -78,7 +77,7 @@ void Worker::thread_fn() { it->second.begin(), it->second.end(), std::back_inserter(tasks)); } } - worker_tasks_.erase(worker_tasks_.begin(), end); + worker_tasks.erase(worker_tasks.begin(), end); } // Make sure tasks are cleared before the next wait for (int i = 0; i < tasks.size(); ++i) { diff --git a/mlx/backend/cuda/worker.h b/mlx/backend/cuda/worker.h index 54a4456ace..451cbf3b9c 100644 --- a/mlx/backend/cuda/worker.h +++ b/mlx/backend/cuda/worker.h @@ -9,12 +9,11 @@ #include #include #include -#include namespace mlx::core::cu { // Run tasks in worker thread, synchronized with cuda stream. -class Worker : public std::enable_shared_from_this { +class Worker { public: explicit Worker(Device& d); ~Worker(); @@ -22,9 +21,6 @@ class Worker : public std::enable_shared_from_this { Worker(const Worker&) = delete; Worker& operator=(const Worker&) = delete; - void start(); - void stop(); - // Add a pending |task| that will run when consumed or commited. void add_task(std::function task); @@ -33,27 +29,31 @@ class Worker : public std::enable_shared_from_this { void commit(cudaStream_t stream); private: - static void signal(void*); + using Tasks = std::vector>; + + // State shared with the detached worker thread, which may outlive the Worker + // and free the state after the CUDA runtime is gone. + struct State { + static void signal(void* data); + void thread_fn(); - void thread_fn(); - std::mutex mtx_; - std::condition_variable cond_; + std::mutex mtx; + std::condition_variable cond; + uint64_t signaled_batch{0}; + bool stop{false}; + std::map worker_tasks; + }; uint64_t committed_batch_{0}; - uint64_t signaled_batch_{0}; // Cuda stream and event for signaling kernel completion. CudaStream signal_stream_; CudaEvent signal_event_; - bool stop_{false}; - // Tasks are put in |pending_tasks_| first, and then moved to - // |worker_tasks_| when end_batch() is called. - using Tasks = std::vector>; + // |state_->worker_tasks| when commit() is called. Tasks pending_tasks_; - std::map worker_tasks_; - std::thread worker_; + std::shared_ptr state_; }; } // namespace mlx::core::cu diff --git a/mlx/compile.cpp b/mlx/compile.cpp index bb17c44962..c3ccd737d1 100644 --- a/mlx/compile.cpp +++ b/mlx/compile.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -302,9 +303,10 @@ std::uintptr_t get_function_address(const std::function& fun) { class CompileCache { public: struct CacheEntry { - CacheEntry(Stream stream, bool shapeless) + CacheEntry(std::optional stream, bool shapeless) : stream(stream), shapeless(shapeless) {}; - Stream stream; + // The default stream when the function was traced, if there was one. + std::optional stream; bool shapeless; std::vector inputs; std::vector outputs; @@ -370,7 +372,8 @@ class CompileCache { // Loop over entries and check: // - Default stream and device match the entry's default stream // - Inputs match i.e. shapes and types must be equal. - auto stream = default_stream(default_device()); + // The default stream is only created when the function actually uses it. + auto stream = default_stream_if_exists(default_device()); for (CacheEntry& entry : entries) { // Check that the default stream and device match if (entry.stream != stream) { @@ -1152,6 +1155,8 @@ ArrayFnWithExtra compile( // Trace to build the graph std::tie(entry.inputs, entry.outputs, entry.extra) = compile_trace(fun, inputs, shapeless); + // Tracing may have created the default stream. + entry.stream = default_stream_if_exists(default_device()); // DFS the graph and get a tape, and a map of array id to (parent, // position in parent inputs) diff --git a/mlx/stream.cpp b/mlx/stream.cpp index b78ee67d67..08f4d6830f 100644 --- a/mlx/stream.cpp +++ b/mlx/stream.cpp @@ -50,6 +50,10 @@ Stream default_stream(Device d) { return s.value(); } +std::optional default_stream_if_exists(Device d) { + return default_stream_storage(d); +} + void set_default_stream(Stream s) { if (!gpu::is_available() && s.device == Device::gpu) { throw std::invalid_argument( diff --git a/mlx/stream.h b/mlx/stream.h index 3099ee2682..bae55751cc 100644 --- a/mlx/stream.h +++ b/mlx/stream.h @@ -2,6 +2,7 @@ #pragma once +#include #include #include @@ -29,6 +30,9 @@ struct MLX_API ThreadLocalStream : public Stream { /** Get the default stream of current thread for the given device. */ MLX_API Stream default_stream(Device d); +/** Get the default stream of current thread if it exists. */ +std::optional default_stream_if_exists(Device d); + /** Make the stream the default for its device on current thread. */ MLX_API void set_default_stream(Stream s); diff --git a/mlx/transforms.cpp b/mlx/transforms.cpp index 694bca53d8..6d2e46f2ab 100644 --- a/mlx/transforms.cpp +++ b/mlx/transforms.cpp @@ -80,14 +80,16 @@ thread_local int detail::RetainGraph::tracing_counter{0}; array eval_impl(std::vector outputs, bool async) { std::deque tape; - // Make an effort to choose a good output stream - Stream stream = default_stream(default_device()); - for (auto& o : outputs) { - if (o.status() == array::Status::unscheduled && o.has_primitive()) { - stream = o.primitive().stream(); - break; + // Make an effort to choose a good output stream, and only create the default + // stream when there is no other choice. + Stream stream = [&outputs]() { + for (auto& o : outputs) { + if (o.status() == array::Status::unscheduled && o.has_primitive()) { + return o.primitive().stream(); + } } - } + return default_stream(default_device()); + }(); // Map of array id that needs fence and stream it's computed on std::unordered_map> needs_fence; diff --git a/tests/scheduler_tests.cpp b/tests/scheduler_tests.cpp index 8a98d35eb9..cc2228e0f3 100644 --- a/tests/scheduler_tests.cpp +++ b/tests/scheduler_tests.cpp @@ -114,6 +114,35 @@ TEST_CASE("test thread unsafe stream") { CHECK_EQ(expected, actual); } +TEST_CASE("test eval does not create default stream") { + auto s = new_thread_unsafe_stream(default_device()); + size_t num_streams = get_streams().size(); + + std::thread t([&] { + async_eval(arange(10, s)); + eval(arange(10, s)); + }); + t.join(); + + CHECK_EQ(get_streams().size(), num_streams); +} + +TEST_CASE("test compile does not create default stream") { + auto s = new_thread_unsafe_stream(default_device()); + size_t num_streams = get_streams().size(); + + std::function(const std::vector&)> fun = + [s](const std::vector& inputs) { + return std::vector{abs(inputs[0], s)}; + }; + auto cfun = compile(fun); + + std::thread t([&] { eval(cfun({array({-1, 2})})); }); + t.join(); + + CHECK_EQ(get_streams().size(), num_streams); +} + TEST_CASE("test thread local stream") { auto s = new_thread_local_stream(default_device()); int result = sum(arange(10, s)).item();