RFC: Improve the performance and usability of linear algebra on CUDA devices

[xwang233]

🚀 The feature, motivation and pitch

Currently, the torch.linalg (https://pytorch.org/docs/stable/linalg.html) package provides linear algebra functionalities in pytorch. The CUDA backend is supported by cuSOLVER and MAGMA libraries.

For now, linear algebra operators in pytorch are implemented in either cuSOLVER or MAGMA, or both. Users can use

torch.backends.cuda.preferred_linalg_library(backend='cusolver')

to prefer one of the two backends. Available options (python str) are default (using heuristics), cusolver, or magma. See doc for details https://pytorch.org/docs/stable/backends.html#torch.backends.cuda.preferred_linalg_library.

However, libraries have limitations and heuristics can't be perfect on all devices, library versions, input batch sizes, and input shapes. We'd like to collect user feedbacks and feature requests on the performance and usability of pytorch linear algebra on CUDA devices. Please leave a comment if you have any suggestions. Thank you!

Alternatives

No response

Additional context

No response

cc @jianyuh @nikitaved @pearu @mruberry @walterddr @IvanYashchuk @xwang233 @lezcano @ptrblck @ngimel

[vadimkantorov]

There were some recent related discussions on API design in #76440. My opinion: for the finest-grain level control allow ops to take op-level execution hints (in the most general case, a dict or namedtuple)

[Balandat]

cc @gpleiss, @jacobrgardner