Use batched library routines when available.

[jpbrodrick89]

I recently performed some profiling comparing the tridiagonal solvers in jax (jax.lax.linalg.tridiagonal_solve) with lineax. I observed that the jax call to the cusparse gtsv2 routine on GPU was much faster than the partially unrolled lineax implementation of the Thomas algorithm when unbatched or for small batch sizes. However, if I increase the batch size there is a crossover point at which lineax is faster because jax.lax.linalg scales ~linearly with batch-size and lineax scales sub-linearly/in constant time (and if I manually increase the unroll in lineax to fully unrolled then lineax is faster when any batch size).

I was looking at making a PR in lineax which used the best of both worlds, but then I started to wonder why the jax implementation did not scale as well. It turns out that the jax implementation does not use the cusparse batch implementations (i.e. cusparse<t>gtsv2StridedBatch() and cusparse<t>gtsvInterleavedBatch()). Instead it just performs a for loop over the batch dimensions.

My immediate objective is to get this addressed sooner rather than later to enable clean completion of my lineax PR that I don't need to change at a later date while improving or maintaining current performance. Note my current use case does not include batch solves (I am doing an implicit solve of the 1D heat equation). In the future I may perhaps want to use the batch solve for ADI, multi-group diffusion, or batched simulations.

I would be more than happy to help but I begin to lose track at the Python/C++ barrier when trawling through the source code.

I have worded the issue title in a very general way, if we just want it to refer to tridiagonal solves happy to rename.

[dfm]

Now that #28696 is merged, JAX lowers to the batched implementation when possible (batch size > 1, and vector RHS). This should be available in the next nightly jaxlib release for experimentation! I haven't done a huge amount of benchmarking, so I'd love to hear when you get new results!

[jpbrodrick89]

Thanks so much for looking into this @dfm , I just tried to install from main (can confirm my venv has the recent changes) but the performance still seems similar (linear complexity when solving for multiple matrices when sub-linear should be possible due to parallelisation up to a point):

Is the problem that I need to wait for the nghtly so I have updated jaxlib as well as jax?

Benchmarks for jax only here for reference (for completeness I tried using the API as is and also explicitly vmapping. As expected, the results are identical). For context, the problem is time-stepping a heat equation implicitly for a plasma but the details are not important:

from functools import partial
from timeit import Timer

import jax
import jax.numpy as jnp
import numpy as np
from jax.lax.linalg import tridiagonal_solve

jax.config.update("jax_enable_x64", True)
dt = 0.05
jax_timeit = {}
jax_vmapped_timeit = {}


def get_diags(Te):
    Te_face = 0.5 * (Te[1:] + Te[:-1])
    ksh_face = -(Te_face**2.5)

    diag = 1 / dt - (ksh_face[1:] + ksh_face[:-1]) / dx**2

    lower = ksh_face[:-1] / dx**2
    diag = diag.at[0].add(lower[0])
    lower = lower.at[0].set(0.0)

    upper = ksh_face[1:] / dx**2
    diag = diag.at[-1].add(upper[-1])
    upper = upper.at[-1].set(0.0)

    return upper, diag, lower


jax_solve_jitted = jax.jit(triadagonal_solve)
vmapped_jax_solve = jax.jit(jax.vmap(lambda lower, diag, upper, rhs: jax_solve(lower, diag, upper, rhs)))

key = jax.random.PRNGKey(42)

# Set up base problem
nx = 200
dx = 2.0 / nx
Te = 2.0 + jnp.tanh(5.0 * jnp.arange(-1 - dx / 2, 1 + dx, dx))
Te = Te.at[0].set(Te[1])
Te = Te.at[-1].set(Te[-2])
upper, diag, lower = get_diags(Te)

batch_sizes = [1, 2, 5, 10, 20, 50, 100, 200, 400, 800, 1600, 3200, 6400]
for batch_size in batch_sizes: 
    # Add random noise just to make double-sure compiler doesn't recognise as multiple copies of the same matrix
    upper_batched = jnp.repeat(upper[None, :], batch_size, axis=0) + jax.random.normal(key, (batch_size, nx)) * 1e-6
    diag_batched = jnp.repeat(diag[None, :], batch_size, axis=0) + jax.random.normal(key, (batch_size, nx)) * 1e-6
    lower_batched = jnp.repeat(lower[None, :], batch_size, axis=0) + jax.random.normal(key, (batch_size, nx)) * 1e-6
    rhs = Te[1:-1, None] / dt
    rhs_batched = jnp.repeat(rhs[None, :], batch_size, axis=0) + jax.random.normal(key, (batch_size, nx)) * 1e-6

    # Solve using batched API   
    jax_solve_jitted(lower_batched, diag_batched, upper_batched, rhs_batched).block_until_ready()
    jax_timer = Timer(
        lambda: jax_solve_jitted(lower_batched, diag_batched, upper_batched, rhs_batched).block_until_ready()
    )
    iters, t = jax_timer.autorange(lambda n, t: print(f"JAX: {t} s (n={n})"))
    jax_timeit[batch_size] = t / iters

    # Solve using vmap explicitly
    vmapped_jax_solve(lower_batched, diag_batched, upper_batched, rhs_batched).block_until_ready()
    jax_vmapped_timer = Timer(
        lambda: vmapped_jax_solve(lower_batched, diag_batched, upper_batched, rhs_batched).block_until_ready()
    )
    iters, t = jax_vmapped_timer.autorange(lambda n, t: print(f"JAX: {t} s (n={n})"))
    jax_vmapped_timeit[batch_size] = t / iters

[dfm]

Yeah, the relevant changes require a new jaxlib (and CUDA plugin) build, not just JAX so you'll need to build jaxlib from source or wait for the next nightly!

[jpbrodrick89]

Thanks for the tip, I will wait for the morning (I can't find a feed informing when nightly's are actually released), but if I need to do this in the future I did try to build from source but think I did something wrong as it still didn't work, I followed these steps:

pyenv virtualenv 3.13.3 newenv
cd ~/.pyenv/versions/3.13.3/envs/newenv/src/jax
python build/build.py build --clang_path /path/to/gpu/llvm/18.1.0/bin/clang++ --wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt
pip install -e .[cuda]

If I did it again, I'd try use git clone instead of pip installing from git but not sure that is actually the issue.