Can I profile jitted CPU code?

[mvsoom]

Profiling jitted CPU code with the jax profiler shows a single call to the compiled code. What other tools are there to accomplish this? (Related to jax issue #774.)

[jakevdp]

There is some information on profiling here: https://jax.readthedocs.io/en/latest/profiling.html

[tgbrooks]

I have the same question as mvsoom and have read and followed the profiling guide (i.e., I'm using jax.profiler.start_profile() and then running my code while collecting a trace with TensorBoard). I still only get a single entry per jitted function with no way of knowing how the time is spent within that function. For example, I see on TensorBoard a slew of "PjitFunction(_objective)" function calls but without any indication of where within my _objective function the time is spent. I tried adding @jax.profiler.annotate_function to my sub functions without any effect.

I am using the CPU only so I haven't used nvprof as suggested in #774 . I understand that profiling optimized compiled code can be tricky in any case but I was hoping to have a little more granularity than I'm getting. Perhaps the instructions in the documentation work better when using a GPU?

[tgbrooks]

I tried again on a PC with a GPU and found that profiling GPU code provides much richer information than when profiling CPU code. Is that known/intentional? It doesn't seem to be documented.

[acertain]

Is it possible to get perf's jit profiling working with jax? a la tensorflow/tensorflow@a45a58b

[joaospinto]

I'm also interested in profiling CPU workflows. Is the current recommended way of doing so to manually add custom tracing events to my code? Some general support for sampling-based profiling (instead of just tracing) would be super useful. CC @jakevdp

[cevheck]

Hi all,

Any updates here? I am trying to optimize a jax.lax.scan function. I see a big discrepancy between the time needed for the first call to the scan and the second one (probably due to caching). I would like to optimize and therefore profile this second (or later) call.

I'm also working on CPU.

[jakevdp]

For doing micro-benchmarks like this, you should keep in mind the effects of jit-compilation, asynchronous dispatch, and data placement: https://docs.jax.dev/en/latest/faq.html#benchmarking-jax-code. Some combination of those is probably leading to the slow first call.