JAX/JIT vs Std Numpy performance: where I am wrong?

Question

Here a simple exercise with a Simpson integration code that I have cooked to accept several functions to integrate over a set of bounds

import numpy as np
def simps(f, a, b, N):
    #N should be even
    dx = (b - a) / N
    x = np.linspace(a, b, N + 1)
    y = f(x)
    w = np.ones_like(y)
    w[2:-1:2] = 2.
    w[1::2]   = 4.
    S = dx / 3 * np.einsum("i...,i...",w,y)
    return S

def funcN(x):
    return np.stack([x**(i/10) * np.exp(-x) for i in range(200)],axis=1)

a = np.arange(0,10,0.1)
b = a+0.05

I am on a CPU device, then I get a 200 x 100 numbers array corresponding to the Int(f_i, a_j,b_j) i:0-199 and j:0-99

%timeit simps(funcN,a,b, 512)

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

Now consider the following JAX/JIT version

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap
from functools import partial
from jax.config import config
config.update("jax_enable_x64", True)   #numpy by default is in double precision

@partial(jit, static_argnums=(0,3))
def jax_simps(f, a,b, N):
    dx = (b - a) / N
    x = jnp.linspace(a, b, N + 1)
    y = f(x)
    w = jnp.ones_like(y)
    w = w.at[2:-1:2].set(2.)
    w = w.at[1::2].set(4.)
    S = dx / 3. * jnp.einsum('i...,i...',w,y)
    return S

@jit
def jax_funcN(x):
    return jnp.stack([x**(i/10) * jnp.exp(-x) for i in range(200)],axis=1)

ja = jnp.arange(0,10,0.1)
jb = ja+0.05

#warm up
jax_simps(jax_funcN,ja,jb, 512).block_until_ready() 

%timeit jax_simps(jax_funcN,ja,jb, 512).block_until_ready()

I have verified that the two code (pure Numpy and JAX/JIT) give the same results as the maximum relative error is of the order of 8. 10^-16.

Now, I got the following timing 933 ms ± 51.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

which is very close to the pure Numpy. Does I have make a very efficient pure Numpy code by chance??? or does I have coded JAX/JIT in a wrong way ?

(nb. using Google collab K80 GPU the timing of JAX/JIT drop to 7.19 ms per loop, keeping the pure Numpy at the level of 1s/loop)

user_2937831 · Answer

From your numbers, it looks like JAX JIT gives a 20% speedup over NumPy on CPU. For CPU execution, NumPy is already pretty optimal: leaving things like autodiff aside, for short sequences of NumPy-like operations JAX's main advantage on CPU is XLA's ability to fuse operations to avoid allocation of temporary arrays for intermediate results, and for this relatively short sequence of operations, it looks like that only buys about 20% improvement.

Now, JAX has other advantages, including autodiff, batching, and (as you mentioned) the ability to lower to accelerators without changing the code. But for execution of a short sequence of vectorized operations on CPU, you often can't do substantially better than NumPy alone.

As a side note: you can speed up both the NumPy and JAX versions by 40-50% by replacing stack with broadcasted operations; for example:

def funcN(x):
  x = x[:, None, :]
  i = np.arange(200)[:, None]
  return x**(i/10) * np.exp(-x)

JAX/JIT vs Std Numpy performance: where I am wrong?

1 Answers