A generic way / tracer to estimate flops of JAX operations?

[adonath]

I was recently looking into estimating the model flops utilization (MFU) for some model I coded in JAX. This gave me the idea that the tracing mechanism in JAX should in principle allow for (at least approximately) estimating the flops for a given model / sequence of operations.

What I would be looking for is something like eval_shape(), but eval_flops(), which takes an example input dtype and shape and returns the estimated flops for execution this function per forward and backward evaluation.

Does something like this already exist? Do you think it would it be possible and useful to interpret a trace this way? Or are the actual flops (after compilation) to different from the naive "execution"?

[yashk2810]

Try using this roofline library: https://cs.opensource.google/jax/jax/+/main:jax/experimental/roofline/roofline.py;l=1?q=roofline&sq=&ss=jax%2Fjax

tests for usage here: https://cs.opensource.google/jax/jax/+/main:tests/roofline_test.py;l=1?q=roofline&ss=jax%2Fjax

[adonath]

Thanks, this looks very promising! I think I will avoid the experimental feature right now and just go with the simpler solution proposed by @carlosgmartin.

[carlosgmartin]

Like cost_analysis?

https://docs.jax.dev/en/latest/aot.html#debug-information-and-analyses-when-available

https://docs.jax.dev/en/latest/jax.stages.html#jax.stages.Lowered.cost_analysis

https://docs.jax.dev/en/latest/jax.stages.html#jax.stages.Compiled.cost_analysis

[adonath]

Thanks @carlosgmartin, this is basically what I was looking for! I guess the backward flops can be included, just by including jax.grad() in the traced function.