Non-symmetric eigendecomposition with CUDA

[jeremyiwk]

There appears to be a non-symmetric eigendecomposition available in JAX using cuSOLVER on the backend - see documentation for jax.lax.linalg.eig.

Is there any plan to incorporate this into CUDA.jl?

[kshyatt]

Is this not covered by https://github.com/JuliaGPU/CUDA.jl/blob/master/lib/cusolver/linalg.jl#L157 ?

[jeremyiwk]

@kshyatt Thank you - yes that appears to cover it as of yesterday! I'll grab at master and play around with it.

Do you have any idea when those changes will be released?

[kshyatt]

Git blame shows that code has been in CUDA.jl for 3 years so you should be able to use it directly. You might need to add a using CUDA.CUSOLVER to your code.

[matteosecli]

I think @jeremyiwk was likely referring to the non-symmetric extension of eigen, that was indeed merged 2 days ago in #2787 and not part of a release yet.
That was possible only thanks to the recent introduction of Xgeev in CUSOLVER by CUDA, which is way faster than the MAGMA implementation JAX has. PyTorch has also wrapped it quite recently, if I recall correctly.

[jeremyiwk]

@matteosecli correct. Thanks for clarifying! I was looking for Xgeev. Indeed, I played with the JAX implementation a bit and found it shockingly slow even compared to the LAPACK version. Could have been user error though.

Very excited about this functionality in CUDA.jl though - thank you!