list of JAX/XLA compiler optimizations

[jeremiecoullon]

Hello!

I regularly get surprised by how JAX (or perhaps rather XLA) seems to do some clever optimizations. For example if you write a function that multiplies a matrix by the inverse of another matrix, the resulting jitted function is as fast as if you hard-code the numerical trick of using a Choleksy factorization. So I'm guessing XLA does that optimization automatically.

I've also noticed that XLA seems to recognise when a term is computed many times in a function; as a result the compiler seems to only compute it once.

However, I can't find a list of these compiler optimizations on the JAX or XLA website. Is this documented somewhere?

[PhilipVinc]

I do agree with you that it would be interesting to know all optimisations performed by Jax/XLA. However it is my understanding that there are different levels where those optimisations happen.
Below is some of my understanding on the XLA compiler toolchain, but I'd also be interested to hear more on it.

For example if you write a function that multiplies a matrix by the inverse of another matrix

That's most likely not a compiler optimisation but just the fact that the inverse of a matrix A is computed by calling lax.solve(A,identity) which in turn uses a LU decomposition and a triangular solve. I call this lowering (of an operation to more fundamental_primitives_).
This does not hit the same logic as the cholesky decomposition, but its performance characteristics are very similar (and XLA might end up generating the same code at the end), which would explain your observed behaviour. So this is simply a consequence of how the inverse of a matrix is lowered to the Jax IR.

I've also noticed that XLA seems to recognise when a term is computed many times in a function; as a result the compiler seems to only compute it once.

In my experience CSE (Common subexpression elimination) is the largest XLA-level optimisation. I think it can eliminate 100% of repeated terms within an expression boundaries (boundaries are for example jax.remat or lax.fori/while loops.

I think XLA might also performs some mathematical transformations like promoting operations like A*v_1, A*v_2 to A*[v1,v2] to leverage better cache locality but I'd be curious to know what heuristics is used to decide when to do this promotion.
I'd also be curious to know if XLA can make some optimisations like replacing ((A1*A2)*v) to the more efficient A1*(A2*v)...

Once XLA is done, I believe it will output LLVM IR on which LLVM will perform several additional optimisations. Probably another CSE pass among them. But since LLVM passes can be configured, it might be interesting to know which passes are enabled and which are disabled...

[soraros]

According to this, it seems JAX doesn't optimise matmul orderings.

And I noticed there is a config flag jax_disable_most_optimizations with which one can disable the optimizations.

[jeremiecoullon]

Thanks for your reply! this was very useful