Jax scans are slower than expected

[dionhaefner]

I am implementing the tridiagonal matrix algorithm (TDMA) to solve many tridiagonal systems of the same shape in two sweeps (one forward and one backward pass).

The shape of each diagonal is something like (100_000, 100), and I vectorize over the leading axis, so this should be reasonably efficient.

In pure NumPy, I would do it like this:

def tdma_naive(a, b, c, d):
    """
    Solves many tridiagonal matrix systems with diagonals a, b, c and RHS vectors d.
    """
    assert a.shape == b.shape and a.shape == c.shape and a.shape == d.shape

    n = a.shape[-1]

    for i in range(1, n):
        w = a[..., i] / b[..., i - 1]
        b[..., i] += -w * c[..., i - 1]
        d[..., i] += -w * d[..., i - 1]

    out = np.empty_like(a)
    out[..., -1] = d[..., -1] / b[..., -1]

    for i in range(n - 2, -1, -1):
        out[..., i] = (d[..., i] - c[..., i] * out[..., i + 1]) / b[..., i]

    return out

The JAX implementation looks like this:

def tdma_jax_kernel(a, b, c, d):
    def compute_primes(last_primes, x):
        last_cp, last_dp = last_primes
        a, b, c, d = x

        denom = 1. / (b - a * last_cp)
        cp = c * denom
        dp = (d - a * last_dp) * denom

        new_primes = (cp, dp)
        return new_primes, new_primes

    diags = (a.T, b.T, c.T, d.T)
    init = jnp.zeros((a.shape[1], a.shape[0]))
    _, (cp, dp) = jax.lax.scan(compute_primes, (init, init), diags)

    def backsubstitution(last_x, x):
        cp, dp = x
        new_x = dp - cp * last_x
        return new_x, new_x

    _, sol = jax.lax.scan(backsubstitution, init, (cp[::-1], dp[::-1]))

    return sol[::-1].T

I implemented the algorithm in a handful of backends (including a sloppily written CUDA kernel). You can see the results in this Gist:

https://gist.github.com/dionhaefner/a97ef80b77e02b36e4b248bb97541161

The executive summary is that Jax is 2.5x slower than Numba on CPU, and 3x slower than my amateurish CUDA kernel on GPU (but is on par with Numba here).

If I eliminate the tranposes from the Jax implementation and transpose the inputs beforehand, the implementation gains a factor 2 of performance on GPU, so it would be nice if scan supported scanning over arbitrary axes.

Is this behavior something that is expected, and is there something else I can do to make the Jax implementation more efficient?

[mattjj]

Thanks for raising this, and for the beautifully clear report! Let's make it faster.

I wouldn't be surprised if XLA:TPU eliminates these transposes (though I admit I'm not certain if it does). I pinged the XLA:CPU/GPU folks for some initial thoughts on this issue.

It sounds like we can solve this on the JAX side. Knowing that manual transposing helps gives us a lot of information; as you suggest, we can make scan work over arbitrary axes, though the bookkeeping may get annoying.

Do you need to differentiate through this? I'm guessing for now at least the answer may be "no", in which case perhaps you can try using a fori_loop or while_loop with indexing instead of a scan? That would also make it a bit closer to your NumPy code IIUC. WDYT?

[shoyer]

+1 for trying fori/while_loop here. For gradients, you probably want to use the implicit function theorem rather than differentiating through the solve — take a look at lax.custom_linear_solve.

[dionhaefner]

Well, here goes my fori_loop implementation:

def tdma_jax_kernel_fori(a, b, c, d):
    def compute_primes(i, val):
        b, d = val
        w = a[..., i] / b[..., i - 1]
        b = jax.ops.index_add(b, jax.ops.index[..., i], -w * c[..., i - 1])
        d = jax.ops.index_add(d, jax.ops.index[..., i], -w * d[..., i - 1])
        return (b, d)

    n = a.shape[2]
    
    b, d = jax.lax.fori_loop(1, n, compute_primes, (b, d))

    def backsubstitution(ir, sol):
        i = n - ir - 1
        sol = jax.ops.index_update(sol, jax.ops.index[..., i], (d[..., i] - c[..., i] * sol[..., i+1]) / b[..., i])
        return sol

    sol = jnp.empty_like(a)
    sol = jax.ops.index_update(sol, jax.ops.index[..., -1], d[..., -1] / b[..., -1])
    sol = jax.lax.fori_loop(1, n, backsubstitution, sol)

    return sol


tdma_jax_fori = jax.jit(tdma_jax_kernel_fori, backend='cpu')
tdma_jax_cuda_fori = jax.jit(tdma_jax_kernel_fori, backend='gpu')

