jnp.split: use dynamic rather than static slices for speed

[jakevdp]

Fixes #12999. See also #9445

This is very similar to the issue & fix in #12219

main branch:

In [1]: import jax.numpy as jnp 
   ...: x = jnp.ones(5000) 
   ...: %time _ = jnp.split(x, jnp.arange(0, 5000, 10))                                                                   
CPU times: user 4.77 s, sys: 118 ms, total: 4.89 s
Wall time: 4.96 s

this branch:

In [1]: import jax.numpy as jnp 
   ...: x = jnp.ones(5000) 
   ...: %time _ = jnp.split(x, jnp.arange(0, 5000, 10))                                                                   
CPU times: user 342 ms, sys: 8.52 ms, total: 351 ms
Wall time: 360 ms

[froystig]

Do we want benchmarks here, like those you added in #12219?

Does this change affect performance for a small number of splits?

[jakevdp]

Good question - for repeated operations after the initial run, we get something like this:

main branch:

In [2]: %timeit jax.block_until_ready(jnp.split(x, jnp.arange(0, 5000, 10)))                                          
40.5 ms ± 4.22 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

this branch:

In [2]: %timeit jax.block_until_ready(jnp.split(x, jnp.arange(0, 5000, 10)))                                          
116 ms ± 4.01 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

So a slowdown of ~3x on repeated calls is the tradeoff for saving ~150x on the initial call.

What do you think?

[jakevdp]

Another option here: we could replace this with

return [ary[int(start): int(end)] for start, end in zip(split_indices[:-1], split_indices[1:])]

Due to #12219, this would now use dynamic rather than static slices and result in the same performance characteristics. I kind of like the idea of delegating the performance question to existing code. What do you think?

[froystig]

From conversation with @hawkinsp: we may want to profile in order to understand why the dynamic slice approach is 3x slower, since it isn't clear that it ought to be (?)

[KeAWang]

Why is dynamic_slice faster than static slice in the first place?

[jakevdp]

Why is dynamic_slice faster than static slice in the first place?

dynamic_slice is slower than static slice if you call it once, because it does not specialize its code on index values, and so at runtime XLA has to perform some logic regarding the value of the start index. But static slice is specialized on static start indices, so each call with different start indices incurs a small overhead in XLA at compile time, and when you accumulate this small overhead thousands of times, it is slower than dynamic_slice, which does not have such overhead because it is not specialized on the index values.

[KeAWang]

Thanks! That explains the performance issues I've been trying to debug with slice vs dynamic_slice vs numpy slice.

[jakevdp]

I can no longer reproduce these performance issues, despite this still lowering to static slice.