Repository: google/jax

Issue:
Slow compilation when slicing
I'm trying to reduce the compilation time of a big program (currently 30s-60s). Here are three minimal examples of code that might be the reason for some of the overhead.

#### Snippet 1:
 `jacrev` for slicing takes a lot longer than for `jacfwd`. Is this expected? Without slicing this effect doesn't occur.
```python
import time

import jax.numpy as jnp
from jax import jacfwd, jacrev, jit


def f1(x):
    return x


def f2(x):
    return x[1:]


def jitting(f, *args):
    s = time.time()
    f(*args).block_until_ready()
    print(f"{f.__name__} jit time: {time.time() - s}")
    s = time.time()
    f(*args).block_until_ready()
    print(f"{f.__name__} cached time: {time.time() - s}")


x = jnp.ones([1000])

print("jit(f1)")
jitting(jit(f1), x)
print("jit(f2)")
jitting(jit(f2), x)

print("jit(jacrev(f1))")
jitting(jit(jacrev(f1)), x)
print("jit(jacrev(f2))")
jitting(jit(jacrev(f2)), x)

print("jit(jacfwd(f1))")
jitting(jit(jacfwd(f1)), x)
print("jit(jacfwd(f2))")
jitting(jit(jacfwd(f2)), x)
```

```
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
jit(f1)
f1 jit time: 0.0003452301025390625
f1 cached time: 8.082389831542969e-05
jit(f2)
f2 jit time: 0.02623915672302246
f2 cached time: 0.00027489662170410156
jit(jacrev(f1))
jacfun jit time: 0.0028450489044189453
jacfun cached time: 0.00012683868408203125
jit(jacrev(f2))
jacfun jit time: 3.8852109909057617
jacfun cached time: 0.0007808208465576172
jit(jacfwd(f1))
jacfun jit time: 0.0038161277770996094
jacfun cached time: 0.0013599395751953125
jit(jacfwd(f2))
jacfun jit time: 0.08190584182739258
jacfun cached time: 0.001171112060546875
```

#### Snippet 2:
 `vjp_fun` with `vmap` compilation time depends heavily on matrix that is mapped over. Is this expected?

