Update internal code for at::_lu_with_info (#56612)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/56612
The goal of this refactoring is to make the `torch.linalg.solve`
to be a composition of calls to `lu_stub` and `lu_solve_stub`.
Once `lu_stub` and `lu_solve_stub` have cuSOLVER-based codepath,
`torch.linalg.solve` will have it as well.
Replaced `lu_with_info_{cpu, cuda}` with one function that calls
to `lu_stub`.
Split MAGMA-based `apply_lu` into `apply_lu_looped_magma`
and `apply_lu_batched_magma`. This simplifies the future switch to
cuSOLVER and cuBLAS libraries.
Test Plan: Imported from OSS
Reviewed By: albanD
Differential Revision: D28248756
Pulled By: mruberry
fbshipit-source-id: 40e02b5be4ff5f78885bcc95685aba581043e096