JAX matmul (100x100) seems ~16x slower than numpy because no asynchronous behavior?!

[RezaRob]

Here's the code:
https://gist.github.com/RezaRob/eb30041cce7e582ac7228dce2310482c

It seems like JAX isn't using multithreading (not truly asynchronous behavior apparently). I just looked at the 'top' command in Linux.

Jax basically seems about 16x slower than numpy on my ~16-threaded machine when "small" 100x100 matrices are multiplied. But that's not really a tiny computation! Since JAX is asynchronous, shouldn't it basically be using multi-threads for this, just like numpy?

The results are the same whether you use JAX_CPU = True or False.

Can somebody please glance at my code and explain what's going on?

And this is the output:

--------------- JAX very large matmul (10000x10000) --------------------
jax.numpy.matmul._cache_size():  0
time: 1.8283798694610596 
 val:  -50.969013
jax.numpy.matmul._cache_size():  1

--------------- JAX very large matmul (10000x10000) --------------------
jax.numpy.matmul._cache_size():  1
time: 0.29285097122192383 
 val:  -12.026682
jax.numpy.matmul._cache_size():  1

-------- Many small (10000 x 100x100) matmul's placed on CPU using  JAX -------------
TFRT_CPU_0
mm._cache_size():  0
a.shape:  (100, 100, 100)
time: 17.741828203201294 
 val:  1313.3522
mm._cache_size():  1

-------- Many small (10000 x 100x100) matmul's placed on CPU using  JAX -------------
TFRT_CPU_0
mm._cache_size():  1
a.shape:  (100, 100, 100)
time: 17.217226028442383 
 val:  132.08763
mm._cache_size():  1

-------- Many small (10000 x 100x100) matmul's placed on CPU using  (standard)numpy  -------------
a.shape:  (100, 100, 100)
time: 1.0805106163024902 
 val:  1313.3482467798458

-------- Many small (10000 x 100x100) matmul's placed on CPU using  (standard)numpy  -------------
a.shape:  (100, 100, 100)
time: 1.0969195365905762 
 val:  132.08913831901737

Thank you.

[RezaRob]

I'm almost leaning towards retracting the question right now. I do know that vmap and jit can be used in this particular example (vmap is so super cool by the way!!), but of course that wasn't the point of the question. I will leave it in the spirit of perhaps getting some interesting discussion. I'm still interested why JAX doesn't go multi-thread at that scale. It picks up multi-threading at 1000x1000 matrix size and becomes more competitive. I suppose that some of the most interesting cases for multi-threading are hampered by the fact that python itself isn't multi-threaded, but if you need variable-sized arrays for some reason, then that's an obvious case where vmap won't do. Perhaps in practice that's not a concern? The other thing is dynamic graphs. Anyway, I'll leave it here, in case anyone wants to comment. Thanks.

[jakevdp]

Hi - thanks for the question! There's some general information about this kind of question in the FAQ: https://jax.readthedocs.io/en/latest/faq.html#is-jax-faster-than-numpy

In particular, when you run linear algebra operations on CPU, they are lowered to whatever BLAS/LAPACK libraries the package was compiled against, and this statement is true for both JAX and NumPy. If NumPy is significantly faster, then I suspect you're using a numpy installation built against a faster BLAS (for example, numpy installed via conda is often built against MKL).

[RezaRob]

Thanks for your comment. I don't see anything in that (FAQ) link that explains why JAX isn't using multiple cores as effectively as numpy for "smaller" operations. This is a very specific issue. Can you please address it directly? Just using 'top' command in Linux shows a huge difference. Numpy maxes out all cores. Jax, at smaller sizes(100x100) doesn't. Nothing to do with "dispatch overhead."

I'm also guessing it's unrelated to BLAS/library versions. I have standard Ubuntu 22.04, python 3, Jax installed through pip (I think?). Nothing weird! 100x100 matrices aren't exactly "tiny and unimportant." JAX makes better use of cores for 1000x1000 matmul ops. So, I'd appreciate a more specific explanation of the issue. I've provided executable code.

[jakevdp]

I don't see anything in that (FAQ) link that explains why JAX isn't using multiple cores as effectively as numpy for "smaller" operations.

If I understand your question, I think addressed that above: for smaller operations, JAX uses whatever BLAS implementation it is compiled against, just as NumPy uses whatever BLAS implementation it is compiled against. If you're seeing vastly different performance between JAX and NumPy on single CPU matmuls, the first thing I'd look at is differences in how jaxlib and numpy were built, and what BLAS kernels they're calling.

[jakevdp]

Looking more closely at your code, though, I think I spoke too soon: I'd now guess that the bulk of the difference is not coming from matmul, but rather from overhead related to array indexing. When you write something like a[x], in NumPy it's a virtually instantaneous creation of a view of the original buffer. In JAX, outside JIT, it allocates a new buffer and copies the data. Doing this over and over outside of a JIT context is slow, and I think that's the main thing your code demonstrates. (to be clear, though, inside JIT these copies are optimized away, so it would not be an issue).

[jakevdp]

For example, if I modify your code like this to avoid timing __getitem___:

        a_slices = list(a)

        t = time.time()
        b = [mm(a_slices[x],a_slices[y]) for x in range(a.shape[0]) 
                                         for y in range(a.shape[0])]
        if not use_numpy:
          b = jax.block_until_ready(b)
        print("time:",time.time()-t)

and then run it on a Colab CPU runtime, I see the following:

-------- Many small (10000 x 100x100) matmul's placed on CPU using  JAX -------------
TFRT_CPU_0
mm._cache_size():  0
a.shape:  (100, 100, 100)
time: 1.0032134056091309
mm._cache_size():  1

-------- Many small (10000 x 100x100) matmul's placed on CPU using  JAX -------------
TFRT_CPU_0
mm._cache_size():  1
a.shape:  (100, 100, 100)
time: 0.5546126365661621
mm._cache_size():  1

-------- Many small (10000 x 100x100) matmul's placed on CPU using  (standard)numpy  -------------
a.shape:  (100, 100, 100)
time: 0.668367862701416

-------- Many small (10000 x 100x100) matmul's placed on CPU using  (standard)numpy  -------------
a.shape:  (100, 100, 100)
time: 0.6782197952270508

[RezaRob]

That's awesome. Thank you so much. Sorry I missed that. The 'top' output is a distraction. Strangely, despite the fact that 'top' still clearly shows numpy maxing all the cores and JAX apparently NOT doing that, with your suggestion, and after JAX "warm-up"/compilations are done, JAX is actually substantially faster than numpy (more than twice). However, that's only timing the matmul on ready-to-go slices:
https://gist.github.com/RezaRob/83059398ab25471de8df8810c02746cf

So, JAX figures out how to use the hardware without running extra threads. I guess it spends more time finding optimizations, but then performs better. Thank you.