[MAGMA][CUDA] eig: deprecate MAGMA and dispatch to cuSOLVER unconditionally (pytorch#173510)

As per title.

Benchmark script:

```python
import torch
import torch.utils.benchmark as benchmark

from itertools import product

results = []

batches = [(), (16,), (64,)]
sizes = [16, 128, 512,  2048]

for b, n in product(batches, sizes):
    shape = b + (n, n)
    print(f"Testing shape={shape}")
    label = "torch.eig"
    sub_label = f"{shape}"
    x = torch.rand(*shape, device="cuda")
    x = x + x.mH
    stmt = "torch.linalg.eig(x)"
    for backend in ("magma", "cusolver"):
        torch.backends.cuda.preferred_linalg_library(backend)
        # warm-up
        for _ in range(5):
            exec(stmt)

        results.append(benchmark.Timer(
            stmt=stmt,
            globals={'x': x},
            label=label,
            sub_label=sub_label,
            description=backend,
        ).blocked_autorange(min_run_time=1))

compare = benchmark.Compare(results)
compare.print()
```

Benchmark results (H100):
```
[-------------------- torch.eig -------------------]
                        |     magma     |   cusolver  | cusolver speedup
1 threads: ------------------------------------------
      (16, 16)          |      67215.3  |       893.3 | 75.24
      (128, 128)        |     390101.4  |     10526.9 | 37.05
      (512, 512)        |    1614706.1  |     61347.5 | 26.32
      (2048, 2048)      |    6447301.5  |    356181.2 | 18.10
      (16, 16, 16)      |     660036.1  |     13095.2 | 50.40
      (16, 128, 128)    |    6530718.7  |    166647.0 | 39.18
      (16, 512, 512)    |   20375827.4  |    994115.1 | 20.49
      (16, 2048, 2048)  |   98335490.6  |   5717112.5 | 17.20
      (64, 16, 16)      |    2167358.2  |     51977.5 | 41.69
      (64, 128, 128)    |   25925259.8  |    664574.7 | 39.01
      (64, 512, 512)    |   84731703.1  |   3946917.0 | 21.46
      (64, 2048, 2048)  |  380878661.3  |  23008593.5 | 16.55

Times are in microseconds (us).
```

Pull Request resolved: pytorch#173510
Approved by: https://github.com/Skylion007

Author: nikitaved

aten/src/ATen/native/cuda/linalg/BatchLinearAlgebra.cpp

@@ -150,13 +150,15 @@ void magmaSyevd(
     value_t* w, scalar_t* wA, magma_int_t ldwa, scalar_t* work, magma_int_t lwork, value_t* rwork,
     magma_int_t lrwork, magma_int_t* iwork, magma_int_t liwork, magma_int_t* info);
 
+#ifdef USE_ROCM
 template<class scalar_t, class value_t=scalar_t>
 void magmaEig(
     magma_vec_t jobvl, magma_vec_t jobvr, magma_int_t n, scalar_t *A, magma_int_t lda,
     scalar_t *w, scalar_t *VL, magma_int_t ldvl,
     scalar_t *VR, magma_int_t ldvr, scalar_t *work, magma_int_t lwork,
     value_t *rwork,
     magma_int_t *info);
+#endif
 
 template<class scalar_t>
 void magmaLuSolve(
@@ -721,6 +723,7 @@ void magmaSyevd<c10::complex<float>, float>(
   AT_CUDA_CHECK(cudaGetLastError());
 }
 
+#ifdef USE_ROCM
 template<>
 void magmaEig<double>(
     magma_vec_t jobvl, magma_vec_t jobvr, magma_int_t n,
@@ -800,6 +803,7 @@ void magmaEig<c10::complex<float>, float>(
          rwork, info);
   AT_CUDA_CHECK(cudaGetLastError());
 }
+#endif
 
 template<>
 void magmaLuSolve<double>(
@@ -971,14 +975,17 @@ magma_trans_t to_magma(TransposeType trans) {
 
 namespace {
 
-void _warn_once_magma_deprecation(const std::string& op_name) {
+void _warn_once_magma_deprecation(const std::string& op_name, bool force_cusolver = true) {
   if (at::globalContext().linalgPreferredBackend() == at::LinalgBackend::Magma) {
+    std::string warn_force_cusolver = force_cusolver
+      ? " " + op_name + " will try dispatching to cuSOLVER instead. " +
+        "If you see any error messages, please, file an issue on GitHub."
+      : "";
     TORCH_WARN_ONCE(
-      op_name, ": ",
-      "MAGMA, as a linear algebra backend, is deprecated and will be removed ",
-      "in future releases. ",
-      op_name, " will try dispatching to cuSOLVER instead. "
-      "If you see any error messages, please, file an issue on GitHub."
+      op_name, ": "
+      "MAGMA, as a linear algebra backend, is deprecated and will be removed "
+      "in future releases.",
+      warn_force_cusolver
     );
   }
 }
@@ -1968,6 +1975,7 @@ This is an in-place routine, content of 'input', 'values', 'vectors' is overwrit
 'infos' is an int Tensor containing error codes for each matrix in the batched input.
 For more information see MAGMA's documentation for GEEV routine.
 */
+#ifdef USE_ROCM
 template <typename scalar_t>
 void apply_magma_eig(Tensor& values, Tensor& vectors, Tensor& input, Tensor& infos, bool compute_eigenvectors) {
 #if !AT_MAGMA_ENABLED()
@@ -2045,25 +2053,22 @@ void linalg_eig_magma(Tensor& eigenvalues, Tensor& eigenvectors, Tensor& infos,
   eigenvectors.copy_(eigenvectors_cpu);
   infos.copy_(infos_cpu);
 }
+#endif // USE_ROCM
+
 void linalg_eig_kernel(Tensor& eigenvalues, Tensor& eigenvectors, Tensor& infos, const Tensor& input, bool compute_eigenvectors) {
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(input.is_cuda());
   // This function calculates the non-symmetric eigendecomposition in-place
   // tensors should be in batched column major memory format
   // the content of eigenvalues, eigenvectors and infos is overwritten by 'linalg_eig_magma' or
   // 'linalg_eig_cusolver_xgeev' both geev routines modify the provided input matrix in-place, therefore we need a copy
-
-  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(input.is_cuda());
-#if defined(CUSOLVER_VERSION) && (CUSOLVER_VERSION >= 11702)
-  auto preferred_backend = at::globalContext().linalgPreferredBackend();
-  switch (preferred_backend) {
-    case at::LinalgBackend::Cusolver:
-    default:
-      linalg_eig_cusolver_xgeev(eigenvalues, eigenvectors, input, infos, compute_eigenvectors);
-      return;
-    case at::LinalgBackend::Magma:
-      break; // MAGMA path handled below
-  }
-#endif
+#ifndef USE_ROCM
+  _warn_once_magma_deprecation("linalg.eig");
+  linalg_eig_cusolver_xgeev(eigenvalues, eigenvectors, input, infos, compute_eigenvectors);
+#else
+  // hipSolver does not have `geev`
+  _warn_once_magma_deprecation("linalg.eig", /*force_cusolver=*/false);
   linalg_eig_magma(eigenvalues, eigenvectors, infos, input, compute_eigenvectors);
+#endif
 }
 
 REGISTER_CUDA_DISPATCH(linalg_eig_stub, &linalg_eig_kernel)

aten/src/ATen/native/cuda/linalg/BatchLinearAlgebraLib.h

@@ -72,7 +72,7 @@ void linalg_eigh_cusolver(const Tensor& eigenvalues,
                           bool upper,
                           bool compute_eigenvectors);
 
-#if defined(CUSOLVER_VERSION) && (CUSOLVER_VERSION >= 11702)
+#if defined(CUSOLVER_VERSION)
 void linalg_eig_cusolver_xgeev(const Tensor& eigenvalues,
                                const Tensor& eigenvectors,
                                const Tensor& input,

test/test_linalg.py

@@ -32,6 +32,7 @@
      runOnRocmArch, MI200_ARCH, MI300_ARCH, MI350_ARCH, NAVI_ARCH, TEST_CUDA)
 from torch.testing._internal.common_device_type import \
     (instantiate_device_type_tests, dtypes, has_cusolver, onlyCPU, skipIf, skipCUDAIfNoMagma, skipCPUIfNoLapack, precisionOverride,
+     skipCUDAIfNoCusolverROCMIfNoMagma,
      skipCUDAIfNoCusolver, skipCUDAIfNoMagmaAndNoCusolver, skipCUDAIfRocm, onlyNativeDeviceTypes, dtypesIfCUDA,
      onlyCUDA, skipMeta, skipCUDAIfNotRocm, dtypesIfMPS, largeTensorTest)
 from torch.testing import make_tensor
@@ -2255,7 +2256,7 @@ def test_norm_fastpaths(self, device):
         self.assertEqual(result, expected)
 
     @skipCPUIfNoLapack
-    @skipCUDAIfNoMagma
+    @skipCUDAIfNoCusolverROCMIfNoMagma
     # NumPy computes only in float64 and complex128 precisions
     # for float32 or complex64 results might be very different from float64 or complex128
     @dtypes(torch.float64, torch.complex128)
@@ -2304,7 +2305,7 @@ def run_test(shape, *, symmetric=False):
             run_test(shape, symmetric=True)
 
     @onlyCUDA
-    @skipCUDAIfNoMagma
+    @skipCUDAIfNoCusolverROCMIfNoMagma
     @dtypes(*floating_and_complex_types())
     def test_eig_identity(self, device, dtype):
 
@@ -2414,7 +2415,7 @@ def run_test(shape, *, symmetric=False):
 
 
     @onlyCUDA
-    @skipCUDAIfNoMagmaAndNoCusolver
+    @skipCUDAIfNoCusolverROCMIfNoMagma
     @dtypes(*floating_and_complex_types())
     def test_eig_out_variants(self, device, dtype):
         from torch.testing._internal.common_utils import random_symmetric_matrix
@@ -2466,19 +2467,8 @@ def run_test(shape, *, symmetric=False):
             run_test(shape, symmetric=True)
 
 
-    @slowTest
-    @onlyCUDA
-    @skipCUDAIfNoMagma
-    @dtypes(torch.float32)
-    def test_eig_check_magma(self, device, dtype):
-        # For CUDA inputs only matrices of size larger than 2048x2048 actually call MAGMA library
-        shape = (2049, 2049)
-        a = make_tensor(shape, dtype=dtype, device=device)
-        w, v = torch.linalg.eig(a)
-        # check correctness using eigendecomposition identity
-        self.assertEqual(a.to(v.dtype) @ v, w * v, atol=1e-3, rtol=1e-3)
-
     @onlyCUDA
+    @skipCUDAIfNoCusolverROCMIfNoMagma
     @dtypes(torch.float32, torch.float64)
     def test_eig_cuda_complex_eigenvectors(self, device, dtype):
         """Test CUDA eigenvector decoding with known ground truth, including batching."""
@@ -2563,8 +2553,8 @@ def test_eig_cuda_complex_eigenvectors(self, device, dtype):
         rhs = vals_batch.unsqueeze(-2) * vecs_batch
         self.assertEqual(lhs, rhs, atol=1e-5, rtol=1e-5)
 
-    @skipCUDAIfNoMagma
     @skipCPUIfNoLapack
+    @skipCUDAIfNoCusolverROCMIfNoMagma
     @dtypes(*floating_and_complex_types())
     def test_eig_errors_and_warnings(self, device, dtype):
         # eig requires the input to be at least 2 dimensional tensor
@@ -2626,7 +2616,7 @@ def test_eig_errors_and_warnings(self, device, dtype):
                 torch.linalg.eig(a, out=(out_w, out_v))
 
     @skipCPUIfNoLapack
-    @skipCUDAIfNoMagma
+    @skipCUDAIfNoCusolverROCMIfNoMagma
     @dtypes(*floating_and_complex_types())
     def test_eig_with_nan(self, device, dtype):
         for val in [np.inf, np.nan]:
@@ -3110,8 +3100,8 @@ def mul_svd_factors(U, S, Vh):
                 S_s = torch.svd(A, compute_uv=False).S
                 self.assertEqual(S_s, S)
 
-    @skipCUDAIfNoMagmaAndNoCusolver
     @skipCPUIfNoLapack
+    @skipCUDAIfNoCusolverROCMIfNoMagma
     @dtypes(torch.complex128)
     def test_invariance_error_spectral_decompositions(self, device, dtype):
         make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=True)

torch/testing/_internal/common_device_type.py

@@ -1828,6 +1828,15 @@ def has_hipsolver():
     return rocm_version >= (5, 3)
 
 
+# Skips a test on CUDA if cuSOLVER is not available,
+# and on ROCm if MAGMA is not available.
+def skipCUDAIfNoCusolverROCMIfNoMagma(fn):
+    if TEST_WITH_ROCM:
+        return skipCUDAIfNoMagma(fn)
+    else:
+        return skipCUDAIfNoCusolver(fn)
+
+
 # Skips a test on CUDA/ROCM if cuSOLVER/hipSOLVER is not available
 def skipCUDAIfNoCusolver(fn):
     return skipCUDAIf(