How to print torch.compile'd graph?

torch.compile
1

The FX graphs can be printed with .code / .graph: torch.fx — PyTorch 2.0 documentation

Can torch.compile’d graphs be inspected (in order to see if there’re any graph breaks)? Should I set TORCH_COMPILE_DEBUG=1?

2

Yup that’s the best flag it’ll show you the fx and inductor IRs

3

Is it possible to get the fx IR by accessing some compiled model attribute? I also saw somehwere some _dynamo.explain(...) method, but not sure if it will print the IR…

4

This tutorial has some good examples of how to get the relevant fx and aten IR from a torch program Google Colab

As far as getting the actual IRs that inductor they should be printed here if running with debug flag pytorch/compile_fx.py at main · pytorch/pytorch · GitHub

5

Thank you for your answers! Btw how does one print Triton compiled code? Basically trying to understand what got fused in the graph and what did not…

6

afk now but last I remember TORCH_COMPILE_DEBUG=1 python train.py should spew out a bunch of logs, look for a line that looks like tmp/inductor/somelongstringsetc if you open that folder you should be able to inspect the actual generated triton code

Also note that you can figure out what got fused or not by looking at the inductor IR pre and post fusion again using the same debug flag but I personally found the triton code to be more readable

And since you’re probably curious you can also inspect the actual PTX code as well triton-ptx.py · GitHub