2x performance degradation after upgrade to 0.4.9

[bevector]

Hi,
After upgrading from 0.4.8 to 0.4.9 performance drops more than 2 times.
Here is simple code to reproduce:

import jax
import jaxlib
import jax.numpy as jnp
import numpy as np
from timeit import repeat

jax.config.update("jax_platform_name", "cpu")


jax_version = jax.__version__
jaxlib_version = jaxlib.__version__

data = jnp.zeros(shape=[100, 4096])

def iterate():
    for i in range(len(data)):
        data[i]

result = repeat(iterate, repeat=10, number=30)
print(f"{jax_version=} {jaxlib_version=}: {np.mean(result)=:.3} {np.std(result)=:.3e}")

- jax_version='0.4.8' jaxlib_version='0.4.7' np.mean(result)=0.949    np.std(result)=1.217e-02
- jax_version='0.4.9' jaxlib_version='0.4.9' np.mean(result)=2.2      np.std(result)=2.117e-02
- numpy:                                     np.mean(result)=0.000282 np.std(result)=6.424e-06

[jakevdp]

Thanks for the report! I suspect this change is due to #15377

That said, this is trace-time overhead and you should find that runtime performance is improved due to the more efficient lowering path. If you have a real-world use-case where you're iterating over many integer indexing operations, the best approach would be to use vectorized indexing instead so the overhead only happens once. Your example function here doesn't do anything, but if you have an example of the actual function that's causing problems, we may be able to recommend a better approach.

[jakevdp]

That's consistent with the 2x trace time you reported above. If you're interested in runtime benchmarking of JAX code, you can start here: https://jax.readthedocs.io/en/latest/faq.html#benchmarking-jax-code

[bevector]

Thanks, I've seen the link but can't figure out how it is related. Here are two screenshots from the profiler showing the time I am talking about. Am I doing something wrong? The code in my first post is just the simplest reproducible example that highlights the degradation.

1st image from 0.4.8 2nd from 0.4.9

[jakevdp]

Which profiler is this?

[bevector]

Which profiler is this?
https://ui.perfetto.dev/

[jakevdp]

Thanks - so are you concerned about trace time or runtime? Your initial report and your profiler results seem to reflect tracing time. I think we knew in #15377 that the trace time would get slower (it adds additional trace-time logic to decide whether to lower to a more efficient dynamic_slice or a less efficient but more general gather)

The reason I linked to the benchmarking section is because it seemed from your initial question that you were doing microbenchmarks, and I thought it would be helpful to link to resources on how to benchmark runtime as opposed to tracing time.

In your perfetto profiler, have you looked at post-tracing runtime cost, or only tracing cost? If tracing cost dominates your overall computational cost, it may be that you would benefit from writing your code with vectorized operations.