Add option to dump intermediate artifacts for jax-triton compilation

[copybara-service[bot]]

Add option to dump intermediate artifacts for jax-triton compilation

Users can set the environment variable JAX_TRITON_DUMP_DIR to dump the intermediate ttir, ttgir, llir and ptx generated during jax-triton compilation. The folder looks something like this:

config                                                                                    multi_head_attention_0d1d2d34d5d6d7c8d9d10d11c12d13d14d15c16d17d18d19c2021d.llir  multi_head_attention_0d1d2d34d5d6d7c8d9d10d11c12d13d14d15c16d17d18d19c2021d.ttgir
multi_head_attention_0d1d2d34d5d6d7c8d9d10d11c12d13d14d15c16d17d18d19c2021d.compile_info  multi_head_attention_0d1d2d34d5d6d7c8d9d10d11c12d13d14d15c16d17d18d19c2021d.ptx   multi_head_attention_0d1d2d34d5d6d7c8d9d10d11c12d13d14d15c16d17d18d19c2021d.ttir

[google-cla[bot]]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.