How do you access XLA's flop estimate for a jitted program?

[avital]

For now, here is how you do it:

In [1]: import jax, jax.numpy as jnp
In [2]: m = jax.xla_computation(lambda x, y: x @ y)(jnp.ones((1000, 1000)), jnp.ones((1000,1000))).as_hlo_module()
In [3]: client = jax.lib.xla_bridge.get_backend()
In [4]: jax.lib.xla_client._xla.hlo_module_cost_analysis(client, m)
Out[4]:
{'bytes accessed': 12000008.0,
'bytes accessed operand 0 {}': 4000000.0,
'bytes accessed operand 1 {}': 4000000.0,
'bytes accessed output {}': 4000008.0,
'flops': 2000000000.0,
'optimal_seconds': 0.0}

(This may become a JAX API later)

[avital]

(This code snippet came from @hawkinsp )

[avital]

Let's wait to add more HOWTOs until we switch over the HOWTO system on to the Linen examples

[jonny-so]

I gather this probably isn't intended to be part of the public API, but I noticed on CPU the estimated flop count for jax.numpy.linalg.cholesky seems to grow only O(D^2). I can't find the code to see how these flop estimates are actually computed - is the problem that any calls out to lapack (on cpu at least) are essentially ignored?

[jheek]

The HLO cost analysis should be taken with a grain of salt indeed. External calls are definitely out of scope.

[JunwenBai]

tf.profiler provides the flops breakdown for conv flops, mul flops, etc. Can jax do sth similar, besides just a total flops?

[cgarciae]

JAX now has this cost_analysis API for jit functions that returns flops info:

import jax

@jax.jit
def f(x, y): 
  return 2 * x + y

# Query for cost analysis, print FLOP estimate
x, y = 3, 4
compiled = f.lower(x, y).compile()
flops = compiled.cost_analysis()[0]['flops']

print(flops) # 2.0

Ref: https://jax.readthedocs.io/en/latest/aot.html

[JunwenBai]

@cgarciae Thx for your reply! I think my question was not very clear. I wonder besides having a total flops (which is 2.0 in your example), can jax provide flops breakdowns? In other words, Conv Flops: 30%, MatMul Flops: 20%, ... etc.
The total flops is useful, but the breakdown flops are even more useful. Do you know if jax can support such breakdowns?