diff --git a/mlx/backend/cuda/cuda_utils.h b/mlx/backend/cuda/cuda_utils.h index 7bae911d26..f8a234ee65 100644 --- a/mlx/backend/cuda/cuda_utils.h +++ b/mlx/backend/cuda/cuda_utils.h @@ -50,6 +50,12 @@ class CudaHandle { } } + Handle release() { + Handle handle = handle_; + handle_ = nullptr; + return handle; + } + operator Handle() const { return handle_; } diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp index 30248f5568..472b3d99fb 100644 --- a/mlx/backend/cuda/device.cpp +++ b/mlx/backend/cuda/device.cpp @@ -462,6 +462,42 @@ bool CommandEncoder::needs_commit() { void CommandEncoder::commit() { nvtx3::scoped_range r("CommandEncoder::commit"); + try { + commit_impl(); + } catch (...) { + // Clear pending CUDA error first. + cudaGetLastError(); + // Clear states. + clear_graph_state(); + node_count_ = 0; + bytes_in_graph_ = 0; + // Clear graph. + try { + graph_.reset(); + } catch (...) { + // Destroying could fail. + graph_.release(); + } + try { + graph_ = CudaGraph(device_); + } catch (...) { + // Keep the original error. + } + // Re-throw the error. + throw; + } +} + +void CommandEncoder::synchronize() { + CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_)); + auto p = std::make_shared>(); + std::future f = p->get_future(); + add_completed_handler([p = std::move(p)]() { p->set_value(); }); + commit(); + f.wait(); +} + +void CommandEncoder::commit_impl() { if (!temporaries_.empty()) { add_completed_handler([temporaries = std::move(temporaries_)]() {}); } @@ -520,13 +556,8 @@ void CommandEncoder::commit() { } // Reset state - from_nodes_.clear(); - to_nodes_.clear(); - graph_deps_key_.clear(); - graph_nodes_key_.clear(); - node_map_.clear(); + clear_graph_state(); graph_ = CudaGraph(device_); - is_graph_updatable_ = true; } // Put completion handlers in a batch. @@ -535,13 +566,16 @@ void CommandEncoder::commit() { bytes_in_graph_ = 0; } -void CommandEncoder::synchronize() { - CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_)); - auto p = std::make_shared>(); - std::future f = p->get_future(); - add_completed_handler([p = std::move(p)]() { p->set_value(); }); - commit(); - f.wait(); +void CommandEncoder::clear_graph_state() { + from_nodes_.clear(); + to_nodes_.clear(); + graph_deps_key_.clear(); + graph_nodes_key_.clear(); + node_map_.clear(); + active_deps_.clear(); + active_outputs_.clear(); + concurrent_nodes_.clear(); + is_graph_updatable_ = true; } Device& device(int cuda_device) { diff --git a/mlx/backend/cuda/device.h b/mlx/backend/cuda/device.h index 15d75082e9..198f0b5ad8 100644 --- a/mlx/backend/cuda/device.h +++ b/mlx/backend/cuda/device.h @@ -138,6 +138,8 @@ class CommandEncoder { std::string id; }; + void commit_impl(); + void clear_graph_state(); void insert_graph_dependencies(GraphNode node); void insert_graph_dependencies(std::vector nodes);