Cusolver and magma timings for linalg.eigh() (symmetric case)

KFrank: <p>Hello Forum (and Johannes)!</p>
<p>The short story:</p>
<p>cusolver outperforms magma for <code>linalg.eigh()</code> (for my test case).</p>
<p>cusolver saturates my gpu with light-to-moderate cpu usage, while magma only shows<br>
moderate gpu usage and heavy (but not saturated) cpu usage.</p>
<p>(I monitor gpu usage with “nvidia-smi -l 1” and cpu usage with the ubuntu / gnome<br>
graphical “System Monitor.”)</p>
<p>I became curious about the <code>linalg</code> “backends” because of this thread:</p>
<aside class="quote quote-modified" data-post="1" data-topic="223386">
  <div class="title">
    <div class="quote-controls"></div>
    <img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/johannes99/48/78328_2.png" class="avatar">
    <div class="quote-title__text-content">
      <a href="https://discuss.pytorch.org/t/torch-linalg-eig-parallelisation/223386">Torch.linalg.eig parallelisation</a> 
    </div>
  </div>
  <blockquote>
    Hello, 
I am using torch for some computations, but am struggeling with improving the performance of my code. I need to get the eigenvalues of some matrizes in the order of magnitude of 6000x6000. It works fine with torch.linalg.eig, but is rather slow, compared to the rest of my computations. If I understand correctly, this is largely because part of the computations need to be performed on CPU. 
Now what I don’t understand is, that the .eig function does not seem to parallelize at all. If I ha…
  </blockquote>
</aside>

<p>It appears that pytorch only supports magma for <code>linalg.eig()</code> (the general<br>
asymmetric case), but poking around in the code Johannes showed me indicated<br>
that both cusolver and magma are supported for <code>linalg.eigh()</code> (the symmetric<br>
and hermitian case), so I ran some comparative timings.</p>
<p>Here is the timing script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)
print (torch.version.cuda)
print (torch.cuda.get_device_name())

import concurrent.futures
from time import time

torch.backends.cuda.preferred_linalg_library ('cusolver')   # to mask subsequent warning

_ = torch.manual_seed (2025)

n = 6000
nb = 4

nWarm = 3
nTime = 10

def batchEigh (tb):
    return  torch.linalg.eigh (tb)

def poolEigh (tb):
    with concurrent.futures.ThreadPoolExecutor(max_workers = 8) as pool:
        futures = [pool.submit (torch.linalg.eigh, t)  for t in tb]
        return  [fut.result()  for fut in futures]

print ('nb:', nb)
print ('n: ', n)
for  dev in ('cpu', 'cuda'):
# for  dev in ('cuda',):
    print (dev, 'timings:')
    if  dev == 'cuda':  backends = ('cusolver', 'magma')
    else:               backends = ('default',)             # just a placeholder
    for  be in backends:
        torch.backends.cuda.preferred_linalg_library (backend = be)
        if  dev == 'cuda':  print (be, '(preferred backend):')  
        tBatch = torch.randn (nb, n, n, device = dev, requires_grad = True)   # batch of nb nxn matrices
        tBatch = (tBatch + tBatch.mT) / 2                                     # make them symmetric
        for  eFunc in (batchEigh, poolEigh):
            for  i in range (nWarm):
                eigh = eFunc (tBatch)
            if  dev == 'cuda':  torch.cuda.synchronize()
            t0 = time()
            for  i in range (nTime):
                eigh = eFunc (tBatch)
            if  dev == 'cuda':  torch.cuda.synchronize()
            t1 = time()
            print ('{:10s} {:8.2f} sec'.format (eFunc.__name__ + ':', (t1 - t0) / nTime))
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.8.0+cu129
12.9
NVIDIA RTX 3000 Ada Generation Laptop GPU
[W1006 12:43:32.703334536 Context.cpp:320] Warning: torch.backends.cuda.preferred_linalg_library is an experimental feature. If you see any error or unexpected behavior when this flag is set please file an issue on GitHub. (function operator())
nb: 4
n:  6000
cpu timings:
batchEigh:    24.04 sec
poolEigh:     29.87 sec
cuda timings:
cusolver (preferred backend):
batchEigh:     7.42 sec
poolEigh:      6.65 sec
magma (preferred backend):
batchEigh:    10.87 sec
poolEigh:      8.18 sec
</code></pre>
<p>Symmetric eigendecomposition is an easier problem (although algorithms may also<br>
be more mature and efficient), as evidenced by the cpu timings.</p>
<p>The cusolver implementation significantly – but not hugely – outperforms magma.</p>
<p>This adds to the argument that pytorch should use cusolver by default for the<br>
general (asymmetric) <code>linagl.eig()</code>.</p>
<p>(Based on Johannes’s reported timings for a non-pytorch cusolver general<br>
eigendecompostion test vs. the pytorch / magma version, it appears that  the<br>
relative performance improvement of cusolver over magma would be even larger<br>
for <code>eig()</code> (general) than for <code>eigh()</code> (symmetric).)</p>
<p>A couple of comments:</p>
<p>Again, parallelizing the batch dimension with a python thread pool increased<br>
magma’s performance significantly (although not by as much as it did for <code>eig()</code>.)</p>
<p>magma again used the cpu heavily, but did use the gpu more heavily (but far<br>
from saturated) than it did for <code>eig()</code>.</p>
<p>I don’t know how to inquire which backend is actually being used, but because<br>
the timings and usage pattern differ depending on which backend is specified<br>
in the <code>torch.backends.cuda.preferred_linalg_library()</code> call, I deduce that<br>
cusolver and magma are being used according to which is chosen as “preferred.”</p>
<p>(Not relevant to cuda, but using the thread pool hurt the cpu performance for<br>
<code>eigh()</code>, rather than helping it significantly as it did for <code>eig()</code>.)</p>
<p>Of course, both the absolute and relative timings will depend on the hardware<br>
and size of the test tensor.</p>
<p>Best.</p>
<p>K. Frank</p>