```python
import time

import jax.numpy as jnp
from jax import jit, jvp, vjp, vmap
from matplotlib import pyplot as plt


def f(x):
    return x + x


def jvp_vmap(x, v):
    y, vjp_fun = vjp(f, x)
    J, = vmap(vjp_fun)(v)
    return J


x = jnp.ones([250])
timings = []
for i in range(1, 50):
    V = jnp.ones([i, x.size])
    s = time.time()
    jit(jvp_vmap)(x, V).block_until_ready()
    timings.append(time.time() - s)
plt.plot(timings)
plt.show()
```
![image](https://user-images.githubusercontent.com/40408671/133804843-bbf3062c-fc02-4a0b-ae25-84396a96377f.png)


#### Snippet 3:
Similar for a gauss-newton vector product
```python
import time

import jax.numpy as jnp
from jax import jit, jvp, vjp, vmap
from matplotlib import pyplot as plt

def f1(x):
    return x[1:]


def gnvp_fun(
    x: jnp.ndarray, vs: jnp.ndarray
) -> jnp.ndarray:
    jvp_fun = lambda input: jvp(f1, (x,), (input,))[1]
    y, vjp_fun = vjp(f1, x)
    return vmap(lambda input: vjp_fun(jvp_fun(input))[0], in_axes=1)(vs).T


x = jnp.ones([100])
timings = []
for i in range(1, 50):
    V = jnp.ones([x.size, i])
    s = time.time()
    jit(gnvp_fun)(x, V).block_until_ready()
    timings.append(time.time() - s)
plt.plot(timings)
plt.show()
```
![image](https://user-images.githubusercontent.com/40408671/133804901-8a9db7c0-a73c-4dea-8e3c-df148f393452.png)


Additionally I have some loops of size 2-4 which are unrolled and increase compilation time. When using `jax.vmap` instead the compilation time reduces but the runtime (CPU) increases significantly. Is this common?


Discussion:
Seems like I am not alone to complain that [`vmap` is slower than `lax.scan` loop](https://github.com/google/jax/issues/7449#issuecomment-922076745), maybe `vmap` is not made for doing auto-batching for small batch size?
> maybe vmap is not made for doing auto-batching for small batch size?

I think that's correct: `vmap` does not really have anything to do with auto-batching, it's more about auto-vectorization.
@tetterl thanks for the report. The shape-dependent compilation time is strange, but note that it has nothing to do with vmap or vjp; you can get the same behavior with the unadorned function `f`:
```python
import time

import jax.numpy as jnp
from jax import jit
from matplotlib import pyplot as plt

def f(x):
    return x + x

timings = []
for i in range(1, 50):
    V = jnp.ones([i, 250])
    s = time.time()
    jit(f)(V).block_until_ready()
    timings.append(time.time() - s)
plt.plot(timings)
plt.show()
```
![download](https://user-images.githubusercontent.com/781659/134089939-a815e71d-f520-4814-a2b5-a6c8fab8ddac.png)

This is quite surprising – I'm not certain what might be causing it, but I'll ask around to see if anyone has ideas.
@mattjj points out that it may be due to some kind of shape-dependent loop unrolling within LLVM.
@jakevdp thanks for the info. I guess that won't be a quick fix then.

Do you have any clue why Snippet 1 is that slow? I still haven't found a way to speed up compilation for code where I use such slicing.
I looked into that – it seems the issue is that for `jacrev` on a slice of size `N`, the HLO embeds a literal encoding of an `N x N` matrix. For example:
```python
import jax
import jax.numpy as jnp

def make_hlo(f, optimize=False, metadata=False, platform=None):
  client = jax.lib.xla_bridge.get_backend(platform)
  print_opts = jax.lib.xla_client._xla.HloPrintOptions.short_parsable()
  print_opts.print_metadata = metadata
  def wrapped_fn(*args, **kwargs):
    c = jax.xla_computation(f)(*args, **kwargs)
    if optimize:
      out = client.compile(c).hlo_modules()[0].to_string(print_opts)
    else:
      out = c.as_hlo_module().to_string(print_opts)
    return out.strip()
  return wrapped_fn

f = jax.jacrev(lambda x: x[1:])
x = jnp.arange(10.0)

print(make_hlo(f)(x))
```
```
HloModule xla_computation_jacfun__8.25

primitive_computation_add.9 {
  parameter.10 = f32[] parameter(0)
  parameter.11 = f32[] parameter(1)
  ROOT add.12 = f32[] add(parameter.10, parameter.11)
}

scatter_add_reducer__11.18 {
  parameter.19 = f32[] parameter(0)
  parameter.20 = f32[] parameter(1)
  ROOT add.21 = f32[] add(parameter.19, parameter.20)
}

ENTRY xla_computation_jacfun__8.25 {
  constant.3 = pred[] constant(false)
  parameter.2 = f32[10]{0} parameter(0)
  constant.4 = s32[] constant(1)
  broadcast.5 = s32[1]{0} broadcast(constant.4), dimensions={}
  gather.6 = f32[9]{0} gather(parameter.2, broadcast.5), offset_dims={0}, collapsed_slice_dims={}, start_index_map={0}, index_vector_dim=0, slice_sizes={9}, indices_are_sorted=true
  broadcast.7 = f32[9]{0} broadcast(gather.6), dimensions={0}
  constant.15 = f32[] constant(0)
  broadcast.16 = f32[10]{0} broadcast(constant.15), dimensions={}
  broadcast.17 = f32[9,10]{1,0} broadcast(broadcast.16), dimensions={1}
  constant.1 = f32[9,9]{1,0} constant({ { 1, 0, 0, 0, 0, 0, 0, 0, 0 }, { 0, 1, 0, 0, 0, 0, 0, 0, 0 }, { 0, 0, 1, 0, 0, 0, 0, 0, 0 }, { 0, 0, 0, 1, 0, 0, 0, 0, 0 }, { 0, 0, 0, 0, 1, 0, 0, 0, 0 }, { 0, 0, 0, 0, 0, 1, 0, 0, 0 }, { 0, 0, 0, 0, 0, 0, 1, 0, 0 }, { 0, 0, 0, 0, 0, 0, 0, 1, 0 }, { 0, 0, 0, 0, 0, 0, 0, 0, 1 } })
  constant.8 = f32[] constant(0)
  reduce.13 = f32[9,9]{1,0} reduce(constant.1, constant.8), dimensions={}, to_apply=primitive_computation_add.9
  broadcast.14 = f32[9,9]{1,0} broadcast(reduce.13), dimensions={0,1}
  scatter.22 = f32[9,10]{1,0} scatter(broadcast.17, broadcast.5, broadcast.14), update_window_dims={0,1}, inserted_window_dims={}, scatter_dims_to_operand_dims={1}, index_vector_dim=0, indices_are_sorted=true, unique_indices=true, to_apply=scatter_add_reducer__11.18
  slice.23 = f32[9,10]{1,0} slice(scatter.22), slice={[0:9], [0:10]}
  ROOT tuple.24 = (f32[9,10]{1,0}) tuple(slice.23)
}
```
In your case, `N` is approaching 1000, so the size of the generated HLO becomes quite large. I suspect that is what's causing the slow compilation.

I'm not entirely sure where this large constant would be coming from, but I think this could probably be considered a bug.
@jakevdp Thanks for looking into this and for providing an example how I can look at the HLO by myself :)

Haven't looked at the inner workings of JAX yet but I suspect that's the matrix that is vmapped over (`jnp.eye(10)[1:, 1:]`) in the implementation of `jacrev`: `J, = vmap(vjp_fun, in_axes=0)(jnp.eye(len(y)))`?

When using `jacfwd` we get: 

```
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
HloModule xla_computation_jacfun.13

ENTRY xla_computation_jacfun.13 {
  constant.3 = pred[] constant(false)
  parameter.2 = f32[10]{0} parameter(0)
  constant.4 = s32[] constant(1)
  broadcast.5 = s32[1]{0} broadcast(constant.4), dimensions={}
  gather.6 = f32[9]{0} gather(parameter.2, broadcast.5), offset_dims={0}, collapsed_slice_dims={}, start_index_map={0}, index_vector_dim=0, slice_sizes={9}, indices_are_sorted=true
  broadcast.8 = f32[9]{0} broadcast(gather.6), dimensions={0}
  constant.1 = f32[10,10]{1,0} constant({ { 1, 0, 0, 0, 0, 0, 0, 0, 0, 0 }, { 0, 1, 0, 0, 0, 0, 0, 0, 0, 0 }, { 0, 0, 1, 0, 0, 0, 0, 0, 0, 0 }, { 0, 0, 0, 1, 0, 0, 0, 0, 0, 0 }, { 0, 0, 0, 0, 1, 0, 0, 0, 0, 0 }, { 0, 0, 0, 0, 0, 1, 0, 0, 0, 0 }, { 0, 0, 0, 0, 0, 0, 1, 0, 0, 0 }, { 0, 0, 0, 0, 0, 0, 0, 1, 0, 0 }, { 0, 0, 0, 0, 0, 0, 0, 0, 1, 0 }, { 0, 0, 0, 0, 0, 0, 0, 0, 0, 1 } })
  gather.7 = f32[10,9]{1,0} gather(constant.1, broadcast.5), offset_dims={0,1}, collapsed_slice_dims={}, start_index_map={1}, index_vector_dim=0, slice_sizes={10,9}, indices_are_sorted=true
  broadcast.9 = f32[10,9]{1,0} broadcast(gather.7), dimensions={0,1}
  transpose.10 = f32[9,10]{0,1} transpose(broadcast.9), dimensions={1,0}
  slice.11 = f32[9,10]{1,0} slice(transpose.10), slice={[0:9], [0:10]}
  ROOT tuple.12 = (f32[9,10]{1,0}) tuple(slice.11)
}
```
So the size of the HLO is also getting large with `jacfwd` but compilation time is still fast. So this doesn't seem to be the main issue? But I agree that embedding a constant like an identity matrix doesn't seem to be optimal, especially since knowledge about the fact that it's an identity matrix could help optimizing.

PS: The constant is also embedded without the slicing.
OK, the slicing issue should be fixed by #7983 - thanks for the report!

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1118,9 +1118,8 @@ def hessian(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
 def _std_basis(pytree):
   leaves, _ = tree_flatten(pytree)
   ndim = sum(map(np.size, leaves))
-  # TODO(mattjj): use a symbolic identity matrix here
   dtype = dtypes.result_type(*leaves)
-  flat_basis = np.eye(ndim, dtype=dtype)
+  flat_basis = jax.numpy.eye(ndim, dtype=dtype)
   return _unravel_array_into_pytree(pytree, 1, flat_basis)
 
 def _unravel_array_into_pytree(pytree, axis, arr):


Test patch: