Is jax convolution really slow on Google Colaboratory?

[shailesh1729]

Hi,

Here is my sample run with CPU runtime on Google Colab.

The basic convolution seems way too slow compared to NumPy.

With the GPU runtime (see below), JAX appears to be just about 30-40% faster compared to NumPy. I am assuming that NumPy is still running on CPU and JAX is running on GPU when the runtime is GPU.

Is there something wrong I am doing? Do we need to configure something so that convolution can be comparable to NumPy on CPU and much faster on GPU?

[jakevdp]

Hi - indeed it looks like JAX's convolution is quite a bit slower than numpy's for this case on CPU. I think there are probably two reasons for this:

1. JAX convolution lowers to the ConvWithPadding operation in XLA, which is designed and optimized for the much more complicated batched multi-dimensional convolutions common in convolutional neural networks. I suspect that not much effort has been put into optimizing the simpler case of 1D convolutions.
2. In general, there has been much more focus on optimizing operations in the XLA GPU backend than in the XLA CPU backend, so you occasionally come across examples where CPU execution is less optimized than we would like.

It would be worth opening a bug about this, because I think we could probably do better in this case.

[jakevdp]

Another note on this: it turns out you can get about a 10X speedup on CPU by introducing a size-1 dimension and doing a 2D convolution:

from jax.scipy import signal
import numpy as np
import jax.numpy as jnp
from jax import jit

p = 100
n = 10000
h = jnp.ones(p) / p
x = jnp.ones(n)

jnp_convolve = jit(jnp.convolve)

@jit
def convolve1(x, h):
  return signal.convolve(x, h)

@jit
def convolve2(x, h):
  return signal.convolve(x[None], h[None])[0]

assert jnp.allclose(jnp_convolve(x, h), convolve1(x, h))
assert jnp.allclose(jnp_convolve(x, h), convolve2(x, h))

%timeit np.convolve(x, h)
# 1000 loops, best of 5: 293 µs per loop

%timeit jnp_convolve(x, h)
# 100 loops, best of 5: 16.4 ms per loop

%timeit convolve1(x, h)
# 100 loops, best of 5: 16.3 ms per loop

%timeit convolve2(x, h)
# 1000 loops, best of 5: 1.48 ms per loop

This should probably be fixed in the XLA:CPU implementation, but if that proves difficult we could use this kind of dummy dimension in the jnp.convolve implementation to make things a bit better.

[shailesh1729]

Thank you for the idea. I will try it out and report if it helps. I have been porting PyWavelets and PyLops into JAX.

• You can see the partial PyWavelets port in cr.sparse.wt
• You can see a slightly different approach to implementing linear operators with JAX in cr.sparse.lop

In this implementation, DWT works on 1D convolutions with (increasingly large) quadrature mirror filters. I was struggling in being able to match the performance of PyWavelets on the CPU.

I will try out this idea of dummy dimension to see if I am able to improve performance or not.

In a linear deblurring example, I wrote which uses a 2D convolution operator, a 2D wavelet basis operator, and LSQR and FISTA algorithms for deblurring the image, I found that the major source of lack of performance was the large convolution operator for blurring.

[shailesh1729]

On a MacBook Pro, the convolve2 proposed by you is actually beating NumPy (although close)

%timeit np.convolve(x, h)
# 242 µs ± 718 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit jnp_convolve(x, h)
# 2.66 ms ± 136 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit convolve1(x, h)
# 2.48 ms ± 10.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit convolve2(x, h)
221 µs ± 9.08 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)