Beef up AOTAutograd logging with aot_id and input descriptions (#89710)
A few things in this PR, that I found useful while debugging some
recent issues:
- We now allocate an aot_id to each aot_function/aot_module invocation,
and print it whenever we report error messages and graph output
logging. Check the comment for why this sort of thing is useful,
and also why it's different from nth_graph. This number is now
incorporated into aot_graph_name
- I noticed that nth_graph only gets incremented when backwards is
compiled. Because backwards is compiled lazily, this means that
multiple forward graphs would have gotten the same ID! I change
nth_graph to always increment to avoid confusion here.
- I added a simple describe_input function, which makes use of
num_params_buffers to tell the user if the input index they're
looking at is a param/buffer or an input. With the help of
https://github.com/pytorch/pytorch/pull/89709 we could give
even more detailed information about inputs (we could also
easily give detailed information about parameters if we stored
a mapping of index to parameter name, but I didn't need this
when debugging so I'll let someone else add it if they need
it.)
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/89710
Approved by: https://github.com/bdhirsh
