Fix the execution of ANISymmetryFunctions on a non-default CUDA stream

[raimis]

ANISymmetryFunctions give incorrect result if executed on a non-default stream.

• Implement CudaANISymmetryFunctions::setStream and CudaANISymmetryFunctions::getStream
• Set the CUDA stream in NNPOps::ANISymmetryFunctions::Holder
• Add tests
• Document

[raimis]

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).

[raimis]

I think, I found a solution.

[raimis]

@peastman Do you have a better suggestion how to implement CudaANISymmetryFunctions::setStream?

[peastman]

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 backward(). Probably forward() and backward() will always be called with the same stream current, but I don't know if that's guaranteed, so it's better not to assume.

[peastman]

One other potential issue. By calling getCurrentCUDAStream(positions.device().index()) does that assume the positions are on the GPU? That doesn't need to be the case. CudaANISymmetryFunctions allows the input and output arrays to be on either the CPU or GPU.

[raimis]

One other potential issue. By calling getCurrentCUDAStream(positions.device().index()) does that assume the positions are on the GPU? That doesn't need to be the case. CudaANISymmetryFunctions allows the input and output arrays to be on either the CPU or GPU.

It is guaranteed. The tensor is copied to GPU before.

        const Tensor positions = positions_.to(tensorOptions);

[raimis]

@peastman this is ready.

[peastman]

Looks good!