pytorch
bcdcf347 - Add cusolver potrs and potrsBatched to the backend of torch.cholesky_solve (#54315)

Commit
3 years ago
Add cusolver potrs and potrsBatched to the backend of torch.cholesky_solve (#54315) Summary: This PR adds cusolver potrs and potrsBatched to the backend of torch.cholesky_solve and torch.linalg.cholesky_solve. `cholesky_solve` heuristics: - If magma is not installed, or batch_size is 1: - If batch_size > 1 and nrhs == 1, dispatch to `cusolverDn<T>potrsBatched`, - Otherwise, dispatch to `cusolverDnXpotrs` (64 bit) and `cusolverDn<T>potrs` (legacy). - Otherwise, use magma. Note: `cusolverDn<T>potrsBatched` only supports `nrhs == 1`. It is used for `nrhs==1` batched matrix if magma is **not** installed. See also https://github.com/pytorch/pytorch/issues/42666 #47953 Todo: - [x] benchmark and heuristic Pull Request resolved: https://github.com/pytorch/pytorch/pull/54315 Reviewed By: ngimel Differential Revision: D27562225 Pulled By: mruberry fbshipit-source-id: 323e5d60610abbbdc8369f5eb112d9fa01da40f6
Author
Parents
Loading