-
Notifications
You must be signed in to change notification settings - Fork 17
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Fix the execution of ANISymmetryFunctions on a non-default CUDA stream #37
Conversation
This is important for the capture of the CUDA graphs, as it has to be done on a non-default stream (https://pytorch.org/docs/1.10.0/notes/cuda.html#cuda-graphs). |
I think, I found a solution. |
@peastman Do you have a better suggestion how to implement |
This looks like a reasonable approach. One very minor change I would suggest is not to cache the stream in a field of the object. Instead of stream = torch::cuda::getCurrentCUDAStream(positions.device().index()).stream();
cudaSymFunc->setStream(stream); I would write cudaSymFunc->setStream(torch::cuda::getCurrentCUDAStream(positions.device().index()).stream()); And then the same in |
One other potential issue. By calling |
It is guaranteed. The tensor is copied to GPU before. const Tensor positions = positions_.to(tensorOptions);
|
@peastman this is ready. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks good!
ANISymmetryFunctions
give incorrect result if executed on a non-default stream.CudaANISymmetryFunctions::setStream
andCudaANISymmetryFunctions::getStream
NNPOps::ANISymmetryFunctions::Holder