Gpytorch: Make use of pytorch's batched symeig

[Balandat]

batched symeig is supported since https://github.com/pytorch/pytorch/pull/21858

We should update batch_symeig to use this new feature (we still want to manually offload to cpu for small matrices for now, until pytorch does it. See discussion here: https://github.com/pytorch/pytorch/issues/22573

We can keep things backward-compatible by catching the exception emitted by the old implementation.

[jacobrgardner]

We can, but to be honest I'm not convinced of the value of using their batching of symeig since it just does a for loop over the individual symeig calls anyways.

What we really want down is support for something like cuSolver's syevjBatched which (a) supports batch mode as a single function, and (b) uses Jacobi instead of implicit QR which will make better use of GPU hardware for small matrices.

[Balandat]

Sure, makes sense. I guess if the number of batches is large then looping in C++ and doing the masking logic in batch mode instead of a python loop could matter thogh.

[vishwakftw]

Just to chime in here: while it is still a for loop internally, the batched symeig implementation saves a lot of time by performing just one clone to get column major matrices, while a naive implementation using a loop would perform as many clones as there are batches. There’s also the reshaping cost (but I’m not sure how high the overhead would be).

On the CPU, I can see a significant improvement in computation times for matrix sizes <= 32 between batched symeig and the naive implementation batched symeig.

As discussed on Slack, a cuSolver based implementation would be really helpful, but it could be beyond the scope of the next release.

[gpleiss]

I have a branch that creates a C wrapper for cusolver's batch symeig (from 4 months ago). I can try to revive it.

[calincru]

@gpleiss @jacobrgardner Do you have any idea whether the cuSolver implementation could be faster than offloading to CPU?

FWIW, I opened an issue some time ago with TF too (see https://github.com/tensorflow/tensorflow/issues/30466) as it seems that their eigh is still slower on GPU despite them using cuSolver.

[gpleiss]

@calincru yeah unfortunately eigen-solving is not really designed for GPUs. The only time when cuSolver is truly faster (and when the CPU is really a bottleneck) is when we have to do batch eigen-solving on, say, 100+ 32x32 matrices. This happens when you use compute the log likelihood of a batch GP with a batch dimension of 10+

[calincru]

I'm working on something right now that deals with many very small matrices (think 100k+ 2x2/3x3 matrices). But as I said, neither PyTorch nor TF seem to show improvements on GPU -- on the contrary, they are extremely slow.

What I'm doing for now is to compute the eigenvalues manually for 2x2 matrices (via formula). This should technically be possible for 3x3 and 4x4 matrices too, but the algorithm for 3x3 matrices is already numerically unstable (see Wikipedia).

Suggestions on what I could do are very welcome. Thanks.

[gpleiss]

Looks like this happened in #928