Timing CPU:

4.56 s ± 147 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

(~15x slower than before)

Timing GPU:

801 ms ± 68.5 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

(~60x slower than before)

Did I do anything wrong?

[mattjj]

@dionhaefner are the output values the same between the new implementation and the old?

There must be something weird going on because scan lowers to a loop that looks just like that, modulo using lax.dynamic_slice and lax.dynamic_update_slice (through convenience wrappers like lax.dynamic_index_in_dim lax.dynamic_update_index_in_dim) rather than the lax.scatter that jax.ops.index_update generates. We've seen cases where lax.dynamic_slice and lax.dynamic_update_slice are faster but this seems pretty extreme. It's the only thing I can think of, though...

@hawkinsp any guesses?

[dionhaefner]

Yes, outputs are the same (and this is almost verbatim the NumPy implementation).

I double-checked my code, and found a small problem when calling the GPU version of the function. The official slowdown on GPU is now a factor of 20, not 80:

272 ms ± 15.7 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

This is also more in line with the CPU slowdown.

[mattjj]

I'm currently wrestling with the fallout of at least one big bug in #2026, so unfortunately I can't dig into this right away. One quick thing I might try next is switching to using dynamic_index_in_dim / dynaminc_update_index_in_dim to see if that closes the gap with scan. (If it did, I'd consider it a bug in index_update.) Then I'd want to follow up on some ideas the XLA folks suggested, and then consider the longer-term solution of #2509.

[bnwebcode]

I'm looking into lax.scan performance in another context and I came across this issue.

I wanted to repeat the benchmarking in this issue on Google Colab to understand the issue better. I used the code from @dionhaefner original gist. However I find that on Colab, the JAX GPU version of the tridiagonal solver is about 6 times slower than measurements by @dionhaefner . The CPU performance measurements are similar.

My notebook is at: https://colab.research.google.com/drive/1-Vblr_7qQd1SR7QkPdMC3Z0jHyv1OJUB?usp=sharing

Probably just user error (Colab is not suitable for this benchmarking?) or has there been a performance regression? Is there maybe a Jax performance regression test suite?

[dionhaefner]

That doesn’t really surprise me; I tested on a Tesla P100 which is quite a bit beefier than Colab’s K80.

[jakevdp]

Colab will give you one of several GPU types depending on availability; you can run !cat /var/colab/hostname to quickly see the type of GPU backend you were assigned (P100 is one of the options available).

[bnwebcode]

OK, thanks for that information. I reconnected a few times and hit upon a P100 (I guess I should sign-up for colab pro...), and I measure 64ms, so around 4.5 times slower than in @dionhaefner's original gist. There does seem to be a discrepancy but I'm not sure if it isn't user error on my part or an actual performance regression.

[Justin-Tan]

Observing similar behaviour with jax.lax.scan when trying to implement the logarithm of the cumulative sum of exponentials using a prefix scan. Here's a MWI:

import jax
import jax.numpy as jnp
from jax import random
from jax import jit, 
from jax.scipy.special import logsumexp

import os, time

def lax_logcumsumexp(x):
    def _logaddexpcarry(carry, x):
        out = jnp.logaddexp(carry, x)
        return out, out
    cumsum_final, _logcumsumexp = jax.lax.scan(_logaddexpcarry, init=-jnp.inf, xs=x)
    return _logcumsumexp

def logcumsumexp(x):
    x_max = jnp.sort(x)[-1] 
    return x_max + jnp.log(jnp.cumsum(jnp.exp(x-x_max)))

# Running below block multiple times - not just first JITing
test_big = random.normal(random.PRNGKey(int(time.time())), (100000,))
%time jit(lax_logcumsumexp)(test_big).block_until_ready()
# About 8-9 ms on CPU, 1.5-2s on GPU
%time jit(logcumsumexp)(test_big).block_until_ready()
# About 75-90 ms on CPU, 600-700 µs on GPU

It's strange how lax.scan is significantly faster on CPU to GPU even with jit - I'm not sure if I'm doing something verboten here. If it matters, I was using a P100 for the GPU (non-Colab). I also tested it in Colab, with a similar Slowdown.

[shoyer]

XLA GPU currently always executes dynamic control flow on the CPU. So small loop iterations (like what you have here) end up much slower, due to the need to frequently synchronize between the CPU/GPU.

You can find a similar example of this sort of slow down in #3076

[clemisch]

I'm sorry to chime in on this old issue; maybe this is better suited as a new one.

I was under the impression that using jax.lax.scan eliminates "dynamic" control flow because everything is known at compile time. In this repro, what has to be transfered to CPU?