[CudaGraph] Handle exceptions thrown while capturing cuda graph #17113
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 a `ScopedCUDAStream` class. - Call to `cudaStreamEndCapture` was previously skipped, end of cuda graph capture now performed as part of RAII-style destructor for `CUDACaptureStream` class. - Restoration of `CUDAThreadEntry::ThreadLocal()->stream` was previously skipped, now restored as part of RAII-style destructor for `CUDACaptureStream` class. - Previously, an error raised from `cudaGraphInstantiate` would leave the `capture_cache_` in an ill-formed state. Now, the `capture_cache_` is only updated after a valid `CUDAGraphCapturedState` has been fully constructed.
vinx13
reviewed
Jun 25, 2024
| CUDA_CALL(cudaStreamBeginCapture(capture_stream_, cudaStreamCaptureModeGlobal)); | ||
| } | ||
| ~CUDACaptureStream() { | ||
| cudaStreamEndCapture(capture_stream_, output_graph_); |
Member
There was a problem hiding this comment.
this should be enclosed with CUDA_CALL to check return code
Contributor
Author
There was a problem hiding this comment.
I tested it with CUDA_CALL to verify, and using the CUDA_CALL would re-introduce the same bug that this change is intended to fix. If the stack is unwinding due to a thrown exception, then throwing another exception would result in std::terminate being called. To avoid this, destructors shouldn't throw exceptions (stackoverflow link).
There are some ways to use std::uncaught_exceptions to determine whether an exception is being unwound, and to conditionally throw an exception if it isn't already the case. However, those tend to be pretty context-dependent, and probably aren't worth using in this case.
vinx13
approved these changes
Jun 27, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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::terminatebeing 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
cudaStreamDestroywas previously skipped, now moved to a RAII-style destructor in aScopedCUDAStreamclass.Call to
cudaStreamEndCapturewas previously skipped, end of cuda graph capture now performed as part of RAII-style destructor forCUDACaptureStreamclass.Restoration of
CUDAThreadEntry::ThreadLocal()->streamwas previously skipped, now restored as part of RAII-style destructor forCUDACaptureStreamclass.Previously, an error raised from
cudaGraphInstantiatewould leave thecapture_cache_in an ill-formed state. Now, thecapture_cache_is only updated after a validCUDAGraphCapturedStatehas been fully constructed.