Best practices for debugging

[hr0nix]

I'm wondering how I can better understand into what a particular kernel gets compiled. For instance,

• Is there a way to see the intermediate representation built by the compiler? It gets printed when there's an error, but is there a way to see it for a successfully compiled kernel?
• What are the interpret and debug parameters to pallas_call and how do I use them?

[sharadmv]

What are the interpret and debug parameters to pallas_call and how do I use them?

• interpret: when this is True, the kernel is compiled using HLO/XLA, not Triton. Specifically we run a loop over program ids and don't parallelize their execution. This emulates how the kernel will be run using Triton.

• debug: when this is True, we print out the IRs used during compilation

My suggestion for debugging practices are:

• if you run into a weird exception or crash, run with debug=True and upload the resulting IRs and exception/crash in a GH issue and we can investigate
• if you're debugging numerics/correctness, run with interpret=True and use jax.debug.print to inspect values
• if you're debugging performance, use xprof/Nvidia performance tools and/or post a benchmark as a GH issue and we can investigate