[Bug] Fantasizing is very slow on GPU due to torch.qr

[danielrjiang]

🐛 Bug

torch.qr is causing things to be very slow on the GPU, when using fantasies. This is likely caused by pytorch/pytorch#22573.

To reproduce

See the following notebook example.
KG_GPU_investigate.ipynb.txt

** Stack trace/error message **
See the following profile results.


kg_cpu.cprofile.txt
kg_gpu.cprofile.txt

Expected Behavior

Running on CUDA should show speedups compared to CPU.

System information

Please complete the following information:

• latest master
• 1.6.0a0
• Mac OS X Catalina 10.15.4

[jacobrgardner]

Maybe we should put in some logic to use the Woodbury formula for small enough matrices? QR is better stability but likely slower for small matrices. @Balandat @gpleiss

[Balandat]

Yeah that makes sense. Also, looking through the code above that section, we know that the new_root here is diagonal if we use cholesky for the root decomposition (i.e. if L is lower triangualr). So we don't have to do any work, we can just do a triangular solve against eye to the get the inverse. So we should probably have two code paths, one where L is diagonal and one where it's not. For the general case of tall non-triangular L we can probably do woodbury.

[jacobrgardner]

Fixed for now by #1224