Performance of broadcast matrix solvers

nedlrichards: <p>TLDR: Computing an inverse and multiplying is much faster than using a solver when <code>A</code> is (1, 10, 10) and <code>b</code> is (5000, 1, 10). Using a solver is fastest when <code>A</code> is (5000, 10, 10) and <code>b</code> is (1, 1, 10). This difference is less pronounced on cuda than cpu. However, when there are multiple leading dimensions on A and B, it seems like the inverse-multiply is still the fastest way to compute a solution in torch on either device.</p>
<p>I have started to use torch in the last year, mainly to replace numpy and scipy in linear algebra calculations. The broadcasting capabilities of torch are really eye opening, it’s cool how many operations there are that can generalize to N dimensions!</p>
<p>I am confused about the performance of broadcasting with torch solvers, however. I am trying to convince my team that they should be using matrix solvers, instead of calculating an inverse matrix and then multiplying. At issue here is a performance critical inverse, where <code>A</code> is an hermitian matrix and <code>b</code> is a set of vectors, with a few shared and broadcast dimensions outside of the final dimensions.</p>
<p>I believe the theory here is clear, we should use a Cholesky decomposition and solver. In practice this does not play out, which may be related to the <a href="https://www.pugetsystems.com/labs/hpc/PyTorch-for-Scientific-Computing---Quantum-Mechanics-Example-Part-3-Code-Optimizations---Batched-Matrix-Operations-Cholesky-Decomposition-and-Inverse-1225/" rel="noopener nofollow ugc">small matrix sizes</a> (matrices in A are 10 x 10). However, the issue persists when I tried torch.linalg.solve and an explicit LU decomposition.</p>
<p>The problem may be in the broadcasting itself. It seems like the solvers do not perform well when the decomposition result is used many times, <em>i.e.</em> when <code>A</code> has a dimension broadcast up to a large dimension in <code>b</code>.</p>
<pre><code class="lang-auto">import torch
import numpy as np
import time

dim_0 = 10
# Switch dimensions 1&amp;2 to change performance
dim_1 = 1
dim_2 = 5000
device = torch.device('cpu')

def make_array(device):    
    B = torch.randn((dim_1, dim_0), dtype=torch.complex64, device=device)
    snapshots = torch.randn((dim_2, dim_0), dtype=torch.complex64, device=device)
    A = torch.einsum('...i, ...j-&gt;...ij', snapshots, snapshots.conj())
    # make sure everything is contiguous
    A = A.contiguous()

    # Create a diagonal matrix with the loading value
    loading = 5e-1 * torch.eye(A.shape[-1], dtype=torch.complex64, device=device)
    # Reshape the diagonal matrix to broadcast with the CSDM tensor
    loading_shape = (A.ndim - 2) * (1,) + (A.shape[-1], A.shape[-1])
    loading = loading.reshape(loading_shape)
    A = A + loading
    return A, B

def time_solve(solver):
    num_loops = 20
    total_time = 0.
    for i in range(num_loops):
        A, B = make_array(device)
        torch.cuda.synchronize()
        tmp = time.time()
        solver(B, A)
        torch.cuda.synchronize()
        if i &gt; 0:
          total_time += time.time() - tmp
    return total_time * 1e6 / num_loops

def solve_cholesky(steering_vectors, csdm):
    csdm_chol = torch.linalg.cholesky(csdm)
    beams_chol = torch.cholesky_solve(steering_vectors.unsqueeze(-1), csdm_chol).squeeze(-1)

def solve(steering_vectors, csdm):
    beams = torch.linalg.solve(csdm, steering_vectors.unsqueeze(-1)).squeeze(-1)

def solve_inverse(steering_vectors, csdm):
    csdm_inverted = torch.linalg.inv(csdm)
    beams_inv = torch.einsum("...ij,...j-&gt;...i", csdm_inverted, steering_vectors,)

print(time_solve(solve_cholesky))
print(time_solve(solve))
print(time_solve(solve_inverse))

</code></pre>

nedlrichards: <p>After more digging, it seems like there are times when computing a matrix inverse is <a href="https://scicomp.stackexchange.com/questions/26423/practical-example-of-why-it-is-not-good-to-invert-a-matrix/26428#26428" rel="noopener nofollow ugc">preferred</a> over a decomposition. Repeated calculations of small systems of equations with the same A seems to fit cleanly into this use case.</p>