[CudaGraph] Handle exceptions thrown while capturing cuda graph #17113
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Prior to this commit, an exception thrown during the capture of a cuda graph would result in
std::terminate
being called. This commit updates the implementation of"vm.builtin.cuda_graph.run_or_capture"
such that a thrown exception can be recovered from, and does not cause any changes to the state of TVM's cuda graph cache.Call to
cudaStreamDestroy
was previously skipped, now moved to a RAII-style destructor in aScopedCUDAStream
class.Call to
cudaStreamEndCapture
was previously skipped, end of cuda graph capture now performed as part of RAII-style destructor forCUDACaptureStream
class.Restoration of
CUDAThreadEntry::ThreadLocal()->stream
was previously skipped, now restored as part of RAII-style destructor forCUDACaptureStream
class.Previously, an error raised from
cudaGraphInstantiate
would leave thecapture_cache_
in an ill-formed state. Now, thecapture_cache_
is only updated after a validCUDAGraphCapturedState
has been fully constructed.